diff --git a/RELEASE.md b/RELEASE.md
index b660edf249e541..575ac36ab4b3b6 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -34,6 +34,13 @@ In `tensorflow/c/experimental/filesystem/filesystem_interface.h`, removed `TF_Tr
     * Exports `__new__` in public API golden files for subclasses of `tuple` (like `tf.io.FixedLenFeature`) to fix false positives during static type checking.>
 * `tf.data`
     * Fixes a bug in `tf.data.Dataset.scan` where the shape of the state returned by `scan_func` was not strictly validated against the initial state.
+*   `tf.image.adjust_contrast`
+
+    *   Registers the missing Python gradient for the `AdjustContrastv2` op, so
+        `tf.image.adjust_contrast` can now be differentiated with
+        `GradientTape`. Fixes
+        [#126083](https://github.com/tensorflow/tensorflow/issues/126083).
+
 *   `tf.experimental.numpy`
 
     *   `tf.experimental.numpy.isclose` and `tf.experimental.numpy.allclose` now
diff --git a/tensorflow/cc/framework/cc_op_gen.cc b/tensorflow/cc/framework/cc_op_gen.cc
index 813cba37dcce43..21cfebf8ff4c0a 100644
--- a/tensorflow/cc/framework/cc_op_gen.cc
+++ b/tensorflow/cc/framework/cc_op_gen.cc
@@ -119,7 +119,12 @@ void WriteClassDecl(const OpInfo& op_info, WritableFile* h) {
       const auto entry = AttrTypeName(attr.type());
       const auto attr_type_name = entry.first;
       const bool use_const = entry.second;
-      const std::string camel_case_name = ToCamelCase(api_def_attr.rename_to());
+      // See cc_op_gen_util.cc's GetOpAttrStruct: attr.name() and
+      // api_def_attr.rename_to() are both unvalidated at OpDef-registration
+      // time and both spliced as a raw C++ identifier below.
+      const std::string safe_attr_name =
+          SafeRenameTo(attr.name(), api_def_attr.rename_to());
+      const std::string camel_case_name = ToCamelCase(safe_attr_name);
       const std::string suffix =
           (camel_case_name == op_info.op_name || camel_case_name == "Attrs")
               ? "_"
@@ -212,10 +217,19 @@ std::string GetConstructorBody(const OpInfo& op_info) {
   for (int i = 0; i < op_info.graph_op_def.input_arg_size(); ++i) {
     const auto& arg(op_info.graph_op_def.input_arg(i));
     const auto& api_def_arg(op_info.api_def.in_arg(i));
+    // See cc_op_gen_util.cc's GetOpAttrStruct comment: arg.name() and
+    // api_def_arg.rename_to() are both unvalidated at OpDef-registration
+    // time. Spliced twice below -- once as the local variable name
+    // (`_<name>`), once via AvoidCPPKeywords as the argument reference --
+    // so both must agree with each other and with the corresponding
+    // .Input(_<name>) reference generated further down from the same
+    // (arg.name(), rename_to()) pair.
+    const std::string safe_arg_name =
+        SafeRenameTo(arg.name(), api_def_arg.rename_to());
     strings::StrAppend(
-        &body, "  auto _", api_def_arg.rename_to(), " = ::tensorflow::ops::",
+        &body, "  auto _", safe_arg_name, " = ::tensorflow::ops::",
         ArgIsList(arg) ? "AsNodeOutList" : "AsNodeOut", "(", scope_str, ", ",
-        AvoidCPPKeywords(api_def_arg.rename_to()), ");\n");
+        AvoidCPPKeywords(safe_arg_name), ");\n");
     absl::StrAppend(&body, "  ", return_on_error, "\n");
   }
 
@@ -228,7 +242,11 @@ std::string GetConstructorBody(const OpInfo& op_info) {
   const std::string spaces = "                     ";
   for (int i = 0; i < op_info.api_def.in_arg_size(); ++i) {
     const auto& arg(op_info.api_def.in_arg(i));
-    absl::StrAppend(&body, spaces, ".Input(_", arg.rename_to(), ")\n");
+    // Same (name, rename_to) pair as the declaration loop above, for the
+    // same index -- SafeRenameTo is pure, so this independently produces
+    // the identical identifier the declaration above already emitted.
+    const std::string safe_arg_name = SafeRenameTo(arg.name(), arg.rename_to());
+    absl::StrAppend(&body, spaces, ".Input(_", safe_arg_name, ")\n");
   }
   for (int i = 0; i < op_info.api_def.attr_size(); ++i) {
     const auto& graph_attr(op_info.graph_op_def.attr(i));
@@ -237,12 +255,18 @@ std::string GetConstructorBody(const OpInfo& op_info) {
         op_info.inferred_input_attrs.end()) {
       continue;
     }
+    // See above: graph_attr.name() and api_def_attr.rename_to() are both
+    // unvalidated, and both are spliced as a raw C++ identifier/field
+    // access below.
+    const std::string safe_attr_name =
+        SafeRenameTo(graph_attr.name(), api_def_attr.rename_to());
     const std::string attr_name =
         api_def_attr.has_default_value()
-            ? absl::StrCat("attrs.", api_def_attr.rename_to(), "_")
-            : AvoidCPPKeywords(api_def_attr.rename_to());
-    strings::StrAppend(&body, spaces, ".Attr(\"", graph_attr.name(), "\", ",
-                       attr_name, ")\n");
+            ? absl::StrCat("attrs.", safe_attr_name, "_")
+            : AvoidCPPKeywords(safe_attr_name);
+    strings::StrAppend(&body, spaces, ".Attr(\"",
+                       absl::CEscape(graph_attr.name()), "\", ", attr_name,
+                       ")\n");
   }
   absl::StrAppend(&body, "  ;\n");
   absl::StrAppend(&body, "  ", scope_str, ".UpdateBuilder(&builder);\n");
diff --git a/tensorflow/cc/framework/cc_op_gen_util.cc b/tensorflow/cc/framework/cc_op_gen_util.cc
index 3186f29ef2794f..3dec3a1e5e6791 100644
--- a/tensorflow/cc/framework/cc_op_gen_util.cc
+++ b/tensorflow/cc/framework/cc_op_gen_util.cc
@@ -495,6 +495,16 @@ std::string AvoidCPPKeywords(absl::string_view name) {
   return std::string(name);
 }
 
+std::string SafeRenameTo(absl::string_view name, absl::string_view rename_to) {
+  if (tensorflow::IsValidAttrOrArgName(rename_to)) {
+    return std::string(rename_to);
+  }
+  if (tensorflow::IsValidAttrOrArgName(name)) {
+    return std::string(name);
+  }
+  return tensorflow::SanitizeToIdentifier(name);
+}
+
 void InferArgAttributes(
     const OpDef::ArgDef& arg,
     std::unordered_map<std::string, std::string>* inferred_attrs) {
@@ -567,14 +577,24 @@ OpInfo::OpInfo(const OpDef& graph_op_def, const ApiDef& api_def,
     const auto& api_def_arg = *FindInputArg(api_def.arg_order(i), api_def);
     arg_types.push_back(
         absl::StrCat("::tensorflow::", ArgIsList(arg) ? "InputList" : "Input"));
-    arg_names.push_back(AvoidCPPKeywords(api_def_arg.rename_to()));
+    // rename_to() comes from ApiDef, a separate message from the OpDef
+    // arg/attr names. Both are spliced as a raw C++ identifier (parameter
+    // name) below: prefer rename_to() when it's a safe identifier, fall
+    // back to the original arg name when THAT is safe, and only sanitize
+    // as a last resort, since IsValidAttrOrArgName is not enforced at
+    // OpDef registration and a legitimately-registered op's argument name
+    // is not guaranteed to already be a safe identifier (e.g.
+    // TFLite_Detection_PostProcess's "raw_outputs/box_encodings").
+    const std::string safe_input_name =
+        SafeRenameTo(arg.name(), api_def_arg.rename_to());
+    arg_names.push_back(AvoidCPPKeywords(safe_input_name));
 
     // TODO(keveman): Include input type information.
     absl::string_view description = api_def_arg.description();
     if (!description.empty()) {
       ConsumeEquals(&description);
-      absl::StrAppend(&comment, "* ", AvoidCPPKeywords(api_def_arg.rename_to()),
-                      ": ", api_def_arg.description(), "\n");
+      absl::StrAppend(&comment, "* ", AvoidCPPKeywords(safe_input_name), ": ",
+                      api_def_arg.description(), "\n");
     }
   }
 
@@ -593,7 +613,14 @@ OpInfo::OpInfo(const OpDef& graph_op_def, const ApiDef& api_def,
     const auto entry = AttrTypeName(attr.type());
     const auto attr_type_name = entry.first;
     const bool use_const = entry.second;
-    std::string attr_name = AvoidCPPKeywords(api_def_attr.rename_to());
+    // See the safe_input_name comment above: rename_to() and the original
+    // attr name are both unvalidated at OpDef-registration time and both
+    // get spliced as a raw C++ identifier here, so prefer rename_to() when
+    // safe, fall back to the original name when THAT is safe, and only
+    // sanitize as a last resort.
+    const std::string safe_attr_name =
+        SafeRenameTo(attr.name(), api_def_attr.rename_to());
+    std::string attr_name = AvoidCPPKeywords(safe_attr_name);
 
     std::string attr_comment;
     if (!api_def_attr.description().empty()) {
@@ -630,7 +657,10 @@ OpInfo::OpInfo(const OpDef& graph_op_def, const ApiDef& api_def,
     bool is_list = ArgIsList(arg);
     output_types.push_back(
         absl::StrCat("::tensorflow::", is_list ? "OutputList" : "Output"));
-    output_names.push_back(AvoidCPPKeywords(api_def_arg.rename_to()));
+    // See the safe_input_name comment above: the output arg's name and
+    // rename_to() need the same three-tier fallback as the input case.
+    output_names.push_back(
+        AvoidCPPKeywords(SafeRenameTo(arg.name(), api_def_arg.rename_to())));
     is_list_output.push_back(is_list);
   }
 
@@ -694,7 +724,15 @@ std::string OpInfo::GetOpAttrStruct() const {
     const auto entry = AttrTypeName(attr.type());
     const auto attr_type_name = entry.first;
     const bool use_const = entry.second;
-    const std::string camel_case_name = ToCamelCase(api_def_attr.rename_to());
+    // See the safe_input_name comment in OpInfo::OpInfo: attr.name() and
+    // api_def_attr.rename_to() are both unvalidated at OpDef-registration
+    // time, and both are spliced as a raw C++ identifier multiple times
+    // below (the setter name, the field access, the static defaults
+    // function name, and the field declaration) -- computed once here so
+    // every splice site below agrees on the same identifier.
+    const std::string safe_attr_name =
+        SafeRenameTo(attr.name(), api_def_attr.rename_to());
+    const std::string camel_case_name = ToCamelCase(safe_attr_name);
     const std::string suffix =
         (camel_case_name == op_name || camel_case_name == "Attrs") ? "_" : "";
     const std::string attr_func_def =
@@ -713,8 +751,7 @@ std::string OpInfo::GetOpAttrStruct() const {
     absl::StrAppend(&setters, "    TF_MUST_USE_RESULT Attrs ", attr_func_def,
                     " x) {\n");
     absl::StrAppend(&setters, "      Attrs ret = *this;\n");
-    absl::StrAppend(&setters, "      ret.", api_def_attr.rename_to(),
-                    "_ = x;\n");
+    absl::StrAppend(&setters, "      ret.", safe_attr_name, "_ = x;\n");
     absl::StrAppend(&setters, "      return ret;\n    }\n\n");
 
     std::string field_initiliazer;
@@ -724,7 +761,7 @@ std::string OpInfo::GetOpAttrStruct() const {
       // Non-empty lists need static storage for their defaults. Define a
       // function with static local variable that stores the array.
       absl::StrAppend(&defaults_static_storage, "    static ", attr_type_name,
-                      " Default_", api_def_attr.rename_to(), "() {\n");
+                      " Default_", safe_attr_name, "() {\n");
       absl::StrAppend(
           &defaults_static_storage, "      static const ",
           ListElementTypeName(attr.type()), " kStorage[] = ",
@@ -733,14 +770,13 @@ std::string OpInfo::GetOpAttrStruct() const {
       absl::StrAppend(&defaults_static_storage, "      return ", attr_type_name,
                       "(kStorage);\n    }\n");
       // Set the field_initializer to call the defined function.
-      absl::StrAppend(&field_initiliazer, "Default_", api_def_attr.rename_to(),
-                      "()");
+      absl::StrAppend(&field_initiliazer, "Default_", safe_attr_name, "()");
     } else {
       field_initiliazer =
           PrintAttrValue(graph_op_def.name(), api_def_attr.default_value());
     }
-    absl::StrAppend(&struct_fields, "    ", attr_type_name, " ",
-                    api_def_attr.rename_to(), "_ = ", field_initiliazer, ";\n");
+    absl::StrAppend(&struct_fields, "    ", attr_type_name, " ", safe_attr_name,
+                    "_ = ", field_initiliazer, ";\n");
   }
 
   if (struct_fields.empty()) {
diff --git a/tensorflow/cc/framework/cc_op_gen_util.h b/tensorflow/cc/framework/cc_op_gen_util.h
index de954fd3f9fc42..b516792d36c7d6 100644
--- a/tensorflow/cc/framework/cc_op_gen_util.h
+++ b/tensorflow/cc/framework/cc_op_gen_util.h
@@ -106,6 +106,16 @@ bool IsCPPKeyword(absl::string_view name);
 
 std::string AvoidCPPKeywords(absl::string_view name);
 
+// Returns a name safe to splice as a raw C++ identifier: `rename_to` if it
+// is already a safe identifier (see IsValidAttrOrArgName), else `name` if
+// THAT is safe, else a sanitized fallback (see SanitizeToIdentifier).
+// `rename_to` comes from ApiDef, a separate message from the OpDef arg/attr
+// name `name` is drawn from; neither is validated at OpDef-registration
+// time (see op_def_util.cc for why), so both must be checked here rather
+// than trusted, at every place either is spliced as a raw C++ identifier
+// into generated source.
+std::string SafeRenameTo(absl::string_view name, absl::string_view rename_to);
+
 void InferArgAttributes(
     const OpDef::ArgDef& arg,
     std::unordered_map<std::string, std::string>* inferred_attrs);
diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD
index 27dbdb1f739634..dba908f9d31d57 100644
--- a/tensorflow/compiler/jit/BUILD
+++ b/tensorflow/compiler/jit/BUILD
@@ -1139,7 +1139,6 @@ cc_library(
         "cluster_scoping_pass.cc",
         "deadness_analysis.cc",
         "deadness_analysis_internal.h",
-        "disable_functional_ops_lowering_for_xla_pass.cc",
         "encapsulate_subgraphs_pass.cc",
         "encapsulate_xla_computations_pass.cc",
         "extract_outside_compilation_pass.cc",
@@ -1155,7 +1154,6 @@ cc_library(
         "clone_constants_for_better_clustering.h",
         "cluster_scoping_pass.h",
         "deadness_analysis.h",
-        "disable_functional_ops_lowering_for_xla_pass.h",
         "encapsulate_subgraphs_pass.h",
         "encapsulate_xla_computations_pass.h",
         "extract_outside_compilation_pass.h",
@@ -1332,7 +1330,6 @@ tf_cc_test(
         "build_xla_ops_pass_test.cc",
         "clone_constants_for_better_clustering_test.cc",
         "cluster_scoping_pass_test.cc",
-        "disable_functional_ops_lowering_for_xla_pass_test.cc",
         "encapsulate_subgraphs_pass_test.cc",
         "encapsulate_xla_computations_pass_test.cc",
         "extract_outside_compilation_pass_test.cc",
diff --git a/tensorflow/compiler/jit/compilability_check_util.cc b/tensorflow/compiler/jit/compilability_check_util.cc
index 8da8b2055c6c2b..c01514f6dcae1a 100644
--- a/tensorflow/compiler/jit/compilability_check_util.cc
+++ b/tensorflow/compiler/jit/compilability_check_util.cc
@@ -354,6 +354,7 @@ bool RecursiveCompilabilityChecker::OpIsInaccurate(const Node& node) const {
 bool RecursiveCompilabilityChecker::OpIsSlow(const Node& node) const {
   // b/128001705: SelfAdjointEigV2 and Svd performance issues.
   // b/135640736: MatrixInverse performance issues.
+  // MatrixDeterminant and LogMatrixDeterminant use the same QR path as Inverse.
   // b/111271662: MatrixSolve performance issues.
   // https://github.com/tensorflow/tensorflow/pull/31012:
   //    ResizeNearestNeighbor, ResizeBilinear, and ResizeBilinearGrad sometimes
@@ -364,6 +365,8 @@ bool RecursiveCompilabilityChecker::OpIsSlow(const Node& node) const {
   return node.type_string() == "SelfAdjointEigV2" ||
          node.type_string() == "Svd" || node.type_string() == "Qr" ||
          node.type_string() == "MatrixInverse" ||
+         node.type_string() == "MatrixDeterminant" ||
+         node.type_string() == "LogMatrixDeterminant" ||
          node.type_string() == "MatrixSolve" ||
          node.type_string() == "ResizeBilinearGrad" ||
          node.type_string() == "NonMaxSuppressionV3" ||
diff --git a/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.cc b/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.cc
deleted file mode 100644
index e41d765276bc0b..00000000000000
--- a/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.cc
+++ /dev/null
@@ -1,76 +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 "tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.h"
-
-#include "tensorflow/compiler/jit/flags.h"
-#include "tensorflow/compiler/jit/xla_cluster_util.h"
-#include "tensorflow/core/framework/node_def_util.h"
-#include "tensorflow/core/graph/graph.h"
-#include "tensorflow/core/protobuf/config.pb.h"
-
-namespace tensorflow {
-
-namespace {
-
-// Returns the effective global JIT level for `options`. Unlike
-// `GetGlobalJitLevelForGraph`, this also works when `options.session_options`
-// is null, which can happen when auto-clustering is enabled purely via the
-// TF_XLA_FLAGS=--tf_xla_auto_jit=N environment variable and the execution
-// path never populates a SessionOptions. In that case we combine just the
-// flag-provided auto-jit level with whether the graph is a single-GPU graph,
-// mirroring what `GetGlobalJitLevelForGraph` does for a DEFAULT ConfigProto
-// setting.
-OptimizerOptions::GlobalJitLevel GetEffectiveGlobalJitLevel(
-    const GraphOptimizationPassOptions& options, const Graph& graph) {
-  if (options.session_options != nullptr) {
-    return GetGlobalJitLevelForGraph(options);
-  }
-
-  const XlaAutoJitFlag& auto_jit_flag =
-      GetMarkForCompilationPassFlags()->xla_auto_jit_flag;
-  auto level_or_off = [](int32_t level) {
-    return level == OptimizerOptions::DEFAULT
-               ? OptimizerOptions::OFF
-               : static_cast<OptimizerOptions::GlobalJitLevel>(level);
-  };
-  return IsSingleGpuGraph(graph)
-             ? level_or_off(auto_jit_flag.optimization_level_single_gpu)
-             : level_or_off(auto_jit_flag.optimization_level_general);
-}
-
-}  // namespace
-
-absl::Status DisableFunctionalOpsLoweringForXlaPass::Run(
-    const GraphOptimizationPassOptions& options) {
-  if (options.graph == nullptr || options.graph->get() == nullptr) {
-    return absl::OkStatus();
-  }
-  Graph* graph = options.graph->get();
-  if (GetEffectiveGlobalJitLevel(options, *graph) < OptimizerOptions::ON_1) {
-    return absl::OkStatus();
-  }
-  for (Node* n : graph->op_nodes()) {
-    if (!n->IsIfNode() && !n->IsCaseNode() && !n->IsWhileNode()) continue;
-    bool lower = false;
-    if (TryGetNodeAttr(n->attrs(), "_lower_using_switch_merge", &lower) &&
-        lower) {
-      n->ClearAttr("_lower_using_switch_merge");
-    }
-  }
-  return absl::OkStatus();
-}
-
-}  // namespace tensorflow
diff --git a/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.h b/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.h
deleted file mode 100644
index cee7eb28187c11..00000000000000
--- a/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.h
+++ /dev/null
@@ -1,41 +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.
-==============================================================================*/
-
-#ifndef TENSORFLOW_COMPILER_JIT_DISABLE_FUNCTIONAL_OPS_LOWERING_FOR_XLA_PASS_H_
-#define TENSORFLOW_COMPILER_JIT_DISABLE_FUNCTIONAL_OPS_LOWERING_FOR_XLA_PASS_H_
-
-#include "tensorflow/core/common_runtime/optimization_registry.h"
-
-namespace tensorflow {
-
-// Auto-clustering lowers functional control flow ops (While/If/Case) into
-// Switch/Merge before MarkForCompilationPass runs. That can cause only
-// fragments of a loop/branch body to end up in a cluster instead of the whole
-// op, which silently corrupts results for some patterns. jit_compile=True
-// never lowers these ops in the first place, so it doesn't hit this.
-//
-// This pass clears the `_lower_using_switch_merge` attribute on functional
-// control flow nodes whenever global JIT is enabled, so auto-clustering skips
-// the lowering the same way jit_compile=True already does.
-class DisableFunctionalOpsLoweringForXlaPass : public GraphOptimizationPass {
- public:
-  DisableFunctionalOpsLoweringForXlaPass() = default;
-
-  absl::Status Run(const GraphOptimizationPassOptions& options) override;
-};
-
-}  // namespace tensorflow
-
-#endif  // TENSORFLOW_COMPILER_JIT_DISABLE_FUNCTIONAL_OPS_LOWERING_FOR_XLA_PASS_H_
diff --git a/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass_test.cc b/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass_test.cc
deleted file mode 100644
index a0adc449288967..00000000000000
--- a/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass_test.cc
+++ /dev/null
@@ -1,193 +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 "tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.h"
-
-#include <memory>
-#include <vector>
-
-#include "tensorflow/compiler/jit/flags.h"
-#include "tensorflow/core/framework/attr_value.pb.h"
-#include "tensorflow/core/framework/node_def_builder.h"
-#include "tensorflow/core/framework/node_def_util.h"
-#include "tensorflow/core/framework/op.h"
-#include "tensorflow/core/framework/types.h"
-#include "tensorflow/core/graph/graph.h"
-#include "tensorflow/core/graph/node_builder.h"
-#include "tensorflow/core/lib/core/status_test_util.h"
-#include "tensorflow/core/platform/test.h"
-#include "tensorflow/core/protobuf/config.pb.h"
-#include "tensorflow/core/public/session_options.h"
-
-namespace tensorflow {
-namespace {
-
-absl::Status AddWhileNode(Graph* graph, absl::string_view name, Node** node) {
-  return NodeBuilder(name, "While")
-      .Input(std::vector<NodeBuilder::NodeOut>{})
-      .Attr("_lower_using_switch_merge", true)
-      .Attr("cond", NameAttrList())
-      .Attr("body", NameAttrList())
-      .Attr("T", DataTypeVector{})
-      .Finalize(graph, node);
-}
-
-absl::Status AddPlaceholder(Graph* graph, absl::string_view name,
-                            DataType dtype, Node** node) {
-  return NodeBuilder(name, "Placeholder")
-      .Attr("dtype", dtype)
-      .Finalize(graph, node);
-}
-
-absl::Status AddIfNode(Graph* graph, absl::string_view name, Node* cond,
-                       Node** node) {
-  return NodeBuilder(name, "If")
-      .Input(cond, 0)
-      .Input(std::vector<NodeBuilder::NodeOut>{})
-      .Attr("_lower_using_switch_merge", true)
-      .Attr("Tcond", DT_BOOL)
-      .Attr("Tin", DataTypeVector{})
-      .Attr("Tout", DataTypeVector{})
-      .Attr("then_branch", NameAttrList())
-      .Attr("else_branch", NameAttrList())
-      .Finalize(graph, node);
-}
-
-absl::Status AddCaseNode(Graph* graph, absl::string_view name,
-                         Node* branch_index, Node** node) {
-  return NodeBuilder(name, "Case")
-      .Input(branch_index, 0)
-      .Input(std::vector<NodeBuilder::NodeOut>{})
-      .Attr("_lower_using_switch_merge", true)
-      .Attr("Tin", DataTypeVector{})
-      .Attr("Tout", DataTypeVector{})
-      .Attr("branches", std::vector<NameAttrList>{NameAttrList()})
-      .Finalize(graph, node);
-}
-
-bool HasLowerUsingSwitchMergeAttr(const Node* n) {
-  bool lower = false;
-  return TryGetNodeAttr(n->attrs(), "_lower_using_switch_merge", &lower) &&
-         lower;
-}
-
-absl::Status RunPass(std::unique_ptr<Graph>* graph,
-                     SessionOptions* session_options) {
-  GraphOptimizationPassOptions options;
-  options.graph = graph;
-  options.session_options = session_options;
-  DisableFunctionalOpsLoweringForXlaPass pass;
-  return pass.Run(options);
-}
-
-TEST(DisableFunctionalOpsLoweringForXlaPassTest,
-     ClearsAttrOnFunctionalNodesWhenGlobalJitEnabled) {
-  auto graph = std::make_unique<Graph>(OpRegistry::Global());
-  Node* while_node;
-  TF_ASSERT_OK(AddWhileNode(graph.get(), "my_while", &while_node));
-  Node* cond;
-  TF_ASSERT_OK(AddPlaceholder(graph.get(), "cond", DT_BOOL, &cond));
-  Node* if_node;
-  TF_ASSERT_OK(AddIfNode(graph.get(), "my_if", cond, &if_node));
-  Node* branch_index;
-  TF_ASSERT_OK(
-      AddPlaceholder(graph.get(), "branch_index", DT_INT32, &branch_index));
-  Node* case_node;
-  TF_ASSERT_OK(AddCaseNode(graph.get(), "my_case", branch_index, &case_node));
-
-  SessionOptions session_options;
-  session_options.config.mutable_graph_options()
-      ->mutable_optimizer_options()
-      ->set_global_jit_level(OptimizerOptions::ON_1);
-
-  TF_ASSERT_OK(RunPass(&graph, &session_options));
-
-  EXPECT_FALSE(HasLowerUsingSwitchMergeAttr(while_node));
-  EXPECT_FALSE(HasLowerUsingSwitchMergeAttr(if_node));
-  EXPECT_FALSE(HasLowerUsingSwitchMergeAttr(case_node));
-}
-
-TEST(DisableFunctionalOpsLoweringForXlaPassTest,
-     PreservesAttrWhenGlobalJitDisabled) {
-  auto graph = std::make_unique<Graph>(OpRegistry::Global());
-  Node* while_node;
-  TF_ASSERT_OK(AddWhileNode(graph.get(), "my_while", &while_node));
-
-  SessionOptions session_options;
-  session_options.config.mutable_graph_options()
-      ->mutable_optimizer_options()
-      ->set_global_jit_level(OptimizerOptions::OFF);
-
-  TF_ASSERT_OK(RunPass(&graph, &session_options));
-
-  EXPECT_TRUE(HasLowerUsingSwitchMergeAttr(while_node));
-}
-
-TEST(DisableFunctionalOpsLoweringForXlaPassTest, NoOpWhenGraphIsMissing) {
-  SessionOptions session_options;
-  session_options.config.mutable_graph_options()
-      ->mutable_optimizer_options()
-      ->set_global_jit_level(OptimizerOptions::ON_2);
-
-  GraphOptimizationPassOptions options;
-  options.session_options = &session_options;
-  DisableFunctionalOpsLoweringForXlaPass pass;
-  TF_ASSERT_OK(pass.Run(options));
-}
-
-class DisableFunctionalOpsLoweringForXlaPassFlagsTest : public ::testing::Test {
- protected:
-  void SetUp() override {
-    flags_ = GetMarkForCompilationPassFlags();
-    original_ = flags_->xla_auto_jit_flag;
-  }
-  void TearDown() override { flags_->xla_auto_jit_flag = original_; }
-
-  MarkForCompilationPassFlags* flags_;
-  XlaAutoJitFlag original_;
-};
-
-TEST_F(DisableFunctionalOpsLoweringForXlaPassFlagsTest,
-       FallsBackToFlagsWhenSessionOptionsIsNull) {
-  auto graph = std::make_unique<Graph>(OpRegistry::Global());
-  Node* while_node;
-  TF_ASSERT_OK(AddWhileNode(graph.get(), "my_while", &while_node));
-
-  flags_->xla_auto_jit_flag.optimization_level_single_gpu =
-      OptimizerOptions::ON_2;
-  flags_->xla_auto_jit_flag.optimization_level_general = OptimizerOptions::ON_2;
-
-  TF_ASSERT_OK(RunPass(&graph, /*session_options=*/nullptr));
-
-  EXPECT_FALSE(HasLowerUsingSwitchMergeAttr(while_node));
-}
-
-TEST_F(DisableFunctionalOpsLoweringForXlaPassFlagsTest,
-       PreservesAttrWhenSessionOptionsIsNullAndFlagsAreOff) {
-  auto graph = std::make_unique<Graph>(OpRegistry::Global());
-  Node* while_node;
-  TF_ASSERT_OK(AddWhileNode(graph.get(), "my_while", &while_node));
-
-  flags_->xla_auto_jit_flag.optimization_level_single_gpu =
-      OptimizerOptions::OFF;
-  flags_->xla_auto_jit_flag.optimization_level_general = OptimizerOptions::OFF;
-
-  TF_ASSERT_OK(RunPass(&graph, /*session_options=*/nullptr));
-
-  EXPECT_TRUE(HasLowerUsingSwitchMergeAttr(while_node));
-}
-
-}  // namespace
-}  // namespace tensorflow
diff --git a/tensorflow/compiler/jit/jit_compilation_pass_registration.cc b/tensorflow/compiler/jit/jit_compilation_pass_registration.cc
index 2abdc413e8725e..a4a9c7d7e11863 100644
--- a/tensorflow/compiler/jit/jit_compilation_pass_registration.cc
+++ b/tensorflow/compiler/jit/jit_compilation_pass_registration.cc
@@ -16,7 +16,6 @@ limitations under the License.
 #include "tensorflow/compiler/jit/build_xla_ops_pass.h"
 #include "tensorflow/compiler/jit/clone_constants_for_better_clustering.h"
 #include "tensorflow/compiler/jit/cluster_scoping_pass.h"
-#include "tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.h"
 #include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h"
 #include "tensorflow/compiler/jit/encapsulate_xla_computations_pass.h"
 #include "tensorflow/compiler/jit/force_xla_constants_on_host_pass.h"
@@ -30,9 +29,6 @@ namespace tensorflow {
 
 // PRE_PLACEMENT passes:
 
-REGISTER_OPTIMIZATION(OptimizationPassRegistry::PRE_PLACEMENT, 5,
-                      DisableFunctionalOpsLoweringForXlaPass);
-
 // EncapsulateXlaComputationsPass rewrites computations generated by the
 // xla.compile() Python code into XlaLaunch nodes.
 REGISTER_OPTIMIZATION(OptimizationPassRegistry::PRE_PLACEMENT, 36,
diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc
index 261ebd3560825a..2fee4832134ae9 100644
--- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc
+++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc
@@ -2152,6 +2152,7 @@ absl::flat_hash_set<std::string> GetKnownXLAAllowlistOp() {
       "LowerBound",
       "MatMul",
       "MatrixBandPart",
+      "MatrixDeterminant",
       "MatrixDiag",
       "MatrixDiagPart",
       "MatrixDiagPartV2",
diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index a0bfc154565e20..5251a2953e7092 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -172,6 +172,7 @@ tf_xla_combined_py_test(
     ],
     tests = [
         # go/keep-sorted start
+        ":matrix_determinant_op_test_lib",
         ":matrix_inverse_op_test_lib",
         ":matrix_solve_op_test_lib",
         ":momentum_test_lib",
@@ -603,6 +604,26 @@ tf_xla_py_strict_test(
     ],
 )
 
+tf_xla_py_strict_test(
+    name = "matrix_determinant_op_test",
+    size = "small",
+    timeout = "moderate",
+    srcs = ["matrix_determinant_op_test.py"],
+    tags = [
+        "no_pip",  # TODO(b/149738646): fix pip install so these tests run on kokoro pip
+        "notap",
+    ],
+    deps = [
+        ":xla_test",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:linalg_ops",
+        "//tensorflow/python/ops:linalg_ops_gen",
+        "//tensorflow/python/platform:test",
+        "//third_party/py/numpy",
+    ],
+)
+
 tf_xla_py_strict_test(
     name = "matrix_inverse_op_test",
     size = "small",
diff --git a/tensorflow/compiler/tests/jit_test.py b/tensorflow/compiler/tests/jit_test.py
index 80f8bdd57c51a9..32787da0b9ecc7 100644
--- a/tensorflow/compiler/tests/jit_test.py
+++ b/tensorflow/compiler/tests/jit_test.py
@@ -642,34 +642,6 @@ def CompiledFunction(x):
       self.assertTrue(InLabels(RunMetadataLabels(run_metadata), "_XlaRun"))
 
 
-class AutoClusteringWhileLoopTest(test.TestCase):
-  """Regression test for auto-clustering silently dropping while_loop output.
-
-  Auto-clustering used to lower functional control flow ops (While/If/Case)
-  to Switch/Merge before MarkForCompilationPass ran, so only fragments of a
-  loop body ended up in a cluster instead of the whole While op. That could
-  silently produce wrong results for some loop bodies, unlike
-  jit_compile=True which never lowers these ops in the first place.
-  """
-
-  def testWhileLoopOutputIsCorrectUnderAutoClustering(self):
-    config = NoRewriteSessionConfig()
-    config.graph_options.optimizer_options.global_jit_level = (
-        config_pb2.OptimizerOptions.ON_2
-    )
-
-    with session_lib.Session(config=config) as sess:
-      x = array_ops.placeholder(dtypes.float32)
-      c = lambda i, _: math_ops.less(i, 5)
-      b = lambda i, x: (i + 1, x + 1.0)
-      _, result = while_loop.while_loop(c, b, (constant_op.constant(0), x))
-
-      output = test_utils.RunWithWarmup(
-          sess, result, {x: np.array([1.0, 2.0, 3.0], dtype=np.float32)}
-      )
-      self.assertAllClose(output, np.array([6.0, 7.0, 8.0], dtype=np.float32))
-
-
 if __name__ == "__main__":
   os.environ["TF_XLA_FLAGS"] = ("--tf_xla_enable_lazy_compilation=true " +
                                 os.environ.get("TF_XLA_FLAGS", ""))
diff --git a/tensorflow/compiler/tests/matrix_determinant_op_test.py b/tensorflow/compiler/tests/matrix_determinant_op_test.py
new file mode 100644
index 00000000000000..66c9b4bbe3cbfc
--- /dev/null
+++ b/tensorflow/compiler/tests/matrix_determinant_op_test.py
@@ -0,0 +1,129 @@
+# 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 XLA implementations of matrix determinant ops."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_linalg_ops
+from tensorflow.python.ops import linalg_ops
+from tensorflow.python.platform import googletest
+
+
+class DeterminantOpTest(xla_test.XLATestCase):
+
+  def _verifyDeterminant(self, x, np_type):
+    y = x.astype(np_type)
+    if y.shape[-1] == 0 and y.shape[-2] == 0:
+      np_det = np.ones(y.shape[:-2], dtype=np_type)
+      np_sign = np.ones(y.shape[:-2], dtype=np_type)
+      np_log_abs = np.zeros(y.shape[:-2], dtype=np_type)
+    else:
+      np_det = np.array(np.linalg.det(y)).astype(np_type)
+      np_sign, np_log_abs = np.linalg.slogdet(y)
+      np_sign = np.array(np_sign).astype(np_type)
+      np_log_abs = np.array(np_log_abs).astype(np_type)
+
+    with self.session() as sess:
+      p = array_ops.placeholder(dtypes.as_dtype(y.dtype), y.shape, name="x")
+      with self.test_scope():
+        det = linalg_ops.matrix_determinant(p)
+        sign, log_abs = gen_linalg_ops.log_matrix_determinant(p)
+      det_out, sign_out, log_abs_out = sess.run(
+          [det, sign, log_abs], feed_dict={p: y}
+      )
+
+    self.assertAllClose(np_det, det_out, rtol=1e-3, atol=1e-3)
+    self.assertShapeEqual(np_det, det)
+    # Compare reconstructed determinants so a QR-vs-LU split of sign/log does
+    # not fail the test when the product still matches. Guard exp() so a
+    # singular matrix's -inf log-abs-det does not warn or overflow.
+    with np.errstate(over="ignore", invalid="ignore"):
+      np_recon = np_sign * np.exp(np_log_abs)
+      tf_recon = sign_out * np.exp(log_abs_out)
+    self.assertAllClose(np_recon, tf_recon, rtol=1e-3, atol=1e-3)
+    self.assertShapeEqual(np_sign, sign)
+    self.assertShapeEqual(np_log_abs, log_abs)
+    # Householder QR may return sign ±1 with a large negative log-abs-det for
+    # singular matrices (NumPy slogdet uses sign 0). Only compare components
+    # where NumPy reports a nonzero sign.
+    non_singular = np.abs(np_sign) > 0
+    if np.any(non_singular):
+      self.assertAllClose(
+          np.where(non_singular, sign_out, 0),
+          np.where(non_singular, np_sign, 0),
+          rtol=1e-3,
+          atol=1e-3,
+      )
+      self.assertAllClose(
+          np.where(non_singular, log_abs_out, 0),
+          np.where(non_singular, np_log_abs, 0),
+          rtol=1e-3,
+          atol=1e-3,
+      )
+
+  def _verifyDeterminantReal(self, x):
+    for np_type in self.float_types & {np.float32, np.float64}:
+      self._verifyDeterminant(x, np_type)
+
+  def testBasic(self):
+    # 1x1
+    self._verifyDeterminantReal(np.array([[7.0]]))
+    # 2x2 with negative determinant: det([[1, 2], [3, 4]]) == -2.
+    # This is the case xla::LogDet() gets wrong (returns NaN).
+    self._verifyDeterminantReal(np.array([[1.0, 2.0], [3.0, 4.0]]))
+    # 2x2 with positive determinant (the motivating jit_compile example).
+    self._verifyDeterminantReal(np.array([[4.0, 7.0], [2.0, 6.0]]))
+    # Rank-deficient but not all-zero. An all-zero matrix hits division by
+    # zero in Householder QR on TPU (SLogDet returns NaN instead of 0).
+    self._verifyDeterminantReal(np.array([[1.0, 2.0], [2.0, 4.0]]))
+    # 3x3 with negative determinant.
+    self._verifyDeterminantReal(
+        np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, -1.0]])
+    )
+    # Well-conditioned 5x5 triangular matrix; det = 2*3*4*5*6 = 720.
+    self._verifyDeterminantReal(
+        np.array([
+            [2.0, 0.0, 0.0, 0.0, 0.0],
+            [1.0, 3.0, 0.0, 0.0, 0.0],
+            [0.0, 1.0, 4.0, 0.0, 0.0],
+            [0.0, 0.0, 1.0, 5.0, 0.0],
+            [0.0, 0.0, 0.0, 1.0, 6.0],
+        ])
+    )
+
+  def testBatch(self):
+    # Mixed signs in the batch: dets are -2 and 3.
+    self._verifyDeterminantReal(
+        np.array([[[1.0, 2.0], [3.0, 4.0]], [[2.0, 1.0], [1.0, 2.0]]])
+    )
+    matrix1 = np.array([[1.0, 2.0], [3.0, 4.0]])
+    matrix2 = np.array([[1.0, 3.0], [3.0, 5.0]])
+    batch = np.concatenate(
+        [np.expand_dims(matrix1, 0), np.expand_dims(matrix2, 0)]
+    )
+    batch = np.tile(batch, [2, 3, 1, 1])
+    self._verifyDeterminantReal(batch)
+
+  def testEmpty(self):
+    self._verifyDeterminantReal(np.empty([0, 0]))
+    self._verifyDeterminantReal(np.empty([0, 2, 2]))
+    self._verifyDeterminantReal(np.empty([2, 0, 0]))
+
+
+if __name__ == "__main__":
+  googletest.main()
diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD
index a617d5a1af9af2..df541ca511272c 100644
--- a/tensorflow/compiler/tf2xla/kernels/BUILD
+++ b/tensorflow/compiler/tf2xla/kernels/BUILD
@@ -112,6 +112,7 @@ cc_library(
         ":lrn_ops",
         ":matmul_op",
         ":matrix_band_part_op",
+        ":matrix_determinant_op",
         ":matrix_diag_ops",
         ":matrix_inverse_op",
         ":matrix_solve_op",
@@ -1027,6 +1028,25 @@ tf_kernel_library(
     ],
 )
 
+tf_kernel_library(
+    name = "matrix_determinant_op",
+    srcs = ["matrix_determinant_op.cc"],
+    deps = [
+        "//tensorflow/compiler/tf2xla:xla_compilation_device",
+        "//tensorflow/compiler/tf2xla:xla_compiler",
+        "//tensorflow/compiler/tf2xla:xla_context",
+        "//tensorflow/compiler/tf2xla:xla_op_registry",
+        "//tensorflow/compiler/tf2xla:xla_resource",
+        "//tensorflow/compiler/tf2xla/ops:xla_ops",
+        "//tensorflow/core:framework",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@xla//xla/hlo/builder:xla_builder",
+        "@xla//xla/hlo/builder/lib:constants",
+        "@xla//xla/hlo/builder/lib:logdet",
+    ],
+)
+
 tf_kernel_library(
     name = "matrix_inverse_op",
     srcs = ["matrix_inverse_op.cc"],
diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_determinant_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_determinant_op.cc
new file mode 100644
index 00000000000000..5d3c8892a51e25
--- /dev/null
+++ b/tensorflow/compiler/tf2xla/kernels/matrix_determinant_op.cc
@@ -0,0 +1,117 @@
+/* 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 "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/lib/logdet.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+
+namespace tensorflow {
+namespace {
+
+absl::Status CheckSquareMatrix(const TensorShape& input_shape) {
+  const int64_t ndims = input_shape.dims();
+  if (ndims < 2) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Input must have rank >= 2, got ", ndims));
+  }
+  if (input_shape.dim_size(ndims - 2) != input_shape.dim_size(ndims - 1)) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Input matrices must be square, got ", input_shape.dim_size(ndims - 2),
+        " != ", input_shape.dim_size(ndims - 1)));
+  }
+  return absl::OkStatus();
+}
+
+// Broadcasts a scalar to the batch shape of a [..., n, n] matrix input.
+xla::XlaOp BroadcastScalarToBatch(xla::XlaOp scalar,
+                                  const TensorShape& input_shape) {
+  std::vector<int64_t> batch_dims(input_shape.dims() - 2);
+  for (int i = 0; i < input_shape.dims() - 2; ++i) {
+    batch_dims[i] = input_shape.dim_size(i);
+  }
+  return xla::Broadcast(scalar, batch_dims);
+}
+
+// slogdet(A) = (sign, log|det|). For n == 0, det is defined to be 1.
+// xla::SLogDet cannot handle that case: it slices Householder taus to n-1.
+xla::SignAndLogDet SLogDetOrEmpty(XlaOpKernelContext* ctx,
+                                  const TensorShape& input_shape) {
+  const int64_t n = input_shape.dim_size(input_shape.dims() - 1);
+  if (n == 0) {
+    const xla::PrimitiveType type = ctx->input_xla_type(0);
+    return xla::SignAndLogDet{
+        BroadcastScalarToBatch(xla::One(ctx->builder(), type), input_shape),
+        BroadcastScalarToBatch(xla::Zero(ctx->builder(), type), input_shape)};
+  }
+  return xla::SLogDet(ctx->Input(0));
+}
+
+class MatrixDeterminantOp : public XlaOpKernel {
+ public:
+  explicit MatrixDeterminantOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape input_shape = ctx->InputShape(0);
+    OP_REQUIRES_OK(ctx, CheckSquareMatrix(input_shape));
+
+    // det = sign * exp(log|det|). Do not use xla::LogDet(): that returns NaN
+    // for matrices with a negative determinant.
+    const xla::SignAndLogDet slogdet = SLogDetOrEmpty(ctx, input_shape);
+    ctx->SetOutput(0, slogdet.sign * xla::Exp(slogdet.logdet));
+  }
+
+ private:
+  MatrixDeterminantOp(const MatrixDeterminantOp&) = delete;
+  void operator=(const MatrixDeterminantOp&) = delete;
+};
+
+class LogMatrixDeterminantOp : public XlaOpKernel {
+ public:
+  explicit LogMatrixDeterminantOp(OpKernelConstruction* ctx)
+      : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape input_shape = ctx->InputShape(0);
+    OP_REQUIRES_OK(ctx, CheckSquareMatrix(input_shape));
+
+    const xla::SignAndLogDet slogdet = SLogDetOrEmpty(ctx, input_shape);
+    ctx->SetOutput(0, slogdet.sign);
+    ctx->SetOutput(1, slogdet.logdet);
+  }
+
+ private:
+  LogMatrixDeterminantOp(const LogMatrixDeterminantOp&) = delete;
+  void operator=(const LogMatrixDeterminantOp&) = delete;
+};
+
+// TODO(b/135640736): Allow complex types once XLA QR/SLogDet is validated for
+// them, matching MatrixInverse.
+REGISTER_XLA_OP(Name("MatrixDeterminant").TypeConstraint("T", kFloatTypes),
+                MatrixDeterminantOp);
+REGISTER_XLA_OP(Name("LogMatrixDeterminant").TypeConstraint("T", kFloatTypes),
+                LogMatrixDeterminantOp);
+
+}  // namespace
+}  // namespace tensorflow
diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD
index 01e7ba21fa66f0..35785f02ee3461 100644
--- a/tensorflow/core/framework/BUILD
+++ b/tensorflow/core/framework/BUILD
@@ -1040,6 +1040,7 @@ cc_library(
         "//tensorflow/core/platform:types",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/strings",
     ],
 )
 
diff --git a/tensorflow/core/framework/op_def_util.cc b/tensorflow/core/framework/op_def_util.cc
index 0a3018a3863bb3..0e61098cc0fd2e 100644
--- a/tensorflow/core/framework/op_def_util.cc
+++ b/tensorflow/core/framework/op_def_util.cc
@@ -21,6 +21,7 @@ limitations under the License.
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/strings/ascii.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/core/framework/attr_value_util.h"
 #include "tensorflow/core/framework/op_def.pb.h"
@@ -267,6 +268,64 @@ bool IsValidOpName(absl::string_view sp) {
   }
 }
 
+// Attribute and argument names are not restricted to CamelCase like op
+// names. Some code generators, such as cc_op_gen.cc and python_op_gen.cc,
+// splice a name verbatim (unescaped) into generated C++ or Python wrapper
+// source as a raw identifier; this function tells such a call site whether
+// a given name is safe to use that way as-is. It is not, and must not be
+// used as, a runtime validity check on OpDef registration: legitimate,
+// already-registered ops can and do use argument names outside this safe
+// set for reasons unrelated to code generation (see SanitizeToIdentifier's
+// comment below for a concrete example), and rejecting them at
+// registration time breaks that legitimate usage.
+bool IsValidAttrOrArgName(absl::string_view sp) {
+  if (sp.empty()) return false;
+  if (sp[0] != '_' && !absl::ascii_isalpha(sp[0])) {
+    return false;
+  }
+  for (size_t i = 1; i < sp.size(); ++i) {
+    char c = sp[i];
+    if (c != '_' && !absl::ascii_isalnum(c)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+// Converts `sp` into a safe identifier by replacing any character outside
+// IsValidAttrOrArgName's safe set with an underscore, and prefixing with an
+// underscore if the result would otherwise start with a digit or be empty.
+//
+// Some legitimate, already-registered OpDefs use argument names outside
+// that safe set for reasons unrelated to code generation -- for example
+// TFLite_Detection_PostProcess registers inputs like
+// "raw_outputs/box_encodings" and outputs like
+// "TFLite_Detection_PostProcess:1". Rejecting such an OpDef at
+// registration/validation time (as ValidateOpDef/ValidateArg once did)
+// breaks that legitimate usage; the actual splice-into-generated-source
+// risk IsValidAttrOrArgName exists for only applies where a name is used
+// as a raw identifier by a code generator such as cc_op_gen.cc or
+// python_op_gen.cc. This function gives those call sites a fallback that
+// is always a syntactically safe identifier, so they can sanitize instead
+// of rejecting -- callers should first check IsValidAttrOrArgName and use
+// the name as-is when it already passes, calling this only for a name
+// that doesn't.
+std::string SanitizeToIdentifier(absl::string_view sp) {
+  std::string result;
+  result.reserve(sp.size() + 1);
+  for (char c : sp) {
+    if (c == '_' || absl::ascii_isalnum(c)) {
+      result.push_back(c);
+    } else {
+      result.push_back('_');
+    }
+  }
+  if (result.empty() || (!absl::ascii_isalpha(result[0]) && result[0] != '_')) {
+    result.insert(result.begin(), '_');
+  }
+  return result;
+}
+
 absl::Status ValidateOpDef(const OpDef& op_def) {
   if (!absl::StartsWith(op_def.name(), "_")) {
     VALIDATE(IsValidOpName(op_def.name()), "Invalid name: ", op_def.name(),
diff --git a/tensorflow/core/framework/op_def_util.h b/tensorflow/core/framework/op_def_util.h
index abaaeefb03c9a8..7bf554a21a934f 100644
--- a/tensorflow/core/framework/op_def_util.h
+++ b/tensorflow/core/framework/op_def_util.h
@@ -29,6 +29,21 @@ limitations under the License.
 
 namespace tensorflow {
 
+// Returns true iff sp is a safe identifier: starts with a letter or
+// underscore, and contains only letters, digits, and underscores. Used to
+// validate any name (OpDef attr/arg names, or ApiDef rename_to values)
+// that downstream code generators (cc_op_gen.cc, python_op_gen.cc) splice
+// into generated C++/Python source, either as identifiers or into string
+// literals. Not a runtime validity check on OpDef registration -- see the
+// definition in op_def_util.cc for why.
+bool IsValidAttrOrArgName(absl::string_view sp);
+
+// Converts `sp` into a safe identifier (see IsValidAttrOrArgName) by
+// replacing unsafe characters with underscores. A code generator that
+// needs a raw identifier from a name that fails IsValidAttrOrArgName
+// should use this rather than the name itself.
+std::string SanitizeToIdentifier(absl::string_view sp);
+
 // Performs a consistency check across the fields of the op_def.
 absl::Status ValidateOpDef(const OpDef& op_def);
 
diff --git a/tensorflow/core/framework/op_def_util_test.cc b/tensorflow/core/framework/op_def_util_test.cc
index 41fd90d4e79fcf..df9b64bd86ee61 100644
--- a/tensorflow/core/framework/op_def_util_test.cc
+++ b/tensorflow/core/framework/op_def_util_test.cc
@@ -68,6 +68,62 @@ void ExpectFailure(const absl::Status& status, const std::string& message) {
 }
 }  // namespace
 
+// IsValidAttrOrArgName and SanitizeToIdentifier are code-generator-facing
+// utilities, not a runtime validity constraint on OpDef registration --
+// ValidateOpDef/ValidateArg deliberately do not call them (see
+// OpDefValidAcceptsNonIdentifierNames below and the comment on
+// IsValidAttrOrArgName's definition for why).
+TEST(IsValidAttrOrArgNameTest, AcceptsSafeIdentifiers) {
+  EXPECT_TRUE(IsValidAttrOrArgName("dtype"));
+  EXPECT_TRUE(IsValidAttrOrArgName("T"));
+  EXPECT_TRUE(IsValidAttrOrArgName("num_threads_2"));
+  EXPECT_TRUE(IsValidAttrOrArgName("_internal"));
+}
+
+TEST(IsValidAttrOrArgNameTest, RejectsUnsafeNames) {
+  // The exact malicious payload from GHSA-2xp6-8g4h-qw72.
+  EXPECT_FALSE(
+      IsValidAttrOrArgName("evil\"); system(\"touch /tmp/PWNED\"); //"));
+  // A legitimate, already-registered op's argument name that is not a
+  // safe identifier for reasons unrelated to code generation.
+  EXPECT_FALSE(IsValidAttrOrArgName("raw_outputs/box_encodings"));
+  EXPECT_FALSE(IsValidAttrOrArgName("TFLite_Detection_PostProcess:1"));
+  EXPECT_FALSE(IsValidAttrOrArgName(""));
+  EXPECT_FALSE(IsValidAttrOrArgName("1abc"));
+}
+
+TEST(SanitizeToIdentifierTest, ProducesASafeIdentifierForAnyInput) {
+  // Every unsafe character becomes an underscore; the malicious payload
+  // that IsValidAttrOrArgName exists to reject becomes an inert, safe
+  // identifier string rather than being spliced verbatim.
+  EXPECT_EQ(SanitizeToIdentifier("evil\"); system(\"touch /tmp/PWNED\"); //"),
+            "evil____system__touch__tmp_PWNED______");
+  EXPECT_EQ(SanitizeToIdentifier("raw_outputs/box_encodings"),
+            "raw_outputs_box_encodings");
+  EXPECT_EQ(SanitizeToIdentifier("TFLite_Detection_PostProcess:1"),
+            "TFLite_Detection_PostProcess_1");
+  // A result that would start with a digit gets a leading underscore
+  // rather than becoming a syntactically invalid identifier.
+  EXPECT_EQ(SanitizeToIdentifier("123_invalid"), "_123_invalid");
+  EXPECT_EQ(SanitizeToIdentifier(""), "_");
+  // A name that is already safe is unaffected other than round-tripping
+  // through the same character-by-character pass.
+  EXPECT_EQ(SanitizeToIdentifier("already_safe"), "already_safe");
+}
+
+// The regression this fix closes: ValidateOpDef must accept an OpDef using
+// argument names outside IsValidAttrOrArgName's safe set, matching a real,
+// already-registered op (TFLite_Detection_PostProcess, registered in
+// tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc) that
+// broke under the runtime rejection this fix removes.
+TEST_F(ValidateOpDefTest, OpDefValidAcceptsNonIdentifierNames) {
+  TF_EXPECT_OK(TestProto(
+      "name: 'TFLite_Detection_PostProcess' "
+      "input_arg { name: 'raw_outputs/box_encodings' type: DT_FLOAT } "
+      "input_arg { name: 'raw_outputs/class_predictions' type: DT_FLOAT } "
+      "output_arg { name: 'TFLite_Detection_PostProcess:1' type: DT_FLOAT }"));
+}
+
 TEST_F(ValidateOpDefTest, OpDefValid) {
   TF_EXPECT_OK(TestBuilder(OpDefBuilder("X").Attr("a: int")));
   TF_EXPECT_OK(TestBuilder(OpDefBuilder("X").Input("a: int32")));
diff --git a/tensorflow/lite/delegates/gpu/delegate.cc b/tensorflow/lite/delegates/gpu/delegate.cc
index e7445905aacced..e472ebb1db4539 100644
--- a/tensorflow/lite/delegates/gpu/delegate.cc
+++ b/tensorflow/lite/delegates/gpu/delegate.cc
@@ -156,54 +156,54 @@ bool ParseOptions(const char* const* options_keys,
                   const char* const* options_values, size_t num_options,
                   TfLiteGpuDelegateOptionsV2* options) {
   for (size_t i = 0; i < num_options; ++i) {
-    if (strcmp(options_keys[i], "is_precision_loss_allowed")) {
+    if (strcmp(options_keys[i], "is_precision_loss_allowed") == 0) {
       if (!absl::SimpleAtoi(options_values[i],
                             &options->is_precision_loss_allowed)) {
         TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: malformed option %s.",
                    options_keys[i]);
         return false;
       }
-    } else if (strcmp(options_keys[i], "inference_preference")) {
+    } else if (strcmp(options_keys[i], "inference_preference") == 0) {
       if (!absl::SimpleAtoi(options_values[i],
                             &options->inference_preference)) {
         TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: malformed option %s.",
                    options_keys[i]);
         return false;
       }
-    } else if (strcmp(options_keys[i], "inference_priority1")) {
+    } else if (strcmp(options_keys[i], "inference_priority1") == 0) {
       if (!absl::SimpleAtoi(options_values[i], &options->inference_priority1)) {
         TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: malformed option %s.",
                    options_keys[i]);
         return false;
       }
-    } else if (strcmp(options_keys[i], "inference_priority2")) {
+    } else if (strcmp(options_keys[i], "inference_priority2") == 0) {
       if (!absl::SimpleAtoi(options_values[i], &options->inference_priority2)) {
         TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: malformed option %s.",
                    options_keys[i]);
         return false;
       }
-    } else if (strcmp(options_keys[i], "inference_priority3")) {
+    } else if (strcmp(options_keys[i], "inference_priority3") == 0) {
       if (!absl::SimpleAtoi(options_values[i], &options->inference_priority3)) {
         TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: malformed option %s.",
                    options_keys[i]);
         return false;
       }
-    } else if (strcmp(options_keys[i], "experimental_flags")) {
+    } else if (strcmp(options_keys[i], "experimental_flags") == 0) {
       if (!absl::SimpleAtoi(options_values[i], &options->experimental_flags)) {
         TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: malformed option %s.",
                    options_keys[i]);
         return false;
       }
-    } else if (strcmp(options_keys[i], "max_delegated_partitions")) {
+    } else if (strcmp(options_keys[i], "max_delegated_partitions") == 0) {
       if (!absl::SimpleAtoi(options_values[i],
                             &options->max_delegated_partitions)) {
         TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: malformed option %s.",
                    options_keys[i]);
         return false;
       }
-    } else if (strcmp(options_keys[i], "serialization_dir")) {
+    } else if (strcmp(options_keys[i], "serialization_dir") == 0) {
       options->serialization_dir = options_values[i];
-    } else if (strcmp(options_keys[i], "model_token")) {
+    } else if (strcmp(options_keys[i], "model_token") == 0) {
       options->model_token = options_values[i];
     } else {
       TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: unknown option %s.",
diff --git a/tensorflow/lite/kernels/BUILD b/tensorflow/lite/kernels/BUILD
index 4e1f229bf90c79..76158c739db190 100644
--- a/tensorflow/lite/kernels/BUILD
+++ b/tensorflow/lite/kernels/BUILD
@@ -2309,6 +2309,7 @@ cc_test(
     deps = [
         ":test_main",
         ":test_util",
+        "//tensorflow/lite/c:c_api_types",
         "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
         "@flatbuffers",
@@ -2478,6 +2479,7 @@ cc_test(
     deps = [
         ":test_main",
         ":test_util",
+        "//tensorflow/lite/c:c_api_types",
         "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
         "@flatbuffers",
@@ -2491,6 +2493,7 @@ cc_test(
     deps = [
         ":test_main",
         ":test_util",
+        "//tensorflow/lite/c:c_api_types",
         "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
         "@flatbuffers",
@@ -2771,6 +2774,7 @@ cc_test(
     deps = [
         ":test_main",
         ":test_util",
+        "//tensorflow/lite/c:c_api_types",
         "//tensorflow/lite/schema:schema_fbs",
         "//tensorflow/lite/types:half",
         "@com_google_googletest//:gtest",
diff --git a/tensorflow/lite/kernels/activations.cc b/tensorflow/lite/kernels/activations.cc
index 616134df602bc0..20ffc244a31225 100644
--- a/tensorflow/lite/kernels/activations.cc
+++ b/tensorflow/lite/kernels/activations.cc
@@ -267,6 +267,10 @@ TfLiteStatus HardSwishPrepare(TfLiteContext* context, TfLiteNode* node) {
     HardSwishParams* params = &data->params;
     const TfLiteTensor* input;
     TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
+    TF_LITE_ENSURE(context, input->params.zero_point >= INT16_MIN);
+    TF_LITE_ENSURE(context, input->params.zero_point <= INT16_MAX);
+    TF_LITE_ENSURE(context, output->params.zero_point >= INT16_MIN);
+    TF_LITE_ENSURE(context, output->params.zero_point <= INT16_MAX);
     params->input_zero_point = input->params.zero_point;
     params->output_zero_point = output->params.zero_point;
     const float input_scale = input->params.scale;
diff --git a/tensorflow/lite/kernels/activations_test.cc b/tensorflow/lite/kernels/activations_test.cc
index 7ce04226d9196c..533c2eff4dbe10 100644
--- a/tensorflow/lite/kernels/activations_test.cc
+++ b/tensorflow/lite/kernels/activations_test.cc
@@ -240,6 +240,24 @@ class QuantizedActivationsOpModel : public BaseActivationsOpModel {
   }
 };
 
+class PrepareOnlyHardSwishOpModel : public SingleOpModel {
+ public:
+  PrepareOnlyHardSwishOpModel(const TensorData& input,
+                              const TensorData& output) {
+    input_ = AddInput(input);
+    output_ = AddOutput(output);
+    SetBuiltinOp(BuiltinOperator_HARD_SWISH, BuiltinOptions_NONE, 0);
+    BuildInterpreter({GetShape(input_)}, /*num_threads=*/1,
+                     /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/false,
+                     /*allocate_and_delegate=*/false);
+  }
+
+ private:
+  int input_;
+  int output_;
+};
+
 const auto kTanhKernelMap = new std::map<string, TfLiteRegistration*>({
     {"Reference", ops::builtin::Register_TANH_REF()},
     {"GenericOptimized", ops::builtin::Register_TANH_GENERIC_OPT()},
@@ -555,6 +573,28 @@ TEST(QuantizedActivationsOpTest, HardSwish) {
   }
 }
 
+TEST(QuantizedActivationsOpTest, HardSwishRejectsZeroPointOutsideInt16Range) {
+  PrepareOnlyHardSwishOpModel input_model(
+      /*input=*/{TensorType_INT8,
+                 {1},
+                 0.0f,
+                 0.0f,
+                 1.0f,
+                 std::numeric_limits<int16_t>::max() + 1},
+      /*output=*/{TensorType_INT8, {}, 0.0f, 0.0f, 1.0f, 0});
+  EXPECT_EQ(input_model.AllocateTensors(), kTfLiteError);
+
+  PrepareOnlyHardSwishOpModel output_model(
+      /*input=*/{TensorType_INT8, {1}, 0.0f, 0.0f, 1.0f, 0},
+      /*output=*/{TensorType_INT8,
+                  {},
+                  0.0f,
+                  0.0f,
+                  1.0f,
+                  std::numeric_limits<int16_t>::min() - 1});
+  EXPECT_EQ(output_model.AllocateTensors(), kTfLiteError);
+}
+
 // See the comment in the reference implementation of quantized HardSwish:
 // A numerical issue significantly affecting ImageNet classification accuracy
 // with MobileNet v3 is only observable at the scale of HardSwish unit tests
diff --git a/tensorflow/lite/kernels/conv.cc b/tensorflow/lite/kernels/conv.cc
index e85ec46fb7dd75..6ffb218a269735 100644
--- a/tensorflow/lite/kernels/conv.cc
+++ b/tensorflow/lite/kernels/conv.cc
@@ -394,11 +394,15 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context,
 
   // Validate stride values
   TF_LITE_ENSURE(context, params->stride_height > 0);
+  TF_LITE_ENSURE(context, params->stride_height <= INT16_MAX);
   TF_LITE_ENSURE(context, params->stride_width > 0);
+  TF_LITE_ENSURE(context, params->stride_width <= INT16_MAX);
 
   // Validate dilation values
   TF_LITE_ENSURE(context, params->dilation_height_factor > 0);
+  TF_LITE_ENSURE(context, params->dilation_height_factor <= INT16_MAX);
   TF_LITE_ENSURE(context, params->dilation_width_factor > 0);
+  TF_LITE_ENSURE(context, params->dilation_width_factor <= INT16_MAX);
 
   const TfLiteTensor* bias = nullptr;
 
@@ -494,6 +498,7 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context,
           params->dilation_height_factor, params->dilation_width_factor,
           input_height, input_width, filter_height, filter_width, padding,
           &out_height, &out_width, &data->padding));
+  TF_LITE_ENSURE_STATUS(ValidatePaddingValuesForInt16(data->padding));
 
   int output_spatial_elements = 0;
   TF_LITE_ENSURE_MSG(context,
diff --git a/tensorflow/lite/kernels/conv_test.cc b/tensorflow/lite/kernels/conv_test.cc
index b3fee38a3a7df8..15722e5c9a36e7 100644
--- a/tensorflow/lite/kernels/conv_test.cc
+++ b/tensorflow/lite/kernels/conv_test.cc
@@ -362,6 +362,36 @@ TEST(ConvolutionPrepareSecurityTest, RejectsPaddingOverflow) {
   EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
 }
 
+TEST(ConvolutionPrepareSecurityTest, RejectsParametersOutsideInt16Range) {
+  constexpr int kTooLarge = std::numeric_limits<int16_t>::max() + 1;
+  PrepareOnlyConvolutionOpModel<float> stride_model(
+      ops::builtin::Register_CONVOLUTION_GENERIC_OPT(),
+      {TensorType_FLOAT32, {1, 1, 1, 1}}, {TensorType_FLOAT32, {1, 1, 1, 1}},
+      {TensorType_FLOAT32, {}},
+      /*stride_width=*/kTooLarge, /*stride_height=*/1);
+  EXPECT_EQ(stride_model.AllocateTensors(), kTfLiteError);
+
+  PrepareOnlyConvolutionOpModel<float> dilation_model(
+      ops::builtin::Register_CONVOLUTION_GENERIC_OPT(),
+      {TensorType_FLOAT32, {1, 1, 1, 1}}, {TensorType_FLOAT32, {1, 1, 1, 1}},
+      {TensorType_FLOAT32, {}},
+      /*stride_width=*/1, /*stride_height=*/1, Padding_VALID,
+      ActivationFunctionType_NONE, /*dilation_width_factor=*/kTooLarge);
+  EXPECT_EQ(dilation_model.AllocateTensors(), kTfLiteError);
+}
+
+TEST(ConvolutionPrepareSecurityTest, RejectsPaddingOutsideInt16Range) {
+  constexpr int kFilterWidth =
+      2 * (std::numeric_limits<int16_t>::max() + 1) + 1;
+  PrepareOnlyConvolutionOpModel<float> m(
+      ops::builtin::Register_CONVOLUTION_GENERIC_OPT(),
+      {TensorType_FLOAT32, {1, 1, 1, 1}},
+      {TensorType_FLOAT32, {1, 1, kFilterWidth, 1}}, {TensorType_FLOAT32, {}},
+      /*stride_width=*/1, /*stride_height=*/1, Padding_SAME);
+
+  EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
+}
+
 TEST_P(ConvolutionOpTest, SimpleTestFloat32) {
   ConvolutionOpModel m(GetRegistration(), {TensorType_FLOAT32, {2, 2, 4, 1}},
                        {TensorType_FLOAT32, {3, 2, 2, 1}},
diff --git a/tensorflow/lite/kernels/depthwise_conv.cc b/tensorflow/lite/kernels/depthwise_conv.cc
index bcbf9aebcd6847..3ab42e12aca1c1 100644
--- a/tensorflow/lite/kernels/depthwise_conv.cc
+++ b/tensorflow/lite/kernels/depthwise_conv.cc
@@ -166,6 +166,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   const int input_channels = SizeOfDimension(input, 3);
   const int output_channels = SizeOfDimension(filter, 3);
   TF_LITE_ENSURE_EQ(context, output_channels % input_channels, 0);
+  TF_LITE_ENSURE(context, output_channels / input_channels <= INT16_MAX);
 
   if (has_bias) {
     TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kBiasTensor, &bias));
@@ -199,6 +200,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                                  params->dilation_width_factor, height, width,
                                  filter_height, filter_width, padding,
                                  &out_height, &out_width, &data->padding));
+  TF_LITE_ENSURE_STATUS(ValidatePaddingValuesForInt16(data->padding));
 
   // Note that quantized inference requires that all tensors have their
   // parameters set. This is usually done during quantized training or
@@ -331,7 +333,9 @@ TfLiteStatus ComputeDepthMultiplier(TfLiteContext* context,
   int num_input_channels = SizeOfDimension(input, 3);
   TF_LITE_ENSURE(context, num_input_channels != 0);
   TF_LITE_ENSURE_EQ(context, num_filter_channels % num_input_channels, 0);
-  *depth_multiplier = num_filter_channels / num_input_channels;
+  const int multiplier = num_filter_channels / num_input_channels;
+  TF_LITE_ENSURE(context, multiplier <= INT16_MAX);
+  *depth_multiplier = static_cast<int16_t>(multiplier);
   return kTfLiteOk;
 }
 
@@ -504,8 +508,8 @@ TfLiteStatus EvalQuantizedPerChannel(TfLiteContext* context, TfLiteNode* node,
 }
 
 TfLiteStatus EvalQuantizedPerChannel16x8(
-    const TfLiteDepthwiseConvParams* params, const OpData* data,
-    const TfLiteTensor* input, const TfLiteTensor* filter,
+    TfLiteContext* context, const TfLiteDepthwiseConvParams* params,
+    const OpData* data, const TfLiteTensor* input, const TfLiteTensor* filter,
     const TfLiteTensor* bias, TfLiteTensor* output) {
   DepthwiseParams op_params;
   op_params.padding_type = PaddingType::kSame;
@@ -515,7 +519,8 @@ TfLiteStatus EvalQuantizedPerChannel16x8(
   op_params.stride_height = params->stride_height;
   op_params.dilation_width_factor = params->dilation_width_factor;
   op_params.dilation_height_factor = params->dilation_height_factor;
-  op_params.depth_multiplier = params->depth_multiplier;
+  TF_LITE_ENSURE_STATUS(ComputeDepthMultiplier(context, input, filter,
+                                               &op_params.depth_multiplier));
   op_params.weights_offset = 0;
   op_params.quantized_activation_min = data->output_activation_min;
   op_params.quantized_activation_max = data->output_activation_max;
@@ -576,7 +581,8 @@ TfLiteStatus EvalHybridPerChannel(TfLiteContext* context, TfLiteNode* node,
   op_params.stride_height = params->stride_height;
   op_params.dilation_width_factor = params->dilation_width_factor;
   op_params.dilation_height_factor = params->dilation_height_factor;
-  op_params.depth_multiplier = params->depth_multiplier;
+  TF_LITE_ENSURE_STATUS(ComputeDepthMultiplier(context, input, filter,
+                                               &op_params.depth_multiplier));
 
   op_params.weights_offset = 0;
   op_params.float_activation_min = output_activation_min;
@@ -647,8 +653,8 @@ TfLiteStatus EvalImpl(TfLiteContext* context, TfLiteNode* node) {
                                                   input, filter, bias, output);
       break;
     case kTfLiteInt16:
-      return EvalQuantizedPerChannel16x8(params, data, input, filter, bias,
-                                         output);
+      return EvalQuantizedPerChannel16x8(context, params, data, input, filter,
+                                         bias, output);
       break;
     default:
       TF_LITE_KERNEL_LOG(context, "Type %d not currently supported.",
diff --git a/tensorflow/lite/kernels/depthwise_conv_test.cc b/tensorflow/lite/kernels/depthwise_conv_test.cc
index 5cd798c024360a..75f08878b7a35d 100644
--- a/tensorflow/lite/kernels/depthwise_conv_test.cc
+++ b/tensorflow/lite/kernels/depthwise_conv_test.cc
@@ -259,6 +259,30 @@ TEST(DepthwiseConvolutionPrepareSecurityTest,
   EXPECT_EQ(stride_model.AllocateTensors(), kTfLiteError);
 }
 
+TEST(DepthwiseConvolutionPrepareSecurityTest, RejectsPaddingOutsideInt16Range) {
+  constexpr int kFilterWidth =
+      2 * (std::numeric_limits<int16_t>::max() + 1) + 1;
+  PrepareOnlyDepthwiseConvolutionOpModel m(
+      ops::builtin::Register_DEPTHWISE_CONVOLUTION_GENERIC_OPT(),
+      {TensorType_FLOAT32, {1, 1, 1, 1}},
+      {TensorType_FLOAT32, {1, 1, kFilterWidth, 1}}, {TensorType_FLOAT32, {}},
+      Padding_SAME);
+
+  EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
+}
+
+TEST(DepthwiseConvolutionPrepareSecurityTest,
+     RejectsDepthMultiplierOutsideInt16Range) {
+  constexpr int kTooLarge = std::numeric_limits<int16_t>::max() + 1;
+  PrepareOnlyDepthwiseConvolutionOpModel m(
+      ops::builtin::Register_DEPTHWISE_CONVOLUTION_GENERIC_OPT(),
+      {TensorType_FLOAT32, {1, 1, 1, 1}},
+      {TensorType_FLOAT32, {1, 1, 1, kTooLarge}}, {TensorType_FLOAT32, {}},
+      Padding_VALID);
+
+  EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
+}
+
 TEST_P(DepthwiseConvolutionOpTest, ActivationReluTest) {
   DepthwiseConvolutionOpModel m(
       GetRegistration(), {TensorType_FLOAT32, {1, 3, 2, 2}},
diff --git a/tensorflow/lite/kernels/gather.cc b/tensorflow/lite/kernels/gather.cc
index d076449de21390..103949da7002dc 100644
--- a/tensorflow/lite/kernels/gather.cc
+++ b/tensorflow/lite/kernels/gather.cc
@@ -105,12 +105,14 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   }
 
   int axis = params->axis;
+  TF_LITE_ENSURE(context, axis >= INT16_MIN && axis <= INT16_MAX);
   if (axis < 0) {
     axis += NumDimensions(input);
   }
   TF_LITE_ENSURE(context, 0 <= axis && axis < NumDimensions(input));
 
   int batch_dims = params->batch_dims;
+  TF_LITE_ENSURE(context, batch_dims >= INT16_MIN && batch_dims <= INT16_MAX);
   // batch_dims should be in range: [-rank(positions), rank(positions)].
   // Negative batch_dims is added with rank of positions.
   if (batch_dims < 0) {
@@ -163,9 +165,12 @@ TfLiteStatus Gather(TfLiteContext* context, const TfLiteGatherParams& params,
   }
   TF_LITE_ENSURE(context, indices_has_only_positive_elements);
 
+  TF_LITE_ENSURE(context, params.axis >= INT16_MIN && params.axis <= INT16_MAX);
+  TF_LITE_ENSURE(context, params.batch_dims >= INT16_MIN &&
+                              params.batch_dims <= INT16_MAX);
   tflite::GatherParams op_params;
-  op_params.axis = params.axis;
-  op_params.batch_dims = params.batch_dims;
+  op_params.axis = static_cast<int16_t>(params.axis);
+  op_params.batch_dims = static_cast<int16_t>(params.batch_dims);
   return optimized_ops::Gather(
       op_params, GetTensorShape(input), GetTensorData<InputT>(input),
       GetTensorShape(positions), GetTensorData<PositionsT>(positions),
diff --git a/tensorflow/lite/kernels/gather_test.cc b/tensorflow/lite/kernels/gather_test.cc
index f5dd5bda80b3de..40148a57bfca8a 100644
--- a/tensorflow/lite/kernels/gather_test.cc
+++ b/tensorflow/lite/kernels/gather_test.cc
@@ -17,6 +17,7 @@ limitations under the License.
 #include <cstdlib>
 #include <cstring>
 #include <initializer_list>
+#include <limits>
 #include <string>
 #include <type_traits>
 #include <vector>
@@ -136,6 +137,47 @@ class GatherOpModel : public SingleOpModel {
   int output_;
 };
 
+class PrepareOnlyGatherOpModel : public SingleOpModel {
+ public:
+  PrepareOnlyGatherOpModel(int axis, int batch_dims) {
+    input_ = AddInput({TensorType_FLOAT32, {1}});
+    positions_ = AddInput({TensorType_INT32, {1}});
+    output_ = AddOutput(TensorType_FLOAT32);
+    SetBuiltinOp(BuiltinOperator_GATHER, BuiltinOptions_GatherOptions,
+                 CreateGatherOptions(builder_, axis, batch_dims).Union());
+    BuildInterpreter({GetShape(input_), GetShape(positions_)},
+                     /*num_threads=*/1,
+                     /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/false,
+                     /*allocate_and_delegate=*/false);
+  }
+
+ private:
+  int input_;
+  int positions_;
+  int output_;
+};
+
+TEST(GatherPrepareTest, RejectsAxisOutsideInt16Range) {
+  PrepareOnlyGatherOpModel too_large(std::numeric_limits<int16_t>::max() + 1,
+                                     /*batch_dims=*/0);
+  EXPECT_EQ(too_large.AllocateTensors(), kTfLiteError);
+
+  PrepareOnlyGatherOpModel too_small(std::numeric_limits<int16_t>::min() - 1,
+                                     /*batch_dims=*/0);
+  EXPECT_EQ(too_small.AllocateTensors(), kTfLiteError);
+}
+
+TEST(GatherPrepareTest, RejectsBatchDimsOutsideInt16Range) {
+  PrepareOnlyGatherOpModel too_large(
+      /*axis=*/0, std::numeric_limits<int16_t>::max() + 1);
+  EXPECT_EQ(too_large.AllocateTensors(), kTfLiteError);
+
+  PrepareOnlyGatherOpModel too_small(
+      /*axis=*/0, std::numeric_limits<int16_t>::min() - 1);
+  EXPECT_EQ(too_small.AllocateTensors(), kTfLiteError);
+}
+
 struct GatherOpTest : public testing::TestWithParam<bool> {};
 
 INSTANTIATE_TEST_SUITE_P(ConstantTensor, GatherOpTest, testing::Bool());
diff --git a/tensorflow/lite/kernels/padding.h b/tensorflow/lite/kernels/padding.h
index e51d683b70d36c..1c63665291aa8c 100644
--- a/tensorflow/lite/kernels/padding.h
+++ b/tensorflow/lite/kernels/padding.h
@@ -33,6 +33,20 @@ inline TfLiteStatus CheckedNarrowPaddingValue(int64_t value, int* result) {
   return kTfLiteOk;
 }
 
+inline TfLiteStatus ValidatePaddingValuesForInt16(
+    const TfLitePaddingValues& padding_values) {
+  const int min = std::numeric_limits<int16_t>::min();
+  const int max = std::numeric_limits<int16_t>::max();
+  if (padding_values.width < min || padding_values.width > max ||
+      padding_values.height < min || padding_values.height > max ||
+      padding_values.width_offset < min || padding_values.width_offset > max ||
+      padding_values.height_offset < min ||
+      padding_values.height_offset > max) {
+    return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
 inline int64_t ComputeEffectiveFilterSize(int filter_size, int dilation_rate) {
   return (static_cast<int64_t>(filter_size) - 1) * dilation_rate + 1;
 }
diff --git a/tensorflow/lite/kernels/perception/max_pool_with_argmax.cc b/tensorflow/lite/kernels/perception/max_pool_with_argmax.cc
index cb0eb842000821..0db7022a490a94 100644
--- a/tensorflow/lite/kernels/perception/max_pool_with_argmax.cc
+++ b/tensorflow/lite/kernels/perception/max_pool_with_argmax.cc
@@ -191,6 +191,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
       op_data->params.stride_height, op_data->params.stride_width, 1, 1, height,
       width, op_data->params.filter_height, op_data->params.filter_width,
       op_data->params.padding, &out_height, &out_width);
+  TF_LITE_ENSURE_STATUS(
+      ValidatePaddingValuesForInt16(op_data->params.computed.padding));
 
   TfLiteIntArray* output_size = TfLiteIntArrayCreate(4);
   output_size->data[0] = batches;
diff --git a/tensorflow/lite/kernels/perception/max_pool_with_argmax_test.cc b/tensorflow/lite/kernels/perception/max_pool_with_argmax_test.cc
index b87bbd8be4a8f3..39fcee09192755 100644
--- a/tensorflow/lite/kernels/perception/max_pool_with_argmax_test.cc
+++ b/tensorflow/lite/kernels/perception/max_pool_with_argmax_test.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include <cstddef>
 #include <cstdint>
+#include <limits>
 #include <memory>
 #include <vector>
 
@@ -41,7 +42,8 @@ class MaxpoolingWithArgMaxOpModel : public SingleOpModel {
                               int stride_width, int filter_height,
                               int filter_width, TfLitePadding padding,
                               const TensorData& output,
-                              const TensorData& indices) {
+                              const TensorData& indices,
+                              bool allocate_and_delegate = true) {
     input_ = AddInput(input);
     output_ = AddOutput(output);
     indices_ = AddOutput(indices);
@@ -49,7 +51,9 @@ class MaxpoolingWithArgMaxOpModel : public SingleOpModel {
     std::vector<uint8_t> custom_option = CreateCustomOptions(
         stride_height, stride_width, filter_height, filter_width, padding);
     SetCustomOp("MaxPoolWithArgmax", custom_option, RegisterMaxPoolWithArgmax);
-    BuildInterpreter({GetShape(input_)});
+    BuildInterpreter({GetShape(input_)}, /*num_threads=*/-1,
+                     /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/true, allocate_and_delegate);
   }
 
   void SetInput(const std::vector<float>& data) {
@@ -102,6 +106,21 @@ class MaxpoolingWithArgMaxOpModel : public SingleOpModel {
   }
 };
 
+TEST(MaxpoolWithArgMaxPrepareTest, RejectsPaddingOutsideInt16Range) {
+  constexpr int kFilterWidth =
+      2 * (std::numeric_limits<int16_t>::max() + 1) + 1;
+  MaxpoolingWithArgMaxOpModel model(
+      /*input=*/{TensorType_FLOAT32, {1, 1, 1, 1}},
+      /*stride_height=*/1, /*stride_width=*/1,
+      /*filter_height=*/1, /*filter_width=*/kFilterWidth,
+      /*padding=*/kTfLitePaddingSame,
+      /*output=*/{TensorType_FLOAT32, {}},
+      /*indices=*/{TensorType_INT32, {}},
+      /*allocate_and_delegate=*/false);
+
+  EXPECT_EQ(model.AllocateTensors(), kTfLiteError);
+}
+
 TEST(MaxpoolWithArgMaxTest, UnsupportedInt64Test) {
   EXPECT_DEATH_IF_SUPPORTED(MaxpoolingWithArgMaxOpModel model(
                                 /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}},
diff --git a/tensorflow/lite/kernels/pooling.cc b/tensorflow/lite/kernels/pooling.cc
index c6e8452de1e43a..82f46da100fb90 100644
--- a/tensorflow/lite/kernels/pooling.cc
+++ b/tensorflow/lite/kernels/pooling.cc
@@ -95,6 +95,7 @@ TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) {
       params->stride_height, params->stride_width, 1, 1, height, width,
       params->filter_height, params->filter_width, padding, &out_height,
       &out_width);
+  TF_LITE_ENSURE_STATUS(ValidatePaddingValuesForInt16(data->padding));
 
   if (input->type == kTfLiteUInt8 || input->type == kTfLiteInt8) {
     if (pool_type == kAverage || pool_type == kMax) {
diff --git a/tensorflow/lite/kernels/pooling_test.cc b/tensorflow/lite/kernels/pooling_test.cc
index 4634c5ce5e7adf..45b6494c72c5f1 100644
--- a/tensorflow/lite/kernels/pooling_test.cc
+++ b/tensorflow/lite/kernels/pooling_test.cc
@@ -15,11 +15,13 @@ limitations under the License.
 #include <stdint.h>
 
 #include <initializer_list>
+#include <limits>
 #include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/c/c_api_types.h"
 #include "tensorflow/lite/kernels/test_util.h"
 #include "tensorflow/lite/schema/schema_generated.h"
 
@@ -119,6 +121,40 @@ class SymmetricQuantizedPoolingOpModel16 : public BasePoolingOpModel {
   }
 };
 
+class PrepareOnlyPoolingOpModel : public SingleOpModel {
+ public:
+  PrepareOnlyPoolingOpModel(BuiltinOperator type, const TensorData& input,
+                            int filter_width, int filter_height,
+                            Padding padding = Padding_VALID) {
+    input_ = AddInput(input);
+    output_ = AddOutput({input.type, {}});
+    SetBuiltinOp(type, BuiltinOptions_Pool2DOptions,
+                 CreatePool2DOptions(builder_, padding, /*stride_w=*/1,
+                                     /*stride_h=*/1, filter_width,
+                                     filter_height, ActivationFunctionType_NONE)
+                     .Union());
+    BuildInterpreter({GetShape(input_)}, /*num_threads=*/1,
+                     /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/false,
+                     /*allocate_and_delegate=*/false);
+  }
+
+ private:
+  int input_;
+  int output_;
+};
+
+TEST(PoolingPrepareSecurityTest, RejectsPaddingOutsideInt16Range) {
+  constexpr int kFilterWidth =
+      2 * (std::numeric_limits<int16_t>::max() + 1) + 1;
+  PrepareOnlyPoolingOpModel m(BuiltinOperator_AVERAGE_POOL_2D,
+                              /*input=*/{TensorType_FLOAT32, {1, 1, 1, 1}},
+                              kFilterWidth,
+                              /*filter_height=*/1, Padding_SAME);
+
+  EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
+}
+
 // Replicate each entry in a vector n times along depth (innermost dimension).
 // The values are incremented by delta, creating ramps offset by each input
 // value. This is used to create simple and predicatable variation.
diff --git a/tensorflow/lite/kernels/split.cc b/tensorflow/lite/kernels/split.cc
index 4b9d4178fdb507..c1fea0cc766f50 100644
--- a/tensorflow/lite/kernels/split.cc
+++ b/tensorflow/lite/kernels/split.cc
@@ -58,6 +58,7 @@ TfLiteStatus ResizeOutputTensors(TfLiteContext* context, TfLiteNode* node,
 
   TF_LITE_ENSURE(context, axis_value >= 0);
   TF_LITE_ENSURE(context, axis_value < NumDimensions(input));
+  TF_LITE_ENSURE(context, axis_value <= INT16_MAX);
 
   const int input_size = SizeOfDimension(input, axis_value);
   TF_LITE_ENSURE(context, num_splits != 0);
@@ -128,6 +129,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
 
   TF_LITE_ENSURE(context, axis_value >= 0);
   TF_LITE_ENSURE(context, axis_value < NumDimensions(op_context.input));
+  TF_LITE_ENSURE(context, axis_value <= INT16_MAX);
 
   // TODO(b/173221795): Our usage of VectorOfTensors could be optimized by
   // calculating it in Prepare, unless we defer shape calculation.
@@ -138,7 +140,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
     VectorOfTensors<scalar> all_outputs(*context, *node->outputs);    \
     tflite::SplitParams op_params;                                    \
     op_params.num_split = NumOutputs(node);                           \
-    op_params.axis = axis_value;                                      \
+    op_params.axis = static_cast<int16_t>(axis_value);                \
     reference_ops::Split(op_params, GetTensorShape(op_context.input), \
                          GetTensorData<scalar>(op_context.input),     \
                          all_outputs.shapes(), all_outputs.data());   \
diff --git a/tensorflow/lite/kernels/split_test.cc b/tensorflow/lite/kernels/split_test.cc
index b971732ded6c9d..c7c47b66ac356a 100644
--- a/tensorflow/lite/kernels/split_test.cc
+++ b/tensorflow/lite/kernels/split_test.cc
@@ -15,12 +15,13 @@ limitations under the License.
 #include <stdint.h>
 
 #include <initializer_list>
+#include <limits>
 #include <sstream>
 #include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/c/c_api_types.h"
 #include "tensorflow/lite/kernels/test_util.h"
 #include "tensorflow/lite/schema/schema_generated.h"
 
@@ -76,6 +77,38 @@ class SplitOpModel : public SingleOpModel {
   std::vector<int> outputs_;
 };
 
+class PrepareOnlySplitOpModel : public SingleOpModel {
+ public:
+  explicit PrepareOnlySplitOpModel(const std::vector<int>& input_shape,
+                                   int axis) {
+    axis_ = AddConstInput(TensorType_INT32, {axis}, {1});
+    input_ = AddInput({TensorType_FLOAT32, input_shape});
+    AddOutput(TensorType_FLOAT32);
+    SetBuiltinOp(BuiltinOperator_SPLIT, BuiltinOptions_SplitOptions,
+                 CreateSplitOptions(builder_, /*num_splits=*/1).Union());
+    BuildInterpreter({{}, GetShape(input_)}, /*num_threads=*/1,
+                     /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/false,
+                     /*allocate_and_delegate=*/false);
+  }
+
+ private:
+  int input_;
+  int axis_;
+};
+
+TEST(SplitAxisRangeTest, RejectsAxisOutsideInt16Range) {
+  constexpr int kAxis = std::numeric_limits<int16_t>::max() + 1;
+  const std::vector<int> input_shape(kAxis + 1, 1);
+  PrepareOnlySplitOpModel const_axis_model(input_shape, kAxis);
+  EXPECT_EQ(const_axis_model.AllocateTensors(), kTfLiteError);
+
+  SplitOpModel dynamic_axis_model({TensorType_FLOAT32, input_shape},
+                                  /*num_splits=*/1);
+  dynamic_axis_model.SetAxis(/*axis=*/-1);
+  EXPECT_EQ(dynamic_axis_model.Invoke(), kTfLiteError);
+}
+
 template <typename T>
 void Check(TestType test_type, int axis, int num_splits,
            std::initializer_list<int> input_shape,
diff --git a/tensorflow/lite/kernels/split_v.cc b/tensorflow/lite/kernels/split_v.cc
index a86a35f474b671..e7c980e0bd472c 100644
--- a/tensorflow/lite/kernels/split_v.cc
+++ b/tensorflow/lite/kernels/split_v.cc
@@ -99,6 +99,7 @@ TfLiteStatus ResizeOutputTensors(TfLiteContext* context, TfLiteNode* node,
 
   TF_LITE_ENSURE(context, axis_value >= 0);
   TF_LITE_ENSURE(context, axis_value < NumDimensions(input));
+  TF_LITE_ENSURE(context, axis_value <= INT16_MAX);
   const int input_size = SizeOfDimension(input, axis_value);
 
   if (minus_one_index != -1) {
@@ -180,13 +181,19 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   }
 
   int axis_value = GetTensorData<int>(op_context.axis)[0];
+  if (axis_value < 0) {
+    axis_value += NumDimensions(op_context.input);
+  }
+  TF_LITE_ENSURE(context, axis_value >= 0);
+  TF_LITE_ENSURE(context, axis_value < NumDimensions(op_context.input));
+  TF_LITE_ENSURE(context, axis_value <= INT16_MAX);
 
   // Use split function to build the outputs since they share the same logic.
 #define TF_LITE_SPLIT_V(scalar)                                     \
   VectorOfTensors<scalar> all_outputs(*context, *node->outputs);    \
   tflite::SplitParams op_params;                                    \
   op_params.num_split = NumOutputs(node);                           \
-  op_params.axis = axis_value;                                      \
+  op_params.axis = static_cast<int16_t>(axis_value);                \
   reference_ops::Split(op_params, GetTensorShape(op_context.input), \
                        GetTensorData<scalar>(op_context.input),     \
                        all_outputs.shapes(), all_outputs.data());
diff --git a/tensorflow/lite/kernels/split_v_test.cc b/tensorflow/lite/kernels/split_v_test.cc
index c856325b55eb8f..b5f8fb3ac1c3e5 100644
--- a/tensorflow/lite/kernels/split_v_test.cc
+++ b/tensorflow/lite/kernels/split_v_test.cc
@@ -15,11 +15,12 @@ limitations under the License.
 #include <stdint.h>
 
 #include <initializer_list>
+#include <limits>
 #include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/c/c_api_types.h"
 #include "tensorflow/lite/kernels/test_util.h"
 #include "tensorflow/lite/schema/schema_generated.h"
 
@@ -87,6 +88,42 @@ class SplitVOpModel : public SingleOpModel {
   std::vector<int> outputs_;
 };
 
+class PrepareOnlySplitVOpModel : public SingleOpModel {
+ public:
+  explicit PrepareOnlySplitVOpModel(const std::vector<int>& input_shape,
+                                    int axis) {
+    input_ = AddInput({TensorType_FLOAT32, input_shape});
+    size_splits_ = AddConstInput(TensorType_INT32, {1}, {1});
+    axis_ = AddConstInput(TensorType_INT32, {axis}, {1});
+    AddOutput(TensorType_FLOAT32);
+    SetBuiltinOp(BuiltinOperator_SPLIT_V, BuiltinOptions_SplitVOptions,
+                 CreateSplitVOptions(builder_, /*num_splits=*/1).Union());
+    BuildInterpreter({GetShape(input_), {}, {}}, /*num_threads=*/1,
+                     /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/false,
+                     /*allocate_and_delegate=*/false);
+  }
+
+ private:
+  int input_;
+  int size_splits_;
+  int axis_;
+};
+
+TEST(SplitVAxisRangeTest, RejectsAxisOutsideInt16Range) {
+  constexpr int kAxis = std::numeric_limits<int16_t>::max() + 1;
+  const std::vector<int> input_shape(kAxis + 1, 1);
+  PrepareOnlySplitVOpModel const_axis_model(input_shape, kAxis);
+  EXPECT_EQ(const_axis_model.AllocateTensors(), kTfLiteError);
+
+  SplitVOpModel dynamic_axis_model({TensorType_FLOAT32, input_shape},
+                                   {TensorType_INT32, {1}},
+                                   /*num_splits=*/1, kAxisIsATensor, {});
+  dynamic_axis_model.SetSizeSplits({1});
+  dynamic_axis_model.SetAxis(/*axis=*/-1);
+  EXPECT_EQ(dynamic_axis_model.Invoke(), kTfLiteError);
+}
+
 template <typename T>
 void Check(TestType test_type, int axis, std::initializer_list<int> input_shape,
            std::initializer_list<int> size_splits_shape,
diff --git a/tensorflow/lite/kernels/transpose_conv.cc b/tensorflow/lite/kernels/transpose_conv.cc
index 66b18585002cf6..49300e61772483 100644
--- a/tensorflow/lite/kernels/transpose_conv.cc
+++ b/tensorflow/lite/kernels/transpose_conv.cc
@@ -29,9 +29,7 @@ limitations under the License.
 // NOLINTNEXTLINE - This header file shouldn't go to the top.
 #include "tensorflow/lite/kernels/internal/portable_tensor_utils.h"
 #include "tensorflow/lite/kernels/internal/reference/integer_ops/transpose_conv.h"
-#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/lite/kernels/internal/reference/transpose_conv.h"
-#include "tensorflow/lite/kernels/internal/tensor.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/internal/types.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
@@ -291,6 +289,34 @@ TfLiteStatus ResizeAndTransposeWeights(TfLiteContext* context,
   return kTfLiteOk;
 }
 
+TfLiteStatus ComputeAndValidatePadding(TfLiteContext* context,
+                                       const TfLiteTransposeConvParams* params,
+                                       const TfLiteTensor* weights,
+                                       const TfLiteTensor* input,
+                                       const TfLiteTensor* output,
+                                       TfLitePaddingValues* padding) {
+  TF_LITE_ENSURE_EQ(context, NumDimensions(output), 4);
+  TF_LITE_ENSURE_EQ(context, NumDimensions(weights), 4);
+  TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4);
+  TF_LITE_ENSURE(context, SizeOfDimension(output, 1) > 0);
+  TF_LITE_ENSURE(context, SizeOfDimension(output, 2) > 0);
+  const int width = SizeOfDimension(output, 2);
+  const int height = SizeOfDimension(output, 1);
+  const int filter_width = SizeOfDimension(weights, 2);
+  const int filter_height = SizeOfDimension(weights, 1);
+
+  int computed_input_height, computed_input_width;
+  TF_LITE_ENSURE_OK(
+      context, ComputePaddingHeightWidthChecked(
+                   params->stride_height, params->stride_width, 1, 1, height,
+                   width, filter_height, filter_width, params->padding,
+                   &computed_input_height, &computed_input_width, padding));
+  TF_LITE_ENSURE_STATUS(ValidatePaddingValuesForInt16(*padding));
+  TF_LITE_ENSURE_EQ(context, computed_input_height, SizeOfDimension(input, 1));
+  TF_LITE_ENSURE_EQ(context, computed_input_width, SizeOfDimension(input, 2));
+  return kTfLiteOk;
+}
+
 template <KernelType kernel_type>
 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   OpData* data = reinterpret_cast<OpData*>(node->user_data);
@@ -321,6 +347,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   // Tensor sanity checks
   TF_LITE_ENSURE_EQ(context, NumDimensions(output_shape), 1);
+  TF_LITE_ENSURE_EQ(context, NumElements(output_shape), 4);
   TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4);
   TF_LITE_ENSURE_EQ(context, NumDimensions(weights), 4);
   TF_LITE_ENSURE(context, params->stride_height > 0);
@@ -406,6 +433,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     }
   } else {
     TF_LITE_ENSURE_STATUS(ResizeTensor(context, output_shape, output));
+    TF_LITE_ENSURE_STATUS(ComputeAndValidatePadding(
+        context, params, weights, input, output, &data->padding));
     if (data->has_col2im) {
       TF_LITE_ENSURE_STATUS(
           ResizeCol2ImTensor(context, output_shape, weights, input, col2im));
@@ -880,20 +909,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
                                                   weights, input, col2im));
   }
 
-  // Get height and width of the output image.
-  const int width = SizeOfDimension(output, 2);
-  const int height = SizeOfDimension(output, 1);
-  const int filter_width = SizeOfDimension(weights, 2);
-  const int filter_height = SizeOfDimension(weights, 1);
-
-  int computed_input_height, computed_input_width;
-  TF_LITE_ENSURE_OK(context, ComputePaddingHeightWidthChecked(
-                                 params->stride_height, params->stride_width, 1,
-                                 1, height, width, filter_height, filter_width,
-                                 params->padding, &computed_input_height,
-                                 &computed_input_width, &data->padding));
-  TF_LITE_ENSURE_EQ(context, computed_input_height, SizeOfDimension(input, 1));
-  TF_LITE_ENSURE_EQ(context, computed_input_width, SizeOfDimension(input, 2));
+  TF_LITE_ENSURE_STATUS(ComputeAndValidatePadding(
+      context, params, weights, input, output, &data->padding));
 
   // Currently support float32, uint8, int8, int16.
   switch (input->type) {
diff --git a/tensorflow/lite/kernels/transpose_conv_test.cc b/tensorflow/lite/kernels/transpose_conv_test.cc
index 507091d74dd6df..c2e0ee1fbcba07 100644
--- a/tensorflow/lite/kernels/transpose_conv_test.cc
+++ b/tensorflow/lite/kernels/transpose_conv_test.cc
@@ -26,7 +26,6 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include "absl/memory/memory.h"
 #include "tensorflow/lite/core/interpreter.h"
 #include "tensorflow/lite/kernels/test_util.h"
 #include "tensorflow/lite/schema/schema_generated.h"
@@ -90,19 +89,19 @@ class BaseTransposeConvOpModel : public SingleOpModel {
 
     if (test_type == TestType::kDynamic) {
       PopulateTensor<int32_t>(output_shape_, output_shape_data);
-      if (!std::is_same<InputType, int16_t>::value &&
-          !std::is_same<InputType, int8_t>::value) {
+      if (!std::is_same_v<InputType, int16_t> &&
+          !std::is_same_v<InputType, int8_t>) {
         PopulateTensor<InputType>(filter_, filter_data);
       }
     }
   }
 
   void SetInput(std::initializer_list<float> data) {
-    if (std::is_same<InputType, uint8_t>::value) {
+    if (std::is_same_v<InputType, uint8_t>) {
       QuantizeAndPopulate<uint8_t>(input_, data);
-    } else if (std::is_same<InputType, int8_t>::value) {
+    } else if (std::is_same_v<InputType, int8_t>) {
       QuantizeAndPopulate<int8_t>(input_, data);
-    } else if (std::is_same<InputType, int16_t>::value) {
+    } else if (std::is_same_v<InputType, int16_t>) {
       QuantizeAndPopulate<int16_t>(input_, data);
     } else {
       PopulateTensor(input_, data);
@@ -224,6 +223,19 @@ TEST(TransposeConvPrepareSecurityTest, RejectsStrideOutsideInt16Range) {
   EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
 }
 
+TEST(TransposeConvPrepareSecurityTest, RejectsPaddingOutsideInt16Range) {
+  constexpr int kFilterWidth =
+      2 * (std::numeric_limits<int16_t>::max() + 1) + 1;
+  PrepareOnlyTransposeConvOpModel<float> m(
+      ops::builtin::Register_TRANSPOSECONV_GENERIC_OPT(), {1, 1, 1, 1},
+      {TensorType_FLOAT32, {1, 1, kFilterWidth, 1}},
+      {TensorType_FLOAT32, {1, 1, 1, 1}}, {TensorType_FLOAT32, {}},
+      Padding_SAME, /*stride_w=*/1, /*stride_h=*/1,
+      ActivationFunctionType_NONE);
+
+  EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
+}
+
 TEST(TransposeConvPrepareSecurityTest, RejectsMismatchedOutputChannels) {
   PrepareOnlyTransposeConvOpModel<float> m(
       ops::builtin::Register_TRANSPOSECONV_GENERIC_OPT(), {1, 1, 1, 2},
@@ -242,8 +254,45 @@ TEST(TransposeConvPrepareSecurityTest, RejectsInconsistentSpatialShape) {
       {TensorType_FLOAT32, {}}, Padding_SAME, /*stride_w=*/1, /*stride_h=*/1,
       ActivationFunctionType_NONE);
 
-  ASSERT_EQ(m.AllocateTensors(), kTfLiteOk);
-  EXPECT_EQ(m.Invoke(), kTfLiteError);
+  EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
+}
+
+class PrepareOnlyInvalidOutputShapeTransposeConvOpModel : public SingleOpModel {
+ public:
+  explicit PrepareOnlyInvalidOutputShapeTransposeConvOpModel(
+      std::initializer_list<int> output_shape_data,
+      std::initializer_list<int> output_shape_dims) {
+    output_shape_ =
+        AddConstInput(TensorType_INT32, output_shape_data, output_shape_dims);
+    filter_ = AddInput({TensorType_FLOAT32, {1, 1, 1, 1}});
+    input_ = AddInput({TensorType_FLOAT32, {1, 1, 1, 1}});
+    output_ = AddOutput({TensorType_FLOAT32, {}});
+
+    SetBuiltinOp(
+        BuiltinOperator_TRANSPOSE_CONV, BuiltinOptions_TransposeConvOptions,
+        CreateTransposeConvOptions(builder_, Padding_SAME, /*stride_w=*/1,
+                                   /*stride_h=*/1, ActivationFunctionType_NONE)
+            .Union());
+    resolver_ = std::make_unique<SingleOpResolver>(
+        BuiltinOperator_TRANSPOSE_CONV,
+        ops::builtin::Register_TRANSPOSECONV_GENERIC_OPT(), /*version=*/1);
+    BuildInterpreter(
+        {GetShape(output_shape_), GetShape(filter_), GetShape(input_)},
+        /*num_threads=*/1, /*allow_fp32_relax_to_fp16=*/false,
+        /*apply_delegate=*/false,
+        /*allocate_and_delegate=*/false);
+  }
+
+ private:
+  int output_shape_;
+  int filter_;
+  int input_;
+  int output_;
+};
+
+TEST(TransposeConvPrepareSecurityTest, RejectsInvalidOutputShapeTensorLength) {
+  PrepareOnlyInvalidOutputShapeTransposeConvOpModel m({1}, {1});
+  EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
 }
 
 // Test case:
@@ -983,19 +1032,19 @@ class BaseTransposeConvBiasOpModel : public SingleOpModel {
                       GetShape(input_), GetShape(bias_)});
     if (test_type == TestType::kDynamic) {
       PopulateTensor<int32_t>(output_shape_, output_shape_data);
-      if (!std::is_same<InputType, int16_t>::value &&
-          !std::is_same<InputType, int8_t>::value) {
+      if (!std::is_same_v<InputType, int16_t> &&
+          !std::is_same_v<InputType, int8_t>) {
         PopulateTensor<FilterType>(filter_, filter_data);
       }
     }
   }
 
   void SetInput(std::initializer_list<float> data) {
-    if (std::is_same<InputType, uint8_t>::value) {
+    if (std::is_same_v<InputType, uint8_t>) {
       QuantizeAndPopulate<uint8_t>(input_, data);
-    } else if (std::is_same<InputType, int8_t>::value) {
+    } else if (std::is_same_v<InputType, int8_t>) {
       QuantizeAndPopulate<int8_t>(input_, data);
-    } else if (std::is_same<InputType, int16_t>::value) {
+    } else if (std::is_same_v<InputType, int16_t>) {
       QuantizeAndPopulate<int16_t>(input_, data);
     } else {
       PopulateTensor(input_, data);
@@ -1003,9 +1052,9 @@ class BaseTransposeConvBiasOpModel : public SingleOpModel {
   }
 
   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<FilterType, int8_t>::value) {
+    } else if (std::is_same_v<FilterType, int8_t>) {
       PerChannelQuantizeBias(bias_, bias);
     } else {
       PopulateTensor(bias_, bias);
diff --git a/tensorflow/lite/kernels/unpack.cc b/tensorflow/lite/kernels/unpack.cc
index cbbff6fe964fc8..7525dc7473e4e7 100644
--- a/tensorflow/lite/kernels/unpack.cc
+++ b/tensorflow/lite/kernels/unpack.cc
@@ -42,10 +42,12 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
   TF_LITE_ENSURE(context, NumElements(input) > 0);
   int axis = data->axis;
+  TF_LITE_ENSURE(context, axis >= INT16_MIN && axis <= INT16_MAX);
   if (axis < 0) {
     axis += NumDimensions(input);
   }
   TF_LITE_ENSURE(context, 0 <= axis && axis < NumDimensions(input));
+  TF_LITE_ENSURE(context, axis <= INT16_MAX);
   if (input->type != kTfLiteInt32 && input->type != kTfLiteFloat32 &&
       input->type != kTfLiteFloat16 && input->type != kTfLiteBFloat16 &&
 #if defined(TFLITE_ENABLE_EXTRA_REFERENCE_KERNELS)
@@ -92,7 +94,7 @@ template <typename T>
 void UnpackImpl(TfLiteContext* context, TfLiteNode* node,
                 const TfLiteTensor* input, int output_count, int axis) {
   tflite::UnpackParams op_params;
-  op_params.axis = axis;
+  op_params.axis = static_cast<int16_t>(axis);
   op_params.num_split = output_count;
   VectorOfTensors<T> all_outputs(*context, *node->outputs);
   reference_ops::Unpack<T>(op_params, GetTensorShape(input),
@@ -106,21 +108,27 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
 
   const TfLiteTensor* input;
   TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  int axis = data->axis;
+  if (axis < 0) {
+    axis += NumDimensions(input);
+  }
+  TF_LITE_ENSURE(context, axis >= 0 && axis < NumDimensions(input));
+  TF_LITE_ENSURE(context, axis <= INT16_MAX);
   switch (TfLiteTypeGetSizeBits(input->type)) {
     case 8: {
-      UnpackImpl<uint8_t>(context, node, input, data->num, data->axis);
+      UnpackImpl<uint8_t>(context, node, input, data->num, axis);
       break;
     }
     case 16: {
-      UnpackImpl<uint16_t>(context, node, input, data->num, data->axis);
+      UnpackImpl<uint16_t>(context, node, input, data->num, axis);
       break;
     }
     case 32: {
-      UnpackImpl<uint32_t>(context, node, input, data->num, data->axis);
+      UnpackImpl<uint32_t>(context, node, input, data->num, axis);
       break;
     }
     case 64: {
-      UnpackImpl<uint64_t>(context, node, input, data->num, data->axis);
+      UnpackImpl<uint64_t>(context, node, input, data->num, axis);
       break;
     }
     default: {
diff --git a/tensorflow/lite/kernels/unpack_test.cc b/tensorflow/lite/kernels/unpack_test.cc
index 89b13ce48a5baa..395603409d482a 100644
--- a/tensorflow/lite/kernels/unpack_test.cc
+++ b/tensorflow/lite/kernels/unpack_test.cc
@@ -15,12 +15,12 @@ limitations under the License.
 #include <stdint.h>
 
 #include <initializer_list>
-#include <iostream>
-#include <type_traits>
+#include <limits>
 #include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "tensorflow/lite/c/c_api_types.h"
 #include "tensorflow/lite/kernels/test_util.h"
 #include "tensorflow/lite/schema/schema_generated.h"
 #include "tensorflow/lite/types/half.h"
@@ -77,6 +77,38 @@ class UnpackOpModel : public SingleOpModel {
   std::vector<int> outputs_;
 };
 
+class PrepareOnlyUnpackOpModel : public SingleOpModel {
+ public:
+  explicit PrepareOnlyUnpackOpModel(const TensorData& input, int num_outputs,
+                                    int axis) {
+    input_ = AddInput(input);
+    for (int i = 0; i < num_outputs; ++i) {
+      AddOutput(input.type);
+    }
+    SetBuiltinOp(BuiltinOperator_UNPACK, BuiltinOptions_UnpackOptions,
+                 CreateUnpackOptions(builder_, num_outputs, axis).Union());
+    BuildInterpreter({GetShape(input_)}, /*num_threads=*/1,
+                     /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/false,
+                     /*allocate_and_delegate=*/false);
+  }
+
+ private:
+  int input_;
+};
+
+TEST(UnpackPrepareTest, RejectsAxisOutsideInt16Range) {
+  PrepareOnlyUnpackOpModel too_large(
+      /*input=*/{TensorType_FLOAT32, {1}}, /*num_outputs=*/1,
+      std::numeric_limits<int16_t>::max() + 1);
+  EXPECT_EQ(too_large.AllocateTensors(), kTfLiteError);
+
+  PrepareOnlyUnpackOpModel too_small(
+      /*input=*/{TensorType_FLOAT32, {1}}, /*num_outputs=*/1,
+      std::numeric_limits<int16_t>::min() - 1);
+  EXPECT_EQ(too_small.AllocateTensors(), kTfLiteError);
+}
+
 template <typename T>
 void Check(int axis, const std::initializer_list<int>& input_shape,
            const std::initializer_list<int>& input_data,
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index fbf0e48fe9e614..1ca50ef4d5da69 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -1649,7 +1649,6 @@ pywrap_library(
     # @unsorted-dict-items
     common_lib_version_scripts = {
         "tensorflow/tensorflow_cc": select({
-            "//tensorflow:windows": None,
             "//tensorflow:macos": "//tensorflow:tf_exported_symbols.lds",
             "//conditions:default": "//tensorflow:tf_version_script.lds",
         }),
diff --git a/tensorflow/python/eager/pywrap_gradient_exclusions.cc b/tensorflow/python/eager/pywrap_gradient_exclusions.cc
index f58e855263dfce..8c34921427963b 100644
--- a/tensorflow/python/eager/pywrap_gradient_exclusions.cc
+++ b/tensorflow/python/eager/pywrap_gradient_exclusions.cc
@@ -429,13 +429,14 @@ absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedInputIndices(
 
 absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedOutputIndices(
     const tensorflow::string &op_name) {
-  static std::array<OpIndexInfo, 487> a = {{
+  static std::array<OpIndexInfo, 488> a = {{
       {"Abs"},
       {"AccumulateNV2"},
       {"Acos"},
       {"Add"},
       {"AddN"},
       {"AddV2"},
+      {"AdjustContrastv2"},
       {"AllToAll"},
       {"Angle"},
       {"ApproxTopK", 1, {0}},
diff --git a/tensorflow/python/eager/pywrap_tfe_src.cc b/tensorflow/python/eager/pywrap_tfe_src.cc
index c1d4f1089e3f69..14c4432064faf9 100644
--- a/tensorflow/python/eager/pywrap_tfe_src.cc
+++ b/tensorflow/python/eager/pywrap_tfe_src.cc
@@ -1870,14 +1870,15 @@ class AccumulatorSet {
     return true;
   }
 
-  void erase(TFE_Py_ForwardAccumulator* element) {
+  bool erase(TFE_Py_ForwardAccumulator* element) {
     MapType::iterator existing = map_.find(element);
     if (existing == map_.end()) {
-      return;
+      return false;
     }
     ListType::iterator list_position = existing->second;
     map_.erase(existing);
     ordered_.erase(list_position);
+    return true;
   }
 
   bool empty() const { return ordered_.empty(); }
@@ -2379,6 +2380,10 @@ void TFE_Py_TapeVariableAccessed(PyObject* variable) {
 }
 
 void TFE_Py_TapeWatchVariable(PyObject* tape, PyObject* variable) {
+  if (!PyObject_TypeCheck(tape, &TFE_Py_Tape_Type)) {
+    PyErr_SetString(PyExc_TypeError, "Expected a TFE_Py_Tape object");
+    return;
+  }
   if (!CouldBackprop()) {
     return;
   }
@@ -2386,6 +2391,10 @@ void TFE_Py_TapeWatchVariable(PyObject* tape, PyObject* variable) {
 }
 
 PyObject* TFE_Py_TapeWatchedVariables(PyObject* tape) {
+  if (!PyObject_TypeCheck(tape, &TFE_Py_Tape_Type)) {
+    PyErr_SetString(PyExc_TypeError, "Expected a TFE_Py_Tape object");
+    return nullptr;
+  }
   return reinterpret_cast<TFE_Py_Tape*>(tape)->tape->GetVariablesAsPyTuple();
 }
 
@@ -2410,11 +2419,23 @@ PyObject* TFE_Py_VariableWatcherNew() {
 }
 
 void TFE_Py_VariableWatcherRemove(PyObject* variable_watcher) {
+  if (!PyObject_TypeCheck(variable_watcher, &TFE_Py_VariableWatcher_Type)) {
+    PyErr_SetString(PyExc_TypeError,
+                    "Expected a TFE_Py_VariableWatcher object");
+    return;
+  }
   auto* stack = GetVariableWatcherSet();
-  stack->erase(reinterpret_cast<TFE_Py_VariableWatcher*>(variable_watcher));
+  bool erased = false;
+  if (stack != nullptr) {
+    auto* vw = reinterpret_cast<TFE_Py_VariableWatcher*>(variable_watcher);
+    erased = stack->erase(vw) > 0;
+  }
   // We kept a reference to the variable watcher in the set to ensure it
   // wouldn't get deleted under us; cleaning it up here.
-  Py_DECREF(variable_watcher);
+  // We only decref if the variable watcher was actually erased from the set.
+  if (erased) {
+    Py_DECREF(variable_watcher);
+  }
 }
 
 void TFE_Py_VariableWatcherVariableAccessed(PyObject* variable) {
@@ -2424,6 +2445,11 @@ void TFE_Py_VariableWatcherVariableAccessed(PyObject* variable) {
 }
 
 PyObject* TFE_Py_VariableWatcherWatchedVariables(PyObject* variable_watcher) {
+  if (!PyObject_TypeCheck(variable_watcher, &TFE_Py_VariableWatcher_Type)) {
+    PyErr_SetString(PyExc_TypeError,
+                    "Expected a TFE_Py_VariableWatcher object");
+    return nullptr;
+  }
   return reinterpret_cast<TFE_Py_VariableWatcher*>(variable_watcher)
       ->variable_watcher->GetVariablesAsPyTuple();
 }
@@ -2908,6 +2934,10 @@ PyObject* TFE_Py_TapeGradient(PyObject* tape, PyObject* target,
                               PyObject* sources_raw,
                               PyObject* unconnected_gradients,
                               TF_Status* status) {
+  if (!PyObject_TypeCheck(tape, &TFE_Py_Tape_Type)) {
+    PyErr_SetString(PyExc_TypeError, "Expected a TFE_Py_Tape object");
+    return nullptr;
+  }
   TFE_Py_Tape* tape_obj = reinterpret_cast<TFE_Py_Tape*>(tape);
   if (!tape_obj->tape->IsPersistent()) {
     auto* tape_set = GetTapeSet();
@@ -3042,6 +3072,11 @@ PyObject* TFE_Py_ForwardAccumulatorNew(bool use_batch) {
 }
 
 PyObject* TFE_Py_ForwardAccumulatorSetAdd(PyObject* accumulator) {
+  if (!PyObject_TypeCheck(accumulator, &TFE_Py_ForwardAccumulator_Type)) {
+    PyErr_SetString(PyExc_TypeError,
+                    "Expected a TFE_Py_ForwardAccumulator object");
+    return nullptr;
+  }
   TFE_Py_ForwardAccumulator* c_accumulator(
       reinterpret_cast<TFE_Py_ForwardAccumulator*>(accumulator));
   c_accumulator->nesting_id = tape_nesting_id_counter.fetch_add(1);
@@ -3058,16 +3093,29 @@ PyObject* TFE_Py_ForwardAccumulatorSetAdd(PyObject* accumulator) {
 }
 
 void TFE_Py_ForwardAccumulatorSetRemove(PyObject* accumulator) {
+  if (!PyObject_TypeCheck(accumulator, &TFE_Py_ForwardAccumulator_Type)) {
+    PyErr_SetString(PyExc_TypeError,
+                    "Expected a TFE_Py_ForwardAccumulator object");
+    return;
+  }
   auto* accumulator_set = GetAccumulatorSet();
+  bool erased = false;
   if (accumulator_set != nullptr) {
-    accumulator_set->erase(
+    erased = accumulator_set->erase(
         reinterpret_cast<TFE_Py_ForwardAccumulator*>(accumulator));
   }
-  Py_DECREF(accumulator);
+  if (erased) {
+    Py_DECREF(accumulator);
+  }
 }
 
 void TFE_Py_ForwardAccumulatorWatch(PyObject* accumulator, PyObject* tensor,
                                     PyObject* tangent) {
+  if (!PyObject_TypeCheck(accumulator, &TFE_Py_ForwardAccumulator_Type)) {
+    PyErr_SetString(PyExc_TypeError,
+                    "Expected a TFE_Py_ForwardAccumulator object");
+    return;
+  }
   int64_t tensor_id = FastTensorId(tensor);
   reinterpret_cast<TFE_Py_ForwardAccumulator*>(accumulator)
       ->accumulator->Watch(tensor_id, tangent);
@@ -3077,6 +3125,11 @@ void TFE_Py_ForwardAccumulatorWatch(PyObject* accumulator, PyObject* tensor,
 // Returns a new reference to the JVP Tensor.
 PyObject* TFE_Py_ForwardAccumulatorJVP(PyObject* accumulator,
                                        PyObject* tensor) {
+  if (!PyObject_TypeCheck(accumulator, &TFE_Py_ForwardAccumulator_Type)) {
+    PyErr_SetString(PyExc_TypeError,
+                    "Expected a TFE_Py_ForwardAccumulator object");
+    return nullptr;
+  }
   PyObject* jvp = reinterpret_cast<TFE_Py_ForwardAccumulator*>(accumulator)
                       ->accumulator->FetchJVP(FastTensorId(tensor));
   if (jvp == nullptr) {
diff --git a/tensorflow/python/eager/pywrap_tfe_test.py b/tensorflow/python/eager/pywrap_tfe_test.py
index 86beaef2200410..04c0460159902f 100644
--- a/tensorflow/python/eager/pywrap_tfe_test.py
+++ b/tensorflow/python/eager/pywrap_tfe_test.py
@@ -383,6 +383,39 @@ def testIntAttrThatDoesNotFitIn32Bits(self):
                                       shape, minval, maxval,
                                       "seed", seed)
 
+  def testTapeAndWatcherTypeCheckAndRefcount(self):
+    dummy_obj = 100
+    with self.assertRaises(TypeError):
+      pywrap_tfe.TFE_Py_TapeWatchVariable(dummy_obj, None)
+    with self.assertRaises(TypeError):
+      pywrap_tfe.TFE_Py_TapeWatchedVariables(dummy_obj)
+    with self.assertRaises(TypeError):
+      pywrap_tfe.TFE_Py_VariableWatcherRemove(dummy_obj)
+    with self.assertRaises(TypeError):
+      pywrap_tfe.TFE_Py_VariableWatcherWatchedVariables(dummy_obj)
+    with self.assertRaises(TypeError):
+      pywrap_tfe.TFE_Py_ForwardAccumulatorSetAdd(dummy_obj)
+    with self.assertRaises(TypeError):
+      pywrap_tfe.TFE_Py_ForwardAccumulatorSetRemove(dummy_obj)
+    with self.assertRaises(TypeError):
+      pywrap_tfe.TFE_Py_ForwardAccumulatorWatch(dummy_obj, None, None)
+    with self.assertRaises(TypeError):
+      pywrap_tfe.TFE_Py_ForwardAccumulatorJVP(dummy_obj, None)
+
+    # Double remove on a tape, watcher, or accumulator should not double decref
+    t = pywrap_tfe.TFE_Py_TapeSetNew(False, False)
+    pywrap_tfe.TFE_Py_TapeSetRemove(t)
+    pywrap_tfe.TFE_Py_TapeSetRemove(t)
+
+    vw = pywrap_tfe.TFE_Py_VariableWatcherNew()
+    pywrap_tfe.TFE_Py_VariableWatcherRemove(vw)
+    pywrap_tfe.TFE_Py_VariableWatcherRemove(vw)
+
+    acc = pywrap_tfe.TFE_Py_ForwardAccumulatorNew(False)
+    pywrap_tfe.TFE_Py_ForwardAccumulatorSetAdd(acc)
+    pywrap_tfe.TFE_Py_ForwardAccumulatorSetRemove(acc)
+    pywrap_tfe.TFE_Py_ForwardAccumulatorSetRemove(acc)
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc
index 89eaa533e54d8c..b713871cbed36a 100644
--- a/tensorflow/python/framework/python_op_gen.cc
+++ b/tensorflow/python/framework/python_op_gen.cc
@@ -199,9 +199,23 @@ std::string DataTypeToPython(DataType dtype, const std::string& dtype_module);
 class ParamNames {
  public:
   // Create param based on Arg.
+  //
+  // rename_to comes from ApiDef, a separate message from the OpDef
+  // arg/attr names. Both are spliced as a raw Python identifier (keyword
+  // argument name) below: prefer rename_to when it's a safe identifier,
+  // fall back to the original name when THAT is safe, and only sanitize
+  // as a last resort, since IsValidAttrOrArgName is not enforced at OpDef
+  // registration and a legitimately-registered op's argument name is not
+  // guaranteed to already be a safe identifier (e.g.
+  // TFLite_Detection_PostProcess's "raw_outputs/box_encodings").
   ParamNames(const std::string& name, const std::string& rename_to)
       : name_(name) {
-    rename_to_ = AvoidPythonReserved(rename_to);
+    const std::string& safe_rename_to =
+        tensorflow::IsValidAttrOrArgName(rename_to) ? rename_to
+        : tensorflow::IsValidAttrOrArgName(name)
+            ? name
+            : tensorflow::SanitizeToIdentifier(name);
+    rename_to_ = AvoidPythonReserved(safe_rename_to);
   }
 
   // Get original parameter name.
diff --git a/tensorflow/python/ops/image_grad.py b/tensorflow/python/ops/image_grad.py
index 39517bfb132d32..bc2d6b22005ba7 100644
--- a/tensorflow/python/ops/image_grad.py
+++ b/tensorflow/python/ops/image_grad.py
@@ -47,6 +47,46 @@ def _ResizeNearestNeighborGrad(op: ops.Operation, grad):
   return [grads, None]
 
 
+@ops.RegisterGradient("AdjustContrastv2")
+def _AdjustContrastGrad(op: ops.Operation, grad):
+  """The derivatives for `tf.image.adjust_contrast`.
+
+  The kernel computes `(images - mean) * contrast_factor + mean`, where
+  `mean` is taken per batch and channel over the last three dimensions,
+  which are interpreted as [height, width, channels].
+
+  Args:
+    op: The `AdjustContrastv2` `Operation`.
+    grad: The tensor representing the gradient w.r.t. the output.
+
+  Returns:
+    The gradients w.r.t. the images and the contrast factor.
+  """
+  images = op.inputs[0]
+  factor = op.inputs[1]
+  factor_t = math_ops.cast(factor, images.dtype)
+  static_rank = images.shape.rank
+  if static_rank is not None:
+    spatial_axes = list(range(static_rank - 3, static_rank - 1))
+  else:
+    dynamic_rank = array_ops.rank(images)
+    spatial_axes = math_ops.range(dynamic_rank - 3, dynamic_rank - 1)
+  mean = math_ops.reduce_mean(images, axis=spatial_axes, keepdims=True)
+  grad_images = grad * factor_t + (1.0 - factor_t) * math_ops.reduce_mean(
+      grad,
+      axis=spatial_axes,
+      keepdims=True,
+  )
+  grad_factor = math_ops.reduce_sum(
+      math_ops.cast(grad, dtypes.float32)
+      * (
+          math_ops.cast(images, dtypes.float32)
+          - math_ops.cast(mean, dtypes.float32)
+      )
+  )
+  return grad_images, grad_factor
+
+
 @ops.RegisterGradient("ResizeBilinear")
 def _ResizeBilinearGrad(op: ops.Operation, grad):
   """The derivatives for bilinear resizing.
diff --git a/tensorflow/python/ops/image_grad_test.py b/tensorflow/python/ops/image_grad_test.py
index 50636b9beef435..28fd605e628455 100644
--- a/tensorflow/python/ops/image_grad_test.py
+++ b/tensorflow/python/ops/image_grad_test.py
@@ -23,6 +23,7 @@
 ScaleAndTranslateOpTest = test_base.ScaleAndTranslateOpTestBase
 CropAndResizeOpTest = test_base.CropAndResizeOpTestBase
 RGBToHSVOpTest = test_base.RGBToHSVOpTestBase
+AdjustContrastOpTest = test_base.AdjustContrastOpTestBase
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/python/ops/image_grad_test_base.py b/tensorflow/python/ops/image_grad_test_base.py
index fc84f3054f9656..b412b3a9f9a5fc 100644
--- a/tensorflow/python/ops/image_grad_test_base.py
+++ b/tensorflow/python/ops/image_grad_test_base.py
@@ -644,5 +644,38 @@ def f_dummy(x):
     self.assertAllClose(numerical_dummy, numerical, atol=1e-4)
 
 
+class AdjustContrastOpTestBase(test.TestCase):
+  """Tests the gradient of tf.image.adjust_contrast.
+
+  The op is affine in the images and linear in the scalar factor, so
+  gradient_checker_v2 validates both slots against the real forward kernel,
+  which is what pins down the spatial reduction axes.
+  """
+
+  def testGradRank3(self):
+    self._check_gradient([4, 5, 3])
+
+  def testGradRank4(self):
+    self._check_gradient([2, 4, 5, 3])
+
+  def testGradRank5(self):
+    self._check_gradient([2, 3, 4, 5, 3])
+
+  def _check_gradient(self, shape):
+    x = np.linspace(0.05, 0.95, num=int(np.prod(shape))).reshape(shape)
+    x = x.astype(np.float32)
+    factor = constant_op.constant(1.3, dtype=dtypes.float32)
+
+    def f(image_tensor, factor_tensor):
+      return image_ops.adjust_contrast(image_tensor, factor_tensor)
+
+    with self.cached_session():
+      analytical, numerical = gradient_checker_v2.compute_gradient(
+          f, [constant_op.constant(x), factor]
+      )
+    max_error = gradient_checker_v2.max_error(analytical, numerical)
+    self.assertLess(max_error, 1e-4)
+
+
 if __name__ == '__main__':
   test.main()
diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index 8fa7ca81622510..27a0c69ba8de52 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -56,6 +56,15 @@
     __name__, 'inf'
 )
 
+# Floating-point types mapped to an integer type of the same size, so their
+# IEEE-754 sign bit can be checked with an integer comparison.
+_SIGN_BITCAST_DTYPES = {
+    dtypes.bfloat16: dtypes.int16,
+    dtypes.float16: dtypes.int16,
+    dtypes.float32: dtypes.int32,
+    dtypes.float64: dtypes.int64,
+}
+
 
 @tf_export.tf_export('experimental.numpy.dot', v1=[])
 @np_utils.np_doc_only('dot')
@@ -784,7 +793,9 @@ def absolute(x):
 @tf_export.tf_export('experimental.numpy.fabs', v1=[])
 @np_utils.np_doc('fabs')
 def fabs(x):
-  return abs(x)
+  # Unlike `absolute`, `fabs` always produces a floating point result,
+  # so an integer argument has to be promoted first.
+  return _scalar(math_ops.abs, x, True)
 
 
 @tf_export.tf_export('experimental.numpy.ceil', v1=[])
@@ -823,6 +834,11 @@ 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)
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 641fdac10b585b..f6eac507899ce2 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -20,6 +20,7 @@
 import numpy as np
 
 from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor
@@ -645,6 +646,18 @@ def testIsInf(self):
     self.assertFalse(np_math_ops.isneginf(x1))
     self.assertFalse(np_math_ops.isneginf(x2))
 
+  def testSignBit(self):
+    for transform in self.array_transforms:
+      values = transform([-1.5, -0.0, 0.0, 1.5])
+      self.assertAllEqual(
+          np_math_ops.signbit(values), [True, True, False, False]
+      )
+    # The sign bit is set even when the value compares equal to zero or NaN.
+    self.assertAllEqual(np_math_ops.signbit([np.nan, -np.nan]), [False, True])
+    self.assertAllEqual(np_math_ops.signbit([-3, 3]), [True, False])
+    negative_zero = ops.convert_to_tensor([-0.0], dtype=dtypes.bfloat16)
+    self.assertAllEqual(np_math_ops.signbit(negative_zero), [True])
+
   def testIsInfFamilyNonFloatInputs(self):
     # A non-floating input has no infinities, but the result must still be an
     # elementwise boolean array shaped like the input, as numpy returns, and
@@ -686,6 +699,29 @@ def testIsInfFamilyFloatInputs(self):
             tf_fun(arg), np_fun(arg), msg='{}({})'.format(np_fun.__name__, arg)
         )
 
+  def testFabsAlwaysReturnsFloat(self):
+    # `fabs` differs from `absolute` in that its result is always floating
+    # point, so an integer argument is promoted rather than passed through.
+    int_args = [
+        [1, -2, 3],
+        -5,
+        np.array([1, -2, 3], dtype=np.int32),
+        np.array([1, -2, 3], dtype=np.int64),
+        np.array([], dtype=np.int32),
+        np.array([[1, -2], [3, -4]], dtype=np.int32),
+    ]
+    for arg in int_args:
+      self.match(
+          np_math_ops.fabs(arg), np.fabs(arg), msg='fabs({})'.format(arg)
+      )
+
+    # A floating point argument keeps its own dtype.
+    for dtype in [np.float16, np.float32, np.float64]:
+      arg = np.array([1.5, -2.5], dtype=dtype)
+      self.match(
+          np_math_ops.fabs(arg), np.fabs(arg), msg='fabs({})'.format(arg)
+      )
+
 
 if __name__ == '__main__':
   tensor.enable_tensor_equality()
diff --git a/tensorflow/python/tfe_wrapper.cc b/tensorflow/python/tfe_wrapper.cc
index c75bfc16a8f50b..990b4cf6a88f75 100644
--- a/tensorflow/python/tfe_wrapper.cc
+++ b/tensorflow/python/tfe_wrapper.cc
@@ -1292,10 +1292,16 @@ PYBIND11_MODULE(_pywrap_tfe, m) {
     return tensorflow::PyoOrThrow(
         TFE_Py_TapeSetNew(persistent.ptr(), watch_accessed_variables.ptr()));
   });
-  m.def("TFE_Py_TapeSetAdd",
-        [](const py::handle& tape) { TFE_Py_TapeSetAdd(tape.ptr()); });
-  m.def("TFE_Py_TapeSetRemove",
-        [](const py::handle& tape) { TFE_Py_TapeSetRemove(tape.ptr()); });
+  m.def("TFE_Py_TapeSetAdd", [](const py::handle& tape) {
+    bool had_err = (PyErr_Occurred() != nullptr);
+    TFE_Py_TapeSetAdd(tape.ptr());
+    if (!had_err && PyErr_Occurred()) throw py::error_already_set();
+  });
+  m.def("TFE_Py_TapeSetRemove", [](const py::handle& tape) {
+    bool had_err = (PyErr_Occurred() != nullptr);
+    TFE_Py_TapeSetRemove(tape.ptr());
+    if (!had_err && PyErr_Occurred()) throw py::error_already_set();
+  });
   m.def("TFE_Py_TapeSetStopOnThread", &TFE_Py_TapeSetStopOnThread);
   m.def("TFE_Py_TapeSetRestartOnThread", &TFE_Py_TapeSetRestartOnThread);
   m.def("TFE_Py_TapeSetIsStopped",
@@ -1355,11 +1361,15 @@ PYBIND11_MODULE(_pywrap_tfe, m) {
   });
   m.def("TFE_Py_TapeWatch",
         [](const py::handle& tape, const py::handle& tensor) {
+          bool had_err = (PyErr_Occurred() != nullptr);
           TFE_Py_TapeWatch(tape.ptr(), tensor.ptr());
+          if (!had_err && PyErr_Occurred()) throw py::error_already_set();
         });
   m.def("TFE_Py_TapeWatchVariable",
         [](const py::handle& tape, const py::handle& variable) {
+          bool had_err = (PyErr_Occurred() != nullptr);
           TFE_Py_TapeWatchVariable(tape.ptr(), variable.ptr());
+          if (!had_err && PyErr_Occurred()) throw py::error_already_set();
         });
   m.def("TFE_Py_TapeWatchedVariables", [](const py::handle& tape) {
     return tensorflow::PyoOrThrow(TFE_Py_TapeWatchedVariables(tape.ptr()));
@@ -1369,7 +1379,9 @@ PYBIND11_MODULE(_pywrap_tfe, m) {
   m.def("TFE_Py_VariableWatcherNew",
         []() { return tensorflow::PyoOrThrow(TFE_Py_VariableWatcherNew()); });
   m.def("TFE_Py_VariableWatcherRemove", [](const py::handle& variable_watcher) {
+    bool had_err = (PyErr_Occurred() != nullptr);
     TFE_Py_VariableWatcherRemove(variable_watcher.ptr());
+    if (!had_err && PyErr_Occurred()) throw py::error_already_set();
   });
   m.def("TFE_Py_VariableWatcherVariableAccessed",
         [](const py::handle& variable) {
@@ -1392,14 +1404,18 @@ PYBIND11_MODULE(_pywrap_tfe, m) {
   });
   m.def("TFE_Py_ForwardAccumulatorSetRemove",
         [](const py::handle& accumulator) {
+          bool had_err = (PyErr_Occurred() != nullptr);
           TFE_Py_ForwardAccumulatorSetRemove(accumulator.ptr());
+          if (!had_err && PyErr_Occurred()) throw py::error_already_set();
         });
 
   m.def("TFE_Py_ForwardAccumulatorWatch",
         [](const py::handle& accumulator, const py::handle& tensor,
            const py::handle& tangent) {
+          bool had_err = (PyErr_Occurred() != nullptr);
           TFE_Py_ForwardAccumulatorWatch(accumulator.ptr(), tensor.ptr(),
                                          tangent.ptr());
+          if (!had_err && PyErr_Occurred()) throw py::error_already_set();
         });
   m.def("TFE_Py_ForwardAccumulatorJVP",
         [](const py::handle& accumulator, const py::handle& tensor) {
diff --git a/third_party/tf_runtime/pointer_union.patch b/third_party/tf_runtime/pointer_union.patch
new file mode 100644
index 00000000000000..334a96ce722eb0
--- /dev/null
+++ b/third_party/tf_runtime/pointer_union.patch
@@ -0,0 +1,17 @@
+diff --git a/lib/basic_kernels/opdefs/basic_kernels.cc b/lib/basic_kernels/opdefs/basic_kernels.cc
+--- a/lib/basic_kernels/opdefs/basic_kernels.cc
++++ b/lib/basic_kernels/opdefs/basic_kernels.cc
+@@ -137,11 +137,11 @@
+ void CallOp::setCalleeFromCallable(CallInterfaceCallable callee) {
+   // Direct call.
+   if (FlatSymbolRefAttr calleeAttr = getCalleeAttr()) {
+-    auto symRef = callee.get<SymbolRefAttr>();
++    auto symRef = cast<SymbolRefAttr>(callee);
+     return setCalleeAttr(cast<FlatSymbolRefAttr>(symRef));
+   }
+   // Indirect call, callee Value is the first operand.
+-  return setOperand(0, callee.get<Value>());
++  return setOperand(0, cast<Value>(callee));
+ }
+ 
+ //===----------------------------------------------------------------------===//
diff --git a/third_party/tf_runtime/workspace.bzl b/third_party/tf_runtime/workspace.bzl
index 7bc99867ef45b5..33d518ba0ac0db 100644
--- a/third_party/tf_runtime/workspace.bzl
+++ b/third_party/tf_runtime/workspace.bzl
@@ -31,5 +31,8 @@ def repo():
         urls = tf_mirror_urls("https://github.com/tensorflow/runtime/archive/{commit}.tar.gz".format(commit = TFRT_COMMIT)),
         # A patch file can be provided for atomic commits to both TF and TFRT.
         # The job that bumps the TFRT_COMMIT also resets patch_file to 'None'.
-        patch_file = ["//third_party/tf_runtime:f16_attr.patch"],
+        patch_file = [
+            "//third_party/tf_runtime:f16_attr.patch",
+            "//third_party/tf_runtime:pointer_union.patch",
+        ],
     )
diff --git a/third_party/xla/third_party/cuda_tile/BUILD.bazel b/third_party/xla/third_party/cuda_tile/BUILD.bazel
new file mode 100644
index 00000000000000..e69de29bb2d1d6
diff --git a/third_party/xla/third_party/cuda_tile/cuda_tile.BUILD b/third_party/xla/third_party/cuda_tile/cuda_tile.BUILD
new file mode 100644
index 00000000000000..72a4d67e2d0297
--- /dev/null
+++ b/third_party/xla/third_party/cuda_tile/cuda_tile.BUILD
@@ -0,0 +1,388 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+cc_binary(
+    name = "cuda-tile-tblgen",
+    srcs = glob([
+        "tools/cuda-tile-tblgen/*.cpp",
+        "tools/cuda-tile-tblgen/*.h",
+    ]),
+    deps = [
+        "@llvm-project//llvm:Support",
+        "@llvm-project//llvm:TableGen",
+        "@llvm-project//mlir:MlirTableGenMain",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TableGen",
+    ],
+)
+
+exports_files(["LICENSE"])
+
+td_library(
+    name = "CudaTileTdFiles",
+    srcs = glob(["include/cuda_tile/Dialect/CudaTile/IR/*.td"]),
+    includes = ["include"],
+    deps = [
+        "@llvm-project//mlir:BuiltinDialectTdFiles",
+        "@llvm-project//mlir:ControlFlowInterfacesTdFiles",
+        "@llvm-project//mlir:FunctionInterfacesTdFiles",
+        "@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
+        "@llvm-project//mlir:OpBaseTdFiles",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+        "@llvm-project//mlir:ViewLikeInterfaceTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "CudaTileDialectIncGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-dialect-decls",
+                "-dialect=cuda_tile",
+            ],
+            "include/cuda_tile/Dialect/CudaTile/IR/Dialect.h.inc",
+        ),
+        (
+            [
+                "-gen-dialect-defs",
+                "-dialect=cuda_tile",
+            ],
+            "include/cuda_tile/Dialect/CudaTile/IR/Dialect.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/cuda_tile/Dialect/CudaTile/IR/Dialect.td",
+    deps = [":CudaTileTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "CudaTileInterfacesIncGen",
+    tbl_outs = [
+        (
+            ["-gen-attr-interface-decls"],
+            "include/cuda_tile/Dialect/CudaTile/IR/AttrInterfaces.h.inc",
+        ),
+        (
+            ["-gen-attr-interface-defs"],
+            "include/cuda_tile/Dialect/CudaTile/IR/AttrInterfaces.cpp.inc",
+        ),
+        (
+            ["-gen-type-interface-decls"],
+            "include/cuda_tile/Dialect/CudaTile/IR/TypeInterfaces.h.inc",
+        ),
+        (
+            ["-gen-type-interface-defs"],
+            "include/cuda_tile/Dialect/CudaTile/IR/TypeInterfaces.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/cuda_tile/Dialect/CudaTile/IR/Interfaces.td",
+    deps = [":CudaTileTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "CudaTileTypesIncGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-typedef-decls",
+                "-typedefs-dialect=cuda_tile",
+            ],
+            "include/cuda_tile/Dialect/CudaTile/IR/Types.h.inc",
+        ),
+        (
+            [
+                "-gen-typedef-defs",
+                "-typedefs-dialect=cuda_tile",
+            ],
+            "include/cuda_tile/Dialect/CudaTile/IR/Types.cpp.inc",
+        ),
+        (
+            [
+                "-gen-type-constraint-decls",
+                "-dialect=cuda_tile",
+            ],
+            "include/cuda_tile/Dialect/CudaTile/IR/TypeConstraints.h.inc",
+        ),
+        (
+            [
+                "-gen-type-constraint-defs",
+                "-dialect=cuda_tile",
+            ],
+            "include/cuda_tile/Dialect/CudaTile/IR/TypeConstraints.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/cuda_tile/Dialect/CudaTile/IR/Types.td",
+    deps = [":CudaTileTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "CudaTileAttrDefsIncGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-attrdef-decls",
+                "-attrdefs-dialect=cuda_tile",
+            ],
+            "include/cuda_tile/Dialect/CudaTile/IR/AttrDefs.h.inc",
+        ),
+        (
+            [
+                "-gen-attrdef-defs",
+                "-attrdefs-dialect=cuda_tile",
+            ],
+            "include/cuda_tile/Dialect/CudaTile/IR/AttrDefs.cpp.inc",
+        ),
+        (
+            [
+                "-gen-enum-decls",
+                "-dialect=cuda_tile",
+            ],
+            "include/cuda_tile/Dialect/CudaTile/IR/Enums.h.inc",
+        ),
+        (
+            [
+                "-gen-enum-defs",
+                "-dialect=cuda_tile",
+            ],
+            "include/cuda_tile/Dialect/CudaTile/IR/Enums.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/cuda_tile/Dialect/CudaTile/IR/AttrDefs.td",
+    deps = [":CudaTileTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "CudaTileRemarksIncGen",
+    tbl_outs = [
+        (
+            ["-gen-enum-decls"],
+            "include/cuda_tile/Dialect/CudaTile/IR/TileIRRemarks.h.inc",
+        ),
+        (
+            ["-gen-enum-defs"],
+            "include/cuda_tile/Dialect/CudaTile/IR/TileIRRemarks.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/cuda_tile/Dialect/CudaTile/IR/Remarks.td",
+    deps = [":CudaTileTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "CudaTileOpsIncGen",
+    tbl_outs = [
+        (
+            ["-gen-op-decls"],
+            "include/cuda_tile/Dialect/CudaTile/IR/Ops.h.inc",
+        ),
+        (
+            ["-gen-op-defs"],
+            "include/cuda_tile/Dialect/CudaTile/IR/Ops.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/cuda_tile/Dialect/CudaTile/IR/Ops.td",
+    deps = [":CudaTileTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "CudaTileOpsCanonicalizationIncGen",
+    tbl_outs = [
+        (
+            ["-gen-rewriters"],
+            "lib/Dialect/CudaTile/IR/OpsCanonicalization.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "lib/Dialect/CudaTile/IR/OpsCanonicalization.td",
+    deps = [
+        ":CudaTileTdFiles",
+        "@llvm-project//mlir:OpBaseTdFiles",
+    ],
+)
+
+cc_library(
+    name = "CudaTileDialect",
+    srcs = glob(
+        ["lib/Dialect/CudaTile/IR/*.cpp"],
+        exclude = ["lib/Dialect/CudaTile/IR/CudaTileTesting.cpp"],
+    ),
+    hdrs = glob(["include/cuda_tile/Dialect/CudaTile/IR/*.h"]),
+    includes = [
+        "include",
+        "lib/Dialect/CudaTile/IR",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":CudaTileAttrDefsIncGen",
+        ":CudaTileDialectIncGen",
+        ":CudaTileInterfacesIncGen",
+        ":CudaTileOpsCanonicalizationIncGen",
+        ":CudaTileOpsIncGen",
+        ":CudaTileTypesIncGen",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:BytecodeOpInterface",
+        "@llvm-project//mlir:ControlFlowInterfaces",
+        "@llvm-project//mlir:FunctionInterfaces",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InferTypeOpInterface",
+        "@llvm-project//mlir:InliningUtils",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:ViewLikeInterface",
+    ],
+)
+
+td_library(
+    name = "CudaTileTransformsTdFiles",
+    srcs = ["include/cuda_tile/Dialect/CudaTile/Transforms/Passes.td"],
+    includes = ["include"],
+    deps = [
+        "@llvm-project//mlir:PassBaseTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "CudaTileTransformsIncGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-pass-decls",
+                "-name=CudaTile",
+            ],
+            "include/cuda_tile/Dialect/CudaTile/Transforms/Passes.h.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/cuda_tile/Dialect/CudaTile/Transforms/Passes.td",
+    deps = [":CudaTileTransformsTdFiles"],
+)
+
+cc_library(
+    name = "CudaTileTransforms",
+    srcs = glob(["lib/Dialect/CudaTile/Transforms/*.cpp"]),
+    hdrs = glob(["include/cuda_tile/Dialect/CudaTile/Transforms/*.h"]),
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":CudaTileDialect",
+        ":CudaTileTransformsIncGen",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)
+
+gentbl_cc_library(
+    name = "CudaTileBytecodeOpsIncGen",
+    tbl_outs = [
+        (
+            ["-gen-cuda-tile-bytecode"],
+            "include/cuda_tile/Bytecode/Writer/Bytecode.inc",
+        ),
+        (
+            ["-gen-cuda-tile-bytecode-reader"],
+            "include/cuda_tile/Bytecode/Reader/BytecodeReader.inc",
+        ),
+    ],
+    tblgen = ":cuda-tile-tblgen",
+    td_file = "include/cuda_tile/Dialect/CudaTile/IR/Ops.td",
+    deps = [":CudaTileTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "CudaTileBytecodeOpcodesIncGen",
+    tbl_outs = [
+        (
+            ["-gen-cuda-tile-opcodes"],
+            "include/cuda_tile/Bytecode/Common/StaticOpcodes.inc",
+        ),
+    ],
+    tblgen = ":cuda-tile-tblgen",
+    td_file = "include/cuda_tile/Dialect/CudaTile/IR/BytecodeOpcodes.td",
+    deps = [":CudaTileTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "CudaTileBytecodeTypeIncGen",
+    tbl_outs = [
+        (
+            ["-gen-cuda-tile-type-bytecode"],
+            "include/cuda_tile/Bytecode/Writer/TypeBytecode.inc",
+        ),
+        (
+            ["-gen-cuda-tile-type-bytecode-reader"],
+            "include/cuda_tile/Bytecode/Reader/TypeBytecodeReader.inc",
+        ),
+    ],
+    tblgen = ":cuda-tile-tblgen",
+    td_file = "include/cuda_tile/Dialect/CudaTile/IR/BytecodeTypeOpcodes.td",
+    deps = [":CudaTileTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "CudaTileBytecodeAttrIncGen",
+    tbl_outs = [
+        (
+            ["-gen-cuda-tile-attr-bytecode"],
+            "include/cuda_tile/Bytecode/Writer/AttrBytecode.inc",
+        ),
+    ],
+    tblgen = ":cuda-tile-tblgen",
+    td_file = "include/cuda_tile/Dialect/CudaTile/IR/BytecodeAttrOpcodes.td",
+    deps = [":CudaTileTdFiles"],
+)
+
+cc_library(
+    name = "CudaTileBytecode",
+    srcs = glob([
+        "lib/Bytecode/**/*.cpp",
+        "lib/Bytecode/**/*.h",
+    ]),
+    hdrs = glob([
+        "include/cuda_tile/Bytecode/**/*.h",
+    ]),
+    includes = [
+        "include",
+        "include/cuda_tile/Bytecode/Common",
+        "include/cuda_tile/Bytecode/Reader",
+        "include/cuda_tile/Bytecode/Writer",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":CudaTileBytecodeAttrIncGen",
+        ":CudaTileBytecodeOpcodesIncGen",
+        ":CudaTileBytecodeOpsIncGen",
+        ":CudaTileBytecodeTypeIncGen",
+        ":CudaTileDialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:TranslateLib",
+    ],
+)
+
+cc_library(
+    name = "CudaTileOptimizer",
+    srcs = ["lib/Dialect/CudaTile/Optimizer/CudaTileOptimizer.cpp"],
+    hdrs = ["include/cuda_tile/Dialect/CudaTile/Optimizer/CudaTileOptimizer.h"],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":CudaTileBytecode",
+        ":CudaTileDialect",
+        ":CudaTileTransforms",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+    ],
+)
diff --git a/third_party/xla/third_party/cuda_tile/workspace.bzl b/third_party/xla/third_party/cuda_tile/workspace.bzl
new file mode 100644
index 00000000000000..7dcb4e0c479335
--- /dev/null
+++ b/third_party/xla/third_party/cuda_tile/workspace.bzl
@@ -0,0 +1,16 @@
+"""Provides the repository macro to import CUDA Tile IR."""
+
+load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
+
+def repo():
+    """Imports CUDA Tile IR."""
+    CUDA_TILE_COMMIT = "af2417041cc939b87ef56d92cfdcf61737c5457e"
+    CUDA_TILE_SHA256 = "81597e49469171bf8fa7319fbd44ebe133001521f484589e3dd3fb3fad282dc0"
+
+    tf_http_archive(
+        name = "cuda_tile",
+        build_file = "//third_party/cuda_tile:cuda_tile.BUILD",
+        sha256 = CUDA_TILE_SHA256,
+        strip_prefix = "cuda-tile-{}".format(CUDA_TILE_COMMIT),
+        urls = tf_mirror_urls("https://github.com/NVIDIA/cuda-tile/archive/{}.tar.gz".format(CUDA_TILE_COMMIT)),
+    )
diff --git a/third_party/xla/third_party/llvm/platforms_config.patch b/third_party/xla/third_party/llvm/platforms_config.patch
index 40af483204777c..995624e06d820e 100644
--- a/third_party/xla/third_party/llvm/platforms_config.patch
+++ b/third_party/xla/third_party/llvm/platforms_config.patch
@@ -15,8 +15,9 @@
 diff --git a/utils/bazel/llvm-project-overlay/llvm/config.bzl b/utils/bazel/llvm-project-overlay/llvm/config.bzl
 --- a/utils/bazel/llvm-project-overlay/llvm/config.bzl
 +++ b/utils/bazel/llvm-project-overlay/llvm/config.bzl
-@@ -49,13 +49,12 @@
+@@ -69,14 +69,13 @@
  backtrace_defines = select({
+     "@platforms//os:emscripten": [],
      "@platforms//os:windows": [],
 -    "@llvm//platforms/config:musl": [],
      "//conditions:default": [
diff --git a/third_party/xla/third_party/tensor_ir/BUILD.bazel b/third_party/xla/third_party/tensor_ir/BUILD.bazel
new file mode 100644
index 00000000000000..e69de29bb2d1d6
diff --git a/third_party/xla/third_party/tensor_ir/tensor_ir.BUILD b/third_party/xla/third_party/tensor_ir/tensor_ir.BUILD
new file mode 100644
index 00000000000000..3d591efd26df72
--- /dev/null
+++ b/third_party/xla/third_party/tensor_ir/tensor_ir.BUILD
@@ -0,0 +1,480 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+exports_files(["LICENSE"])
+
+td_library(
+    name = "TensorIRTdFiles",
+    srcs = glob(["include/tensor_ir/Dialect/*.td"]),
+    includes = ["include"],
+    deps = [
+        "@llvm-project//mlir:BuiltinDialectTdFiles",
+        "@llvm-project//mlir:ControlFlowInterfacesTdFiles",
+        "@llvm-project//mlir:FunctionInterfacesTdFiles",
+        "@llvm-project//mlir:GPUOpsTdFiles",
+        "@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
+        "@llvm-project//mlir:OpBaseTdFiles",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "TensorIRDialectIncGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-dialect-decls",
+                "-dialect=nv_tensor_ir",
+            ],
+            "include/tensor_ir/Dialect/TensorDialect.h.inc",
+        ),
+        (
+            [
+                "-gen-dialect-defs",
+                "-dialect=nv_tensor_ir",
+            ],
+            "include/tensor_ir/Dialect/TensorDialect.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/tensor_ir/Dialect/TensorDialect.td",
+    deps = [":TensorIRTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "TensorIROpsIncGen",
+    tbl_outs = [
+        (
+            ["-gen-op-decls"],
+            "include/tensor_ir/Dialect/TensorOps.h.inc",
+        ),
+        (
+            ["-gen-op-defs"],
+            "include/tensor_ir/Dialect/TensorOps.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/tensor_ir/Dialect/TensorOps.td",
+    deps = [":TensorIRTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "TensorIREnumsIncGen",
+    tbl_outs = [
+        (
+            ["-gen-enum-decls"],
+            "include/tensor_ir/Dialect/TensorEnums.h.inc",
+        ),
+        (
+            ["-gen-enum-defs"],
+            "include/tensor_ir/Dialect/TensorEnums.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/tensor_ir/Dialect/TensorEnums.td",
+    deps = [":TensorIRTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "TensorIROpInterfacesIncGen",
+    tbl_outs = [
+        (
+            ["-gen-op-interface-decls"],
+            "include/tensor_ir/Dialect/TensorOpInterfaces.h.inc",
+        ),
+        (
+            ["-gen-op-interface-defs"],
+            "include/tensor_ir/Dialect/TensorOpInterfaces.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/tensor_ir/Dialect/TensorInterfaces.td",
+    deps = [":TensorIRTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "TensorIRAttrInterfacesIncGen",
+    tbl_outs = [
+        (
+            ["-gen-attr-interface-decls"],
+            "include/tensor_ir/Dialect/TensorAttrInterfaces.h.inc",
+        ),
+        (
+            ["-gen-attr-interface-defs"],
+            "include/tensor_ir/Dialect/TensorAttrInterfaces.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/tensor_ir/Dialect/TensorAttrInterfaces.td",
+    deps = [":TensorIRTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "TensorIRAttrsIncGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-attrdef-decls",
+                "-attrdefs-dialect=nv_tensor_ir",
+            ],
+            "include/tensor_ir/Dialect/TensorAttrs.h.inc",
+        ),
+        (
+            [
+                "-gen-attrdef-defs",
+                "-attrdefs-dialect=nv_tensor_ir",
+            ],
+            "include/tensor_ir/Dialect/TensorAttrs.cpp.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/tensor_ir/Dialect/TensorAttrs.td",
+    deps = [":TensorIRTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "TensorIROpsCanonicalizationIncGen",
+    strip_include_prefix = "lib/Dialect",
+    tbl_outs = [
+        (
+            ["-gen-rewriters"],
+            "lib/Dialect/TensorOpsCanonicalization.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "lib/Dialect/TensorOpsCanonicalization.td",
+    deps = [
+        ":TensorIRTdFiles",
+        "@llvm-project//mlir:OpBaseTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "TensorIRTransformPassesIncGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-pass-decls",
+                "-name=NVTensorIRTransform",
+            ],
+            "include/tensor_ir/Transform/Passes.h.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/tensor_ir/Transform/Passes.td",
+    deps = [
+        ":TensorIRTdFiles",
+        "@llvm-project//mlir:PassBaseTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "TensorToCudaTileConversionPassIncGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-pass-decls",
+                "-name=TensorToCudaTileConversion",
+            ],
+            "include/tensor_ir/Conversion/TensorToCudaTile/Passes.h.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "include/tensor_ir/Conversion/TensorToCudaTile/Passes.td",
+    deps = [
+        ":TensorIRTdFiles",
+        "@llvm-project//mlir:PassBaseTdFiles",
+    ],
+)
+
+cc_library(
+    name = "NVTensorIRSupport",
+    srcs = ["lib/Support/TCutegen.cpp"],
+    hdrs = [
+        "include/tensor_ir/Support/Macros.h",
+        "include/tensor_ir/Support/Status.h",
+        "include/tensor_ir/Support/TCutegen.h",
+    ],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_library(
+    name = "NVTensorIRCudaApi",
+    srcs = ["lib/Support/CudaApi.cpp"],
+    hdrs = ["include/tensor_ir/Support/CudaApi.h"],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":NVTensorIRSupport",
+        "@llvm-project//llvm:Support",
+        "@local_config_cuda//cuda:cuda_headers",
+    ],
+)
+
+cc_library(
+    name = "NVTensorIRDialect",
+    srcs = [
+        "lib/Dialect/Canonicalization.cpp",
+        "lib/Dialect/TensorAttrs.cpp",
+        "lib/Dialect/TensorDialect.cpp",
+        "lib/Dialect/TensorOps.cpp",
+    ],
+    hdrs = [
+        "include/tensor_ir/Dialect/TensorIR.h",
+        "include/tensor_ir/Dialect/TensorIRAttrs.h",
+    ],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":NVTensorIRSupport",
+        ":TensorIRAttrInterfacesIncGen",
+        ":TensorIRAttrsIncGen",
+        ":TensorIRDialectIncGen",
+        ":TensorIREnumsIncGen",
+        ":TensorIROpInterfacesIncGen",
+        ":TensorIROpsCanonicalizationIncGen",
+        ":TensorIROpsIncGen",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:BytecodeOpInterface",
+        "@llvm-project//mlir:ControlFlowInterfaces",
+        "@llvm-project//mlir:FunctionInterfaces",
+        "@llvm-project//mlir:GPUDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InferTypeOpInterface",
+        "@llvm-project//mlir:SideEffectInterfaces",
+    ],
+)
+
+cc_library(
+    name = "NVTensorIRUtils",
+    srcs = [
+        "lib/Utils/ComputeCapability.cpp",
+        "lib/Utils/Utils.cpp",
+    ],
+    hdrs = [
+        "include/tensor_ir/Utils/ComputeCapability.h",
+        "include/tensor_ir/Utils/Utils.h",
+    ],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":NVTensorIRDialect",
+        ":NVTensorIRSupport",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:BytecodeWriter",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_library(
+    name = "NVTensorIRAnalysis",
+    srcs = [
+        "lib/Analysis/KernelArgLayout.cpp",
+        "lib/Analysis/TileAnalyzer.cpp",
+        "lib/Analysis/TileCandidateGenerator.cpp",
+    ],
+    hdrs = [
+        "include/tensor_ir/Analysis/TileAnalyzer.h",
+        "include/tensor_ir/Analysis/TileCandidateGenerator.h",
+        "include/tensor_ir/Compiler/CudaTile/KernelArgLayout.h",
+    ],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":NVTensorIRDialect",
+        ":NVTensorIRRuntime",
+        ":NVTensorIRUtils",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
+cc_library(
+    name = "NVTensorIRTransform",
+    srcs = glob(["lib/Transform/*.cpp"]),
+    hdrs = ["include/tensor_ir/Transform/Passes.h"],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":NVTensorIRAnalysis",
+        ":NVTensorIRDialect",
+        ":NVTensorIRSupport",
+        ":NVTensorIRUtils",
+        ":TensorIRTransformPassesIncGen",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_library(
+    name = "NVTensorIRToCudaTileConversion",
+    srcs = glob(["lib/Conversion/TensorToCudaTile/*.cpp"]),
+    hdrs = [
+        "include/tensor_ir/Conversion/TensorToCudaTile/Options.h",
+        "include/tensor_ir/Conversion/TensorToCudaTile/TensorToCudaTile.h",
+        "include/tensor_ir/Conversion/TensorToCudaTile/TensorToCudaTileInternal.h",
+    ],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":NVTensorIRAnalysis",
+        ":NVTensorIRDialect",
+        ":NVTensorIRSupport",
+        ":NVTensorIRUtils",
+        ":TensorToCudaTileConversionPassIncGen",
+        "@cuda_tile//:CudaTileDialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:DialectUtils",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)
+
+cc_library(
+    name = "NVTensorIRCudaTilePipelines",
+    srcs = ["lib/Compiler/CudaTile/Pipelines.cpp"],
+    hdrs = ["include/tensor_ir/Compiler/CudaTile/Pipelines.h"],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":NVTensorIRDialect",
+        ":NVTensorIRToCudaTileConversion",
+        ":NVTensorIRTransform",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Transforms",
+    ],
+)
+
+cc_library(
+    name = "NVTensorIRRuntime",
+    srcs = ["lib/Runtime/CudaTileRuntimeKernel.cpp"],
+    hdrs = [
+        "include/tensor_ir/Runtime/CudaTile/CudaTileRuntimeKernel.h",
+        "include/tensor_ir/Runtime/CudaTile/KernelArgLayout.h",
+        "include/tensor_ir/Runtime/CudaTile/KernelLaunchHelpers.h",
+        "include/tensor_ir/Runtime/CudaTile/RuntimeOperandAccessor.h",
+        "include/tensor_ir/Runtime/IRuntimeKernel.h",
+        "include/tensor_ir/Runtime/Types.h",
+    ],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":NVTensorIRCudaApi",
+        ":NVTensorIRSupport",
+        "@cuda_tile//:CudaTileBytecode",
+        "@cuda_tile//:CudaTileDialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@local_config_cuda//cuda:cuda_headers",
+    ],
+)
+
+cc_library(
+    name = "NVTensorIRReference",
+    srcs = glob(["lib/Reference/*.cpp"]),
+    hdrs = [
+        "include/tensor_ir/Reference/reference_graph.h",
+        "include/tensor_ir/Reference/reference_node.h",
+        "include/tensor_ir/Reference/simplified_tensor.h",
+        "include/tensor_ir/Reference/tensor_memory.h",
+        "lib/Reference/constant_utils.h",
+    ],
+    includes = [
+        "include",
+        "lib/Reference",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":NVTensorIRCudaApi",
+        ":NVTensorIRDialect",
+        ":NVTensorIRRuntime",
+        ":NVTensorIRSupport",
+        ":NVTensorIRUtils",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@local_config_cuda//cuda:cuda_headers",
+    ],
+)
+
+cc_library(
+    name = "NVTensorIRRegistration",
+    srcs = ["lib/Registration/Registration.cpp"],
+    hdrs = ["include/tensor_ir/Registration/Registration.h"],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":NVTensorIRDialect",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FuncExtensions",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
+cc_library(
+    name = "NVTensorIRCompiler",
+    srcs = [
+        "lib/Compiler/Compiler.cpp",
+        "lib/Compiler/CudaTile/CudaTileCompiler.cpp",
+        "lib/Compiler/CudaTile/CudaTileFrontend.cpp",
+        "lib/Compiler/CudaTile/TileIRAssembly.cpp",
+    ],
+    hdrs = [
+        "include/tensor_ir/Compiler/CompileOptions.h",
+        "include/tensor_ir/Compiler/Compiler.h",
+        "include/tensor_ir/Compiler/CudaTile/CudaTileCompiler.h",
+        "include/tensor_ir/Compiler/CudaTile/CudaTileFrontend.h",
+        "include/tensor_ir/Compiler/CudaTile/TileIRAssembly.h",
+    ],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":NVTensorIRAnalysis",
+        ":NVTensorIRCudaTilePipelines",
+        ":NVTensorIRDialect",
+        ":NVTensorIRRegistration",
+        ":NVTensorIRRuntime",
+        ":NVTensorIRSupport",
+        ":NVTensorIRToCudaTileConversion",
+        ":NVTensorIRUtils",
+        "@cuda_tile//:CudaTileBytecode",
+        "@cuda_tile//:CudaTileDialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_library(
+    name = "NVTensorIRCAPI",
+    srcs = ["lib/CAPI/TensorIR.cpp"],
+    hdrs = ["include/tensor_ir-c/TensorIR.h"],
+    includes = ["include"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":NVTensorIRCompiler",
+        ":NVTensorIRDialect",
+        ":NVTensorIRRegistration",
+        ":NVTensorIRRuntime",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:CAPIIRHeaders",
+    ],
+)
diff --git a/third_party/xla/third_party/tensor_ir/workspace.bzl b/third_party/xla/third_party/tensor_ir/workspace.bzl
new file mode 100644
index 00000000000000..089aacd7ceafe8
--- /dev/null
+++ b/third_party/xla/third_party/tensor_ir/workspace.bzl
@@ -0,0 +1,16 @@
+"""Provides the repository macro to import Tensor IR."""
+
+load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
+
+def repo():
+    """Imports Tensor IR."""
+    TENSOR_IR_COMMIT = "63692d79629e6f32a1d8757695590a59e0adbafd"
+    TENSOR_IR_SHA256 = "b80794d7c2bfb1bc1ca432d892977becba8d35ec6c18c586acbd648ccc8074dd"
+
+    tf_http_archive(
+        name = "tensor_ir",
+        build_file = "//third_party/tensor_ir:tensor_ir.BUILD",
+        sha256 = TENSOR_IR_SHA256,
+        strip_prefix = "tensor-ir-{}".format(TENSOR_IR_COMMIT),
+        urls = tf_mirror_urls("https://github.com/NVIDIA/tensor-ir/archive/{}.tar.gz".format(TENSOR_IR_COMMIT)),
+    )
diff --git a/third_party/xla/workspace2.bzl b/third_party/xla/workspace2.bzl
index 923de0b6bda95d..0f49a1db57c284 100644
--- a/third_party/xla/workspace2.bzl
+++ b/third_party/xla/workspace2.bzl
@@ -31,6 +31,7 @@ load("//third_party/brotli:workspace.bzl", brotli = "repo")
 load("//third_party/clang_toolchain:cc_configure_clang.bzl", "cc_download_clang_toolchain")
 load("//third_party/compute_library:workspace.bzl", compute_library = "repo")
 load("//third_party/cpuinfo:workspace.bzl", cpuinfo = "repo")
+load("//third_party/cuda_tile:workspace.bzl", cuda_tile = "repo")
 load("//third_party/cudnn_frontend:workspace.bzl", cudnn_frontend = "repo")
 load("//third_party/cutlass:workspace.bzl", cutlass = "repo")
 load("//third_party/cutlass_cutedsl_runtime:workspace.bzl", cutlass_cutedsl_runtime = "repo")
@@ -81,6 +82,7 @@ load("//third_party/spdlog:workspace.bzl", spdlog = "repo")
 load("//third_party/sqlite:workspace.bzl", sqlite = "repo")
 load("//third_party/stablehlo:workspace.bzl", stablehlo = "repo")
 load("//third_party/system_libpci:workspace.bzl", system_libpci = "repo")
+load("//third_party/tensor_ir:workspace.bzl", tensor_ir = "repo")
 load("//third_party/tensorrt:tensorrt_configure.bzl", "tensorrt_configure")
 load("//third_party/tensorrt:workspace.bzl", tensorrt = "repo")
 load("//third_party/transformer_engine:workspace.bzl", transformer_engine = "repo")
@@ -147,11 +149,13 @@ def _initialize_third_party():
     spdlog()
     sqlite()
     stablehlo()
+    tensor_ir()
     tensorrt()
     transformer_engine()
     triton()
     uv()
     xnnpack()
+    cuda_tile()
     cutlass()
     cutlass_cutedsl_runtime()
 
diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD
index f7ee78cd9db44b..864d7e0fe5126b 100644
--- a/third_party/xla/xla/backends/gpu/transforms/BUILD
+++ b/third_party/xla/xla/backends/gpu/transforms/BUILD
@@ -182,45 +182,6 @@ xla_cc_test(
     ],
 )
 
-cc_library(
-    name = "alias_in_place_outputs",
-    srcs = ["alias_in_place_outputs.cc"],
-    hdrs = ["alias_in_place_outputs.h"],
-    deps = [
-        "//xla:shape_util",
-        "//xla/hlo/analysis:hlo_reachability",
-        "//xla/hlo/analysis:indexing_analysis",
-        "//xla/hlo/analysis:symbolic_map",
-        "//xla/hlo/ir:hlo",
-        "//xla/hlo/pass:hlo_pass",
-        "//xla/hlo/utils:hlo_traversal",
-        "//xla/service/gpu:backend_configs_cc",
-        "//xla/service/gpu:cublas_cudnn",
-        "//xla/service/gpu:ir_emission_utils",
-        "@com_google_absl//absl/algorithm:container",
-        "@com_google_absl//absl/container:flat_hash_set",
-        "@com_google_absl//absl/log",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings:string_view",
-        "@llvm-project//mlir:IR",
-    ],
-)
-
-xla_cc_test(
-    name = "alias_in_place_outputs_test",
-    srcs = ["alias_in_place_outputs_test.cc"],
-    deps = [
-        ":alias_in_place_outputs",
-        "//xla:shape_util",
-        "//xla/hlo/ir:hlo",
-        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
-        "@com_google_absl//absl/log",
-        "@com_google_absl//absl/strings:string_view",
-        "@com_google_googletest//:gtest_main",
-        "@llvm-project//mlir:IR",
-    ],
-)
-
 cc_library(
     name = "async_wrapper",
     srcs = ["async_wrapper.cc"],
diff --git a/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs.cc b/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs.cc
deleted file mode 100644
index 2d87ff77af0b93..00000000000000
--- a/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs.cc
+++ /dev/null
@@ -1,208 +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/transforms/alias_in_place_outputs.h"
-
-#include <cstdint>
-#include <memory>
-#include <vector>
-
-#include "absl/algorithm/container.h"
-#include "absl/container/flat_hash_set.h"
-#include "absl/log/log.h"
-#include "absl/status/statusor.h"
-#include "absl/strings/string_view.h"
-#include "mlir/IR/MLIRContext.h"
-#include "xla/hlo/analysis/hlo_reachability.h"
-#include "xla/hlo/analysis/indexing_analysis.h"
-#include "xla/hlo/analysis/indexing_map.h"
-#include "xla/hlo/analysis/symbolic_expr.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/hlo/utils/hlo_traversal.h"
-#include "xla/service/gpu/backend_configs.pb.h"
-#include "xla/service/gpu/cublas_cudnn.h"
-#include "xla/service/gpu/ir_emission_utils.h"
-#include "xla/shape.h"
-#include "xla/shape_util.h"
-
-namespace xla::gpu {
-namespace {
-
-bool IsTritonFusion(const HloFusionInstruction& fusion) {
-  auto config = fusion.backend_config<GpuBackendConfig>();
-  if (!config.ok()) {
-    return false;
-  }
-  absl::string_view kind = config->fusion_backend_config().kind();
-  return kind == kTritonFusionKind || kind == kTritonGemmFusionKind ||
-         kind == kTritonNestedGemmFusionKind;
-}
-
-bool IsSupportedCublasLtMatmul(const HloInstruction& hlo) {
-  return IsCublasLtMatmul(hlo) || IsCublasLtMatmulF8(hlo) ||
-         IsCublasLtMatmulMx(hlo) || IsCublasLtGroupedMatmul(hlo);
-}
-
-// Returns true if the composed output→input indexing map for the given operand
-// of the fusion is an identity, i.e. output element [i,j,...] reads input
-// element [i,j,...]. This covers pure elementwise chains and bitcasts that
-// cancel. Operands read through multiple data-flow paths are rejected.
-bool HasIdentityIndexing(const GroupedByOpIndexing& grouped,
-                         const HloInstruction* operand) {
-  auto it = grouped.find(operand);
-  if (it == grouped.end() || it->second.size() != 1) {
-    return false;
-  }
-  const IndexingMap& map = it->second.begin()->map();
-  return map.GetSymbolicMap().IsIdentity() && map.GetRangeVars().empty() &&
-         map.GetSymbolicConstraints().empty();
-}
-
-// Returns true if overwriting `operand` in place avoids rather than forces a
-// copy: `operand` must be a writable intermediate (not a parameter, constant,
-// or root) that `user` reads exactly once, and every other user of `operand`
-// must precede `user` in the graph (e.g. a residual read before the op
-// overwrites it).
-bool BeneficialToAlias(const HloInstruction* operand,
-                       const HloInstruction* user,
-                       const HloReachabilityMap& reachability) {
-  if (operand->opcode() == HloOpcode::kParameter ||
-      operand->opcode() == HloOpcode::kConstant || operand->IsRoot()) {
-    return false;
-  }
-  if (absl::c_count(user->operands(), operand) != 1) {
-    // Don't alias if the same buffer is used multiple times in the user.
-    return false;
-  }
-  return absl::c_all_of(operand->users(), [&](const HloInstruction* other) {
-    return other == user || reachability.IsReachable(other, user);
-  });
-}
-
-bool AliasFusion(HloFusionInstruction& fusion,
-                 const HloReachabilityMap& reachability,
-                 mlir::MLIRContext& mlir_ctx) {
-  if (!fusion.shape().IsArray() ||
-      !fusion.output_to_operand_aliasing().empty()) {
-    return false;
-  }
-  std::unique_ptr<HloFusionAdaptor> fusion_adaptor =
-      HloFusionAdaptor::ForInstruction(&fusion);
-  GroupedByOpIndexing grouped = ComputeGroupedOutputToInputIndexing(
-      *fusion_adaptor, fusion_adaptor->GetRoots()[0], &mlir_ctx);
-  const HloInstruction* root = fusion.fused_expression_root();
-  for (int64_t i = 0; i < fusion.operand_count(); ++i) {
-    const HloInstruction* operand = fusion.operand(i);
-    // Aliasing requires matching element type and dimensions (layouts may
-    // differ), as enforced by the HLO verifier.
-    if (!ShapeUtil::Compatible(operand->shape(), root->shape())) {
-      continue;
-    }
-    if (!HasIdentityIndexing(grouped, operand)) {
-      continue;
-    }
-    if (!BeneficialToAlias(operand, &fusion, reachability)) {
-      continue;
-    }
-    VLOG(2) << "Aliasing output of " << fusion.name() << " to operand " << i
-            << " (" << operand->name() << ")";
-    fusion.set_output_to_operand_aliasing(
-        {{/*output_index=*/{}, {/*operand_number=*/i, /*operand_index=*/{}}}});
-    return true;
-  }
-  return false;
-}
-
-bool AliasCublasLtMatmul(HloCustomCallInstruction& custom_call,
-                         const HloReachabilityMap& reachability) {
-  if (!custom_call.output_to_operand_aliasing().empty()) {
-    return false;
-  }
-
-  const int64_t bias_idx = IsCublasLtGroupedMatmul(custom_call) ? 3 : 2;
-  if (custom_call.operand_count() <= bias_idx) {
-    return false;
-  }
-
-  auto config = custom_call.backend_config<GpuBackendConfig>();
-  if (!config.ok() || config->gemm_backend_config().beta() == 0.0) {
-    return false;
-  }
-
-  // cuBLASLt returns either the result array or a (result, scratch) tuple.
-  const Shape& shape = custom_call.shape();
-  ShapeIndex output_index;
-  const Shape* output_shape = &shape;
-  if (shape.IsTuple()) {
-    if (shape.tuple_shapes().empty()) {
-      return false;
-    }
-    output_index = ShapeIndex{0};
-    output_shape = &shape.tuple_shapes(0);
-  } else if (!shape.IsArray()) {
-    return false;
-  }
-
-  const HloInstruction* bias = custom_call.operand(bias_idx);
-  if (!ShapeUtil::Equal(bias->shape(), *output_shape)) {
-    return false;
-  }
-  if (!BeneficialToAlias(bias, &custom_call, reachability)) {
-    return false;
-  }
-
-  VLOG(2) << "Aliasing output " << output_index.ToString() << " of "
-          << custom_call.name() << " to operand " << bias_idx << " ("
-          << bias->name() << ")";
-  custom_call.set_output_to_operand_aliasing(
-      {{output_index, {/*operand_number=*/bias_idx, /*operand_index=*/{}}}});
-  return true;
-}
-
-}  // namespace
-
-absl::StatusOr<bool> AliasInPlaceOutputs::RunImpl(
-    HloModule* module,
-    const absl::flat_hash_set<absl::string_view>& execution_threads) {
-  bool changed = false;
-  RegisterSymbolicExprStorage(mlir_context_);
-  for (HloComputation* computation :
-       module->MakeNonfusionComputations(execution_threads)) {
-    std::unique_ptr<HloReachabilityMap> reachability;
-    // Only build the reachability map if needed, reuse if already built.
-    const auto reachability_of = [&]() -> const HloReachabilityMap& {
-      if (reachability == nullptr) {
-        reachability = HloReachabilityMap::Build(computation);
-      }
-      return *reachability;
-    };
-    for (HloInstruction* instr : computation->instructions()) {
-      if (auto* fusion = DynCast<HloFusionInstruction>(instr);
-          fusion != nullptr && IsTritonFusion(*fusion)) {
-        changed |= AliasFusion(*fusion, reachability_of(), *mlir_context_);
-      } else if (auto* call = DynCast<HloCustomCallInstruction>(instr);
-                 call != nullptr && IsSupportedCublasLtMatmul(*call)) {
-        changed |= AliasCublasLtMatmul(*call, reachability_of());
-      }
-    }
-  }
-  return changed;
-}
-
-}  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs.h b/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs.h
deleted file mode 100644
index f833678195ae27..00000000000000
--- a/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs.h
+++ /dev/null
@@ -1,65 +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_TRANSFORMS_ALIAS_IN_PLACE_OUTPUTS_H_
-#define XLA_BACKENDS_GPU_TRANSFORMS_ALIAS_IN_PLACE_OUTPUTS_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 mlir {
-class MLIRContext;
-}  // namespace mlir
-
-namespace xla::gpu {
-
-// Annotates GPU operations whose output buffer can overwrite an operand buffer
-// in place with `output_to_operand_aliasing`, eliminating copies.
-//
-// Supported operations:
-//   * Triton fusions (GEMMs, elementwise): an operand is aliased when it
-//     reaches the fusion root through a single data-flow path consisting of
-//     pure elementwise ops or bitcasts that cancel, and its shape is
-//     compatible with the root.
-//   * cuBLASLt custom calls (`__cublas$lt$matmul`, `__cublas$lt$matmul$f8`,
-//     `__cublas$lt$matmul$mx`, `__cublas$lt$groupedMatmul`): the bias/C
-//     operand is aliased when `beta != 0` and its shape matches the output.
-//
-// In both cases, the aliased operand must be a writable intermediate whose
-// only other users (if any) precede this operation in the dataflow graph.
-//
-// Must run before copy insertion.
-class AliasInPlaceOutputs : public HloModulePass {
- public:
-  explicit AliasInPlaceOutputs(mlir::MLIRContext* mlir_context)
-      : mlir_context_(mlir_context) {}
-
-  absl::string_view name() const override { return "alias_in_place_outputs"; }
-
- protected:
-  absl::StatusOr<bool> RunImpl(
-      HloModule* module,
-      const absl::flat_hash_set<absl::string_view>& execution_threads) override;
-
- private:
-  mlir::MLIRContext* mlir_context_;
-};
-
-}  // namespace xla::gpu
-
-#endif  // XLA_BACKENDS_GPU_TRANSFORMS_ALIAS_IN_PLACE_OUTPUTS_H_
diff --git a/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs_test.cc b/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs_test.cc
deleted file mode 100644
index 00c67074d1bc22..00000000000000
--- a/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs_test.cc
+++ /dev/null
@@ -1,434 +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/transforms/alias_in_place_outputs.h"
-
-#include <cstdint>
-#include <memory>
-#include <utility>
-#include <vector>
-
-#include <gmock/gmock.h>
-#include <gtest/gtest.h>
-#include "absl/log/log.h"
-#include "absl/strings/string_view.h"
-#include "mlir/IR/MLIRContext.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/testlib/hlo_hardware_independent_test_base.h"
-#include "xla/shape_util.h"
-
-namespace xla::gpu {
-namespace {
-
-using ::testing::ElementsAre;
-using ::testing::Pair;
-
-class AliasInPlaceOutputsTest : public HloHardwareIndependentTestBase {
- protected:
-  using AliasingList =
-      std::vector<std::pair<ShapeIndex, std::pair<int64_t, ShapeIndex>>>;
-
-  // Output-to-operand aliasing of the single candidate instruction
-  // (Triton fusion or cuBLASLt custom call) in the entry computation.
-  const AliasingList& GetAliasing(const HloModule& module) {
-    for (const HloInstruction* instr :
-         module.entry_computation()->instructions()) {
-      if (const auto* fusion = DynCast<HloFusionInstruction>(instr)) {
-        return fusion->output_to_operand_aliasing();
-      }
-      if (const auto* call = DynCast<HloCustomCallInstruction>(instr)) {
-        return call->output_to_operand_aliasing();
-      }
-    }
-    LOG(FATAL) << "no candidate instruction in entry computation";
-  }
-
-  void ExpectAlias(absl::string_view hlo, int64_t operand,
-                   ShapeIndex output = {}) {
-    ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo));
-    ASSERT_OK_AND_ASSIGN(bool changed,
-                         AliasInPlaceOutputs(&mlir_context_).Run(module.get()));
-    EXPECT_TRUE(changed);
-    EXPECT_THAT(GetAliasing(*module),
-                ElementsAre(Pair(output, Pair(operand, ShapeIndex{}))));
-  }
-
-  void ExpectNoAlias(absl::string_view hlo) {
-    ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo));
-    ASSERT_OK_AND_ASSIGN(bool changed,
-                         AliasInPlaceOutputs(&mlir_context_).Run(module.get()));
-    EXPECT_FALSE(changed);
-    EXPECT_TRUE(GetAliasing(*module).empty());
-  }
-
-  mlir::MLIRContext mlir_context_;
-};
-
-// --- Triton Fusion Tests ---
-
-// dot(a, b) + x with x a dead intermediate => x is aliased to the output.
-TEST_F(AliasInPlaceOutputsTest, TritonAliasesDeadBiasOperand) {
-  ExpectAlias(R"(
-    HloModule m
-
-    triton {
-      a = f32[128,64] parameter(0)
-      b = f32[64,128] parameter(1)
-      x = f32[128,128] parameter(2)
-      dot = f32[128,128] dot(a, b),
-        lhs_contracting_dims={1}, rhs_contracting_dims={0}
-      ROOT add = f32[128,128] add(dot, x)
-    }
-
-    ENTRY e {
-      a = f32[128,64] parameter(0)
-      b = f32[64,128] parameter(1)
-      c = f32[128,128] parameter(2)
-      x = f32[128,128] add(c, c)
-      ROOT fusion = f32[128,128] fusion(a, b, x), kind=kCustom, calls=triton,
-        backend_config={"fusion_backend_config":{"kind":"__triton_gemm"}}
-    })",
-              /*operand=*/2);
-}
-
-// Post-autotuning Triton GEMMs are converted to __triton_nested_gemm_fusion.
-TEST_F(AliasInPlaceOutputsTest, TritonAliasesPostAutotuningNestedGemmFusion) {
-  ExpectAlias(R"(
-    HloModule m
-
-    triton {
-      a = f32[128,64] parameter(0)
-      b = f32[64,128] parameter(1)
-      x = f32[128,128] parameter(2)
-      dot = f32[128,128] dot(a, b),
-        lhs_contracting_dims={1}, rhs_contracting_dims={0}
-      ROOT add = f32[128,128] add(dot, x)
-    }
-
-    ENTRY e {
-      a = f32[128,64] parameter(0)
-      b = f32[64,128] parameter(1)
-      c = f32[128,128] parameter(2)
-      x = f32[128,128] add(c, c)
-      ROOT fusion = f32[128,128] fusion(a, b, x), kind=kCustom, calls=triton,
-        backend_config={"fusion_backend_config":{"kind":"__triton_nested_gemm_fusion"}}
-    })",
-              /*operand=*/2);
-}
-
-// The bias reaches the root through a non-identity access (transpose), so it is
-// not safe to alias even though shapes match.
-TEST_F(AliasInPlaceOutputsTest, TritonDoesNotAliasNonIdentityAccess) {
-  ExpectNoAlias(R"(
-    HloModule m
-
-    triton {
-      a = f32[128,64] parameter(0)
-      b = f32[64,128] parameter(1)
-      x = f32[128,128] parameter(2)
-      xt = f32[128,128] transpose(x), dimensions={1,0}
-      dot = f32[128,128] dot(a, b),
-        lhs_contracting_dims={1}, rhs_contracting_dims={0}
-      ROOT add = f32[128,128] add(dot, xt)
-    }
-
-    ENTRY e {
-      a = f32[128,64] parameter(0)
-      b = f32[64,128] parameter(1)
-      c = f32[128,128] parameter(2)
-      x = f32[128,128] add(c, c)
-      ROOT fusion = f32[128,128] fusion(a, b, x), kind=kCustom, calls=triton,
-        backend_config={"fusion_backend_config":{"kind":"__triton_gemm"}}
-    })");
-}
-
-// The bias is a plain entry parameter (read-only input) => not beneficial.
-TEST_F(AliasInPlaceOutputsTest, TritonDoesNotAliasParameterOperand) {
-  ExpectNoAlias(R"(
-    HloModule m
-
-    triton {
-      a = f32[128,64] parameter(0)
-      b = f32[64,128] parameter(1)
-      x = f32[128,128] parameter(2)
-      dot = f32[128,128] dot(a, b),
-        lhs_contracting_dims={1}, rhs_contracting_dims={0}
-      ROOT add = f32[128,128] add(dot, x)
-    }
-
-    ENTRY e {
-      a = f32[128,64] parameter(0)
-      b = f32[64,128] parameter(1)
-      x = f32[128,128] parameter(2)
-      ROOT fusion = f32[128,128] fusion(a, b, x), kind=kCustom, calls=triton,
-        backend_config={"fusion_backend_config":{"kind":"__triton_gemm"}}
-    })");
-}
-
-// The bias has another user that is independent => not safe to alias.
-TEST_F(AliasInPlaceOutputsTest, TritonDoesNotAliasIndependentLiveOperand) {
-  ExpectNoAlias(R"(
-    HloModule m
-
-    triton {
-      a = f32[128,64] parameter(0)
-      b = f32[64,128] parameter(1)
-      x = f32[128,128] parameter(2)
-      dot = f32[128,128] dot(a, b),
-        lhs_contracting_dims={1}, rhs_contracting_dims={0}
-      ROOT add = f32[128,128] add(dot, x)
-    }
-
-    ENTRY e {
-      a = f32[128,64] parameter(0)
-      b = f32[64,128] parameter(1)
-      c = f32[128,128] parameter(2)
-      x = f32[128,128] add(c, c)
-      fusion = f32[128,128] fusion(a, b, x), kind=kCustom, calls=triton,
-        backend_config={"fusion_backend_config":{"kind":"__triton_gemm"}}
-      ROOT t = (f32[128,128], f32[128,128]) tuple(fusion, x)
-    })");
-}
-
-// The bias has another user, but that user is a predecessor of the fusion
-// (residual pattern: a = norm(x); fusion = dot(a, b) + x). Safe to alias!
-TEST_F(AliasInPlaceOutputsTest,
-       TritonAliasesResidualOperandWithPredecessorUser) {
-  ExpectAlias(R"(
-    HloModule m
-
-    triton {
-      a = f32[128,64] parameter(0)
-      b = f32[64,128] parameter(1)
-      x = f32[128,128] parameter(2)
-      dot = f32[128,128] dot(a, b),
-        lhs_contracting_dims={1}, rhs_contracting_dims={0}
-      ROOT add = f32[128,128] add(dot, x)
-    }
-
-    ENTRY e {
-      c = f32[128,128] parameter(0)
-      b = f32[64,128] parameter(1)
-      x = f32[128,128] add(c, c)
-      slice = f32[128,64] slice(x), slice={[0:128], [0:64]}
-      ROOT fusion = f32[128,128] fusion(slice, b, x), kind=kCustom, calls=triton,
-        backend_config={"fusion_backend_config":{"kind":"__triton_gemm"}}
-    })",
-              /*operand=*/2);
-}
-
-// --- cuBLASLt Custom Call Tests ---
-
-// __cublas$lt$matmul with beta=1 and dead bias operand => aliases output {0} to
-// operand {2, {}}.
-TEST_F(AliasInPlaceOutputsTest, CublasLtAliasesDeadBiasOperand) {
-  ExpectAlias(R"(
-    HloModule m
-
-    ENTRY e {
-      p0 = f32[128,64] parameter(0)
-      p1 = f32[64,128] parameter(1)
-      p2 = f32[128,128] parameter(2)
-      bias = f32[128,128] add(p2, p2)
-      ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, bias),
-        custom_call_target="__cublas$lt$matmul",
-        backend_config={"gemm_backend_config":{"beta":1.0}}
-    })",
-              /*operand=*/2, /*output=*/ShapeIndex{0});
-}
-
-// __cublas$lt$matmul with beta=0 => no aliasing.
-TEST_F(AliasInPlaceOutputsTest, CublasLtDoesNotAliasWhenBetaIsZero) {
-  ExpectNoAlias(R"(
-    HloModule m
-
-    ENTRY e {
-      p0 = f32[128,64] parameter(0)
-      p1 = f32[64,128] parameter(1)
-      p2 = f32[128,128] parameter(2)
-      bias = f32[128,128] add(p2, p2)
-      ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, bias),
-        custom_call_target="__cublas$lt$matmul",
-        backend_config={"gemm_backend_config":{"beta":0.0}}
-    })");
-}
-
-// __cublas$lt$matmul where bias is a parameter => no aliasing.
-TEST_F(AliasInPlaceOutputsTest, CublasLtDoesNotAliasParameterOperand) {
-  ExpectNoAlias(R"(
-    HloModule m
-
-    ENTRY e {
-      p0 = f32[128,64] parameter(0)
-      p1 = f32[64,128] parameter(1)
-      p2 = f32[128,128] parameter(2)
-      ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, p2),
-        custom_call_target="__cublas$lt$matmul",
-        backend_config={"gemm_backend_config":{"beta":1.0}}
-    })");
-}
-
-// cuBLASLt in a residual chain (AlphaFold3 pattern):
-// bias is used by an intermediate layer (e.g. norm/slice) that produces LHS of
-// matmul. Because the other user precedes matmul, aliasing is safe!
-TEST_F(AliasInPlaceOutputsTest, CublasLtAliasesResidualChain) {
-  ExpectAlias(R"(
-    HloModule m
-
-    ENTRY e {
-      p0 = f32[128,128] parameter(0)
-      p1 = f32[64,128] parameter(1)
-      bias = f32[128,128] add(p0, p0)
-      lhs = f32[128,64] slice(bias), slice={[0:128], [0:64]}
-      matmul = (f32[128,128], s8[4096]) custom-call(lhs, p1, bias),
-        custom_call_target="__cublas$lt$matmul",
-        backend_config={"gemm_backend_config":{"beta":1.0}}
-      ROOT gte = f32[128,128] get-tuple-element(matmul), index=0
-    })",
-              /*operand=*/2, /*output=*/ShapeIndex{0});
-}
-
-// cuBLASLt where other user does NOT precede matmul => no aliasing.
-TEST_F(AliasInPlaceOutputsTest, CublasLtDoesNotAliasIndependentUser) {
-  ExpectNoAlias(R"(
-    HloModule m
-
-    ENTRY e {
-      p0 = f32[128,64] parameter(0)
-      p1 = f32[64,128] parameter(1)
-      p2 = f32[128,128] parameter(2)
-      bias = f32[128,128] add(p2, p2)
-      matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, bias),
-        custom_call_target="__cublas$lt$matmul",
-        backend_config={"gemm_backend_config":{"beta":1.0}}
-      ROOT t = ((f32[128,128], s8[4096]), f32[128,128]) tuple(matmul, bias)
-    })");
-}
-
-// cuBLASLt FP8 (__cublas$lt$matmul$f8) with beta=1.
-TEST_F(AliasInPlaceOutputsTest, CublasLtF8AliasesDeadBiasOperand) {
-  ExpectAlias(R"(
-    HloModule m
-
-    ENTRY e {
-      p0 = f8e4m3fn[128,64] parameter(0)
-      p1 = f8e4m3fn[64,128] parameter(1)
-      p2 = f32[128,128] parameter(2)
-      bias = f32[128,128] add(p2, p2)
-      ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, bias),
-        custom_call_target="__cublas$lt$matmul$f8",
-        backend_config={"gemm_backend_config":{"beta":1.0}}
-    })",
-              /*operand=*/2, /*output=*/ShapeIndex{0});
-}
-
-// Grouped GEMM (__cublas$lt$groupedMatmul) with beta=1 (bias at operand index
-// 3).
-TEST_F(AliasInPlaceOutputsTest, CublasLtGroupedMatmulAliasesBiasOperand) {
-  ExpectAlias(R"(
-    HloModule m
-
-    ENTRY e {
-      p0 = f32[128,64] parameter(0)
-      p1 = f32[64,128] parameter(1)
-      sizes = s32[2] parameter(2)
-      p3 = f32[128,128] parameter(3)
-      bias = f32[128,128] add(p3, p3)
-      ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, sizes, bias),
-        custom_call_target="__cublas$lt$groupedMatmul",
-        backend_config={"gemm_backend_config":{"beta":1.0}}
-    })",
-              /*operand=*/3, /*output=*/ShapeIndex{0});
-}
-
-// Block-scaled MX matmul (__cublas$lt$matmul$mx) with beta=1.
-TEST_F(AliasInPlaceOutputsTest, CublasLtMatmulMxAliasesDeadBiasOperand) {
-  ExpectAlias(R"(
-    HloModule m
-
-    ENTRY e {
-      p0 = f8e4m3fn[128,64] parameter(0)
-      p1 = f8e4m3fn[64,128] parameter(1)
-      p2 = f32[128,128] parameter(2)
-      bias = f32[128,128] add(p2, p2)
-      ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, bias),
-        custom_call_target="__cublas$lt$matmul$mx",
-        backend_config={"gemm_backend_config":{"beta":1.0}}
-    })",
-              /*operand=*/2, /*output=*/ShapeIndex{0});
-}
-
-// Triton fusion with tuple output is not supported for aliasing.
-TEST_F(AliasInPlaceOutputsTest, TritonDoesNotAliasTupleFusion) {
-  ExpectNoAlias(R"(
-    HloModule m
-
-    triton {
-      a = f32[128,64] parameter(0)
-      b = f32[64,128] parameter(1)
-      x = f32[128,128] parameter(2)
-      dot = f32[128,128] dot(a, b),
-        lhs_contracting_dims={1}, rhs_contracting_dims={0}
-      add = f32[128,128] add(dot, x)
-      ROOT t = (f32[128,128], f32[128,128]) tuple(add, dot)
-    }
-
-    ENTRY e {
-      a = f32[128,64] parameter(0)
-      b = f32[64,128] parameter(1)
-      c = f32[128,128] parameter(2)
-      x = f32[128,128] add(c, c)
-      ROOT fusion = (f32[128,128], f32[128,128]) fusion(a, b, x), kind=kCustom, calls=triton,
-        backend_config={"fusion_backend_config":{"kind":"__triton_gemm"}}
-    })");
-}
-
-// cuBLASLt with constant bias operand => no aliasing.
-TEST_F(AliasInPlaceOutputsTest, CublasLtDoesNotAliasConstantOperand) {
-  ExpectNoAlias(R"(
-    HloModule m
-
-    ENTRY e {
-      p0 = f32[128,64] parameter(0)
-      p1 = f32[64,128] parameter(1)
-      c = f32[128,128] constant({...})
-      ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, c),
-        custom_call_target="__cublas$lt$matmul",
-        backend_config={"gemm_backend_config":{"beta":1.0}}
-    })");
-}
-
-// cuBLASLt where bias shape does not match output shape (e.g. vector broadcast
-// bias) => no aliasing.
-TEST_F(AliasInPlaceOutputsTest, CublasLtDoesNotAliasMismatchedShapeBias) {
-  ExpectNoAlias(R"(
-    HloModule m
-
-    ENTRY e {
-      p0 = f32[128,64] parameter(0)
-      p1 = f32[64,128] parameter(1)
-      p2 = f32[128] parameter(2)
-      bias = f32[128] add(p2, p2)
-      ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, bias),
-        custom_call_target="__cublas$lt$matmul",
-        backend_config={"gemm_backend_config":{"beta":1.0}}
-    })");
-}
-
-}  // namespace
-}  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc
index 91baddc0f6ab36..59a515813eb372 100644
--- a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc
@@ -2416,6 +2416,42 @@ TEST_F(ConvFusionRewriterIntegrationTest,
       )");
 }
 
+TEST_F(ConvFusionRewriterUnitTest, EpilogueNotFusedOnPreAmpere) {
+  const char* const hlo_string = R"(
+    HloModule Test
+
+    ENTRY Test {
+      input = f16[1,16,16,16] parameter(0)
+      filter = f16[3,3,16,32] parameter(1)
+      conv = f16[1,16,16,32] convolution(input, filter),
+               window={size=3x3 pad=1_1x1_1},
+               dim_labels=b01f_01io->b01f
+      bias = f16[32] parameter(2)
+      bcast = f16[1,16,16,32] broadcast(bias), dimensions={3}
+      ROOT add = f16[1,16,16,32] add(conv, bcast)
+    })";
+
+  se::DeviceDescription volta_device;
+  volta_device.set_gpu_compute_capability(se::CudaComputeCapability::Volta());
+
+  // On Volta (SM70), epilogue should NOT be fused into the custom fusion.
+  RunAndMatch(hlo_string,
+              m::Add(m::Fusion(m::Parameter(0), m::Parameter(1))
+                         .WithFusionKind(HloInstruction::FusionKind::kCustom),
+                     m::Broadcast(m::Parameter(2))),
+              /*run_algebraic_simplifier=*/false, volta_device);
+
+  se::DeviceDescription ampere_device;
+  ampere_device.set_gpu_compute_capability(se::CudaComputeCapability::Ampere());
+
+  // On Ampere+ (SM80+), epilogue SHOULD be fused into the custom fusion.
+  RunAndMatch(hlo_string,
+              m::Fusion(m::Parameter(0), m::Parameter(1), m::Parameter(2))
+                  .WithFusionKind(HloInstruction::FusionKind::kCustom)
+                  .WithShape(F16, {1, 16, 16, 32}),
+              /*run_algebraic_simplifier=*/false, ampere_device);
+}
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc
index 3b5495b800ed89..a60ab8e6ad60fb 100644
--- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc
@@ -87,6 +87,14 @@ bool IsEpilogueOpSupportedByCuDNN(const HloInstruction& hlo,
   if (is_nchw) {
     return false;
   }
+
+  // cuDNN only supports epilogues on Ampere and above.
+  const se::CudaComputeCapability* cuda_cc =
+      device_info.gpu_compute_capability().cuda_compute_capability();
+  if (cuda_cc != nullptr && !cuda_cc->IsAtLeastAmpere()) {
+    return false;
+  }
+
   const HloOpcode opcode = hlo.opcode();
   // Do not fuse chained converts (a convert whose operand is already a
   // convert). Note: Fusing chained converts could steal a convert from an
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 9b1cadc113604e..92dcb8c92f7d13 100644
--- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc
+++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc
@@ -275,6 +275,76 @@ LogicalResult MemRefSliceOp::verify() {
   return success();
 }
 
+std::optional<bool> MemRefSliceOp::sliceStridesAcrossSourceTiles(
+    const int64_t source_size, const int64_t slice_size,
+    const int64_t source_tile_size, const int64_t result_tile_size,
+    Value offset) {
+  CHECK_EQ(source_tile_size % result_tile_size, 0);
+  DCHECK(offset == nullptr || isGuaranteedDivisible(offset, result_tile_size));
+  const std::optional<int64_t> maybe_cst_offset =
+      offset ? getConstantIntValue(offset) : std::nullopt;
+  if (maybe_cst_offset && slice_size != ShapedType::kDynamic) {
+    // Fully static slice
+    return (*maybe_cst_offset + slice_size - 1) / source_tile_size !=
+           *maybe_cst_offset / source_tile_size;
+  }
+  if (slice_size != ShapedType::kDynamic && slice_size <= result_tile_size) {
+    // We never stride at all
+    return false;
+  }
+  if (source_size != ShapedType::kDynamic && source_size <= source_tile_size) {
+    // Source has only one tile
+    return false;
+  }
+  if (slice_size != ShapedType::kDynamic && slice_size > source_tile_size) {
+    // Slice is too big to be contained in a single source tile
+    return true;
+  }
+  // TODO(apaszke,tlongeri): Should we relax the requirement for the shape to
+  // be divisible by the slice size? We need to consider if accessing the last
+  // partial tile is allowed or not.
+  if (slice_size != ShapedType::kDynamic &&
+      source_tile_size % slice_size == 0 &&
+      source_size != ShapedType::kDynamic && source_size % slice_size == 0 &&
+      offset != nullptr && isGuaranteedDivisible(offset, slice_size)) {
+    // Slice is guaranteed to fit in a single source tile.
+    return false;
+  }
+  return std::nullopt;
+}
+
+std::optional<bool> MemRefSliceOp::sliceStridesWithinSourceTiles(
+    const int64_t source_size, const int64_t slice_size,
+    const int64_t source_tile_size, const int64_t result_tile_size,
+    Value offset) {
+  CHECK_EQ(source_tile_size % result_tile_size, 0);
+  DCHECK(offset == nullptr || isGuaranteedDivisible(offset, result_tile_size));
+  if (source_tile_size == result_tile_size) {
+    // The source tile isn't subdivided into result tiles
+    return false;
+  }
+  if (slice_size != ShapedType::kDynamic) {
+    if (slice_size <= result_tile_size) {
+      // We never stride at all
+      return false;
+    }
+    if (slice_size > 2 * result_tile_size) {
+      // We stride more than once. We've checked that the result tile is smaller
+      // than the source tile, so we must stride within source tiles at least
+      // once.
+      return true;
+    }
+    // We stride exactly once. Is it within or across source tiles?
+    if (offset != nullptr) {
+      if (const std::optional<int64_t> rem =
+              getRemainder(offset, source_tile_size)) {
+        return *rem != source_tile_size - result_tile_size;
+      }
+    }
+  }
+  return std::nullopt;
+}
+
 std::optional<std::string> MemRefSliceOp::verifyOffsetAndSizeTileAlignment(
     std::array<int64_t, 2> tc_target_shape, MemRefType source_type_override) {
   mlir::MemRefType source_ty = getMemRef().getType();
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 ca958bb36b3d2f..0c2391d5fb776b 100644
--- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td
+++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td
@@ -1199,6 +1199,19 @@ def TPU_MemRefSliceOp : TPU_Op<"memref_slice", [Pure, AttrSizedOperandSegments]>
     std::optional<std::string> verifyOffsetAndSizeTileAlignment(
         std::array<int64_t, 2> tc_target_shape,
         MemRefType source_type_override = nullptr);
+
+    // Given a source shape size, a slice shape size, a source tile size, a
+    // result tile size and optionally a slice offset, the below functions
+    // returns a tri-state boolean (nullopt means "unknown") indicating whether
+    // the result tiles of the slice stride across or within the source tiles.
+    // The result tile should divide the source tile and the offset, if
+    // provided, should be divisible by the result tile size.
+    static std::optional<bool> sliceStridesAcrossSourceTiles(
+        int64_t source_size, int64_t slice_size, int64_t source_tile_size,
+        int64_t result_tile_size, Value offset);
+    static std::optional<bool> sliceStridesWithinSourceTiles(
+        int64_t source_size, int64_t slice_size, int64_t source_tile_size,
+        int64_t result_tile_size, Value offset);
   }];
   let hasVerifier = 1;
   let hasCanonicalizer = 1;
diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.cc b/third_party/xla/xla/pjrt/common_pjrt_client.cc
index 63e0f87459d866..f3da0d85167f07 100644
--- a/third_party/xla/xla/pjrt/common_pjrt_client.cc
+++ b/third_party/xla/xla/pjrt/common_pjrt_client.cc
@@ -405,8 +405,12 @@ absl::StatusOr<PjRtDeviceEventRef> CommonPjRtClient::LinearizeIntoImpl(
                 HostBufferSemantics::kImmutableOnlyDuringCall,
             device_shape, byte_strides.value_or(absl::Span<const int64_t>()),
             raw_buffer));
+    // TODO(parkers): IsCpuId because the linearization pool shares with the
+    // execute pool on cpu, so we have potential deadlocks.
     if (host_buffer_semantics ==
-        HostBufferSemantics::kImmutableOnlyDuringCall) {
+            HostBufferSemantics::kImmutableOnlyDuringCall ||
+        IsCpuId(platform_id()) ||
+        memory_space->kind_id() == UnpinnedHostMemorySpace::kKindId) {
       if (!linearized.IsAvailable()) {
         tsl::BlockUntilReady(linearized.GetAsyncValue());
       }
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
index ecbc7338db887f..64d780de8066ec 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
@@ -1078,45 +1078,6 @@ bool PjRtCpuClient::BufferFromHostBufferSupportsZeroCopy(
       data, type, dims, byte_strides, shape);
 }
 
-absl::StatusOr<PjRtDeviceEventRef> PjRtCpuClient::LinearizeHostBufferInto(
-    const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
-    std::optional<absl::Span<int64_t const>> byte_strides,
-    HostBufferSemantics host_buffer_semantics,
-    absl::AnyInvocable<void() &&> on_done_with_host_buffer,
-    const xla::Shape& device_shape, PjRtRawBufferRef raw_buffer) {
-  if (device_shape.IsToken()) {
-    return PjRtDeviceEventRef(tsl::MakeAvailableAsyncValueRef<CpuEvent>());
-  }
-  auto* cpp_buf = raw_buffer->down_cast<CpuRawBuffer>();
-  if (cpp_buf == nullptr) {
-    return absl::InvalidArgumentError("Not a CPU raw buffer");
-  }
-  return cpp_buf->CopyFromHostBuffer(
-      data, type, dims, byte_strides, host_buffer_semantics,
-      std::move(on_done_with_host_buffer), device_shape, async_work_runner(),
-      raw_client()->eigen_intraop_pool(),
-      raw_client()->max_transpose_threads());
-}
-
-absl::StatusOr<PjRtDeviceEventRef> PjRtCpuClient::LinearizeInto(
-    const LiteralSlice& literal, const xla::Shape& device_shape,
-    HostBufferSemantics host_buffer_semantics, PjRtRawBufferRef raw_buffer) {
-  if (host_buffer_semantics ==
-      PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall) {
-    return absl::UnimplementedError(
-        "ImmutableOnlyDuringCall semantics is not supported on CPU.");
-  }
-  if (device_shape.IsToken()) {
-    return PjRtDeviceEventRef(tsl::MakeAvailableAsyncValueRef<CpuEvent>());
-  }
-  auto* cpp_buf = raw_buffer->down_cast<CpuRawBuffer>();
-  if (cpp_buf == nullptr) {
-    return absl::InvalidArgumentError("Not a CPU raw buffer");
-  }
-  return cpp_buf->CopyFromLiteral(literal, device_shape.layout(),
-                                  async_work_runner());
-}
-
 absl::StatusOr<CompiledMemoryStats> PjRtCpuExecutable::GetCompiledMemoryStats()
     const {
   const auto& buffer_assignment = cpu_executable_->buffer_assignment();
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h
index 0bd6837ebf2b0a..2276bb60252b63 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.h
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h
@@ -298,18 +298,6 @@ class PjRtCpuClient final : public CommonPjRtClientImpl {
   absl::StatusOr<int> GetMemorySpaceKindForShape(
       const Shape& shape) const override;
 
-  absl::StatusOr<PjRtDeviceEventRef> LinearizeHostBufferInto(
-      const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
-      std::optional<absl::Span<int64_t const>> byte_strides,
-      HostBufferSemantics host_buffer_semantics,
-      absl::AnyInvocable<void() &&> on_done_with_host_buffer,
-      const xla::Shape& device_shape, PjRtRawBufferRef raw_buffer) override;
-
-  absl::StatusOr<PjRtDeviceEventRef> LinearizeInto(
-      const LiteralSlice& literal, const xla::Shape& device_shape,
-      HostBufferSemantics host_buffer_semantics,
-      PjRtRawBufferRef raw_buffer) override;
-
   bool BufferFromHostBufferSupportsZeroCopy(
       const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
       std::optional<absl::Span<int64_t const>> byte_strides, const Shape& shape,
diff --git a/third_party/xla/xla/pjrt/cpu/raw_buffer.cc b/third_party/xla/xla/pjrt/cpu/raw_buffer.cc
index b2c0b29ef5879c..113de407961ac9 100644
--- a/third_party/xla/xla/pjrt/cpu/raw_buffer.cc
+++ b/third_party/xla/xla/pjrt/cpu/raw_buffer.cc
@@ -228,135 +228,6 @@ CpuRawBuffer::CopyRawDeviceToHostAndReturnEvent(
   return PjRtDeviceEventRef(std::move(event));
 }
 
-absl::StatusOr<PjRtDeviceEventRef> CpuRawBuffer::CopyFromLiteral(
-    const LiteralSlice& literal, const xla::Layout& layout,
-    AsyncWorkRunner* async_work_runner) {
-  auto event = tsl::MakeConstructedAsyncValueRef<CpuEvent>();
-  async_work_runner->Execute([literal, layout, event, buffer = buffer_]() {
-    CHECK(buffer.IsConcrete());
-    const xla::Shape& shape = literal.shape();
-    if (shape.IsToken()) {
-    } else if ((!shape.has_layout() &&
-                !xla::LayoutUtil::IsMonotonicWithDim0Major(layout)) ||
-               shape.layout() != layout) {
-      auto shape_copy = xla::ShapeUtil::MakeShape(
-          literal.shape().element_type(), literal.shape().dimensions());
-      shape_copy.mutable_layout()->mutable_minor_to_major()->assign(
-          layout.minor_to_major().begin(), layout.minor_to_major().end());
-
-      xla::Literal literal_copy(shape_copy);
-      CHECK_OK(literal_copy.CopyFrom(literal));
-      PackOrCopy(literal_copy.shape().element_type(), literal_copy,
-                 buffer->untyped_data(), buffer->size_bytes());
-    } else {
-      PackOrCopy(literal.shape().element_type(), literal,
-                 buffer->untyped_data(), buffer->size_bytes());
-    }
-    event.SetStateConcrete();
-  });
-  return PjRtDeviceEventRef(std::move(event));
-}
-
-absl::StatusOr<PjRtDeviceEventRef> CpuRawBuffer::CopyFromHostBuffer(
-    const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
-    std::optional<absl::Span<int64_t const>> byte_strides,
-    PjRtClient::HostBufferSemantics host_buffer_semantics,
-    absl::AnyInvocable<void() &&> on_done_with_host_buffer, const Shape& shape,
-    AsyncWorkRunner* async_work_runner, tsl::thread::ThreadPool* thread_pool,
-    int max_transpose_threads) {
-  CommonPjRtClient* client =
-      absl::down_cast<CommonPjRtClient*>(memory_space()->client());
-  tsl::AsyncValueRef<CpuDeviceMemory> device_buffer = buffer_;
-  bool has_default_layout =
-      !byte_strides || HasMajorToMinorLayout(type, dims, *byte_strides);
-  const int bit_width = primitive_util::BitWidth(type);
-  // Packed arrays are unpacked on host and packed on device.
-  bool is_packed = primitive_util::IsSubByteNonPredType(type);
-
-  size_t byte_size = ShapeUtil::ByteSizeOf(shape);
-  if (is_packed) {
-    byte_size *= 8 / bit_width;
-  }
-  auto dst_data_ptr = device_buffer->untyped_data();
-  if (!has_default_layout ||
-      (shape.has_layout() &&
-       !LayoutUtil::IsMonotonicWithDim0Major(shape.layout())) ||
-      is_packed) {
-    // If the input array does not have a major-to-minor layout or device layout
-    // is not major-to-minor, transpose it into the device layout. Currently we
-    // choose to always do this synchronously.
-    // TODO(phawkins): consider performing the transpose asynchronously.
-    // TODO(phawkins): parallelize the transpose.
-    std::shared_ptr<TransposePlan> transpose;
-    {
-      absl::InlinedVector<int64_t, 4> permutation(dims.size());
-      if (shape.has_layout()) {
-        absl::c_reverse_copy(shape.layout().minor_to_major(),
-                             permutation.begin());
-      } else {
-        absl::c_iota(permutation, 0);
-      }
-      TransposePlan::Options options;
-      options.elem_size_in_bytes = primitive_util::ByteWidth(type);
-      options.dims = dims;
-      options.permutation = permutation;
-      if (is_packed) {
-        options.dest_bits_per_element = bit_width;
-      }
-      if (byte_strides) {
-        options.input_striding = TransposePlan::Striding{*byte_strides};
-      }
-      if (thread_pool) {
-        options.num_threads =
-            std::min(thread_pool->NumThreads(), max_transpose_threads);
-      }
-      ABSL_ASSIGN_OR_RETURN(transpose, client->GetTransposePlan(options));
-    }
-    std::optional<std::function<void(std::function<void(void)>)>> schedule_work;
-    if (thread_pool && max_transpose_threads > 1) {
-      schedule_work = [thread_pool](std::function<void(void)> work) {
-        thread_pool->Schedule(std::move(work));
-      };
-    }
-    transpose->Execute(data, dst_data_ptr, schedule_work);
-    if (on_done_with_host_buffer) {
-      std::move(on_done_with_host_buffer)();
-      on_done_with_host_buffer = nullptr;
-    }
-  } else {
-    bool should_sync_copy =
-        host_buffer_semantics ==
-            PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall ||
-        (byte_size < kSmallDataTransferByteSize);
-    if (should_sync_copy) {
-      std::memcpy(dst_data_ptr, data, byte_size);
-      if (on_done_with_host_buffer) {
-        std::move(on_done_with_host_buffer)();
-        on_done_with_host_buffer = nullptr;
-      }
-    } else {
-      tsl::AsyncValueRef<CpuEvent> copy_event =
-          tsl::MakeConstructedAsyncValueRef<CpuEvent>();
-      auto result = PjRtDeviceEventRef(copy_event.CopyRef());
-      async_work_runner->Execute([device_buffer, dst_data_ptr, data, byte_size,
-                                  copy_event = std::move(copy_event),
-                                  on_done_with_host_buffer = std::move(
-                                      on_done_with_host_buffer)]() mutable {
-        tsl::profiler::TraceMe traceme("H2D Dispatch");
-        std::memcpy(dst_data_ptr, data, byte_size);
-        if (on_done_with_host_buffer) {
-          std::move(on_done_with_host_buffer)();
-          on_done_with_host_buffer = nullptr;
-        }
-        // Signal copy is complete.
-        copy_event.SetStateConcrete();
-      });
-      return result;
-    }
-  }
-  return PjRtDeviceEventRef(tsl::MakeAvailableAsyncValueRef<CpuEvent>());
-}
-
 absl::StatusOr<PjRtDeviceEventRef> CpuRawBuffer::MakeAllocationReadyEvent() {
   return PjRtDeviceEventRef(tsl::MakeAvailableAsyncValueRef<CpuEvent>());
 }
diff --git a/third_party/xla/xla/pjrt/cpu/raw_buffer.h b/third_party/xla/xla/pjrt/cpu/raw_buffer.h
index a9a73a3d6f779d..845f0f4560ce45 100644
--- a/third_party/xla/xla/pjrt/cpu/raw_buffer.h
+++ b/third_party/xla/xla/pjrt/cpu/raw_buffer.h
@@ -125,20 +125,8 @@ class CpuRawBuffer : public CommonPjRtRawBufferImpl {
       void* dst, int64_t offset, int64_t transfer_size,
       PjRtDeviceEventRefVector dependencies) override;
 
-  absl::StatusOr<PjRtDeviceEventRef> CopyFromLiteral(
-      const LiteralSlice& literal, const xla::Layout& layout,
-      AsyncWorkRunner* async_work_runner);
-
   absl::StatusOr<PjRtDeviceEventRef> MakeAllocationReadyEvent() override;
 
-  absl::StatusOr<PjRtDeviceEventRef> CopyFromHostBuffer(
-      const void* data, PrimitiveType type, absl::Span<int64_t const> dims,
-      std::optional<absl::Span<int64_t const>> byte_strides,
-      PjRtClient::HostBufferSemantics host_buffer_semantics,
-      absl::AnyInvocable<void() &&> on_done_with_host_buffer,
-      const Shape& shape, AsyncWorkRunner* async_work_runner,
-      tsl::thread::ThreadPool* thread_pool, int max_transpose_threads);
-
   void CopyTo(
       PjRtRawBufferRef dst_raw_buffer,
       PjRtDeviceEventPromiseRef definition_event_promise,
diff --git a/third_party/xla/xla/pjrt/plugin/xla_cpu/BUILD b/third_party/xla/xla/pjrt/plugin/xla_cpu/BUILD
index f8d22af61a690c..843696e48bef8e 100644
--- a/third_party/xla/xla/pjrt/plugin/xla_cpu/BUILD
+++ b/third_party/xla/xla/pjrt/plugin/xla_cpu/BUILD
@@ -145,6 +145,7 @@ xla_cc_test(
     deps = [
         ":cpu_topology",
         ":cpu_topology_description",
+        "//xla:shape_util",
         "//xla/backends/cpu:target_machine_options",
         "//xla/pjrt:host_memory_spaces",
         "//xla/pjrt:pjrt_common",
diff --git a/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description.cc b/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description.cc
index 7a99f291db2750..088fe3053b9c08 100644
--- a/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description.cc
+++ b/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description.cc
@@ -77,7 +77,15 @@ absl::StatusOr<Layout> CpuTopologyDescription::GetDefaultLayout(
                            PrimitiveType_Name(element_type));
   }
   Shape shape = ShapeUtil::MakeShape(element_type, dims);
-  return LayoutUtil::GetWithDefaultLayout(shape).layout();
+  Layout layout = LayoutUtil::GetWithDefaultLayout(shape).layout();
+  // `GetWithDefaultLayout` returns a padded layout for sub-byte types since the
+  // notion of "default" is context dependent and in this case means the default
+  // for literals for historical reasons. Because of this, we need to manually
+  // populate the `element_size_in_bits` for sub-byte types here.
+  if (primitive_util::IsSubByteNonPredType(element_type)) {
+    layout.set_element_size_in_bits(primitive_util::BitWidth(element_type));
+  }
+  return layout;
 }
 
 absl::StatusOr<int> CpuTopologyDescription::GetMemorySpaceKindForShape(
diff --git a/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description_test.cc b/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description_test.cc
index cb91e98cfca424..8d4512566a3330 100644
--- a/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description_test.cc
+++ b/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description_test.cc
@@ -23,11 +23,13 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "xla/backends/cpu/target_machine_options.h"
+#include "xla/layout.h"
 #include "xla/pjrt/host_memory_spaces.h"
 #include "xla/pjrt/pjrt_common.h"
 #include "xla/pjrt/pjrt_compiler.h"
 #include "xla/pjrt/pjrt_device_dimensions.h"
 #include "xla/pjrt/plugin/xla_cpu/cpu_topology.h"
+#include "xla/shape_util.h"
 #include "xla/tsl/platform/statusor.h"
 
 namespace xla {
@@ -192,5 +194,26 @@ TEST(CpuTopologyDescriptionTest, ProcessIdAndIndexOnProcess) {
                    .ok());
 }
 
+TEST(CpuTopologyDescriptionTest, GetDefaultLayout) {
+  std::vector<CpuTopology::CpuDevice> cpu_devices = {{0, 0}};
+  xla::cpu::TargetMachineOptions target_machine_options(
+      /*triple=*/"triple", /*cpu=*/"cpu", /*features=*/"");
+  CpuTopologyDescription topology(
+      xla::CpuId(), "cpu", "1.0",
+      CpuTopology(cpu_devices, target_machine_options));
+
+  for (PrimitiveType element_type : {S4, S32}) {
+    const Shape shape = ShapeUtil::MakeShape(element_type, {2, 3});
+    ASSERT_OK_AND_ASSIGN(
+        Layout default_layout,
+        topology.GetDefaultLayout(shape.element_type(), shape.dimensions()));
+    ASSERT_OK_AND_ASSIGN(
+        Shape canonical_shape,
+        topology.MakeCanonicalShapeForMemorySpace(CpuDeviceMemorySpace::kKindId,
+                                                  shape, /*layout=*/nullptr));
+    EXPECT_EQ(default_layout, canonical_shape.layout());
+  }
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index ab249ebdbd073f..152db83f4beeb1 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -2558,7 +2558,7 @@ cc_library(
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
-        "//xla/tsl/platform:logging",
+        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/log:check",
@@ -3701,16 +3701,6 @@ xla_cc_test(
     ],
 )
 
-cc_library(
-    name = "computation_placer_hdr",
-    hdrs = ["computation_placer.h"],
-    deps = [
-        ":device_assignment",
-        "//xla/stream_executor:platform_id",
-        "@com_google_absl//absl/status:statusor",
-    ],
-)
-
 cc_library(
     name = "human_readable_profile_builder",
     srcs = ["human_readable_profile_builder.cc"],
diff --git a/third_party/xla/xla/service/batchnorm_expander.cc b/third_party/xla/xla/service/batchnorm_expander.cc
index 554e29df485515..bcdf46bdcb3f28 100644
--- a/third_party/xla/xla/service/batchnorm_expander.cc
+++ b/third_party/xla/xla/service/batchnorm_expander.cc
@@ -38,7 +38,7 @@ limitations under the License.
 #include "xla/literal_util.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
-#include "xla/tsl/platform/logging.h"
+#include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
 
@@ -217,32 +217,34 @@ absl::Status BatchNormExpanderVisitor::HandleBatchNormTraining(
   HloComputation* add_reduce_computation =
       GetOrCreateScalarAddComputation(ptype);
 
+  // X^2.
+  auto operand_squared =
+      add_binary(operand_shape, HloOpcode::kMultiply, operand, operand);
   // Sum[X].
   auto sum = add(HloInstruction::CreateReduce(feature_shape, operand, zero,
                                               dimensions_without_feature,
                                               add_reduce_computation));
 
+  // Sum[X^2].
+  auto squared_sum = add(HloInstruction::CreateReduce(
+      feature_shape, operand_squared, zero, dimensions_without_feature,
+      add_reduce_computation));
+
   // E[X].
   auto mean = add(Mean(elements_per_feature, sum, add));
 
   auto mean_broadcasted = feature_broadcast(mean);
 
-  // X - E[X].
-  auto operand_minus_mean = add_binary(operand_shape, HloOpcode::kSubtract,
-                                       operand, mean_broadcasted);
-
-  // (X - E[X])^2.
-  auto operand_minus_mean_squared =
-      add_binary(operand_shape, HloOpcode::kMultiply, operand_minus_mean,
-                 operand_minus_mean);
+  // E[X^2].
+  auto square_mean = add(Mean(elements_per_feature, squared_sum, add));
 
-  // Sum[(X - E[X])^2].
-  auto squared_diff_sum = add(HloInstruction::CreateReduce(
-      feature_shape, operand_minus_mean_squared, zero,
-      dimensions_without_feature, add_reduce_computation));
+  // E^2[X].
+  auto mean_square =
+      add_binary(feature_shape, HloOpcode::kMultiply, mean, mean);
 
-  // Var[X] = E[(X - E[X])^2].
-  auto var = add(Mean(elements_per_feature, squared_diff_sum, add));
+  // Var[X].
+  auto var =
+      add_binary(feature_shape, HloOpcode::kSubtract, square_mean, mean_square);
 
   auto var_broadcasted = feature_broadcast(var);
 
@@ -253,6 +255,10 @@ absl::Status BatchNormExpanderVisitor::HandleBatchNormTraining(
   // 1 / Sqrt[Var[X] + epsilon].
   auto rsqrt_var_add_epsilon = add(Rsqrt(var_add_epsilon));
 
+  // X - E[X].
+  auto operand_minus_mean = add_binary(operand_shape, HloOpcode::kSubtract,
+                                       operand, mean_broadcasted);
+
   // (X - E[X]) / Sqrt[Var[X] + epsilon].
   auto normalized = add_binary(operand_shape, HloOpcode::kMultiply,
                                operand_minus_mean, rsqrt_var_add_epsilon);
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 156ae91d48e18d..1357ef9030f976 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1743,7 +1743,6 @@ cc_library(
     deps = [
         ":alias_info",
         "//xla:xla_proto_cc",
-        "//xla/backends/gpu/transforms:alias_in_place_outputs",
         "//xla/backends/gpu/transforms:alias_passthrough_params",
         "//xla/backends/gpu/transforms:copy_fusion",
         "//xla/backends/gpu/transforms:sanitize_constant_names",
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 5658324242ab93..0939d1c897195e 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -1992,11 +1992,8 @@ absl::Status GpuCompiler::OptimizeHloModule(
 absl::Status GpuCompiler::RunPreSchedulingCopyInsertion(
     HloModule& hlo_module, const se::DeviceDescription& device_description,
     const GpuAliasInfo* alias_info) {
-  ABSL_ASSIGN_OR_RETURN(BorrowedMlirContext borrowed_context,
-                   mlir_context_pool_.GetOrCreate());
-  mlir::MLIRContext* mlir_context = borrowed_context->get();
   return PreSchedulingCopyInsertionPipeline(hlo_module.config(), alias_info,
-                                            device_description, mlir_context)
+                                            device_description)
       .Run(&hlo_module, {HloInstruction::kMainExecutionThread})
       .status();
 }
diff --git a/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.cc b/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.cc
index 824e37e3e2da53..b8ac06a1546d5b 100644
--- a/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.cc
+++ b/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.cc
@@ -19,10 +19,10 @@ limitations under the License.
 #include <memory>
 #include <utility>
 
-#include "xla/backends/gpu/transforms/alias_in_place_outputs.h"
 #include "xla/backends/gpu/transforms/alias_passthrough_params.h"
 #include "xla/backends/gpu/transforms/copy_fusion.h"
 #include "xla/backends/gpu/transforms/sanitize_constant_names.h"
+#include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/pass/hlo_pass_pipeline.h"
 #include "xla/hlo/transforms/simplifiers/hlo_dce.h"
 #include "xla/service/copy_insertion.h"
@@ -40,8 +40,7 @@ namespace gpu {
 
 HloPassPipeline PreSchedulingCopyInsertionPipeline(
     const HloModuleConfig& config, const GpuAliasInfo* alias_info,
-    const se::DeviceDescription& device_description,
-    mlir::MLIRContext* mlir_context) {
+    const se::DeviceDescription& device_description) {
   const DebugOptions& debug_options = config.debug_options();
 
   // In some cases, we have to place the result of an instruction in a temporary
@@ -68,7 +67,6 @@ HloPassPipeline PreSchedulingCopyInsertionPipeline(
   if (config.alias_passthrough_params()) {
     pipeline.AddPass<AliasPassthroughParams>();
   }
-  pipeline.AddPass<AliasInPlaceOutputs>(mlir_context);
   pipeline.AddPass<LoopScheduleLinearizer>(alias_info);
 
   if (debug_options.xla_gpu_copy_insertion_use_region_analysis()) {
diff --git a/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.h b/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.h
index 7fdf8422aa0639..7b87575e35c0d5 100644
--- a/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.h
+++ b/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.h
@@ -21,10 +21,6 @@ limitations under the License.
 #include "xla/service/gpu/alias_info.h"
 #include "xla/stream_executor/device_description.h"
 
-namespace mlir {
-class MLIRContext;
-}  // namespace mlir
-
 namespace xla {
 namespace gpu {
 
@@ -32,8 +28,7 @@ namespace gpu {
 // This pipeline must run before scheduling to ensure correctness.
 HloPassPipeline PreSchedulingCopyInsertionPipeline(
     const HloModuleConfig& config, const GpuAliasInfo* alias_info,
-    const se::DeviceDescription& device_description,
-    mlir::MLIRContext* mlir_context);
+    const se::DeviceDescription& device_description);
 
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h
index 716f54fc01ec6b..f0d5c78703e8e7 100644
--- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h
+++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h
@@ -166,7 +166,7 @@ Useful logging and error messages
     prefetch a buffer to alternate memory, according to some heuristic and not
     based on limited copy resource.
     * If the CostAnalysisPrefetchIntervalPicker is used, which is the default,
-      live range too long is governed by the picker's
+      live range too short is governed by the picker's
       min_overlap_to_async_copy_ratio argument.
 
   - "Finding allocation for": Magical logging phrase indicating the point in
@@ -206,13 +206,13 @@ Useful logging and error messages
 #include "xla/hlo/analysis/hlo_dataflow_analysis.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/utils/hlo_live_range.h"
+#include "xla/layout.h"
 #include "xla/service/heap_simulator/heap_simulator.h"
 #include "xla/service/hlo.pb.h"
 #include "xla/service/hlo_value.h"
 #include "xla/service/memory_space_assignment/allocation.h"
 #include "xla/service/memory_space_assignment/memory_space_assignment.pb.h"
 #include "xla/service/memory_space_assignment/options.h"
-#include "xla/shape.h"
 #include "xla/util.h"
 
 namespace xla {
diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD
index f2a1919062e564..6b44cb146095dd 100644
--- a/third_party/xla/xla/tests/BUILD
+++ b/third_party/xla/xla/tests/BUILD
@@ -3485,13 +3485,13 @@ xla_test(
     name = "batch_norm_training_test",
     srcs = ["batch_norm_training_test.cc"],
     deps = [
-        ":hlo_test_base",
+        ":hlo_pjrt_test_base",
         ":xla_internal_test_main",  # fixdeps: keep
         "//xla:literal_util",
         "//xla/hlo/testlib:test",
         "//xla/tests:xla_test_backend_predicates",
+        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/status",
-        "@com_google_googletest//:gtest",
     ],
 )
 
diff --git a/third_party/xla/xla/tests/batch_norm_training_test.cc b/third_party/xla/xla/tests/batch_norm_training_test.cc
index 9cc9dac24f3eec..fdc14558fd48f0 100644
--- a/third_party/xla/xla/tests/batch_norm_training_test.cc
+++ b/third_party/xla/xla/tests/batch_norm_training_test.cc
@@ -13,15 +13,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include <cmath>
 #include <utility>
 
 #include "xla/tests/xla_test_backend_predicates.h"
-#include <gmock/gmock.h>
 #include "absl/status/status.h"
 #include "xla/hlo/testlib/test.h"
 #include "xla/literal_util.h"
-#include "xla/tests/hlo_test_base.h"
+#include "xla/tests/hlo_pjrt_test_base.h"
+#include "xla/tsl/platform/statusor.h"
 
 namespace xla {
 namespace {
@@ -41,14 +40,15 @@ ENTRY entry {
 class BatchNormTrainingTest : public HloTestBase {};
 
 TEST_F(BatchNormTrainingTest, CorrectComputation) {
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kModuleStr));
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          ParseAndReturnVerifiedModule(kModuleStr));
 
   auto input = LiteralUtil::CreateR2<float>({{1.0}, {2.0}});
   auto scale = LiteralUtil::CreateR1<float>({0.5});
   auto offset = LiteralUtil::CreateR1<float>({0.1});
 
-  ASSERT_OK_AND_ASSIGN(auto result,
-                       Execute(std::move(module), {&input, &scale, &offset}));
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto result, Execute(std::move(module), {&input, &scale, &offset}));
 
   // Decompose result tuple
   auto result_tuple = result.DecomposeTuple();
@@ -78,50 +78,12 @@ TEST_F(BatchNormTrainingTest, CorrectComputation) {
   }
 }
 
-TEST_F(BatchNormTrainingTest, LargeOffset) {
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kModuleStr));
-
-  auto input =
-      LiteralUtil::CreateR2<float>({{10000.0f + 1.0f}, {10000.0f + 2.0f}});
-  auto scale = LiteralUtil::CreateR1<float>({0.5f});
-  auto offset = LiteralUtil::CreateR1<float>({0.1f});
-
-  ASSERT_OK_AND_ASSIGN(auto result,
-                       Execute(std::move(module), {&input, &scale, &offset}));
-
-  auto result_tuple = result.DecomposeTuple();
-
-  auto expected_output =
-      LiteralUtil::CreateR2<float>({{-0.399003029f}, {0.599003f}});
-  auto expected_batch_mean = LiteralUtil::CreateR1<float>({10000.0f + 1.5f});
-  auto expected_batch_var = LiteralUtil::CreateR1<float>({0.25f});
-
-  const float tolerance = 1e-4f;
-
-  for (int i = 0; i < expected_output.element_count(); ++i) {
-    EXPECT_FALSE(std::isnan(result_tuple[0].data<float>()[i]));
-    EXPECT_NEAR(result_tuple[0].data<float>()[i],
-                expected_output.data<float>()[i], tolerance);
-  }
-
-  for (int i = 0; i < expected_batch_mean.element_count(); ++i) {
-    EXPECT_NEAR(result_tuple[1].data<float>()[i],
-                expected_batch_mean.data<float>()[i], tolerance);
-  }
-
-  for (int i = 0; i < expected_batch_var.element_count(); ++i) {
-    EXPECT_FALSE(std::isnan(result_tuple[2].data<float>()[i]));
-    EXPECT_GE(result_tuple[2].data<float>()[i], 0.0f);
-    EXPECT_NEAR(result_tuple[2].data<float>()[i],
-                expected_batch_var.data<float>()[i], tolerance);
-  }
-}
-
 TEST_F(BatchNormTrainingTest, ReturnsErrorWhenHloPassesDisabled) {
   if (test::DeviceTypeIsOneOf({test::kGpu, test::kInterpreter, test::kTpu})) {
     GTEST_SKIP();
   }
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kModuleStr));
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          ParseAndReturnVerifiedModule(kModuleStr));
 
   auto status_or_result =
       Execute(std::move(module), {}, /*run_hlo_passes=*/false);
diff --git a/third_party/xla/xla/tpu/c_api_conversions.h b/third_party/xla/xla/tpu/c_api_conversions.h
index 4394030b089984..b93c4a461c7520 100644
--- a/third_party/xla/xla/tpu/c_api_conversions.h
+++ b/third_party/xla/xla/tpu/c_api_conversions.h
@@ -104,10 +104,10 @@ void Destroy(XLA_HloModuleConfig* c_config);
 // Helper for managing stack based C -> C++ conversions.
 template <class CType>
 struct StackHelper {
-  explicit StackHelper() {}
+  explicit StackHelper() : value() {}
 
   template <class CppType>
-  explicit StackHelper(const CppType& t) {
+  explicit StackHelper(const CppType& t) : value() {
     ::ApiConverter::ToC(t, &value);
   }
   ~StackHelper() { ::ApiConverter::Destroy(&value); }
@@ -117,7 +117,7 @@ struct StackHelper {
     return ::ApiConverter::FromC(&value);
   }
 
-  mutable CType value;
+  mutable CType value{};
 };
 
 }  // namespace ApiConverter
diff --git a/third_party/xla/xla/tpu/c_api_conversions_test.cc b/third_party/xla/xla/tpu/c_api_conversions_test.cc
index db67de69760887..e2415ca48f9b87 100644
--- a/third_party/xla/xla/tpu/c_api_conversions_test.cc
+++ b/third_party/xla/xla/tpu/c_api_conversions_test.cc
@@ -323,6 +323,27 @@ TEST(ProtoHelper, EmptyProto) {
   stream_executor::tpu::SerializedProto_Free(serialized_proto);
 }
 
+TEST(StackHelper, DefaultConstruct) {
+  {
+    StackHelper<XLA_Shape> stack_shape;
+    EXPECT_EQ(stack_shape.value.dimensions.size, 0);
+  }
+  {
+    StackHelper<XLA_Layout> stack_layout;
+    EXPECT_EQ(stack_layout.value.minor_to_major.size, 0);
+  }
+  {
+    StackHelper<XLA_Tile> stack_tile;
+    EXPECT_EQ(stack_tile.value.dimensions.size, 0);
+  }
+}
+
+TEST(StackHelper, Conversion) {
+  xla::Shape cpp_shape = xla::ShapeUtil::MakeShapeWithType<float>({4, 3});
+  StackHelper<XLA_Shape> stack_shape(cpp_shape);
+  EXPECT_EQ(stack_shape.AsCpp<xla::Shape>(), cpp_shape);
+}
+
 // TODO(b/290654348): SE_DeviceAddressBase, SE_DeviceAddressAllocator,
 // SE_MaybeOwningDeviceAddress