diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements.in b/ci/official/requirements_updater/numpy1_requirements/requirements.in
index 51461aada6f659..5eefeed6e2cc35 100644
--- a/ci/official/requirements_updater/numpy1_requirements/requirements.in
+++ b/ci/official/requirements_updater/numpy1_requirements/requirements.in
@@ -31,7 +31,7 @@ portpicker == 1.6.0
 scipy >= 1.13.0, < 1.15.0
 requests >= 2.33.0
 packaging==23.2
-setuptools==78.1.1
+setuptools==83.0.0
 jax==0.4.7
 zstandard==0.23.0
 # The dependencies below are needed for TF wheel testing.
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 137bc9e338a4e4..e8238d6e0c9d98 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
@@ -953,9 +953,9 @@ zstandard==0.25.0 \
     # via -r ci/official/requirements_updater/requirements.in
 
 # The following packages are considered to be unsafe in a requirements file:
-setuptools==78.1.1 \
-    --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
-    --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
+setuptools==83.0.0 \
+    --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \
+    --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   tb-nightly
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 8b37e5bd8b7bde..9237c82440f566 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
@@ -952,9 +952,9 @@ zstandard==0.25.0 \
     # via -r ci/official/requirements_updater/requirements.in
 
 # The following packages are considered to be unsafe in a requirements file:
-setuptools==78.1.1 \
-    --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
-    --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
+setuptools==83.0.0 \
+    --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \
+    --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   tb-nightly
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 7c0230a0cbe19b..12559bc05817da 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
@@ -952,9 +952,9 @@ zstandard==0.25.0 \
     # via -r ci/official/requirements_updater/requirements.in
 
 # The following packages are considered to be unsafe in a requirements file:
-setuptools==78.1.1 \
-    --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
-    --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
+setuptools==83.0.0 \
+    --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \
+    --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   tb-nightly
diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.cc b/tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.cc
index 17218d4607d5a2..0f5f6dbd4dfa96 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.cc
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.cc
@@ -237,7 +237,7 @@ absl::Status LowerHlotoLoops(mlir::ModuleOp module,
       mlir::kernel_gen::transforms::CreateBufferReusePass());
   // Approximate Tanh using standard operations.
   pm.addNestedPass<FuncOp>(
-      ::mlir::mhlo::createLegalizeTrigonometricToApproximationPass());
+      ::mlir::mhlo::createLegalizeTanhToApproximationPass());
   // Transform the Linalg ops inside of the loop nest into parallel loops.
   pm.addNestedPass<FuncOp>(::mlir::createConvertLinalgToParallelLoopsPass());
 
@@ -295,7 +295,7 @@ absl::Status LowerLoopsToGPU(mlir::ModuleOp module, bool index_64bit,
   pm.addNestedPass<FuncOp>(mlir::bufferization::createPromoteBuffersToStackPass(
       [](Value alloc) { return IsSmallAlloc(alloc); }));
   // Free all temporaries,
-  pm.addNestedPass<FuncOp>(mlir::deallocation::createBufferDeallocationPass());
+  pm.addNestedPass<FuncOp>(mlir::deallocation::createBufferDeallocation());
   pm.addPass(mlir::createCanonicalizerPass());
   pm.addNestedPass<FuncOp>(::mlir::createConvertLinalgToLoopsPass());
 
@@ -359,9 +359,11 @@ absl::Status LowerKernelBodiesToLowLevelIr(mlir::ModuleOp module,
   auto& kernelPm = pm.nest<::mlir::gpu::GPUModuleOp>();
   kernelPm.addPass(::mlir::createSCFToControlFlowPass());
 #if TENSORFLOW_USE_ROCM
-  kernelPm.addPass(mlir::createGpuKernelToRocdlPass(architecture));
+  mlir::GpuKernelToROCDLPassOptions options;
+  options.chipset = architecture;
+  kernelPm.addPass(mlir::createGpuKernelToROCDLPass(options));
 #elif GOOGLE_CUDA
-  kernelPm.addPass(mlir::createGpuKernelToNvvmPass());
+  kernelPm.addPass(mlir::createGpuKernelToNVVMPass());
   kernelPm.addPass(mlir::NVVM::createNVVMOptimizeForTargetPass());
 #endif
   // Remove all location information to prevent a debug build.
diff --git a/tensorflow/core/debug/debug_io_utils.cc b/tensorflow/core/debug/debug_io_utils.cc
index 92ed75e02a9df9..0f77a9529742aa 100644
--- a/tensorflow/core/debug/debug_io_utils.cc
+++ b/tensorflow/core/debug/debug_io_utils.cc
@@ -694,6 +694,14 @@ absl::Status DebugFileIO::DumpTensorToEventFile(
 }
 
 absl::Status DebugFileIO::RecursiveCreateDir(Env* env, const std::string& dir) {
+  if (dir.empty()) {
+    // io::Dirname() returns "" for any path with no '/' component; this is
+    // a fixed point (Dirname("") == ""). It means there is no parent
+    // directory left to create: any remaining path is relative to the
+    // current working directory. Stop here instead of recursing forever.
+    return absl::OkStatus();
+  }
+
   if (env->FileExists(dir).ok() && env->IsDirectory(dir).ok()) {
     // The path already exists as a directory. Return OK right away.
     return absl::OkStatus();
diff --git a/tensorflow/core/debug/debug_io_utils_test.cc b/tensorflow/core/debug/debug_io_utils_test.cc
index fde63f53331cf1..e5355ab4aeb8e8 100644
--- a/tensorflow/core/debug/debug_io_utils_test.cc
+++ b/tensorflow/core/debug/debug_io_utils_test.cc
@@ -15,10 +15,17 @@ limitations under the License.
 
 #include "tensorflow/core/debug/debug_io_utils.h"
 
+#include <cstdio>
 #include <cstdlib>
 #include <memory>
 #include <unordered_set>
 
+#if defined(PLATFORM_WINDOWS)
+#include <direct.h>
+#else
+#include <unistd.h>
+#endif
+
 #include "absl/synchronization/notification.h"
 #include "tensorflow/core/debug/debug_callback_registry.h"
 #include "tensorflow/core/debug/debug_node_key.h"
@@ -32,11 +39,41 @@ limitations under the License.
 #include "tensorflow/core/lib/io/path.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/util/event.pb.h"
 
 namespace tensorflow {
 namespace {
 
+#if defined(PLATFORM_WINDOWS)
+#define TFDBG_GETCWD _getcwd
+#define TFDBG_CHDIR _chdir
+#else
+#define TFDBG_GETCWD getcwd
+#define TFDBG_CHDIR chdir
+#endif
+
+// Temporarily changes the process's current working directory for the
+// lifetime of this object, then restores the original directory. Used by
+// tests that must exercise a bare relative (slash-less) path argument, which
+// otherwise depends on whatever the test runner's ambient CWD happens to be.
+class ScopedChdir {
+ public:
+  explicit ScopedChdir(const std::string& new_dir) {
+    char buf[FILENAME_MAX];
+    CHECK(TFDBG_GETCWD(buf, sizeof(buf)) != nullptr);
+    old_dir_ = buf;
+    CHECK_EQ(0, TFDBG_CHDIR(new_dir.c_str()));
+  }
+  ~ScopedChdir() { (void)TFDBG_CHDIR(old_dir_.c_str()); }
+
+ private:
+  std::string old_dir_;
+};
+
+#undef TFDBG_GETCWD
+#undef TFDBG_CHDIR
+
 class DebugIOUtilsTest : public ::testing::Test {
  public:
   void Initialize() {
@@ -151,6 +188,60 @@ TEST_F(DebugIOUtilsTest, DumpFloatTensorToFileSunnyDay) {
   ASSERT_EQ(0, undeleted_dirs);
 }
 
+TEST_F(DebugIOUtilsTest, DumpTensorToDirWithRelativeDumpRootSunnyDay) {
+  Initialize();
+  // Run inside a directory guaranteed to be writable, so the relative-path
+  // dump below doesn't depend on whatever the test runner's ambient current
+  // working directory happens to be.
+  ScopedChdir scoped_chdir(testing::TmpDir());
+
+  // A relative directory with no path separator at all. io::Dirname() on a
+  // path like this returns "", and DebugFileIO::RecursiveCreateDir used to
+  // recurse on that empty result forever instead of treating it as "no
+  // parent directory left to create" (see GitHub issue #123114).
+  const std::string test_dir = "tfdbg_relative_dump_root_test";
+  if (env_->FileExists(test_dir).ok()) {
+    int64_t undeleted_files = 0;
+    int64_t undeleted_dirs = 0;
+    TF_ASSERT_OK(
+        env_->DeleteRecursively(test_dir, &undeleted_files, &undeleted_dirs));
+  }
+
+  const uint64_t wall_time = env_->NowMicros();
+  const DebugNodeKey kDebugNodeKey("/job:localhost/replica:0/task:0/cpu:0",
+                                   "foo/bar/qux/tensor_a", 0, "DebugIdentity");
+
+  std::string dump_file_path;
+  TF_ASSERT_OK(DebugFileIO::DumpTensorToDir(
+      kDebugNodeKey, *tensor_a_, wall_time, test_dir, &dump_file_path));
+
+  // Read the file into a Event proto.
+  Event event;
+  TF_ASSERT_OK(ReadEventFromFile(dump_file_path, &event));
+
+  ASSERT_GE(wall_time, event.wall_time());
+  ASSERT_EQ(1, event.summary().value().size());
+  ASSERT_EQ(kDebugNodeKey.debug_node_name,
+            event.summary().value(0).node_name());
+
+  Tensor a_prime(DT_FLOAT);
+  ASSERT_TRUE(a_prime.FromProto(event.summary().value(0).tensor()));
+
+  // Verify tensor shape and value.
+  ASSERT_EQ(tensor_a_->shape(), a_prime.shape());
+  for (int i = 0; i < a_prime.flat<float>().size(); ++i) {
+    ASSERT_EQ(tensor_a_->flat<float>()(i), a_prime.flat<float>()(i));
+  }
+
+  // Tear down temporary file and directories.
+  int64_t undeleted_files = 0;
+  int64_t undeleted_dirs = 0;
+  TF_ASSERT_OK(
+      env_->DeleteRecursively(test_dir, &undeleted_files, &undeleted_dirs));
+  ASSERT_EQ(0, undeleted_files);
+  ASSERT_EQ(0, undeleted_dirs);
+}
+
 TEST_F(DebugIOUtilsTest, DumpStringTensorToFileSunnyDay) {
   Initialize();
 
diff --git a/tensorflow/core/util/example_proto_fast_parsing.cc b/tensorflow/core/util/example_proto_fast_parsing.cc
index c9384af72346ae..5776007faf3572 100644
--- a/tensorflow/core/util/example_proto_fast_parsing.cc
+++ b/tensorflow/core/util/example_proto_fast_parsing.cc
@@ -2372,7 +2372,15 @@ absl::Status ParseContextDenseFeatures(
             reinterpret_cast<const uint8_t*>(feature_proto.data()),
             feature_proto.size());
         EnableAliasing(&stream);
-        num_elements += ParseFeature(dtype, &stream, &out, &out_offset);
+        const int num_added = ParseFeature(dtype, &stream, &out, &out_offset);
+        if (num_added < 0) {
+          // This should be unreachable -- we already scanned the feature in
+          // GetContextFeatureLengths, and it hasn't changed since then.
+          return absl::InvalidArgumentError(
+              absl::StrCat("Error in context feature ", c.feature_name,
+                           " in example ", ExampleName(example_names, e)));
+        }
+        num_elements += num_added;
       }
       if (num_elements != data_max_elements) {
         return absl::InvalidArgumentError(
@@ -2422,10 +2430,17 @@ absl::Status ParseContextSparseFeatures(
           reinterpret_cast<const uint8_t*>(feature_proto.data()),
           feature_proto.size());
       EnableAliasing(&stream);
-      size_t num_added =
+      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
+        // GetContextFeatureLengths, and it hasn't changed since then.
+        return absl::InvalidArgumentError(
+            absl::StrCat("Error in context feature ", c.feature_name,
+                         " in example ", ExampleName(example_names, e)));
+      }
       num_elements += num_added;
-      max_num_cols = std::max(max_num_cols, num_added);
+      max_num_cols = std::max(max_num_cols, static_cast<size_t>(num_added));
       for (int i = 0; i < num_added; i++) {
         if (is_batch) *out_indices++ = e;
         *out_indices++ = i;
@@ -2492,8 +2507,15 @@ absl::Status ParseContextRaggedFeatures(
             reinterpret_cast<const uint8_t*>(feature_proto.data()),
             feature_proto.size());
         EnableAliasing(&stream);
-        size_t num_added =
+        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
+          // GetContextFeatureLengths, and it hasn't changed since then.
+          return absl::InvalidArgumentError(
+              absl::StrCat("Error in context feature ", c.feature_name,
+                           " in example ", ExampleName(example_names, e)));
+        }
         split += num_added;
       }
       if (int32_splits) {
diff --git a/tensorflow/python/kernel_tests/nn_ops/atrous_conv2d_test.py b/tensorflow/python/kernel_tests/nn_ops/atrous_conv2d_test.py
index 9ce2873ecdd0ee..9068677a086380 100644
--- a/tensorflow/python/kernel_tests/nn_ops/atrous_conv2d_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/atrous_conv2d_test.py
@@ -169,6 +169,24 @@ def testAtrousConv2DInvalid(self):
             padding="SAME")
         self.evaluate(op)
 
+  def testAtrousConv2DInvalidInputRank(self):
+    with self.assertRaisesRegex(ValueError, "rank 4"):
+      nn_ops.atrous_conv2d(
+          value=np.ones([10]),
+          filters=np.ones([1, 1, 1, 1]),
+          rate=1,
+          padding="SAME",
+      )
+
+  def testAtrousConv2DInvalidFilterRank(self):
+    with self.assertRaisesRegex(ValueError, "rank 4"):
+      nn_ops.atrous_conv2d(
+          value=np.ones([1, 1, 1, 1]),
+          filters=np.ones([10]),
+          rate=1,
+          padding="SAME",
+      )
+
 
 class AtrousConv2DTransposeTest(test.TestCase):
 
@@ -250,7 +268,7 @@ def testAtrousDepthwiseConv2DForward(self):
   def testInvalidInputRank(self):
     value = array_ops.zeros([10], dtype=dtypes.float32)
     filters = array_ops.zeros([5, 5, 1, 8], dtype=dtypes.float32)
-    with self.assertRaisesRegex(ValueError, "rank at least 3"):
+    with self.assertRaisesRegex(ValueError, "rank 4"):
       nn_ops.atrous_conv2d(value, filters, rate=1, padding="VALID")
 
 
diff --git a/tensorflow/python/kernel_tests/signal/fft_ops_test.py b/tensorflow/python/kernel_tests/signal/fft_ops_test.py
index ed338f0e77acc6..e7a38dc47a9025 100644
--- a/tensorflow/python/kernel_tests/signal/fft_ops_test.py
+++ b/tensorflow/python/kernel_tests/signal/fft_ops_test.py
@@ -24,6 +24,7 @@
 from tensorflow.python.eager import context
 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
 from tensorflow.python.ops import gen_spectral_ops
@@ -470,6 +471,16 @@ def test_grad_random(self, rank, extra_dims, np_type):
     self._check_grad_complex(self._tf_ifft_for_rank(rank), re, im,
                              rtol=tol, atol=tol)
 
+  def testNDOpsAcceptNonTensorInput(self):
+    # The axes used to be inferred from the input before it was converted to
+    # a tensor, so a plain Python list raised AttributeError rather than
+    # being accepted like the equivalent tensor. Build the ops in a graph so
+    # this does not depend on an FFTND/IFFTND kernel being registered.
+    x = [[1.0, 2.0], [3.0, 4.0]]
+    with ops.Graph().as_default():
+      self.assertIsNotNone(fft_ops.fftnd(x))
+      self.assertIsNotNone(fft_ops.ifftnd(x))
+
 
 @test_util.run_all_in_graph_and_eager_modes
 class RFFTOpsTest(BaseFFTOpsTest, parameterized.TestCase):
@@ -523,6 +534,16 @@ def testIRFFTWrongLengthRaisesBeforeNegativeLastLength(self):
     with self.assertRaisesRegex(ValueError, "Dimension must be 2 but is 1"):
       fft_ops.irfftnd(x, fft_length=[-5])
 
+  def testNDOpsAcceptNonTensorInput(self):
+    # The axes used to be inferred from the input before it was converted to
+    # a tensor, so a plain Python list raised AttributeError rather than
+    # being accepted like the equivalent tensor. Build the ops in a graph so
+    # this does not depend on an RFFTND/IRFFTND kernel being registered.
+    x = [[1.0, 2.0], [3.0, 4.0]]
+    with ops.Graph().as_default():
+      self.assertIsNotNone(fft_ops.rfftnd(x))
+      self.assertIsNotNone(fft_ops.irfftnd(x))
+
   def _np_fftn(self, x, fft_length=None, axes=None, norm=None):
     return np.fft.rfftn(x, s=fft_length, axes=axes, norm=norm)
 
diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD
index c0fdb0eee4297f..fea13cbd711535 100644
--- a/tensorflow/python/ops/BUILD
+++ b/tensorflow/python/ops/BUILD
@@ -3844,6 +3844,7 @@ py_test(
     deps = [
         ":array_ops",
         ":tensor_array_ops",
+        ":variables",
         "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:constant_op",
         "//tensorflow/python/framework:dtypes",
diff --git a/tensorflow/python/ops/nn_ops.py b/tensorflow/python/ops/nn_ops.py
index 0dcf653ef1811a..149668d9766a3e 100644
--- a/tensorflow/python/ops/nn_ops.py
+++ b/tensorflow/python/ops/nn_ops.py
@@ -1912,8 +1912,9 @@ def atrous_conv2d(value, filters, rate, padding, name=None):
         [batch, height, width, out_channels].
 
   Raises:
-    ValueError: If input/output depth does not match `filters`' shape, or if
-      padding is other than `'VALID'` or `'SAME'`.
+    ValueError: If input/output depth does not match `filters`' shape, if
+      `value` or `filters` is not rank 4, or if padding is other than
+      `'VALID'` or `'SAME'`.
 
   References:
     Multi-Scale Context Aggregation by Dilated Convolutions:
@@ -1932,6 +1933,10 @@ def atrous_conv2d(value, filters, rate, padding, name=None):
       (https://ieeexplore.ieee.org/abstract/document/6738831)
       ([pdf](https://arxiv.org/pdf/1302.1700.pdf))
   """
+  value = ops.convert_to_tensor(value, name="value")
+  filters = ops.convert_to_tensor(filters, name="filters")
+  value.shape.assert_has_rank(4)
+  filters.shape.assert_has_rank(4)
   return convolution(
       input=value,
       filter=filters,
diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index a2ba023e44cf54..42caa5e7d514fd 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -1022,14 +1022,35 @@ def moveaxis(a, source, destination):  # pylint: disable=missing-docstring
   a_rank = np_utils._maybe_static(array_ops.rank(a))  # pylint: disable=protected-access
 
   def _correct_axis(axis, rank):
-    if axis < 0:
-      return axis + rank
-    return axis
+    if isinstance(axis, (int, np.integer)) and isinstance(
+        rank, (int, np.integer)
+    ):
+      axis = int(axis)
+      rank = int(rank)
+      if not (-rank <= axis < rank):
+        raise ValueError(
+            f'Argument `axis` (received axis={axis}) is out of bounds '
+            f'for input {a} of rank {rank}.'
+        )
+      return axis + rank if axis < 0 else axis
+    rank_t = ops.convert_to_tensor(rank)
+    axis_t = ops.convert_to_tensor(axis)
+    control_flow_assert.Assert(
+        math_ops.reduce_all(
+            math_ops.logical_and(axis_t >= -rank_t, axis_t < rank_t)
+        ),
+        ['axis', axis_t, 'is out of bounds for array of dimension', rank_t],
+    )
+    return array_ops.where_v2(axis_t < 0, np_utils.add(axis_t, rank_t), axis_t)
 
   source = tuple(_correct_axis(axis, a_rank) for axis in source)
   destination = tuple(_correct_axis(axis, a_rank) for axis in destination)
 
-  if a.shape.rank is not None:
+  if (
+      isinstance(a_rank, (int, np.integer))
+      and builtins.all(isinstance(x, (int, np.integer)) for x in source)
+      and builtins.all(isinstance(x, (int, np.integer)) for x in destination)
+  ):
     perm = [i for i in range(a_rank) if i not in source]
     for dest, src in sorted(zip(destination, source)):
       assert dest <= len(perm)
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 e9c6327c3c459c..1045b154b4d36e 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1314,6 +1314,11 @@ def _test(*args):
     _test(a, tuple(range(6)), tuple(reversed(range(6))))
     _test(a, (), ())
 
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.moveaxis(a, -8, 0)
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.moveaxis(a, 0, 8)
+
   def testFlip(self):
     np.random.seed(0)
     random_seed.set_seed(0)
diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index 27a0c69ba8de52..b31cc630874afb 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -1169,7 +1169,7 @@ def f(a):
       )
     if n < 0:
       raise ValueError(
-          f'Argument `order` must be a non-negative integer. Received: axis={n}'
+          f'Argument `n` must be a non-negative integer. Received: n={n}'
       )
     slice1 = [slice(None)] * nd
     slice2 = [slice(None)] * nd
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 f6eac507899ce2..cf7e185cfeaf4e 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -423,6 +423,15 @@ def testCrossDynamicUnknownBatchDim(self):
     )
     self.match(dynamic_cross(a, b), np.cross(a, b), check_dtype=False)
 
+  def testDiffErrorMessage(self):
+    # Verify the error message for negative n mentions the parameter correctly.
+    x = np_array_ops.array([1, 2, 3])
+    with self.assertRaisesRegex(
+        ValueError,
+        r'Argument `n` must be a non-negative integer\. Received: n=-1',
+    ):
+      np_math_ops.diff(x, n=-1)
+
   def testAverageWrongShape(self):
     with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError, r''):
       np_math_ops.average(np.ones([2, 3]), weights=np.ones([2, 4]))
diff --git a/tensorflow/python/ops/signal/fft_ops.py b/tensorflow/python/ops/signal/fft_ops.py
index 5c3cc9d59f27cd..85c484f0587f07 100644
--- a/tensorflow/python/ops/signal/fft_ops.py
+++ b/tensorflow/python/ops/signal/fft_ops.py
@@ -233,11 +233,13 @@ def _fftn(input_tensor, fft_length=None, axes=None, norm=None, name=None):
     with _ops.name_scope(
         name, default_name, [input_tensor, fft_length, axes]
     ) as name:
-      axes = _process_empty_axes(input_tensor, axes)
-      fft_rank = axes.shape[0]
+      # Convert first: inferring the axes reads `input_tensor.shape`,
+      # which a list or a scalar does not have.
       input_tensor = _ops.convert_to_tensor(
           input_tensor, preferred_dtype=_dtypes.complex64
       )
+      axes = _process_empty_axes(input_tensor, axes)
+      fft_rank = axes.shape[0]
       input_tensor.shape.with_rank_at_least(fft_rank)
       if fft_length is None:
         fft_length = _infer_fft_length_for_fftn(input_tensor)
@@ -273,11 +275,13 @@ def _ifftn(input_tensor, fft_length=None, axes=None, norm=None, name=None):
     with _ops.name_scope(
         name, default_name, [input_tensor, fft_length, axes]
     ) as name:
-      axes = _process_empty_axes(input_tensor, axes)
-      fft_rank = axes.shape[0]
+      # Convert first: inferring the axes reads `input_tensor.shape`,
+      # which a list or a scalar does not have.
       input_tensor = _ops.convert_to_tensor(
           input_tensor, preferred_dtype=_dtypes.complex64
       )
+      axes = _process_empty_axes(input_tensor, axes)
+      fft_rank = axes.shape[0]
       input_tensor.shape.with_rank_at_least(fft_rank)
       if fft_length is None:
         fft_length = _infer_fft_length_for_fftn(input_tensor)
@@ -313,11 +317,13 @@ def _rfftn(input_tensor, fft_length=None, axes=None, norm=None, name=None):
     with _ops.name_scope(
         name, default_name, [input_tensor, fft_length, axes]
     ) as name:
-      axes = _process_empty_axes(input_tensor, axes)
-      fft_rank = axes.shape[0]
+      # Convert first: inferring the axes reads `input_tensor.shape`,
+      # which a list or a scalar does not have.
       input_tensor = _ops.convert_to_tensor(
           input_tensor, preferred_dtype=_dtypes.float32
       )
+      axes = _process_empty_axes(input_tensor, axes)
+      fft_rank = axes.shape[0]
       if input_tensor.dtype not in (_dtypes.float32, _dtypes.float64):
         raise ValueError(
             "RFFT requires tf.float32 or tf.float64 inputs, got: %s"
@@ -370,11 +376,13 @@ def _irfftn(input_tensor, fft_length=None, axes=None, norm=None, name=None):
     with _ops.name_scope(
         name, default_name, [input_tensor, fft_length]
     ) as name:
-      axes = _process_empty_axes(input_tensor, axes)
-      fft_rank = axes.shape[0]
+      # Convert first: inferring the axes reads `input_tensor.shape`,
+      # which a list or a scalar does not have.
       input_tensor = _ops.convert_to_tensor(
           input_tensor, preferred_dtype=_dtypes.complex64
       )
+      axes = _process_empty_axes(input_tensor, axes)
+      fft_rank = axes.shape[0]
       input_tensor.shape.with_rank_at_least(fft_rank)
       if input_tensor.dtype not in (_dtypes.complex64, _dtypes.complex128):
         raise ValueError(
diff --git a/tensorflow/python/ops/tensor_array_ops.py b/tensorflow/python/ops/tensor_array_ops.py
index f2f74de0c06c04..5e90b10e7cb5fa 100644
--- a/tensorflow/python/ops/tensor_array_ops.py
+++ b/tensorflow/python/ops/tensor_array_ops.py
@@ -710,6 +710,15 @@ def close(self, name=None):
 # pylint: enable=protected-access
 
 
+# Advice shared by the errors raised when an eager-mode TensorArray is used
+# while a graph is being built.
+_EAGER_TENSOR_ARRAY_IN_GRAPH_ADVICE = (
+    "This is not currently supported. You may be attempting to capture a "
+    "TensorArray inside a tf.function or tf.data map function. Instead, "
+    "construct a new TensorArray inside the function."
+)
+
+
 class _EagerTensorArray:
   """Eager-compatible implementation of TensorArray."""
 
@@ -807,6 +816,30 @@ def grad(self, source, flow=None, name=None):
         "gradient implementation does not use/need this function to compute "
         "gradients of operations that use TensorArrays.")
 
+  def _check_symbolic_input(self, value, method_name, argument_name):
+    """Checks that `value` is usable as a Python value.
+
+    An `_EagerTensorArray` is only created while executing eagerly, so a
+    symbolic `value` means an eager-mode TensorArray was captured by a
+    `tf.function` or a `tf.data` map function. Without this check `value` is
+    used as a Python value below, which fails with a confusing error such as
+    "Using a symbolic `tf.Tensor` as a Python `bool` is not allowed".
+
+    Args:
+      value: the argument passed to `method_name`.
+      method_name: name of the `TensorArray` method being called.
+      argument_name: name of the argument `value` was passed as.
+
+    Raises:
+      NotImplementedError: if a graph is being built and `value` is symbolic.
+    """
+    if not context.executing_eagerly() and tensor_util.is_tf_type(value):
+      raise NotImplementedError(
+          "Attempting to call TensorArray.%s() with a symbolic `%s` on an "
+          "eager-mode TensorArray. %s"
+          % (method_name, argument_name, _EAGER_TENSOR_ARRAY_IN_GRAPH_ADVICE)
+      )
+
   def read(self, index, name=None):
     """See TensorArray."""
     del name  # not meaningful when executing eagerly.
@@ -814,6 +847,8 @@ def read(self, index, name=None):
     if isinstance(index, ops.EagerTensor):
       index = index.numpy()
 
+    self._check_symbolic_input(index, "read", "index")
+
     if index < 0:
       raise errors_impl.OutOfRangeError(
           None, None,
@@ -851,11 +886,14 @@ def _write(self, index, value):
       errors_impl.InvalidArgumentError: `value` dtype does not match dtype.
       errors_impl.OutOfRangeError: `index` is out of bounds.
       ValueError: shape of `value` is not consistent with inferred shape.
+      NotImplementedError: a graph is being built and `index` is symbolic.
     """
 
     if isinstance(index, ops.EagerTensor):
       index = index.numpy()
 
+    self._check_symbolic_input(index, "write", "index")
+
     if index < 0:
       raise errors_impl.OutOfRangeError(
           None, None,
@@ -920,6 +958,9 @@ def gather(self, indices, name=None):
     del name  # not meaningful when executing eagerly.
     if isinstance(indices, ops.EagerTensor):
       indices = indices.numpy()
+
+    self._check_symbolic_input(indices, "gather", "indices")
+
     return array_ops_stack.stack([self._maybe_zero(i) for i in indices])
 
   def concat(self, name=None):
@@ -956,12 +997,25 @@ def scatter(self, indices, value, name=None):
     del name  # not meaningful when executing eagerly.
     if isinstance(indices, ops.EagerTensor):
       indices = indices.numpy()
+
+    self._check_symbolic_input(indices, "scatter", "indices")
+
     for index, val in zip(indices, array_ops_stack.unstack(value)):
       self._write(index, val)  # pylint: disable=protected-access
     return self.parent()
 
   def split(self, value, lengths, name=None):
     """See TensorArray."""
+    if not context.executing_eagerly():
+      # `lengths` is converted to a tensor below and then read as a Python
+      # value, which is not possible while a graph is being built. This holds
+      # even for a `lengths` that was passed as a Python value.
+      raise NotImplementedError(
+          "Attempting to call TensorArray.split() on an eager-mode "
+          "TensorArray while building a graph. %s"
+          % _EAGER_TENSOR_ARRAY_IN_GRAPH_ADVICE
+      )
+
     # TODO(b/129870929): Fix after all callers provide proper init dtype.
     value = ops.convert_to_tensor(
         value, preferred_dtype=self._dtype, name="value")
diff --git a/tensorflow/python/ops/tensor_array_ops_test.py b/tensorflow/python/ops/tensor_array_ops_test.py
index 77ecaa990eed4e..c42a37bff0ccb8 100644
--- a/tensorflow/python/ops/tensor_array_ops_test.py
+++ b/tensorflow/python/ops/tensor_array_ops_test.py
@@ -22,6 +22,7 @@
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.ops import variables
 from tensorflow.python.platform import test
 
 
@@ -88,6 +89,130 @@ def test_shape_inference_stack_concat(self):
     self.assertEqual(new_arr.stack().shape, (4, 2, 3))
     self.assertEqual(new_arr.concat().shape, (8, 3))
 
+  @test_util.run_v2_only
+  def test_write_symbolic_index_on_captured_eager_tensor_array(self):
+    values = tensor_array_ops.TensorArray(
+        dtypes.int32, size=0, dynamic_size=True, clear_after_read=False
+    )
+
+    @def_function.function
+    def fn(index):
+      return values.write(index, 1)
+
+    with self.assertRaisesRegex(
+        NotImplementedError, 'construct a new TensorArray inside'
+    ):
+      fn(constant_op.constant(0, dtypes.int32))
+
+  @test_util.run_v2_only
+  def test_write_variable_index_on_captured_eager_tensor_array(self):
+    values = tensor_array_ops.TensorArray(
+        dtypes.int32, size=0, dynamic_size=True, clear_after_read=False
+    )
+    index = variables.Variable(0, dtype=dtypes.int32)
+
+    @def_function.function
+    def fn():
+      return values.write(index, 1)
+
+    with self.assertRaisesRegex(
+        NotImplementedError, 'construct a new TensorArray inside'
+    ):
+      fn()
+
+  @test_util.run_v2_only
+  def test_read_symbolic_index_on_captured_eager_tensor_array(self):
+    values = tensor_array_ops.TensorArray(
+        dtypes.int32, size=2, clear_after_read=False
+    )
+    values = values.write(0, 1).write(1, 2)
+
+    @def_function.function
+    def fn(index):
+      return values.read(index)
+
+    with self.assertRaisesRegex(
+        NotImplementedError, 'construct a new TensorArray inside'
+    ):
+      fn(constant_op.constant(1, dtypes.int32))
+
+  @test_util.run_v2_only
+  def test_read_python_index_on_captured_eager_tensor_array(self):
+    # A concrete index still resolves to a Python value, so this keeps working.
+    values = tensor_array_ops.TensorArray(
+        dtypes.int32, size=2, clear_after_read=False
+    )
+    values = values.write(0, 1).write(1, 2)
+
+    @def_function.function
+    def fn():
+      return values.read(1)
+
+    self.assertAllEqual(fn(), 2)
+
+  @test_util.run_v2_only
+  def test_gather_symbolic_indices_on_captured_eager_tensor_array(self):
+    values = tensor_array_ops.TensorArray(
+        dtypes.int32, size=2, clear_after_read=False
+    )
+    values = values.write(0, 1).write(1, 2)
+
+    @def_function.function
+    def fn(indices):
+      return values.gather(indices)
+
+    with self.assertRaisesRegex(
+        NotImplementedError, 'construct a new TensorArray inside'
+    ):
+      fn(constant_op.constant([0, 1], dtypes.int32))
+
+  @test_util.run_v2_only
+  def test_gather_python_indices_on_captured_eager_tensor_array(self):
+    values = tensor_array_ops.TensorArray(
+        dtypes.int32, size=2, clear_after_read=False
+    )
+    values = values.write(0, 1).write(1, 2)
+
+    @def_function.function
+    def fn():
+      return values.gather([0, 1])
+
+    self.assertAllEqual(fn(), [1, 2])
+
+  @test_util.run_v2_only
+  def test_scatter_symbolic_indices_on_captured_eager_tensor_array(self):
+    values = tensor_array_ops.TensorArray(
+        dtypes.int32, size=2, clear_after_read=False, element_shape=[1]
+    )
+
+    @def_function.function
+    def fn(indices):
+      return values.scatter(
+          indices, constant_op.constant([[3], [4]], dtypes.int32)
+      )
+
+    with self.assertRaisesRegex(
+        NotImplementedError, 'construct a new TensorArray inside'
+    ):
+      fn(constant_op.constant([0, 1], dtypes.int32))
+
+  @test_util.run_v2_only
+  def test_split_on_captured_eager_tensor_array(self):
+    values = tensor_array_ops.TensorArray(
+        dtypes.int32, size=2, clear_after_read=False
+    )
+
+    @def_function.function
+    def fn():
+      return values.split(
+          constant_op.constant([1, 2, 3, 4], dtypes.int32), [2, 2]
+      )
+
+    with self.assertRaisesRegex(
+        NotImplementedError, 'construct a new TensorArray inside'
+    ):
+      fn()
+
 
 if __name__ == '__main__':
   test.main()
diff --git a/tensorflow/tensorflow.default.bzl b/tensorflow/tensorflow.default.bzl
index 3a16e2c4a0ee23..083ed242b31b70 100644
--- a/tensorflow/tensorflow.default.bzl
+++ b/tensorflow/tensorflow.default.bzl
@@ -15,6 +15,7 @@
 
 """Default (OSS) build versions of TensorFlow general-purpose build extensions."""
 
+load("@local_config_syslibs//:build_defs.bzl", "if_system_lib")
 load(
     "@rules_ml_toolchain//py/rules_pywrap:pywrap.default.bzl",
     _pywrap_aware_cc_import = "pywrap_aware_cc_import",
@@ -80,7 +81,6 @@ if_portable = _if_portable
 ADDITIONAL_API_INDEXABLE_SETTINGS = _ADDITIONAL_API_INDEXABLE_SETTINGS
 if_indexing_source_code = _if_indexing_source_code
 pywrap_tensorflow_macro = _pywrap_tensorflow_macro
-tf_cc_shared_library = _tf_cc_shared_library
 pytype_library = _pytype_library
 tf_py_test = _tf_py_test
 tf_py_strict_test = _tf_py_test
@@ -125,7 +125,49 @@ pywrap_aware_tf_cc_shared_object = _pywrap_aware_tf_cc_shared_object
 pywrap_aware_filegroup = _pywrap_aware_filegroup
 pywrap_aware_genrule = _pywrap_aware_genrule
 pywrap_aware_cc_import = _pywrap_aware_cc_import
-pywrap_library = _pywrap_library
 pywrap_common_library = _pywrap_common_library
 stripped_cc_info = _stripped_cc_info
 pywrap_binaries = _pywrap_binaries
+
+def tf_system_libs_linkopts():
+    """Returns linker flags for system libraries configured via TF_SYSTEM_LIBS."""
+    return (
+        if_system_lib("boringssl", ["-lssl", "-lcrypto"]) +
+        if_system_lib("com_github_googlecloudplatform_google_cloud_cpp", ["-lgoogle_cloud_cpp_common", "-lgoogle_cloud_cpp_bigtable", "-lgoogle_cloud_cpp_storage"]) +
+        if_system_lib("com_github_grpc_grpc", ["-lgrpc++", "-lgrpc", "-lgpr"]) +
+        if_system_lib("com_google_protobuf", ["-lprotobuf"]) +
+        if_system_lib("com_googlesource_code_re2", ["-lre2"]) +
+        if_system_lib("curl", ["-lcurl"]) +
+        if_system_lib("flatbuffers", ["-lflatbuffers"]) +
+        if_system_lib("gif", ["-lgif"]) +
+        if_system_lib("hwloc", ["-lhwloc"]) +
+        if_system_lib("icu", ["-licui18n", "-licuuc", "-licudata"]) +
+        if_system_lib("jsoncpp_git", ["-ljsoncpp"]) +
+        if_system_lib("libjpeg_turbo", ["-ljpeg"]) +
+        if_system_lib("org_sqlite", ["-lsqlite3"]) +
+        if_system_lib("png", ["-lpng"]) +
+        if_system_lib("snappy", ["-lsnappy"]) +
+        if_system_lib("zlib", ["-lz"])
+    )
+
+def tf_cc_shared_library(name, linkopts = [], **kwargs):
+    _tf_cc_shared_library(
+        name = name,
+        linkopts = linkopts + tf_system_libs_linkopts(),
+        **kwargs
+    )
+
+def pywrap_library(name, common_lib_linkopts = {}, **kwargs):
+    updated_common_lib_linkopts = dict(common_lib_linkopts)
+    sys_linkopts = tf_system_libs_linkopts()
+    if sys_linkopts:
+        for lib in ["tensorflow/tensorflow_framework", "tensorflow/tensorflow_cc"]:
+            if lib in updated_common_lib_linkopts:
+                updated_common_lib_linkopts[lib] = updated_common_lib_linkopts[lib] + sys_linkopts
+            else:
+                updated_common_lib_linkopts[lib] = sys_linkopts
+    _pywrap_library(
+        name = name,
+        common_lib_linkopts = updated_common_lib_linkopts,
+        **kwargs
+    )
diff --git a/tensorflow/tools/pip_package/setup.py.tpl b/tensorflow/tools/pip_package/setup.py.tpl
index e11f2ac3be0c7e..c506c29134c6bf 100644
--- a/tensorflow/tools/pip_package/setup.py.tpl
+++ b/tensorflow/tools/pip_package/setup.py.tpl
@@ -107,7 +107,6 @@ REQUIRED_PACKAGES = [
     'packaging',
     'protobuf >= 6.31.1, < 8.0.0',
     'requests >= 2.21.0, < 3',
-    'setuptools',
     'six >= 1.12.0',
     'termcolor >= 1.1.0',
     'typing_extensions >= 3.6.6',
diff --git a/third_party/xla/.github/workflows/nightly_benchmarks.yml b/third_party/xla/.github/workflows/nightly_benchmarks.yml
index 2c9dee7b59a524..d6b111fa083746 100644
--- a/third_party/xla/.github/workflows/nightly_benchmarks.yml
+++ b/third_party/xla/.github/workflows/nightly_benchmarks.yml
@@ -30,8 +30,8 @@ concurrency:
 jobs:
   run_benchmarks:
     name: Run XLA Nightly Benchmarks
-    uses: google-ml-infra/bap/.github/workflows/run-benchmarks.yaml@aa854ba7bdb77966a9cc1a013417f1c1a208c4c0 # main
+    uses: google-ml-infra/bap/.github/workflows/run-benchmarks.yaml@887b34acdb5c70bbd2bc93bb11c2d0a103698e7f # main
     with:
       registry_file: "xla/tools/benchmarks/benchmark_registry.pbtxt"
       tag_filter: "scheduled"
-      bap_ref: "aa854ba7bdb77966a9cc1a013417f1c1a208c4c0"
+      bap_ref: "887b34acdb5c70bbd2bc93bb11c2d0a103698e7f"
diff --git a/third_party/xla/.github/workflows/run_benchmarks.yml b/third_party/xla/.github/workflows/run_benchmarks.yml
index bc0e96de2f2c6d..d5130492ba9b06 100644
--- a/third_party/xla/.github/workflows/run_benchmarks.yml
+++ b/third_party/xla/.github/workflows/run_benchmarks.yml
@@ -57,10 +57,10 @@ permissions:
 jobs:
   run_benchmarks:
     name: XLA Benchmarks
-    uses: google-ml-infra/bap/.github/workflows/run-benchmarks.yaml@aa854ba7bdb77966a9cc1a013417f1c1a208c4c0 # main
+    uses: google-ml-infra/bap/.github/workflows/run-benchmarks.yaml@887b34acdb5c70bbd2bc93bb11c2d0a103698e7f # main
     with:
       registry_file: "xla/tools/benchmarks/benchmark_registry.pbtxt"
-      bap_ref: "aa854ba7bdb77966a9cc1a013417f1c1a208c4c0"
+      bap_ref: "887b34acdb5c70bbd2bc93bb11c2d0a103698e7f"
       ab_mode: ${{ inputs.ab_mode }}
       baseline_ref: ${{ inputs.baseline_ref }}
       experiment_ref: ${{ inputs.experiment_ref }}
diff --git a/third_party/xla/xla/backends/cpu/tests/BUILD b/third_party/xla/xla/backends/cpu/tests/BUILD
index 672430c7392d80..79c4519867f9be 100644
--- a/third_party/xla/xla/backends/cpu/tests/BUILD
+++ b/third_party/xla/xla/backends/cpu/tests/BUILD
@@ -46,9 +46,6 @@ xla_test(
 xla_test(
     name = "ynn_fusion_test",
     srcs = ["ynn_fusion_test.cc"],
-    backend_tags = {
-        "cpu": ["broken"],  # b/516277731
-    },
     backends = ["cpu"],
     deps = [
         "//xla:error_spec",
diff --git a/third_party/xla/xla/core/collectives/collectives_registry.h b/third_party/xla/xla/core/collectives/collectives_registry.h
index fbd94e47d7a64a..df8e35362861c5 100644
--- a/third_party/xla/xla/core/collectives/collectives_registry.h
+++ b/third_party/xla/xla/core/collectives/collectives_registry.h
@@ -59,14 +59,13 @@ class CollectivesRegistry {
 #define XLA_COLLECTIVES_REGISTER_(PLATFORM, NAME, PRIORITY, IMPL, N) \
   XLA_COLLECTIVES_REGISTER__(PLATFORM, NAME, PRIORITY, IMPL, N)
 #define XLA_COLLECTIVES_REGISTER__(PLATFORM, NAME, PRIORITY, IMPL, N)         \
-  ABSL_ATTRIBUTE_UNUSED static const bool xla_collectives_##N##_registered_ = \
-      [] {                                                                    \
-        absl::Status status = ::xla::CollectivesRegistry::Register(           \
-            PLATFORM, NAME, PRIORITY, IMPL);                                  \
-        if (!status.ok()) {                                                   \
-          LOG(ERROR) << "Failed to register XLA collectives: " << status;     \
-        }                                                                     \
-        return true;                                                          \
-      }()
+  [[maybe_unused]] static const bool xla_collectives_##N##_registered_ = [] { \
+    absl::Status status =                                                     \
+        ::xla::CollectivesRegistry::Register(PLATFORM, NAME, PRIORITY, IMPL); \
+    if (!status.ok()) {                                                       \
+      LOG(ERROR) << "Failed to register XLA collectives: " << status;         \
+    }                                                                         \
+    return true;                                                              \
+  }()
 
 #endif  // XLA_CORE_COLLECTIVES_COLLECTIVES_REGISTRY_H_
diff --git a/third_party/xla/xla/service/cpu/cpu_float_support.h b/third_party/xla/xla/service/cpu/cpu_float_support.h
index e9502839bb008a..656a386b69ed6f 100644
--- a/third_party/xla/xla/service/cpu/cpu_float_support.h
+++ b/third_party/xla/xla/service/cpu/cpu_float_support.h
@@ -40,11 +40,30 @@ class CpuFloatSupport : public FloatSupport {
 
   bool SupportsLowPrecisionOperand(const HloInstruction& hlo,
                                    int64_t operand_index) const override {
-    if (LowPrecisionType() == BF16 || LowPrecisionType() == F16) {
+    if (LowPrecisionType() == BF16) {
       switch (hlo.opcode()) {
         case HloOpcode::kSort:
         case HloOpcode::kCompare:
         case HloOpcode::kSelect:
+        case HloOpcode::kBroadcast:
+        case HloOpcode::kConcatenate:
+        case HloOpcode::kCopy:
+        case HloOpcode::kDynamicSlice:
+        case HloOpcode::kDynamicUpdateSlice:
+        case HloOpcode::kGather:
+        case HloOpcode::kPad:
+        case HloOpcode::kReshape:
+        case HloOpcode::kReverse:
+        case HloOpcode::kScatter:
+        case HloOpcode::kSlice:
+        case HloOpcode::kTranspose:
+        case HloOpcode::kAbs:
+        case HloOpcode::kNegate:
+        case HloOpcode::kSign:
+        case HloOpcode::kMaximum:
+        case HloOpcode::kMinimum:
+        case HloOpcode::kClamp:
+        case HloOpcode::kSelectAndScatter:
           return true;
         default:
           break;
@@ -54,10 +73,29 @@ class CpuFloatSupport : public FloatSupport {
   }
 
   bool SupportsLowPrecisionOutput(const HloInstruction& hlo) const override {
-    if (LowPrecisionType() == BF16 || LowPrecisionType() == F16) {
+    if (LowPrecisionType() == BF16) {
       switch (hlo.opcode()) {
         case HloOpcode::kSort:
         case HloOpcode::kSelect:
+        case HloOpcode::kBroadcast:
+        case HloOpcode::kConcatenate:
+        case HloOpcode::kCopy:
+        case HloOpcode::kDynamicSlice:
+        case HloOpcode::kDynamicUpdateSlice:
+        case HloOpcode::kGather:
+        case HloOpcode::kPad:
+        case HloOpcode::kReshape:
+        case HloOpcode::kReverse:
+        case HloOpcode::kScatter:
+        case HloOpcode::kSlice:
+        case HloOpcode::kTranspose:
+        case HloOpcode::kAbs:
+        case HloOpcode::kNegate:
+        case HloOpcode::kSign:
+        case HloOpcode::kMaximum:
+        case HloOpcode::kMinimum:
+        case HloOpcode::kClamp:
+        case HloOpcode::kSelectAndScatter:
           return true;
         default:
           break;
diff --git a/third_party/xla/xla/service/cpu/cpu_float_support_test.cc b/third_party/xla/xla/service/cpu/cpu_float_support_test.cc
index 1891860619ff0f..d4a8c53462192f 100644
--- a/third_party/xla/xla/service/cpu/cpu_float_support_test.cc
+++ b/third_party/xla/xla/service/cpu/cpu_float_support_test.cc
@@ -200,46 +200,41 @@ INSTANTIATE_TEST_SUITE_P(SkipInstructionTestSuite, SkipInstructionTest,
                          SkipInstructionTest::Name);
 
 TEST_F(TargetMachineTestBase, SortNotUpcast) {
-  for (absl::string_view type_str : {"bf16", "f16"}) {
-    std::string hlo_text = absl::StrReplaceAll(R"(
+  std::string hlo_text = R"(
 HloModule test_module
 
 compare {
-  p0 = $type$[] parameter(0)
-  p1 = $type$[] parameter(1)
+  p0 = bf16[] parameter(0)
+  p1 = bf16[] parameter(1)
   p2 = s32[] parameter(2)
   p3 = s32[] parameter(3)
   ROOT cmp = pred[] compare(p0, p1), direction=LT
 }
 
 ENTRY main {
-  k = $type$[100] parameter(0)
+  k = bf16[100] parameter(0)
   v = s32[100] parameter(1)
-  ROOT sort = ($type$[100], s32[100]) sort(k, v), dimensions={0}, to_apply=compare, is_stable=true
+  ROOT sort = (bf16[100], s32[100]) sort(k, v), dimensions={0}, to_apply=compare, is_stable=true
 }
-)",
-                                               {{"$type$", type_str}});
-
-    ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text));
+)";
 
-    PrimitiveType low_precision_type = (type_str == "bf16") ? BF16 : F16;
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text));
 
-    CpuFloatSupport cpu_float_support(
-        low_precision_type, [](const HloInstruction&) { return false; });
+  CpuFloatSupport cpu_float_support(
+      BF16, [](const HloInstruction&) { return false; });
 
-    FloatNormalization float_normalization(&cpu_float_support);
-    ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get()));
-    EXPECT_FALSE(upcast);
+  FloatNormalization float_normalization(&cpu_float_support);
+  ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get()));
+  EXPECT_FALSE(upcast);
 
-    HloInstruction* root = module->entry_computation()->root_instruction();
-    EXPECT_EQ(root->opcode(), HloOpcode::kSort);
-    EXPECT_EQ(root->operand(0)->shape().element_type(), low_precision_type);
-    EXPECT_EQ(root->operand(1)->shape().element_type(), S32);
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_EQ(root->opcode(), HloOpcode::kSort);
+  EXPECT_EQ(root->operand(0)->shape().element_type(), BF16);
+  EXPECT_EQ(root->operand(1)->shape().element_type(), S32);
 
-    HloComputation* compare_comp = root->to_apply();
-    EXPECT_EQ(compare_comp->parameter_instruction(0)->shape().element_type(),
-              low_precision_type);
-  }
+  HloComputation* compare_comp = root->to_apply();
+  EXPECT_EQ(compare_comp->parameter_instruction(0)->shape().element_type(),
+            BF16);
 }
 
 TEST_F(TargetMachineTestBase, SortUpcastForUnsupportedType) {
@@ -280,42 +275,107 @@ ENTRY main {
 }
 
 TEST_F(TargetMachineTestBase, CompareAndSelectNotUpcast) {
-  for (absl::string_view type_str : {"bf16", "f16"}) {
-    std::string hlo_text = absl::StrReplaceAll(R"(
+  std::string hlo_text = R"(
 HloModule test_module
 
 ENTRY main {
-  p0 = $type$[100] parameter(0)
-  p1 = $type$[100] parameter(1)
+  p0 = bf16[100] parameter(0)
+  p1 = bf16[100] parameter(1)
   cmp = pred[100] compare(p0, p1), direction=LT
-  ROOT select = $type$[100] select(cmp, p0, p1)
+  ROOT select = bf16[100] select(cmp, p0, p1)
 }
-)",
-                                               {{"$type$", type_str}});
+)";
 
-    ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text));
 
-    PrimitiveType low_precision_type = (type_str == "bf16") ? BF16 : F16;
+  CpuFloatSupport cpu_float_support(
+      BF16, [](const HloInstruction&) { return false; });
 
-    CpuFloatSupport cpu_float_support(
-        low_precision_type, [](const HloInstruction&) { return false; });
+  FloatNormalization float_normalization(&cpu_float_support);
+  ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get()));
+  EXPECT_FALSE(upcast);
+
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_EQ(root->opcode(), HloOpcode::kSelect);
+  EXPECT_EQ(root->shape().element_type(), BF16);
+  EXPECT_EQ(root->operand(1)->shape().element_type(), BF16);
+  EXPECT_EQ(root->operand(2)->shape().element_type(), BF16);
+
+  const HloInstruction* cmp = root->operand(0);
+  EXPECT_EQ(cmp->opcode(), HloOpcode::kCompare);
+  EXPECT_EQ(cmp->operand(0)->shape().element_type(), BF16);
+  EXPECT_EQ(cmp->operand(1)->shape().element_type(), BF16);
+}
+
+TEST_F(TargetMachineTestBase, DataMovementOpsNotUpcast) {
+  std::string hlo_text = R"(
+HloModule test_module
+
+ENTRY main {
+  p0 = bf16[100, 10] parameter(0)
+  transpose = bf16[10, 100] transpose(p0), dimensions={1, 0}
+  reshape = bf16[1000] reshape(transpose)
+  ROOT copy = bf16[1000] copy(reshape)
+}
+)";
 
-    FloatNormalization float_normalization(&cpu_float_support);
-    ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get()));
-    EXPECT_FALSE(upcast);
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text));
 
-    HloInstruction* root = module->entry_computation()->root_instruction();
-    EXPECT_EQ(root->opcode(), HloOpcode::kSelect);
-    EXPECT_EQ(root->shape().element_type(), low_precision_type);
-    EXPECT_EQ(root->operand(1)->shape().element_type(), low_precision_type);
-    EXPECT_EQ(root->operand(2)->shape().element_type(), low_precision_type);
+  CpuFloatSupport cpu_float_support(
+      BF16, [](const HloInstruction&) { return false; });
 
-    const HloInstruction* cmp = root->operand(0);
-    EXPECT_EQ(cmp->opcode(), HloOpcode::kCompare);
-    EXPECT_EQ(cmp->operand(0)->shape().element_type(), low_precision_type);
-    EXPECT_EQ(cmp->operand(1)->shape().element_type(), low_precision_type);
-  }
+  FloatNormalization float_normalization(&cpu_float_support);
+  ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get()));
+  EXPECT_FALSE(upcast);
+
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_EQ(root->opcode(), HloOpcode::kCopy);
+  EXPECT_EQ(root->shape().element_type(), BF16);
+
+  HloInstruction* reshape = root->mutable_operand(0);
+  EXPECT_EQ(reshape->opcode(), HloOpcode::kReshape);
+  EXPECT_EQ(reshape->shape().element_type(), BF16);
+
+  HloInstruction* transpose = reshape->mutable_operand(0);
+  EXPECT_EQ(transpose->opcode(), HloOpcode::kTranspose);
+  EXPECT_EQ(transpose->shape().element_type(), BF16);
+}
+
+TEST_F(TargetMachineTestBase, ElementwiseOpsNotUpcast) {
+  std::string hlo_text = R"(
+HloModule test_module
+
+ENTRY main {
+  p0 = bf16[100] parameter(0)
+  p1 = bf16[100] parameter(1)
+  abs = bf16[100] abs(p0)
+  negate = bf16[100] negate(abs)
+  ROOT max = bf16[100] maximum(negate, p1)
+}
+)";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text));
+
+  CpuFloatSupport cpu_float_support(
+      BF16, [](const HloInstruction&) { return false; });
+
+  FloatNormalization float_normalization(&cpu_float_support);
+  ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get()));
+  EXPECT_FALSE(upcast);
+
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_EQ(root->opcode(), HloOpcode::kMaximum);
+  EXPECT_EQ(root->shape().element_type(), BF16);
+
+  HloInstruction* negate = root->mutable_operand(0);
+  EXPECT_EQ(negate->opcode(), HloOpcode::kNegate);
+  EXPECT_EQ(negate->shape().element_type(), BF16);
+
+  HloInstruction* abs = negate->mutable_operand(0);
+  EXPECT_EQ(abs->opcode(), HloOpcode::kAbs);
+  EXPECT_EQ(abs->shape().element_type(), BF16);
 }
 
 }  // namespace
+
 }  // namespace xla::cpu
diff --git a/third_party/xla/xla/service/cpu/onednn_float_support.cc b/third_party/xla/xla/service/cpu/onednn_float_support.cc
index b749cccbdd49f6..97043196ca1cca 100644
--- a/third_party/xla/xla/service/cpu/onednn_float_support.cc
+++ b/third_party/xla/xla/service/cpu/onednn_float_support.cc
@@ -55,11 +55,18 @@ bool OneDnnFloatSupport::IsSupported(const HloInstruction& hlo) const {
     case HloOpcode::kSelectAndScatter:
     case HloOpcode::kSlice:
     case HloOpcode::kTranspose:
+    // Elementwise ops without precision loss.
+    case HloOpcode::kAbs:
+    case HloOpcode::kNegate:
+    case HloOpcode::kSign:
+    case HloOpcode::kMaximum:
+    case HloOpcode::kMinimum:
+    case HloOpcode::kClamp:
     // Other special ops.
     case HloOpcode::kBitcast:
       return true;
     case HloOpcode::kSort:
-      return LowPrecisionType() == BF16 || LowPrecisionType() == F16;
+      return LowPrecisionType() == BF16;
     default:
       return 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 601f4436beffa6..bfa997822e86de 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
@@ -166,12 +166,40 @@ stream_executor::DeviceDescription TestGpuDeviceInfo::AMDMI210DeviceInfo() {
   b.set_core_count(104);
   b.set_fpus_per_core(128);
   b.set_block_dim_limit_x(2'147'483'647);
-  b.set_block_dim_limit_y(2'147'483'647);
-  b.set_block_dim_limit_z(2'147'483'647);
+  b.set_block_dim_limit_y(65536);
+  b.set_block_dim_limit_z(65536);
   b.set_memory_bandwidth(1'638'400'000'000);
   b.set_l2_cache_size(8 * 1024 * 1024);
   b.set_clock_rate_ghz(1.7);
   b.set_device_memory_size(67'628'957'696);
+  b.set_registers_per_core_limit(131072);
+  b.set_registers_per_block_limit(131072);
+  b.set_runtime_version(stream_executor::SemanticVersion{6, 0, 0});
+  b.set_driver_version(stream_executor::SemanticVersion{6, 0, 0});
+  return b;
+}
+
+stream_executor::DeviceDescription TestGpuDeviceInfo::AMDMI300DeviceInfo() {
+  stream_executor::DeviceDescription b;
+  b.set_gpu_compute_capability(stream_executor::GpuComputeCapability(
+      stream_executor::RocmComputeCapability("gfx942")));
+  b.set_threads_per_block_limit(1024);
+  b.set_threads_per_warp(64);
+  b.set_shared_memory_per_block(64 * 1024);
+  b.set_shared_memory_per_block_optin(64 * 1024);
+  b.set_shared_memory_per_core(64 * 1024);
+  b.set_threads_per_core_limit(2048);
+  b.set_core_count(304);
+  b.set_fpus_per_core(128);
+  b.set_block_dim_limit_x(2'147'483'647);
+  b.set_block_dim_limit_y(65536);
+  b.set_block_dim_limit_z(65536);
+  b.set_memory_bandwidth(5'300'000'000'000);
+  b.set_l2_cache_size(4 * 1024 * 1024);
+  b.set_clock_rate_ghz(2.1);
+  b.set_device_memory_size(int64_t{192} * 1024 * 1024 * 1024);
+  b.set_registers_per_core_limit(131072);
+  b.set_registers_per_block_limit(131072);
   b.set_runtime_version(stream_executor::SemanticVersion{6, 0, 0});
   b.set_driver_version(stream_executor::SemanticVersion{6, 0, 0});
   return b;
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 256704f256a167..b74996d4faecf4 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
@@ -41,6 +41,7 @@ class TestGpuDeviceInfo {
           stream_executor::GpuComputeCapability{
               stream_executor::CudaComputeCapability(10, 0)});
   static stream_executor::DeviceDescription AMDMI210DeviceInfo();
+  static stream_executor::DeviceDescription AMDMI300DeviceInfo();
   static stream_executor::DeviceDescription AMDMI350DeviceInfo();
   static stream_executor::DeviceDescription AMDRX7900DeviceInfo();
   // Returns default RTXA6000 or AMDMI210 device info
diff --git a/third_party/xla/xla/service/gpu/model/collective_interpolator_test.cc b/third_party/xla/xla/service/gpu/model/collective_interpolator_test.cc
index e4fec1d9f95ff5..7cd40cd94ef59c 100644
--- a/third_party/xla/xla/service/gpu/model/collective_interpolator_test.cc
+++ b/third_party/xla/xla/service/gpu/model/collective_interpolator_test.cc
@@ -1083,8 +1083,7 @@ INSTANTIATE_TEST_SUITE_P(
     });
 
 TEST(DefaultCollectivePerfTableTest, EstimatesGfx942DefaultProfile) {
-  se::DeviceDescription device_info = TestGpuDeviceInfo::RTXA6000DeviceInfo();
-  device_info.set_rocm_compute_capability("gfx942");
+  se::DeviceDescription device_info = TestGpuDeviceInfo::AMDMI300DeviceInfo();
   ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<CollectiveInterpolator> interpolator,
       CollectiveInterpolator::Create(kNumGpusPerHost, device_info));
@@ -1133,8 +1132,7 @@ TEST(DefaultCollectivePerfTableTest, EstimatesGfx942DefaultProfile) {
 }
 
 TEST(DefaultCollectivePerfTableTest, EstimatesGfx950DefaultProfile) {
-  se::DeviceDescription device_info = TestGpuDeviceInfo::RTXA6000DeviceInfo();
-  device_info.set_rocm_compute_capability("gfx950");
+  se::DeviceDescription device_info = TestGpuDeviceInfo::AMDMI350DeviceInfo();
   ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<CollectiveInterpolator> interpolator,
       CollectiveInterpolator::Create(kNumGpusPerHost, device_info));
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 6accee94d58dad..d2d7d0f4c041a9 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
@@ -334,15 +334,11 @@ class MatmulInterpolatorDefaultTableTest
   }
 
   std::unique_ptr<MatmulInterpolator> GetMatmulInterpolatorGfx942() {
-    se::DeviceDescription device_info = TestGpuDeviceInfo::AMDMI210DeviceInfo();
-    device_info.set_rocm_compute_capability("gfx942");
-    return GetMatmulInterpolator(device_info);
+    return GetMatmulInterpolator(TestGpuDeviceInfo::AMDMI300DeviceInfo());
   }
 
   std::unique_ptr<MatmulInterpolator> GetMatmulInterpolatorGfx950() {
-    se::DeviceDescription device_info = TestGpuDeviceInfo::AMDMI210DeviceInfo();
-    device_info.set_rocm_compute_capability("gfx950");
-    return GetMatmulInterpolator(device_info);
+    return GetMatmulInterpolator(TestGpuDeviceInfo::AMDMI350DeviceInfo());
   }
 };
 
@@ -739,8 +735,7 @@ INSTANTIATE_TEST_SUITE_P(
            info) { return info.param.test_name; });
 
 TEST(DefaultMatmulPerfTableTest, Gfx942InterpolatesBetweenGridPoints) {
-  se::DeviceDescription device_info = TestGpuDeviceInfo::AMDMI210DeviceInfo();
-  device_info.set_rocm_compute_capability("gfx942");
+  se::DeviceDescription device_info = TestGpuDeviceInfo::AMDMI300DeviceInfo();
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<MatmulInterpolator> interpolator,
                        MatmulInterpolator::Create(device_info));
   ASSERT_OK_AND_ASSIGN(
@@ -838,8 +833,7 @@ INSTANTIATE_TEST_SUITE_P(
            info) { return info.param.test_name; });
 
 TEST(DefaultMatmulPerfTableTest, Gfx950InterpolatesBetweenGridPoints) {
-  se::DeviceDescription device_info = TestGpuDeviceInfo::AMDMI210DeviceInfo();
-  device_info.set_rocm_compute_capability("gfx950");
+  se::DeviceDescription device_info = TestGpuDeviceInfo::AMDMI350DeviceInfo();
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<MatmulInterpolator> interpolator,
                        MatmulInterpolator::Create(device_info));
   ASSERT_OK_AND_ASSIGN(
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 04d0e81405ba0e..f52179e018ad57 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
@@ -980,8 +980,7 @@ TEST_F(IsSolLatencyEstimatorEnabledTest, CreatesEstimatorWithGfx950Profiles) {
     }
   )";
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo));
-  gpu_device_info_ = TestGpuDeviceInfo::AMDMI210DeviceInfo();
-  gpu_device_info_.set_rocm_compute_capability("gfx950");
+  gpu_device_info_ = TestGpuDeviceInfo::AMDMI350DeviceInfo();
 
   SchedulerConfig scheduler_config;
   ASSERT_OK_AND_ASSIGN(
@@ -1030,8 +1029,7 @@ TEST_F(IsSolLatencyEstimatorEnabledTest, CreatesEstimatorWithGfx942Profiles) {
     }
   )";
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo));
-  gpu_device_info_ = TestGpuDeviceInfo::AMDMI210DeviceInfo();
-  gpu_device_info_.set_rocm_compute_capability("gfx942");
+  gpu_device_info_ = TestGpuDeviceInfo::AMDMI300DeviceInfo();
 
   SchedulerConfig scheduler_config;
   ASSERT_OK_AND_ASSIGN(
diff --git a/third_party/xla/xla/stream_executor/host/host_stream_factory.h b/third_party/xla/xla/stream_executor/host/host_stream_factory.h
index b8d908d59e9ab1..d34afc8ef320fa 100644
--- a/third_party/xla/xla/stream_executor/host/host_stream_factory.h
+++ b/third_party/xla/xla/stream_executor/host/host_stream_factory.h
@@ -63,11 +63,11 @@ class HostStreamFactoryRegistrar {
 #define REGISTER_HOST_STREAM_FACTORY(factory, priority) \
   INTERNAL_REGISTER_HOST_STREAM_FACTORY(factory, priority, __COUNTER__)
 
-#define INTERNAL_REGISTER_HOST_STREAM_FACTORY(factory, priority, ctr) \
-  ABSL_ATTRIBUTE_UNUSED static ::stream_executor::host::              \
-      HostStreamFactoryRegistrar<factory>                             \
-      INTERNAL_REGISTER_LOCAL_HOST_STREAM_FACTORY_NAME(ctr) {         \
-    priority                                                          \
+#define INTERNAL_REGISTER_HOST_STREAM_FACTORY(factory, priority, ctr)          \
+  [[maybe_unused]] static ::stream_executor::host::HostStreamFactoryRegistrar< \
+      factory>                                                                 \
+  INTERNAL_REGISTER_LOCAL_HOST_STREAM_FACTORY_NAME(ctr) {                      \
+    priority                                                                   \
   }
 
 // __COUNTER__ must go through another macro to be properly expanded