Merge remote-tracking branch 'upstream/master' into feat/auto-strategy

MODULE.bazel

@@ -219,6 +219,7 @@ python.toolchain(python_version = "3.10")
 python.toolchain(python_version = "3.11")
 python.toolchain(python_version = "3.12")
 python.toolchain(python_version = "3.13")
+python.toolchain(python_version = "3.14")
 
 pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
 
@@ -237,6 +238,7 @@ pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
     "3.11",
     "3.12",
     "3.13",
+    "3.14",
 ]]
 
 use_repo(pip, "pypi_mods", pypi = "tf_pypi")

RELEASE.md

@@ -52,9 +52,9 @@ In `tensorflow/c/experimental/filesystem/filesystem_interface.h`, removed `TF_Tr
 
 ### Major Features and Improvements
 
+*   Support for Python 3.14 has been added.
 * `tf.distribute.experimental`
     * Introduces `tf.distribute.experimental.auto_strategy`, a dynamic factory API that automatically detects available hardware (TPUs, multiple GPUs, or multi-worker clusters via `TF_CONFIG`) and resolves the optimal distribution strategy, safely falling back to single device execution.
-
 * `tf.lite`
     * Adds support for QUI4 (Quantized Unsigned 4-bit) in Dequantize operator.
     * Adds support for FP16 and BF16 in Unpack operator.

ci/official/containers/ml_build/Dockerfile

@@ -71,7 +71,6 @@ RUN if [ "${INSTALL_CMAKE_TOOLCHAIN}" = "true" ]; then \
 # Setup links for TensorFlow to compile.
 # Referenced in devel.usertools/*.bazelrc.
 # Set python3.12 as the default python version.
-# TF does not support python3.13.
 RUN ln -sf /usr/bin/python3.12 /usr/bin/python3
 RUN ln -sf /usr/bin/python3.12 /usr/bin/python
 RUN ln -sf /usr/lib/python3.12 /usr/lib/tf_python

ci/official/containers/ml_build/builder.requirements.txt

@@ -11,7 +11,7 @@ pyyaml
 # For JAX
 build ~= 1.2.2
 # uv is faster than pip for installing Python packages.
-uv ~= 0.5.30
+uv ~= 0.12.23
 
 # For running wheel verification script
 immutabledict ~= 4.2.1

ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt

@@ -764,13 +764,13 @@ typing-extensions==4.14.1 \
     #   -r ci/official/requirements_updater/requirements.in
     #   optree
     #   rich
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.6 \
-    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
-    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.41.3 \
     --hash=sha256:488609bc63a29322326e05560731bf7bfea8e48ad646e1f5e40d366607de0942 \

ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt

@@ -763,13 +763,13 @@ typing-extensions==4.14.1 \
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   optree
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.6 \
-    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
-    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.41.3 \
     --hash=sha256:488609bc63a29322326e05560731bf7bfea8e48ad646e1f5e40d366607de0942 \

ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt

@@ -763,13 +763,13 @@ typing-extensions==4.14.1 \
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   optree
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.6 \
-    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
-    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.41.3 \
     --hash=sha256:488609bc63a29322326e05560731bf7bfea8e48ad646e1f5e40d366607de0942 \

configure.py

@@ -1332,6 +1332,7 @@ def main():
         clang_version = retrieve_clang_version(clang_compiler_path)
         if is_s390x():
           write_to_bazelrc('build --linkopt=-rtlib=compiler-rt')
+          write_to_bazelrc('build --host_linkopt=-rtlib=compiler-rt')
         disable_clang_offsetof_extension(clang_version)
     if is_windows():
       environ_cp['TF_NEED_CLANG'] = str(choose_compiler_Win(environ_cp))

requirements_lock_3_10.txt

@@ -888,13 +888,13 @@ typing-extensions==4.14.1 \
     #   -r ci/official/requirements_updater/requirements.in
     #   optree
     #   rich
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.6 \
-    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
-    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.46.2 \
     --hash=sha256:33ae60725d69eaa249bc1982e739943c23b34b58d51f1cb6253453773aca6e65 \

requirements_lock_3_11.txt

@@ -887,13 +887,13 @@ typing-extensions==4.14.1 \
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   optree
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.6 \
-    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
-    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.46.2 \
     --hash=sha256:33ae60725d69eaa249bc1982e739943c23b34b58d51f1cb6253453773aca6e65 \

requirements_lock_3_12.txt

@@ -887,13 +887,13 @@ typing-extensions==4.14.1 \
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   optree
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.6 \
-    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
-    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.46.2 \
     --hash=sha256:33ae60725d69eaa249bc1982e739943c23b34b58d51f1cb6253453773aca6e65 \

requirements_lock_3_13.txt

@@ -894,13 +894,13 @@ typing-extensions==4.14.1 \
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   optree
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.6 \
-    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
-    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.46.2 \
     --hash=sha256:33ae60725d69eaa249bc1982e739943c23b34b58d51f1cb6253453773aca6e65 \

requirements_lock_3_14.txt

@@ -1128,13 +1128,13 @@ typing-extensions==4.14.1 \
     #   -r ci/official/requirements_updater/requirements.in
     #   grpcio
     #   optree
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.8 \
-    --hash=sha256:63a77fb8892bf28ebc3178683445222aa500e48ebad5ec77b0ad80f8726b1f50 \
-    --hash=sha256:9bad61a4268dac112f1c5cd4630a56ede601b6ed420300677a869083d70a4c44
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.46.2 \
     --hash=sha256:33ae60725d69eaa249bc1982e739943c23b34b58d51f1cb6253453773aca6e65 \

requirements_lock_3_14_freethreaded.txt

@@ -1129,13 +1129,13 @@ typing-extensions==4.14.1 \
     #   -r ci/official/requirements_updater/requirements.in
     #   grpcio
     #   optree
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.8 \
-    --hash=sha256:63a77fb8892bf28ebc3178683445222aa500e48ebad5ec77b0ad80f8726b1f50 \
-    --hash=sha256:9bad61a4268dac112f1c5cd4630a56ede601b6ed420300677a869083d70a4c44
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.46.2 \
     --hash=sha256:33ae60725d69eaa249bc1982e739943c23b34b58d51f1cb6253453773aca6e65 \

tensorflow/c/eager/abstract_context.h

@@ -34,7 +34,7 @@ class AbstractContext {
  protected:
   enum AbstractContextKind { kGraph, kMlir, kEager, kTfrt, kTape, kOpHandler };
   explicit AbstractContext(AbstractContextKind kind) : kind_(kind) {}
-  virtual ~AbstractContext() {}
+  virtual ~AbstractContext() = default;
 
  public:
   AbstractContextKind getKind() const { return kind_; }

tensorflow/c/eager/abstract_operation.h

@@ -40,7 +40,7 @@ class AbstractOperation {
     kOpHandler
   };
   explicit AbstractOperation(AbstractOperationKind kind) : kind_(kind) {}
-  virtual ~AbstractOperation() {}
+  virtual ~AbstractOperation() = default;
 
  public:
   AbstractOperationKind getKind() const { return kind_; }

tensorflow/c/eager/abstract_tensor_handle.h

@@ -31,7 +31,7 @@ class AbstractTensorHandle : public core::RefCounted {
  protected:
   enum AbstractTensorHandleKind { kGraph, kMlir, kEager, kTfrt, kCustomDevice };
   explicit AbstractTensorHandle(AbstractTensorHandleKind kind) : kind_(kind) {}
-  ~AbstractTensorHandle() override {}
+  ~AbstractTensorHandle() override = default;
 
  public:
   // Returns tensor dtype.

tensorflow/c/eager/c_api_distributed_test.cc

@@ -310,7 +310,7 @@ string VariableAddFunction() {
 class GraphErrorInjectionPass : public tensorflow::GraphOptimizationPass {
  public:
   static bool enabled_;
-  GraphErrorInjectionPass() {}
+  GraphErrorInjectionPass() = default;
 
   absl::Status Run(
       const tensorflow::GraphOptimizationPassOptions& options) override {

tensorflow/c/eager/c_api_unified_experimental_graph.cc

@@ -321,7 +321,7 @@ class GraphOperation : public TracingOperation {
   static bool classof(const AbstractOperation* ptr) {
     return ptr->getKind() == kGraph;
   }
-  ~GraphOperation() override {}
+  ~GraphOperation() override = default;
 
  private:
   friend class GraphContext;  // For access to op_.

tensorflow/c/eager/gradients.h

@@ -58,7 +58,7 @@ class GradientFunction {
       AbstractContext* ctx,
       absl::Span<AbstractTensorHandle* const> grad_outputs,
       absl::Span<AbstractTensorHandle*> grad_inputs) = 0;
-  virtual ~GradientFunction() {}
+  virtual ~GradientFunction() = default;
 };
 
 // Metadata from the forward operation that is made available to the

tensorflow/c/eager/graph_function.cc

@@ -26,7 +26,7 @@ namespace graph {
 GraphFunction::GraphFunction(FunctionDef fdef)
     : AbstractFunction(kGraph),
       func_record_(new FunctionRecord(std::move(fdef), {}, true)) {}
-GraphFunction::~GraphFunction() {}
+GraphFunction::~GraphFunction() = default;
 absl::Status GraphFunction::GetFunctionDef(const FunctionDef **fdef) {
   *fdef = &(func_record_->fdef());
   return absl::OkStatus();

tensorflow/c/eager/immediate_execution_context.h

@@ -272,7 +272,7 @@ class ImmediateExecutionContext : public AbstractContext {
  protected:
   explicit ImmediateExecutionContext(AbstractContextKind kind)
       : AbstractContext(kind) {}
-  ~ImmediateExecutionContext() override {}
+  ~ImmediateExecutionContext() override = default;
 };
 
 namespace internal {

tensorflow/c/eager/immediate_execution_distributed_manager.h

@@ -33,7 +33,7 @@ class WorkerCacheInterface;
 
 class ImmediateExecutionDistributedManager {
  public:
-  virtual ~ImmediateExecutionDistributedManager() {}
+  virtual ~ImmediateExecutionDistributedManager() = default;
 
   // Set up distributed execution environment on local and remote tasks.
   // When `reset_context` is true, initialize new cluster context state based

tensorflow/c/eager/immediate_execution_tensor_handle.h

@@ -83,7 +83,7 @@ class ImmediateExecutionTensorHandle : public AbstractTensorHandle {
  protected:
   explicit ImmediateExecutionTensorHandle(AbstractTensorHandleKind kind)
       : AbstractTensorHandle(kind) {}
-  ~ImmediateExecutionTensorHandle() override {}
+  ~ImmediateExecutionTensorHandle() override = default;
 };
 
 namespace internal {

tensorflow/c/eager/tape.h

@@ -94,7 +94,7 @@ using OpTape =
 template <typename Gradient, typename BackwardFunction, typename TapeTensor>
 class VSpace {
  public:
-  virtual ~VSpace() {}
+  virtual ~VSpace() = default;
 
   // Returns the number of elements in the gradient tensor.
   virtual int64_t NumElements(Gradient* tensor) const = 0;

tensorflow/c/experimental/filesystem/plugins/gcs/cleanup.h

@@ -92,14 +92,14 @@ class Cleanup {
   bool is_released() const { return released_; }
 
  private:
-  static_assert(!std::is_reference<F>::value, "F must not be a reference");
+  static_assert(!std::is_reference_v<F>, "F must not be a reference");
 
   bool released_ = false;
   F f_;
 };
 
 template <int&... ExplicitParameterBarrier, typename F,
-          typename DecayF = typename std::decay<F>::type>
+          typename DecayF = std::decay_t<F>>
 Cleanup<DecayF> MakeCleanup(F&& f) {
   return Cleanup<DecayF>(std::forward<F>(f));
 }

tensorflow/c/experimental/gradients/array_grad.cc

@@ -35,7 +35,7 @@ class IdentityNGradientFunction : public GradientFunction {
     }
     return absl::OkStatus();
   }
-  ~IdentityNGradientFunction() override {}
+  ~IdentityNGradientFunction() override = default;
 };
 }  // namespace
 

tensorflow/c/experimental/gradients/nn_grad.cc

@@ -107,7 +107,7 @@ class SparseSoftmaxCrossEntropyWithLogitsGradientFunction
     grad_inputs[1] = nullptr;
     return absl::OkStatus();
   }
-  ~SparseSoftmaxCrossEntropyWithLogitsGradientFunction() override {}
+  ~SparseSoftmaxCrossEntropyWithLogitsGradientFunction() override = default;
 
  private:
   vector<AbstractTensorHandle*> forward_outputs_;
@@ -146,7 +146,7 @@ class BiasAddGradientFunction : public GradientFunction {
 
     return absl::OkStatus();
   }
-  ~BiasAddGradientFunction() override {}
+  ~BiasAddGradientFunction() override = default;
 
  private:
   AttrBuilder forward_attrs_;

tensorflow/c/experimental/stream_executor/stream_executor_test.cc

@@ -150,7 +150,7 @@ TEST(StreamExecutor, UnifiedMemoryAllocateNotSet) {
 /*** StreamExecutor behavior tests ***/
 class StreamExecutorTest : public ::testing::Test {
  protected:
-  StreamExecutorTest() {}
+  StreamExecutorTest() = default;
   void SetUp() override {
     test_util::PopulateDefaultPlatform(&platform_, &platform_fns_);
     test_util::PopulateDefaultDeviceFns(&device_fns_);

tensorflow/cc/client/client_session.cc

@@ -61,7 +61,7 @@ ClientSession::ClientSession(const Scope& scope,
 // Define destructor here so we can forward declare `Impl` in client_session.h.
 // If we define a dtor in the header file or use the default dtor,
 // unique_ptr<Impl> needs the complete type.
-ClientSession::~ClientSession() {}
+ClientSession::~ClientSession() = default;
 
 SessionOptions ClientSession::Impl::MakeDefaultSessionOptions(
     const std::string& target) {

tensorflow/cc/framework/gradient_checker.cc

@@ -69,14 +69,14 @@ SET_BASE_UNITS_FOR_TYPE(complex128, ({{1, 0}, {0, 1}}));
 // When T is complex, it assigns the real and complex values to successive rows
 // or columns in the matrix depending on the expand_by_row parameter
 template <typename T, typename JAC_T>
-typename std::enable_if<std::is_floating_point<T>::value>::type SetJacobian(
+std::enable_if_t<std::is_floating_point_v<T>> SetJacobian(
     typename TTypes<JAC_T>::Matrix* jacobian, const int row, const int col,
     const T& value, const bool expand_by_row) {
   (*jacobian)(row, col) = JAC_T{value};
 }
 
 template <typename T, typename JAC_T>
-typename std::enable_if<is_complex<T>::value>::type SetJacobian(
+std::enable_if_t<is_complex<T>::value> SetJacobian(
     typename TTypes<JAC_T>::Matrix* jacobian, const int row, const int col,
     const T& value, const bool expand_by_row) {
   (*jacobian)(row, col) = JAC_T{value.real()};

tensorflow/cc/framework/ops.h

@@ -112,9 +112,9 @@ class Input {
   struct Initializer {
     /// Construct from a scalar value of an arithmetic type or a type that can
     /// be converted to a string (eg. a string literal).
-    template <typename T, typename = typename std::enable_if<
-                              std::is_arithmetic<T>::value ||
-                              std::is_convertible<T, std::string>::value>::type>
+    template <typename T, typename = std::enable_if_t<
+                              std::is_arithmetic_v<T> ||
+                              std::is_convertible_v<T, std::string>>>
     Initializer(const T& v) {  // NOLINT(runtime/explicit)
       typedef typename RealType<T>::type RealT;
       Tensor t(DataTypeToEnum<RealT>::v(), TensorShape());
@@ -125,9 +125,9 @@ class Input {
     Initializer(const Tensor& t) : tensor(t) {}  // NOLINT(runtime/explicit)
 
     /// Construct from a scalar value and an explicit shape
-    template <typename T, typename = typename std::enable_if<
-                              std::is_arithmetic<T>::value ||
-                              std::is_convertible<T, std::string>::value>::type>
+    template <typename T, typename = std::enable_if_t<
+                              std::is_arithmetic_v<T> ||
+                              std::is_convertible_v<T, std::string>>>
     Initializer(const T& v, const TensorShape& shape) {
       typedef typename RealType<T>::type RealT;
       Tensor t(DataTypeToEnum<RealT>::v(), shape);
@@ -138,9 +138,9 @@ class Input {
     }
 
     /// Construct from a initializer list of scalars (a one-dimensional tensor).
-    template <typename T, typename = typename std::enable_if<
-                              std::is_arithmetic<T>::value ||
-                              std::is_convertible<T, std::string>::value>::type>
+    template <typename T, typename = std::enable_if_t<
+                              std::is_arithmetic_v<T> ||
+                              std::is_convertible_v<T, std::string>>>
     Initializer(
         const std::initializer_list<T>& v) {  // NOLINT(runtime/explicit)
       typedef typename RealType<T>::type RealT;
@@ -151,9 +151,9 @@ class Input {
     }
 
     /// Construct from a initializer list of scalars and an explicit shape.
-    template <typename T, typename = typename std::enable_if<
-                              std::is_arithmetic<T>::value ||
-                              std::is_convertible<T, std::string>::value>::type>
+    template <typename T, typename = std::enable_if_t<
+                              std::is_arithmetic_v<T> ||
+                              std::is_convertible_v<T, std::string>>>
     Initializer(const std::initializer_list<T>& v, const TensorShape& shape) {
       typedef typename RealType<T>::type RealT;
       Tensor t(DataTypeToEnum<RealT>::v(), shape);
@@ -175,7 +175,7 @@ class Input {
     Initializer(const std::initializer_list<Initializer>& v);
 
     // START_SKIP_DOXYGEN
-    template <typename T, bool = std::is_convertible<T, std::string>::value>
+    template <typename T, bool = std::is_convertible_v<T, std::string>>
     struct RealType {
       typedef tstring type;
     };
@@ -210,9 +210,9 @@ class Input {
   /// * A Tensor object.
   Input(const Output& o) : output_(o) {}  // NOLINT(runtime/explicit)
 
-  template <typename T, typename = typename std::enable_if<
-                            std::is_arithmetic<T>::value ||
-                            std::is_convertible<T, std::string>::value>::type>
+  template <typename T,
+            typename = std::enable_if_t<std::is_arithmetic_v<T> ||
+                                        std::is_convertible_v<T, std::string>>>
   Input(const T& v)  // NOLINT(runtime/explicit)
       : Input(Initializer(v)) {}
 

tensorflow/cc/framework/scope.cc

@@ -29,7 +29,7 @@ Scope::Scope(Impl* impl) : impl_(impl) {}
 
 Scope::Scope(const Scope& other) : impl_(new Impl(*other.impl())) {}
 
-Scope::~Scope() {}
+Scope::~Scope() = default;
 
 Scope& Scope::operator=(const Scope& other) {
   // We can't copy Impls because of the const members, use copy ctor instead

tensorflow/cc/saved_model/bundle_v2_test.cc

@@ -41,7 +41,7 @@ constexpr char kTestData[] = "cc/saved_model/testdata";
 
 class BundleV2Test : public ::testing::Test {
  protected:
-  BundleV2Test() {}
+  BundleV2Test() = default;
 
   void RestoreVarsAndVerify(SavedModelV2Bundle* bundle,
                             std::vector<std::string> expected_names) {

tensorflow/cc/saved_model/reader_test.cc

@@ -51,7 +51,7 @@ string NonChunkedSavedModel() {
 
 class ReaderTest : public ::testing::Test {
  protected:
-  ReaderTest() {}
+  ReaderTest() = default;
 
   void CheckMetaGraphDef(const MetaGraphDef& meta_graph_def) {
     const auto& tags = meta_graph_def.meta_info_def().tags();

tensorflow/cc/saved_model/saved_model_bundle_lite_test.cc

@@ -42,7 +42,7 @@ constexpr char kTestDataInitOpV2[] =
 
 class LoaderTest : public ::testing::Test {
  protected:
-  LoaderTest() {}
+  LoaderTest() = default;
 
   string MakeSerializedExample(float x) {
     tensorflow::Example example;

tensorflow/cc/saved_model/saved_model_bundle_test.cc

@@ -57,7 +57,7 @@ constexpr char kV2ModuleSavedModelChecksum[] = "15788619162413586750";
 
 class LoaderTest : public ::testing::Test {
  protected:
-  LoaderTest() {}
+  LoaderTest() = default;
 
   string MakeSerializedExample(float x) {
     tensorflow::Example example;

tensorflow/compiler/jit/BUILD

@@ -1643,6 +1643,7 @@ cc_library(
     hdrs = ["device_compiler_client.h"],
     visibility = [":internal"],
     deps = [
+        ":common",
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/core/util:determinism",
         "@xla//xla/client:executable_build_options",
@@ -1653,8 +1654,11 @@ tf_cc_test(
     name = "device_compiler_client_test",
     srcs = ["device_compiler_client_test.cc"],
     deps = [
+        ":common",
         ":device_compiler_client",
         "@com_google_googletest//:gtest_main",
+        "@xla//xla/hlo/builder:xla_computation",
+        "@xla//xla/service:hlo_proto_cc",
     ],
 )
 

tensorflow/compiler/jit/deadness_analysis.cc

@@ -132,7 +132,7 @@ class Predicate {
   virtual absl::Span<Predicate* const> GetOperands() const = 0;
 
   virtual Kind kind() const = 0;
-  virtual ~Predicate() {}
+  virtual ~Predicate() = default;
 
   // Invokes func on p and on all of its operands recursively.  Does not invoke
   // `func` on the same Predicate instance twice.  Aborts the search if `func`
@@ -1566,7 +1566,7 @@ void DeadnessAnalysisImpl::Print() const {
 
 }  // namespace
 
-DeadnessAnalysis::~DeadnessAnalysis() {}
+DeadnessAnalysis::~DeadnessAnalysis() = default;
 
 /*static*/ absl::Status DeadnessAnalysis::Run(
     const Graph& graph, std::unique_ptr<DeadnessAnalysis>* result) {

tensorflow/compiler/jit/defs.cc

@@ -32,6 +32,8 @@ const char* const kXlaInternalScopeAttr = "_XlaInternalScope";
 
 const char* const kXlaClusterIdAttr = "_xla_compile_id";
 
+const char* const kXlaDeterministicAttr = "_XlaDeterministic";
+
 static std::atomic<bool> xla_devices_creation_required(false);
 
 // Request XLA:GPU and XLA:CPU device creation. Deprecated, only used by XRT

tensorflow/compiler/jit/defs.h

@@ -38,6 +38,9 @@ extern const char* const kXlaInternalScopeAttr;  // "_XlaInternalScope"
 // The id of the compiled cluster.
 extern const char* const kXlaClusterIdAttr;  // "_xla_compile_id"
 
+// Implies per-cluster deterministic XLA ops.
+extern const char* const kXlaDeterministicAttr;  // "_XlaDeterministic"
+
 [[deprecated("XLA:CPU/GPU devices are deprecated")]] void
 RequestXlaDevicesCreation();
 

tensorflow/compiler/jit/device_compiler_client.cc

@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/jit/device_compiler_client.h"
 
+#include "tensorflow/compiler/jit/defs.h"
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
 #include "tensorflow/core/util/determinism.h"
 
@@ -37,10 +38,23 @@ xla::ExecutableBuildOptions GetExecutableBuildOptions(
   build_options.set_alias_passthrough_params(options.alias_passthrough_params);
   build_options.mutable_debug_options()->set_xla_detailed_logging(
       options.detailed_logging);
-  if (tensorflow::OpDeterminismRequired()) {
+  bool cluster_deterministic = false;
+  if (result.computation != nullptr) {
+    const auto& frontend_attrs =
+        result.computation->proto().frontend_attributes().map();
+    auto it = frontend_attrs.find(kXlaDeterministicAttr);
+    if (it != frontend_attrs.end() && it->second == "true") {
+      cluster_deterministic = true;
+    }
+  }
+  if (tensorflow::OpDeterminismRequired() || cluster_deterministic) {
     build_options.mutable_debug_options()
         ->set_xla_gpu_exclude_nondeterministic_ops(true);
   }
+  if (cluster_deterministic) {
+    build_options.mutable_debug_options()
+        ->set_xla_gpu_experimental_cost_model_gemm_tiling_default(true);
+  }
   return build_options;
 }
 

tensorflow/compiler/jit/device_compiler_client_test.cc

@@ -15,7 +15,12 @@ limitations under the License.
 
 #include "tensorflow/compiler/jit/device_compiler_client.h"
 
+#include <memory>
+
 #include <gtest/gtest.h>
+#include "tensorflow/compiler/jit/defs.h"
+#include "xla/hlo/builder/xla_computation.h"
+#include "xla/service/hlo.pb.h"
 
 namespace tensorflow {
 namespace {
@@ -72,5 +77,28 @@ TEST(GetExecutableOptionTest, DumpingWithoutDetailedLogging) {
   EXPECT_TRUE(build_option.debug_options().xla_enable_dumping());
 }
 
+TEST(GetExecutableOptionTest, XlaDeterministicFrontendAttribute) {
+  XlaCompiler::Options options;
+  XlaCompiler::CompilationResult result;
+  result.computation = std::make_shared<xla::XlaComputation>();
+
+  auto default_build_option =
+      GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/-1);
+  EXPECT_FALSE(default_build_option.debug_options()
+                   .xla_gpu_exclude_nondeterministic_ops());
+  EXPECT_FALSE(default_build_option.debug_options()
+                   .xla_gpu_experimental_cost_model_gemm_tiling_default());
+
+  (*result.computation->mutable_proto()
+        ->mutable_frontend_attributes()
+        ->mutable_map())[kXlaDeterministicAttr] = "true";
+  auto deterministic_build_option =
+      GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/-1);
+  EXPECT_TRUE(deterministic_build_option.debug_options()
+                  .xla_gpu_exclude_nondeterministic_ops());
+  EXPECT_TRUE(deterministic_build_option.debug_options()
+                  .xla_gpu_experimental_cost_model_gemm_tiling_default());
+}
+
 }  // namespace
 }  // namespace tensorflow

tensorflow/compiler/jit/encapsulate_util.h

@@ -94,7 +94,7 @@ struct XlaClusterInfo {
   // TODO(klimek): In c++17 we'll be able to initialize host_compute_core
   // without losing aggregate initialization, which allows us to get rid of
   // the constructor definitions again.
-  XlaClusterInfo() {}
+  XlaClusterInfo() = default;
   XlaClusterInfo(const std::string& cluster_name,
                  const NameAttrList& func_name_attrs, Node* node,
                  const std::map<std::string, int>& host_compute_core)

tensorflow/compiler/jit/xla_activity_listener.cc

@@ -100,6 +100,6 @@ void RegisterXlaActivityListener(
 
 void XlaActivityListener::Flush() {}
 
-XlaActivityListener::~XlaActivityListener() {}
+XlaActivityListener::~XlaActivityListener() = default;
 
 }  // namespace tensorflow

tensorflow/compiler/jit/xla_activity_listener_test.cc

@@ -48,7 +48,7 @@ class TestListener : public XlaActivityListener {
     return absl::OkStatus();
   }
 
-  ~TestListener() override {}
+  ~TestListener() override = default;
 
   const XlaAutoClusteringActivity& auto_clustering_activity() const {
     return auto_clustering_activity_;

tensorflow/compiler/mlir/BUILD

@@ -149,6 +149,7 @@ cc_library(
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/log:vlog_is_on",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings:string_view",
         "@llvm-project//llvm:Support",

tensorflow/compiler/mlir/lite/BUILD

@@ -1100,6 +1100,7 @@ cc_library(
         ":pass",
         ":pass_options",
         ":tensorflow_lite_ops",
+        "@com_google_absl//absl/strings",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
@@ -1553,6 +1554,7 @@ cc_library(
     srcs = [
         "transforms/decompose_hybrid_quantization.cc",
         "transforms/default_quant_params.cc",
+        "transforms/fold_fp16_casts_into_fully_connected_pass.cc",
         "transforms/fold_stablehlo_constant_transforms_pass.cc",
         "transforms/fuse_a4w2_drq_fully_connected_pass.cc",
         "transforms/generated_post_quantize.inc",

tensorflow/compiler/mlir/lite/common/tfl_pass_config.h

@@ -118,6 +118,10 @@ struct PassConfig {
 
   // When set to true, enable unsafe single batch rank reduction.
   bool unsafe_single_batch_rank_reduction = false;
+
+  // When set to true, fold f16 <-> f32 casts around float fully_connected ops
+  // into them, producing mixed precision (f16 activation) fully_connected ops.
+  bool fold_fp16_casts_into_fully_connected = false;
 };
 
 inline llvm::raw_ostream& operator<<(llvm::raw_ostream& os,
@@ -153,7 +157,9 @@ inline llvm::raw_ostream& operator<<(llvm::raw_ostream& os,
             << tflite::ConverterFlags::ModelOriginFramework_Name(
                    pass_config.model_origin_framework)
             << "\nunsafe_single_batch_rank_reduction: "
-            << pass_config.unsafe_single_batch_rank_reduction << "\n";
+            << pass_config.unsafe_single_batch_rank_reduction
+            << "\nfold_fp16_casts_into_fully_connected: "
+            << pass_config.fold_fp16_casts_into_fully_connected << "\n";
 }
 
 }  // namespace TFL

tensorflow/compiler/mlir/lite/converter_flags.proto

@@ -39,7 +39,7 @@ enum FileFormat {
 // of as properties of models, instead describing how models are to be
 // processed in the context of the present tooling job.
 //
-// Next ID to use: 74.
+// Next ID to use: 75.
 message ConverterFlags {
   reserved 54, 61;
 
@@ -398,4 +398,11 @@ message ConverterFlags {
   // Directory for debug artifacts and printed IR dumps when debug mode is
   // enabled.
   optional string debug_dir = 73;
+
+  // If true, fold the f16 -> f32 cast feeding a float fully_connected and the
+  // f32 -> f16 cast consuming it into the fully_connected, which then takes
+  // and produces f16 activations while keeping its weights unchanged. The
+  // resulting mixed precision fully_connected requires runtime support.
+  // WARNING: Experimental interface, subject to change.
+  optional bool fold_fp16_casts_into_fully_connected = 74 [default = false];
 }

tensorflow/compiler/mlir/lite/flatbuffer_export.cc

@@ -1584,10 +1584,9 @@ std::optional<BufferOffset<tflite::Buffer>> Translator::BuildBuffer(
                             /*hash=*/0);
 
     std::string buffer_data = buffer.GetData();
-    if (custom_option_alignment_.has_value()) {
-      builder_.ForceVectorAlignment(buffer_data.size(), sizeof(uint8_t),
-                                    custom_option_alignment_.value());
-    }
+    builder_.ForceVectorAlignment(
+        buffer_data.size(), sizeof(uint8_t),
+        custom_option_alignment_.value_or(kFbAlignment));
     ReleaseBufferData(buffer);
     return tflite::CreateBuffer(
         builder_, builder_.CreateVector(
@@ -3969,6 +3968,7 @@ std::optional<BufferOffset<tflite::SubGraph>> Translator::BuildSubGraph(
 BufferOffset<tflite::Metadata> Translator::BuildMetadata(StringRef name,
                                                          StringRef content) {
   auto buffer_index = buffers_.size();
+  builder_.ForceVectorAlignment(content.size(), sizeof(uint8_t), kFbAlignment);
   auto buffer_data = builder_.CreateVector(
       reinterpret_cast<const uint8_t*>(content.data()), content.size());
   buffers_.push_back(tflite::CreateBuffer(builder_, buffer_data));

tensorflow/compiler/mlir/lite/flatbuffer_operator.h

@@ -128,7 +128,7 @@ void BuiltinOptionsToAttributes(
     // NOLINTNEXTLINE
     llvm::SmallVectorImpl<mlir::NamedAttribute>& attributes);
 
-// While the last several tensors could be optional tensors for an tfl op, the
+// While the last several tensors could be optional tensors for a tfl op, the
 // number of input operands could vary. This function gets the min/max number of
 // operands from tflite op name.
 llvm::MinMax OperandNumbersMinMax(llvm::StringRef op_name);

tensorflow/compiler/mlir/lite/ir/tfl_ops.cc

@@ -6018,11 +6018,55 @@ int64_t L2NormalizationOp::GetArithmeticCount(Operation* op) {
 // PadOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult PadOp::fold(FoldAdaptor) {
+OpFoldResult PadOp::fold(FoldAdaptor adaptor) {
   if (InputOutputHasSameShape(getInput().getType(), getOutput().getType()))
     return getInput();
 
-  return {};
+  // Fold a constant input padded by constant amounts. Left unfolded, a
+  // constant-input PAD survives into the flatbuffer, and delegates that
+  // require a runtime input (e.g. ML Drift) reject it.
+  if (!ShouldFoldOperation(this->getOperation())) return {};
+
+  auto input = dyn_cast_or_null<DenseElementsAttr>(adaptor.getInput());
+  if (!input) return {};
+  auto output_type = mlir::dyn_cast<RankedTensorType>(getOutput().getType());
+  if (!output_type || !output_type.hasStaticShape()) return {};
+  // DenseElementsAttr does not support quantized element types.
+  Type element_type = output_type.getElementType();
+  if (!element_type.isSignlessIntOrFloat()) return {};
+
+  auto input_type = mlir::cast<ShapedType>(input.getType());
+  const int64_t rank = input_type.getRank();
+  const llvm::SmallVector<int64_t> paddings =
+      UnpackIndexOperand(adaptor.getPadding());
+  if (paddings.size() != 2 * rank) return {};
+  for (int64_t d = 0; d < rank; ++d) {
+    const int64_t before = paddings[2 * d], after = paddings[2 * d + 1];
+    if (before < 0 || after < 0 ||
+        input_type.getDimSize(d) + before + after !=
+            output_type.getDimSize(d)) {
+      return {};
+    }
+  }
+
+  // TFL PAD fills with zero (PADV2 takes an explicit value).
+  Attribute zero = mlir::isa<FloatType>(element_type)
+                       ? Attribute(FloatAttr::get(element_type, 0.0))
+                       : Attribute(IntegerAttr::get(element_type, 0));
+  std::vector<Attribute> result(output_type.getNumElements(), zero);
+
+  const FlatIndHelper read_inds(input_type);
+  const FlatIndHelper write_inds(output_type);
+  llvm::SmallVector<int64_t> offset(rank);
+  for (int64_t d = 0; d < rank; ++d) offset[d] = paddings[2 * d];
+
+  int64_t read_flat_ind = 0;
+  for (Attribute value : input.getValues<Attribute>()) {
+    llvm::SmallVector<int64_t> ind = read_inds.GetShapedInd(read_flat_ind++);
+    FlatIndHelper::AddOffset(ind, offset);
+    result[write_inds.GetFlatInd(ind)] = value;
+  }
+  return DenseElementsAttr::get(output_type, result);
 }
 
 // When padding amounts are constants, cast them to i32. XNN can only

tensorflow/compiler/mlir/lite/python/BUILD

@@ -117,6 +117,7 @@ cc_library(
         ":conversion_failure_reporter",
         "//tensorflow/compiler/mlir/lite:common",
         "//tensorflow/compiler/mlir/lite:converter_flags_proto_cc",
+        "//tensorflow/compiler/mlir/lite:downcast_x64",
         "//tensorflow/compiler/mlir/lite:flatbuffer_export",
         "//tensorflow/compiler/mlir/lite:optimize_broadcast_like_pass_options",
         "//tensorflow/compiler/mlir/lite:tensorflow_lite",

tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc

@@ -59,6 +59,7 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
 #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h"
 #include "tensorflow/compiler/mlir/lite/transforms/cast_bf16_ops_to_f32_pass.h"
+#include "tensorflow/compiler/mlir/lite/transforms/downcast_x64_pass.h"
 #include "tensorflow/compiler/mlir/lite/transforms/large_constant_fold_pass.h"
 #include "tensorflow/compiler/mlir/lite/transforms/optimize_broadcast_like_pass.h"
 #include "tensorflow/compiler/mlir/lite/transforms/optimize_broadcast_like_pass_options.h"
@@ -140,7 +141,8 @@ std::unique_ptr<mlir::Pass> CreatePruneDeadResourcesPass() {
 }  // namespace
 
 void AddPipelinePasses(mlir::OpPassManager& pass_manager,
-                       const mlir::TFL::PassConfig& pass_config) {
+                       const mlir::TFL::PassConfig& pass_config,
+                       const tflite::ConverterFlags& converter_flags) {
   // =========================================================================
   // 1. Skip-to-TFLite & Pre-Lowering Passes
   // =========================================================================
@@ -303,10 +305,20 @@ void AddPipelinePasses(mlir::OpPassManager& pass_manager,
   pass_manager.addPass(CreatePruneDeadResourcesPass());
   pass_manager.addNestedPass<mlir::func::FuncOp>(
       mlir::TFL::CreateFuseA4W2DRQFullyConnectedPass());
+  if (pass_config.fold_fp16_casts_into_fully_connected) {
+    pass_manager.addNestedPass<mlir::func::FuncOp>(
+        mlir::TFL::CreateFoldFp16CastsIntoFullyConnectedPass());
+  }
   pass_manager.addNestedPass<mlir::func::FuncOp>(
       mlir::createCanonicalizerPass());
   pass_manager.addPass(mlir::createSymbolDCEPass());
   pass_manager.addPass(mlir::TFL::CreateCleanupOptimizationBarrierPass());
+  if (!converter_flags.enable_x64()) {
+    pass_manager.addNestedPass<mlir::func::FuncOp>(
+        mlir::TFL::CreateDowncastX64Pass());
+    pass_manager.addNestedPass<mlir::func::FuncOp>(
+        mlir::createCanonicalizerPass());
+  }
   pass_manager.addPass(mlir::odml::createLegalizeStablehloToVhloPass());
   pass_manager.addPass(mlir::createReconcileUnrealizedCastsPass());
   pass_manager.addPass(CreatePruneDeadResourcesPass());
@@ -426,7 +438,7 @@ absl::Status ConvertStableHloToTFLite(
       tensorflow::InitPassManager(pm, converter_flags.debug_options());
     }
 
-    AddPipelinePasses(pm, pass_config);
+    AddPipelinePasses(pm, pass_config, converter_flags);
 
     mlir::StatusScopedDiagnosticHandler status_handler(context,
                                                        /*propagate=*/false);

tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc

@@ -394,6 +394,8 @@ absl::Status ConvertMlirBytecodeToTFLite(
   pass_config.fold_fp16_resource_casts =
       converter_flags.fold_fp16_resource_casts();
   pass_config.unfold_batch_matmul = converter_flags.unfold_batchmatmul();
+  pass_config.fold_fp16_casts_into_fully_connected =
+      converter_flags.fold_fp16_casts_into_fully_connected();
 
   return mlir::TFL::ConvertStableHloToTFLite(*module, converter_flags,
                                              pass_config, export_stream);

tensorflow/compiler/mlir/lite/quantization/lite/quantize_model_test.cc

@@ -103,6 +103,14 @@ absl::Status QuantizeModel(
 
   auto flatbuffer_model = FlatBufferModelAbslError::BuildFromBuffer(
       output_buffer.data(), output_buffer.size());
+  for (const auto* buffer : *flatbuffer_model->GetModel()->buffers()) {
+    if (buffer != nullptr && buffer->data() != nullptr) {
+      EXPECT_EQ((buffer->data()->Data() -
+                 reinterpret_cast<const uint8_t*>(output_buffer.data())) %
+                    16,
+                0);
+    }
+  }
   *model = UnPackFlatBufferModel(*flatbuffer_model->GetModel());
   return absl::OkStatus();
 }

tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/BUILD

@@ -111,8 +111,12 @@ cc_library(
     ],
     deps = [
         "//tensorflow/compiler/mlir/lite:tensorflow_lite",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite_quantize",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:AsmParser",
         "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:QuantOps",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TransformUtils",
         "@xla//xla/mlir_hlo",

tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/custom_call.cc

@@ -15,16 +15,29 @@ limitations under the License.
 
 #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/custom_call.h"
 
+#include <algorithm>
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
 #include <optional>
 
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
 #include "mlir/AsmParser/AsmParser.h"  // from @llvm-project
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
 #include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
 #include "mlir/IR/OperationSupport.h"  // from @llvm-project
 #include "mlir/IR/PatternMatch.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "mlir/Support/LogicalResult.h"  // from @llvm-project
 #include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h"
 #include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
 
 namespace mlir {
@@ -61,10 +74,242 @@ mlir::DictionaryAttr ParseSerializedTFLOpAttributes(
   return dict_attribute;
 }
 
+// Quantization parameters carried in the backend_config of a `tfl.quantize` /
+// `tfl.dequantize` custom call.
+struct QuantParams {
+  SmallVector<double, 4> scales;
+  SmallVector<int64_t, 4> zero_points;
+  int num_bits = 8;
+  bool is_signed = true;
+  bool narrow_range = false;
+  int32_t quantized_dimension = 0;
+};
+
+QuantParams ParseQuantParams(DictionaryAttr attributes) {
+  QuantParams params;
+
+  if (auto scale_attr = attributes.getAs<mlir::FloatAttr>("scale")) {
+    params.scales.push_back(scale_attr.getValueAsDouble());
+  } else if (auto scale_array = attributes.getAs<mlir::ArrayAttr>("scale")) {
+    for (auto elem : scale_array) {
+      if (auto f_attr = mlir::dyn_cast<mlir::FloatAttr>(elem)) {
+        params.scales.push_back(f_attr.getValueAsDouble());
+      } else if (auto i_attr = mlir::dyn_cast<mlir::IntegerAttr>(elem)) {
+        params.scales.push_back(static_cast<double>(i_attr.getInt()));
+      }
+    }
+  } else if (auto scale_dense =
+                 attributes.getAs<mlir::DenseFPElementsAttr>("scale")) {
+    for (auto val : scale_dense.getValues<APFloat>()) {
+      params.scales.push_back(val.convertToDouble());
+    }
+  } else if (auto scale_int = attributes.getAs<mlir::IntegerAttr>("scale")) {
+    params.scales.push_back(static_cast<double>(scale_int.getInt()));
+  }
+
+  if (auto zp_attr = attributes.getAs<mlir::IntegerAttr>("zero_point")) {
+    params.zero_points.push_back(zp_attr.getInt());
+  } else if (auto zp_array = attributes.getAs<mlir::ArrayAttr>("zero_point")) {
+    for (auto elem : zp_array) {
+      if (auto i_attr = mlir::dyn_cast<mlir::IntegerAttr>(elem)) {
+        params.zero_points.push_back(i_attr.getInt());
+      }
+    }
+  } else if (auto zp_dense =
+                 attributes.getAs<mlir::DenseIntElementsAttr>("zero_point")) {
+    for (auto val : zp_dense.getValues<APInt>()) {
+      params.zero_points.push_back(val.getSExtValue());
+    }
+  }
+
+  if (auto num_bits_attr = attributes.getAs<mlir::IntegerAttr>("num_bits")) {
+    params.num_bits = num_bits_attr.getInt();
+  }
+  if (auto is_signed_attr = attributes.getAs<mlir::BoolAttr>("is_signed")) {
+    params.is_signed = is_signed_attr.getValue();
+  }
+  if (auto narrow_range_attr =
+          attributes.getAs<mlir::BoolAttr>("narrow_range")) {
+    params.narrow_range = narrow_range_attr.getValue();
+  }
+  if (auto dim_attr =
+          attributes.getAs<mlir::IntegerAttr>("quantized_dimension")) {
+    params.quantized_dimension = dim_attr.getInt();
+  }
+
+  if (params.scales.empty()) {
+    params.scales.push_back(1.0);
+  }
+  if (params.zero_points.empty()) {
+    params.zero_points.push_back(0);
+  }
+  return params;
+}
+
+// Builds the quantized tensor type of `shape` described by `params`. Returns
+// null if `params.quantized_dimension` is out of range for `shape`.
+RankedTensorType BuildQuantizedTensorType(Builder& builder,
+                                          const QuantParams& params,
+                                          ArrayRef<int64_t> shape,
+                                          Type expressed_type, Location loc) {
+  Type quantized_element_type;
+  if (params.scales.size() == 1) {
+    quantized_element_type = TFL::GetPerTensorQuantizedTensorType(
+        builder, params.scales[0], params.zero_points[0], expressed_type,
+        params.num_bits, loc, params.narrow_range, params.is_signed);
+  } else {
+    const int64_t rank = static_cast<int64_t>(shape.size());
+    int64_t quant_dim = params.quantized_dimension;
+    if (quant_dim < 0) {
+      quant_dim += rank;
+    }
+    if (quant_dim < 0 || quant_dim >= rank) {
+      return nullptr;
+    }
+    quantized_element_type = TFL::GetPerAxisQuantizedTensorType(
+        builder, params.scales, params.zero_points,
+        static_cast<int32_t>(quant_dim), expressed_type, params.num_bits, loc,
+        params.narrow_range, params.is_signed);
+  }
+  if (!quantized_element_type) {
+    return nullptr;
+  }
+  return RankedTensorType::get(shape, quantized_element_type);
+}
+
+bool ApproxEqual(double a, double b) {
+  return std::abs(a - b) <= 1e-6 * std::max(std::abs(a), std::abs(b));
+}
+
+// Returns true if `qtype` has the scales and zero points in `params`.
+bool MatchesScalesAndZeroPoints(quant::QuantizedType qtype,
+                                const QuantParams& params) {
+  if (auto uniform = mlir::dyn_cast<quant::UniformQuantizedType>(qtype)) {
+    return params.scales.size() == 1 &&
+           ApproxEqual(uniform.getScale(), params.scales[0]) &&
+           uniform.getZeroPoint() == params.zero_points[0];
+  }
+  if (auto per_axis =
+          mlir::dyn_cast<quant::UniformQuantizedPerAxisType>(qtype)) {
+    ArrayRef<double> scales = per_axis.getScales();
+    ArrayRef<int64_t> zero_points = per_axis.getZeroPoints();
+    if (scales.size() != params.scales.size() ||
+        zero_points.size() != params.zero_points.size()) {
+      return false;
+    }
+    for (size_t i = 0; i < scales.size(); ++i) {
+      if (!ApproxEqual(scales[i], params.scales[i]) ||
+          zero_points[i] != params.zero_points[i]) {
+        return false;
+      }
+    }
+    return true;
+  }
+  return false;
+}
+
+// Rewrites a `tfl.quantize` custom call on `input` into `tfl.quantize`.
+LogicalResult RewriteQuantizeCustomCall(mhlo::CustomCallOp op, Value input,
+                                        PatternRewriter& rewriter) {
+  DictionaryAttr attributes =
+      ParseSerializedTFLOpAttributes(op.getBackendConfig(), op.getContext());
+  if (!attributes) {
+    return failure();
+  }
+  QuantParams params = ParseQuantParams(attributes);
+
+  ShapedType input_shaped_type = mlir::cast<ShapedType>(input.getType());
+  RankedTensorType output_type =
+      BuildQuantizedTensorType(rewriter, params, input_shaped_type.getShape(),
+                               input_shaped_type.getElementType(), op.getLoc());
+  if (!output_type) {
+    return op.emitError("invalid quantization parameters for tfl.quantize");
+  }
+  auto quant_op = rewriter.create<TFL::QuantizeOp>(
+      op.getLoc(), output_type, input, TypeAttr::get(output_type));
+  rewriter.replaceOp(op, quant_op.getOutput());
+  return success();
+}
+
+// Rewrites a `tfl.dequantize` custom call on `input` into `tfl.dequantize`.
+// - An already-quantized input is dequantized with its own type, which must
+//   match the scales and zero points in the backend_config.
+// - An integer constant input becomes a `tfl.pseudo_qconst` with the quantized
+//   type described by the backend_config.
+// Any other input (e.g. a runtime integer tensor) can't be given a quantized
+// type and is rejected.
+LogicalResult RewriteDequantizeCustomCall(mhlo::CustomCallOp op, Value input,
+                                          PatternRewriter& rewriter) {
+  DictionaryAttr attributes =
+      ParseSerializedTFLOpAttributes(op.getBackendConfig(), op.getContext());
+  if (!attributes) {
+    return failure();
+  }
+
+  if (auto call_op = input.getDefiningOp<mhlo::CustomCallOp>()) {
+    if (call_op.getCallTargetName() == "tfl.quantize") {
+      return failure();
+    }
+  }
+  ShapedType input_shaped_type = mlir::cast<ShapedType>(input.getType());
+  Type output_type = op.getResultTypes().front();
+  Type expressed_type = mlir::cast<ShapedType>(output_type).getElementType();
+  QuantParams params = ParseQuantParams(attributes);
+
+  Value tfl_dequant_input = input;
+  if (auto qtype = mlir::dyn_cast<quant::QuantizedType>(
+          input_shaped_type.getElementType())) {
+    if (!MatchesScalesAndZeroPoints(qtype, params)) {
+      return op.emitError(
+          "tfl.dequantize scale/zero_point don't match the quantized type of "
+          "its input");
+    }
+  } else {
+    ElementsAttr const_value;
+    Location const_loc = op.getLoc();
+    if (auto const_op = input.getDefiningOp<arith::ConstantOp>()) {
+      const_value = mlir::dyn_cast<ElementsAttr>(const_op.getValue());
+      const_loc = const_op.getLoc();
+    } else if (auto const_op = input.getDefiningOp<mhlo::ConstantOp>()) {
+      const_value = mlir::dyn_cast<ElementsAttr>(const_op.getValue());
+      const_loc = const_op.getLoc();
+    }
+    if (!const_value ||
+        !mlir::isa<mlir::IntegerType>(const_value.getElementType())) {
+      return rewriter.notifyMatchFailure(
+          op, "tfl.dequantize input must be quantized or an integer constant");
+    }
+    RankedTensorType qtensor_type =
+        BuildQuantizedTensorType(rewriter, params, input_shaped_type.getShape(),
+                                 expressed_type, op.getLoc());
+    if (!qtensor_type) {
+      return op.emitError("invalid quantization parameters for tfl.dequantize");
+    }
+    auto qconst_op = rewriter.create<TFL::QConstOp>(
+        const_loc, qtensor_type, TypeAttr::get(qtensor_type), const_value);
+    tfl_dequant_input = qconst_op.getResult();
+  }
+
+  auto dequant_op = rewriter.create<TFL::DequantizeOp>(op.getLoc(), output_type,
+                                                       tfl_dequant_input);
+  rewriter.replaceOp(op, dequant_op.getOutput());
+  return success();
+}
+
 LogicalResult ConvertCustomCallOp::matchAndRewrite(
     mhlo::CustomCallOp mhlo_custom_call, OpAdaptor adaptor,
     ConversionPatternRewriter& rewriter) const {
   auto call_target_name = mhlo_custom_call.getCallTargetName();
+
+  if (call_target_name == "tfl.quantize") {
+    return RewriteQuantizeCustomCall(mhlo_custom_call,
+                                     adaptor.getOperands().front(), rewriter);
+  }
+  if (call_target_name == "tfl.dequantize") {
+    return RewriteDequantizeCustomCall(mhlo_custom_call,
+                                       adaptor.getOperands().front(), rewriter);
+  }
+
   if (call_target_name.starts_with("tfl.")) {
     auto bc = mhlo_custom_call.getBackendConfig();
     if (mlir::DictionaryAttr attributes =
@@ -101,6 +346,34 @@ LogicalResult ConvertCustomCallOp::matchAndRewrite(
   return success();
 }
 
+class RewriteQuantizeCustomCallOp
+    : public OpRewritePattern<mhlo::CustomCallOp> {
+ public:
+  using OpRewritePattern<mhlo::CustomCallOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mhlo::CustomCallOp op,
+                                PatternRewriter& rewriter) const override {
+    if (op.getCallTargetName() != "tfl.quantize") {
+      return failure();
+    }
+    return RewriteQuantizeCustomCall(op, op.getOperand(0), rewriter);
+  }
+};
+
+class RewriteDequantizeCustomCallOp
+    : public OpRewritePattern<mhlo::CustomCallOp> {
+ public:
+  using OpRewritePattern<mhlo::CustomCallOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mhlo::CustomCallOp op,
+                                PatternRewriter& rewriter) const override {
+    if (op.getCallTargetName() != "tfl.dequantize") {
+      return failure();
+    }
+    return RewriteDequantizeCustomCall(op, op.getOperand(0), rewriter);
+  }
+};
+
 // Removes the `mhlo.custom_call @shape_assertion` custom call which represents
 // an assertion that the first operand (`assert_what`) evaluates to `true`.
 // This is a temporary workaround for unblocking dynamic model conversion
@@ -132,6 +405,61 @@ LogicalResult RemoveCustomCallWithShapeAssertion::matchAndRewrite(
   return success();
 }
 
+// TFL ops whose quantized output uses the same scale and zero point as their
+// first operand: they only move, select or replicate elements.
+constexpr llvm::StringLiteral kScalePreservingTflOps[] = {
+    "tfl.broadcast_to", "tfl.depth_to_space", "tfl.expand_dims",
+    "tfl.gather",       "tfl.gather_nd",      "tfl.reshape",
+    "tfl.reverse_v2",   "tfl.slice",          "tfl.space_to_depth",
+    "tfl.squeeze",      "tfl.strided_slice",  "tfl.tile",
+    "tfl.transpose",
+};
+
+// Gives a scale-preserving `tfl.*` custom call the quantized type of its first
+// operand, so that e.g. `reshape(quantize(x))` stays quantized. Ops that
+// compute new values (FC, BMM, add, ...) need their own output scale and are
+// deliberately not handled here.
+class PropagateQuantizedTypeToCustomCallOp
+    : public OpRewritePattern<mhlo::CustomCallOp> {
+ public:
+  using OpRewritePattern<mhlo::CustomCallOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mhlo::CustomCallOp op,
+                                PatternRewriter& rewriter) const override {
+    if (!llvm::is_contained(kScalePreservingTflOps, op.getCallTargetName())) {
+      return failure();
+    }
+    if (op.getNumOperands() == 0 || op.getNumResults() != 1) {
+      return failure();
+    }
+    auto operand_type = mlir::dyn_cast<ShapedType>(op.getOperand(0).getType());
+    if (!operand_type) {
+      return failure();
+    }
+    auto qtype =
+        mlir::dyn_cast<quant::QuantizedType>(operand_type.getElementType());
+    if (!qtype) {
+      return failure();
+    }
+    auto res_shaped_type =
+        mlir::dyn_cast<ShapedType>(op.getResult(0).getType());
+    if (!res_shaped_type) {
+      return failure();
+    }
+    // Only retype results that already hold the storage type, e.g. i8 -> !quant
+    // with i8 storage.
+    auto res_int_type =
+        mlir::dyn_cast<mlir::IntegerType>(res_shaped_type.getElementType());
+    if (!res_int_type ||
+        res_int_type.getWidth() != qtype.getStorageTypeIntegralWidth()) {
+      return failure();
+    }
+    rewriter.modifyOpInPlace(
+        op, [&] { op.getResult(0).setType(res_shaped_type.clone(qtype)); });
+    return success();
+  }
+};
+
 std::optional<bool> IsCustomCallLegal(mhlo::CustomCallOp op) {
   auto call_target_name = op.getCallTargetName();
   if (call_target_name.starts_with("custom_call.")) {
@@ -160,6 +488,9 @@ void PopulateCustomCallPatterns(MLIRContext* ctx, RewritePatternSet& patterns,
 void PopulateCustomCallPreparePatterns(MLIRContext* ctx,
                                        RewritePatternSet& patterns) {
   patterns.add<RemoveCustomCallWithShapeAssertion>(ctx);
+  patterns.add<RewriteQuantizeCustomCallOp>(ctx);
+  patterns.add<RewriteDequantizeCustomCallOp>(ctx);
+  patterns.add<PropagateQuantizedTypeToCustomCallOp>(ctx);
 }
 
 }  // namespace odml

tensorflow/compiler/mlir/lite/tests/cast_bf16_ops_to_f32.mlir

@@ -101,3 +101,15 @@ func.func private @test.impl(%arg0: tensor<2x2xbf16>) -> tensor<2x2xbf16> {
   %0 = "tfl.add"(%arg0, %arg0) <{fused_activation_function = "NONE"}> : (tensor<2x2xbf16>, tensor<2x2xbf16>) -> tensor<2x2xbf16>
   return %0 : tensor<2x2xbf16>
 }
+
+// -----
+
+// CHECK-LABEL: func.func @keep_non_bf16_cast_round_trip
+func.func @keep_non_bf16_cast_round_trip(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
+  // CHECK: %[[NARROW:.*]] = "tfl.cast"(%arg0) : (tensor<2x2xf32>) -> tensor<2x2xf16>
+  // CHECK: %[[WIDEN:.*]] = "tfl.cast"(%[[NARROW]]) : (tensor<2x2xf16>) -> tensor<2x2xf32>
+  // CHECK: return %[[WIDEN]] : tensor<2x2xf32>
+  %0 = "tfl.cast"(%arg0) : (tensor<2x2xf32>) -> tensor<2x2xf16>
+  %1 = "tfl.cast"(%0) : (tensor<2x2xf16>) -> tensor<2x2xf32>
+  return %1 : tensor<2x2xf32>
+}

tensorflow/compiler/mlir/lite/tests/const-fold.mlir

@@ -537,6 +537,37 @@ func.func @ConstFoldPad(%arg0: tensor<15600xf32>) -> tensor<15600xf32> {
   // CHECK:  return %arg0
 }
 
+// CHECK-LABEL: @ConstFoldPadConstInput
+func.func @ConstFoldPadConstInput() -> tensor<3x4xf32> {
+  %0 = "tfl.pseudo_const"() {value = dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf32>} : () -> tensor<2x2xf32>
+  %1 = "tfl.pseudo_const"() {value = dense<[[1, 0], [0, 2]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
+  %2 = "tfl.pad"(%0, %1) : (tensor<2x2xf32>, tensor<2x2xi32>) -> tensor<3x4xf32>
+  func.return %2 : tensor<3x4xf32>
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00], [1.000000e+00, 2.000000e+00, 0.000000e+00, 0.000000e+00], [3.000000e+00, 4.000000e+00, 0.000000e+00, 0.000000e+00]]> : tensor<3x4xf32>
+  // CHECK-NOT: tfl.pad
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstFoldPadConstInputI64Paddings
+func.func @ConstFoldPadConstInputI64Paddings() -> tensor<4xi32> {
+  %0 = "tfl.pseudo_const"() {value = dense<[7, 8]> : tensor<2xi32>} : () -> tensor<2xi32>
+  %1 = "tfl.pseudo_const"() {value = dense<[[1, 1]]> : tensor<1x2xi64>} : () -> tensor<1x2xi64>
+  %2 = "tfl.pad"(%0, %1) : (tensor<2xi32>, tensor<1x2xi64>) -> tensor<4xi32>
+  func.return %2 : tensor<4xi32>
+  // CHECK: %[[CST:.*]] = arith.constant dense<[0, 7, 8, 0]> : tensor<4xi32>
+  // CHECK-NOT: tfl.pad
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @NotFoldPadRuntimeInput
+func.func @NotFoldPadRuntimeInput(%arg0: tensor<2x2xf32>) -> tensor<3x4xf32> {
+  %0 = "tfl.pseudo_const"() {value = dense<[[1, 0], [0, 2]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
+  %1 = "tfl.pad"(%arg0, %0) : (tensor<2x2xf32>, tensor<2x2xi32>) -> tensor<3x4xf32>
+  func.return %1 : tensor<3x4xf32>
+  // CHECK: %[[PAD:.*]] = "tfl.pad"(%arg0
+  // CHECK: return %[[PAD]]
+}
+
 func.func @ConstFoldPadV2(%arg0: tensor<15600xf32>) -> tensor<15600xf32> {
   %0 = "tfl.pseudo_const"() {value = dense<0> : tensor<1x2xi32>} : () -> tensor<1x2xi32>
   %1 = "tfl.pseudo_const"() {value = dense<0.0> : tensor<f32>} : () -> tensor<f32>

tensorflow/compiler/mlir/lite/tests/downcast-x64.mlir

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 // ==============================================================================
-// RUN: litert-opt %s --tfl-downcast-x64 --canonicalize | FileCheck %s
+// RUN: litert-opt %s --tfl-downcast-x64 --canonicalize -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL: testFuncSignature
 // CHECK: (%arg0: tensor<4xi32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32>
@@ -37,10 +37,140 @@ func.func @testI64ConstantDowncast() -> tensor<i64> {
   func.return %0 : tensor<i64>
 }
 
+// INT64_MAX / INT64_MIN sentinels saturate to INT32_MAX / INT32_MIN instead of
+// wrapping to -1 / 0.
+// CHECK-LABEL: testI64ConstantSaturates
+func.func @testI64ConstantSaturates() -> tensor<4xi64> {
+  // CHECK: %[[CST:.*]] = arith.constant dense<[2147483647, -2147483648, 2147483647, -2147483648]> : tensor<4xi32>
+  // CHECK: return %[[CST]]
+  // expected-warning @+1 {{saturating them to [INT32_MIN, INT32_MAX]}}
+  %0 = arith.constant dense<[9223372036854775807, -9223372036854775808, 2147483648, -2147483649]> : tensor<4xi64>
+  func.return %0 : tensor<4xi64>
+}
+
+// CHECK-LABEL: testI64SplatConstantSaturates
+func.func @testI64SplatConstantSaturates() -> tensor<2x3xi64> {
+  // CHECK: %[[CST:.*]] = arith.constant dense<2147483647> : tensor<2x3xi32>
+  // CHECK: return %[[CST]]
+  // expected-warning @+1 {{saturating them to [INT32_MIN, INT32_MAX]}}
+  %0 = arith.constant dense<9223372036854775807> : tensor<2x3xi64>
+  func.return %0 : tensor<2x3xi64>
+}
+
+// CHECK-LABEL: testTflConstSaturates
+func.func @testTflConstSaturates() -> tensor<2xi64> {
+  // CHECK: %[[CST:.*]] = arith.constant dense<[-2147483648, 7]> : tensor<2xi32>
+  // CHECK: return %[[CST]]
+  // expected-warning @+1 {{saturating them to [INT32_MIN, INT32_MAX]}}
+  %0 = "tfl.pseudo_const"() {value = dense<[-9223372036854775808, 7]> : tensor<2xi64>} : () -> tensor<2xi64>
+  func.return %0 : tensor<2xi64>
+}
+
+// Unsigned constants clamp to [0, UINT32_MAX] with unsigned comparisons: a
+// value with the top bit set saturates to UINT32_MAX rather than being read as
+// negative and clamped to the bottom of the range. (arith.constant requires
+// signless types, so these use tfl.pseudo_const.)
+// CHECK-LABEL: testUI64ConstantSaturates
+func.func @testUI64ConstantSaturates() -> tensor<3xui64> {
+  // CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[4294967295, 4294967295, 7]> : tensor<3xui32>}> : () -> tensor<3xui32>
+  // CHECK: return %[[CST]]
+  // expected-warning @+1 {{saturating them to [0, UINT32_MAX]}}
+  %0 = "tfl.pseudo_const"() {value = dense<[18446744073709551615, 4294967296, 7]> : tensor<3xui64>} : () -> tensor<3xui64>
+  func.return %0 : tensor<3xui64>
+}
+
+// Unsigned values that fit in 32 bits but exceed INT32_MAX are kept.
+// CHECK-LABEL: testUI64ConstantInRangeNoWarning
+func.func @testUI64ConstantInRangeNoWarning() -> tensor<2xui64> {
+  // CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[4294967295, 2147483648]> : tensor<2xui32>}> : () -> tensor<2xui32>
+  // CHECK: return %[[CST]]
+  %0 = "tfl.pseudo_const"() {value = dense<[4294967295, 2147483648]> : tensor<2xui64>} : () -> tensor<2xui64>
+  func.return %0 : tensor<2xui64>
+}
+
+// ui64 function arguments narrow to ui32, preserving signedness.
+// CHECK-LABEL: testUI64FuncSignature
+// CHECK: (%arg0: tensor<4xui32>) -> tensor<4xui32>
+func.func @testUI64FuncSignature(%arg0: tensor<4xui64>) -> tensor<4xui64> {
+  // CHECK: return %arg0 : tensor<4xui32>
+  func.return %arg0 : tensor<4xui64>
+}
+
+// In-range values pass through unchanged and emit no warning.
+// CHECK-LABEL: testI64ConstantInRangeNoWarning
+func.func @testI64ConstantInRangeNoWarning() -> tensor<2xi64> {
+  // CHECK: %[[CST:.*]] = arith.constant dense<[2147483647, -2147483648]> : tensor<2xi32>
+  // CHECK: return %[[CST]]
+  %0 = arith.constant dense<[2147483647, -2147483648]> : tensor<2xi64>
+  func.return %0 : tensor<2xi64>
+}
+
 // CHECK-LABEL: testGeneralOpForceConvert
 func.func @testGeneralOpForceConvert(%arg0: tensor<2xf64>) -> tensor<2xf64> {
   // CHECK: %[[ADD:.*]] = tfl.add %arg0, %arg0 {fused_activation_function = "NONE"} : tensor<2xf32>
   // CHECK: return %[[ADD]]
   %0 = "tfl.add"(%arg0, %arg0) {fused_activation_function = "NONE"} : (tensor<2xf64>, tensor<2xf64>) -> tensor<2xf64>
   func.return %0 : tensor<2xf64>
-}
+}
+
+// Block arguments of nested regions (here a region-based tfl.while, which is
+// only outlined into functions after this pass runs) are narrowed along with
+// the operands feeding them.
+// CHECK-LABEL: testWhileRegionBlockArgs
+// CHECK-SAME: (%arg0: tensor<i32>) -> tensor<i32>
+func.func @testWhileRegionBlockArgs(%arg0: tensor<i64>) -> tensor<i64> {
+  // CHECK: "tfl.while"(%arg0) ({
+  // CHECK: ^bb0(%[[CARG:.*]]: tensor<i32>):
+  // CHECK: tfl.less(%[[CARG]], %{{.*}}) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  // CHECK: ^bb0(%[[BARG:.*]]: tensor<i32>):
+  // CHECK: tfl.add %[[BARG]], %{{.*}} {fused_activation_function = "NONE"} : tensor<i32>
+  // CHECK: }) : (tensor<i32>) -> tensor<i32>
+  %0 = "tfl.while"(%arg0) ({
+    ^bb0(%carg: tensor<i64>):
+      %limit = arith.constant dense<10> : tensor<i64>
+      %cond = "tfl.less"(%carg, %limit) : (tensor<i64>, tensor<i64>) -> tensor<i1>
+      "tfl.yield"(%cond) : (tensor<i1>) -> ()
+  }, {
+    ^bb0(%barg: tensor<i64>):
+      %one = arith.constant dense<1> : tensor<i64>
+      %next = "tfl.add"(%barg, %one) {fused_activation_function = "NONE"} : (tensor<i64>, tensor<i64>) -> tensor<i64>
+      "tfl.yield"(%next) : (tensor<i64>) -> ()
+  }) : (tensor<i64>) -> tensor<i64>
+  func.return %0 : tensor<i64>
+}
+
+// Resource-backed constants (what lightweight_conversion produces for large
+// tensors) are narrowed into a new blob with the same saturation rules.
+// res_i64 = [INT64_MAX, INT64_MIN, 5].
+// CHECK-LABEL: testI64ResourceConstantSaturates
+func.func @testI64ResourceConstantSaturates() -> tensor<3xi64> {
+  // CHECK: %[[CST:.*]] = arith.constant dense_resource<res_i64_x32> : tensor<3xi32>
+  // CHECK: return %[[CST]]
+  // expected-warning @+1 {{saturating them to [INT32_MIN, INT32_MAX]}}
+  %0 = arith.constant dense_resource<res_i64> : tensor<3xi64>
+  func.return %0 : tensor<3xi64>
+}
+
+// res_ui64 = [UINT64_MAX, 7].
+// CHECK-LABEL: testUI64ResourceConstantSaturates
+func.func @testUI64ResourceConstantSaturates() -> tensor<2xui64> {
+  // CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense_resource<res_ui64_x32> : tensor<2xui32>}> : () -> tensor<2xui32>
+  // CHECK: return %[[CST]]
+  // expected-warning @+1 {{saturating them to [0, UINT32_MAX]}}
+  %0 = "tfl.pseudo_const"() {value = dense_resource<res_ui64> : tensor<2xui64>} : () -> tensor<2xui64>
+  func.return %0 : tensor<2xui64>
+}
+
+// The narrowed blobs: 64-byte alignment prefix, then
+// [INT32_MAX, INT32_MIN, 5] and [UINT32_MAX, 7] as little-endian 32-bit.
+// CHECK-DAG: res_i64_x32: "0x40000000FFFFFF7F0000008005000000"
+// CHECK-DAG: res_ui64_x32: "0x40000000FFFFFFFF07000000"
+
+{-#
+  dialect_resources: {
+    builtin: {
+      res_i64: "0x08000000FFFFFFFFFFFFFF7F00000000000000800500000000000000",
+      res_ui64: "0x08000000FFFFFFFFFFFFFFFF0700000000000000"
+    }
+  }
+#-}

tensorflow/compiler/mlir/lite/tests/fold-fp16-casts-into-fully-connected.mlir

@@ -0,0 +1,86 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+// RUN: litert-opt %s -tfl-fold-fp16-casts-into-fully-connected | FileCheck %s
+
+// CHECK-LABEL: @FoldCasts
+// CHECK-SAME: (%[[X:.*]]: tensor<4x64xf16>, %[[W:.*]]: tensor<32x64xf32>, %[[B:.*]]: tensor<32xf32>)
+func.func @FoldCasts(%arg0: tensor<4x64xf16>, %w: tensor<32x64xf32>, %b: tensor<32xf32>) -> tensor<4x32xf16> {
+  %0 = "tfl.cast"(%arg0) : (tensor<4x64xf16>) -> tensor<4x64xf32>
+  %1 = "tfl.fully_connected"(%0, %w, %b) {fused_activation_function = "RELU", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<4x64xf32>, tensor<32x64xf32>, tensor<32xf32>) -> tensor<4x32xf32>
+  %2 = "tfl.cast"(%1) : (tensor<4x32xf32>) -> tensor<4x32xf16>
+  func.return %2 : tensor<4x32xf16>
+  // CHECK-NOT: tfl.cast
+  // CHECK: %[[FC:.*]] = "tfl.fully_connected"(%[[X]], %[[W]], %[[B]]) <{fused_activation_function = "RELU", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<4x64xf16>, tensor<32x64xf32>, tensor<32xf32>) -> tensor<4x32xf16>
+  // CHECK-NOT: tfl.cast
+  // CHECK: return %[[FC]]
+}
+
+// The input cast feeding several fully_connected ops (e.g. q/k/v projections)
+// is folded into each of them.
+
+// CHECK-LABEL: @FoldSharedInputCast
+func.func @FoldSharedInputCast(%arg0: tensor<4x64xf16>, %w0: tensor<32x64xf32>, %w1: tensor<16x64xf32>) -> (tensor<4x32xf16>, tensor<4x16xf16>) {
+  %none = "tfl.no_value"() {value} : () -> none
+  %0 = "tfl.cast"(%arg0) : (tensor<4x64xf16>) -> tensor<4x64xf32>
+  %1 = "tfl.fully_connected"(%0, %w0, %none) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<4x64xf32>, tensor<32x64xf32>, none) -> tensor<4x32xf32>
+  %2 = "tfl.cast"(%1) : (tensor<4x32xf32>) -> tensor<4x32xf16>
+  %3 = "tfl.fully_connected"(%0, %w1, %none) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<4x64xf32>, tensor<16x64xf32>, none) -> tensor<4x16xf32>
+  %4 = "tfl.cast"(%3) : (tensor<4x16xf32>) -> tensor<4x16xf16>
+  func.return %2, %4 : tensor<4x32xf16>, tensor<4x16xf16>
+  // CHECK-NOT: tfl.cast
+  // CHECK: "tfl.fully_connected"(%arg0, %arg1, %{{.*}}) {{.*}} -> tensor<4x32xf16>
+  // CHECK: "tfl.fully_connected"(%arg0, %arg2, %{{.*}}) {{.*}} -> tensor<4x16xf16>
+  // CHECK-NOT: tfl.cast
+}
+
+// The f32 result has another user, so the fully_connected must stay f32.
+
+// CHECK-LABEL: @NoFoldOutputHasOtherUsers
+func.func @NoFoldOutputHasOtherUsers(%arg0: tensor<4x64xf16>, %w: tensor<32x64xf32>) -> (tensor<4x32xf16>, tensor<4x32xf32>) {
+  %none = "tfl.no_value"() {value} : () -> none
+  %0 = "tfl.cast"(%arg0) : (tensor<4x64xf16>) -> tensor<4x64xf32>
+  %1 = "tfl.fully_connected"(%0, %w, %none) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<4x64xf32>, tensor<32x64xf32>, none) -> tensor<4x32xf32>
+  %2 = "tfl.cast"(%1) : (tensor<4x32xf32>) -> tensor<4x32xf16>
+  func.return %2, %1 : tensor<4x32xf16>, tensor<4x32xf32>
+  // CHECK: "tfl.fully_connected"({{.*}}) {{.*}} -> tensor<4x32xf32>
+  // CHECK: "tfl.cast"
+}
+
+// The activation is not cast from f16, so there is nothing to fold.
+
+// CHECK-LABEL: @NoFoldInputNotFromF16
+func.func @NoFoldInputNotFromF16(%arg0: tensor<4x64xf32>, %w: tensor<32x64xf32>) -> tensor<4x32xf16> {
+  %none = "tfl.no_value"() {value} : () -> none
+  %0 = "tfl.fully_connected"(%arg0, %w, %none) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<4x64xf32>, tensor<32x64xf32>, none) -> tensor<4x32xf32>
+  %1 = "tfl.cast"(%0) : (tensor<4x32xf32>) -> tensor<4x32xf16>
+  func.return %1 : tensor<4x32xf16>
+  // CHECK: "tfl.fully_connected"({{.*}}) {{.*}} -> tensor<4x32xf32>
+  // CHECK: "tfl.cast"
+}
+
+// A bf16 input cast is a different precision contract and must not be folded.
+
+// CHECK-LABEL: @NoFoldBf16
+func.func @NoFoldBf16(%arg0: tensor<4x64xbf16>, %w: tensor<32x64xf32>) -> tensor<4x32xbf16> {
+  %none = "tfl.no_value"() {value} : () -> none
+  %0 = "tfl.cast"(%arg0) : (tensor<4x64xbf16>) -> tensor<4x64xf32>
+  %1 = "tfl.fully_connected"(%0, %w, %none) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<4x64xf32>, tensor<32x64xf32>, none) -> tensor<4x32xf32>
+  %2 = "tfl.cast"(%1) : (tensor<4x32xf32>) -> tensor<4x32xbf16>
+  func.return %2 : tensor<4x32xbf16>
+  // CHECK: "tfl.cast"
+  // CHECK: "tfl.fully_connected"({{.*}}) {{.*}} -> tensor<4x32xf32>
+  // CHECK: "tfl.cast"
+}

tensorflow/compiler/mlir/lite/transforms/cast_bf16_ops_to_f32_pass.cc

@@ -148,6 +148,12 @@ class RemoveUnneededCastOps : public OpRewritePattern<mlir::TFL::CastOp> {
     if (orig_input.getType() != op.getType()) {
       return failure();
     }
+    // Only fold round trips through bf16, which is what this pass introduces.
+    // Other narrowing round trips (e.g. f32 -> f16 -> f32) are explicit
+    // precision changes in the source program and must be preserved.
+    if (!getElementTypeOrSelf(prev_cast.getType()).isBF16()) {
+      return failure();
+    }
     rewriter.replaceOp(op, orig_input);
     return success();
   }

tensorflow/compiler/mlir/lite/transforms/downcast_x64_pass.cc

@@ -15,13 +15,27 @@ limitations under the License.
 
 #include "tensorflow/compiler/mlir/lite/transforms/downcast_x64_pass.h"
 
+#include <cstdint>
+#include <cstring>
+#include <limits>
 #include <memory>
+#include <string>
+#include <utility>
 
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/AsmState.h"  // from @llvm-project
+#include "mlir/IR/Block.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
 #include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
 #include "mlir/IR/Types.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
 #include "mlir/Pass/Pass.h"  // from @llvm-project
@@ -52,12 +66,152 @@ mlir::Type ConvertX64TypeToX32(mlir::Type type) {
   if (type.isF64()) {
     return mlir::Float32Type::get(type.getContext());
   }
-  if (type.isInteger(64)) {
-    return mlir::IntegerType::get(type.getContext(), 32);
+  // Preserve signedness so ui64 narrows to ui32 rather than signless i32.
+  if (auto int_type = mlir::dyn_cast<mlir::IntegerType>(type);
+      int_type && int_type.getWidth() == 64) {
+    return mlir::IntegerType::get(type.getContext(), 32,
+                                  int_type.getSignedness());
   }
   return type;
 }
 
+// Emits the saturation warning on `op`.
+void EmitSaturationWarning(mlir::Operation* op, bool is_unsigned) {
+  op->emitWarning()
+      << "enable_x64=false: 64-bit integer constant has values outside the "
+         "32-bit range; saturating them to "
+      << (is_unsigned ? "[0, UINT32_MAX]" : "[INT32_MIN, INT32_MAX]")
+      << ". Pass enable_x64=true to keep 64-bit types.";
+}
+
+// Narrows a `DenseIntElementsAttr` of 64-bit integers to 32 bits, clamping
+// out-of-range values. Sets `*clamped` if any value was clamped.
+mlir::DenseIntElementsAttr NarrowDenseI64(mlir::DenseIntElementsAttr value,
+                                          mlir::Type new_element_type,
+                                          bool* clamped) {
+  const bool is_unsigned = value.getElementType().isUnsignedInteger();
+  const APInt min_value =
+      is_unsigned ? APInt::getZero(64) : APInt::getSignedMinValue(32).sext(64);
+  const APInt max_value = is_unsigned ? APInt::getMaxValue(32).zext(64)
+                                      : APInt::getSignedMaxValue(32).sext(64);
+  auto less_than = [is_unsigned](const APInt& a, const APInt& b) {
+    return is_unsigned ? a.ult(b) : a.slt(b);
+  };
+  return mlir::cast<mlir::DenseIntElementsAttr>(
+      value.mapValues(new_element_type, [&](const APInt& v) -> APInt {
+        if (less_than(v, min_value)) {
+          *clamped = true;
+          return min_value.trunc(32);
+        }
+        if (less_than(max_value, v)) {
+          *clamped = true;
+          return max_value.trunc(32);
+        }
+        return v.trunc(32);
+      }));
+}
+
+// Narrows a `DenseResourceElementsAttr` of 64-bit integers to 32 bits by
+// allocating a new blob, clamping out-of-range values. Returns a null
+// attribute if the resource's blob is unavailable. Sets `*clamped` if any
+// value was clamped.
+mlir::DenseResourceElementsAttr NarrowResourceI64(
+    mlir::DenseResourceElementsAttr value, mlir::ShapedType new_type,
+    bool* clamped) {
+  mlir::AsmResourceBlob* blob = value.getRawHandle().getBlob();
+  if (!blob) {
+    if (auto* resource = value.getRawHandle().getResource()) {
+      blob = resource->getBlob();
+    }
+  }
+  if (!blob) return nullptr;
+
+  const bool is_unsigned = value.getElementType().isUnsignedInteger();
+  llvm::ArrayRef<char> src_bytes = blob->getData();
+  const int64_t num_elements = new_type.getNumElements();
+  if (src_bytes.size() != num_elements * sizeof(int64_t)) return nullptr;
+  mlir::AsmResourceBlob new_blob = mlir::HeapAsmResourceBlob::allocate(
+      num_elements * sizeof(int32_t), /*align=*/64, /*dataIsMutable=*/true);
+  llvm::MutableArrayRef<char> dst_bytes = new_blob.getMutableData();
+
+  // Copy element-wise through memcpy: the source blob's alignment is not
+  // guaranteed to be suitable for a direct int64_t view.
+  for (int64_t i = 0; i < num_elements; ++i) {
+    const char* src = src_bytes.data() + i * sizeof(int64_t);
+    char* dst = dst_bytes.data() + i * sizeof(int32_t);
+    if (is_unsigned) {
+      uint64_t v;
+      std::memcpy(&v, src, sizeof(v));
+      if (v > std::numeric_limits<uint32_t>::max()) {
+        *clamped = true;
+        v = std::numeric_limits<uint32_t>::max();
+      }
+      const uint32_t narrowed = static_cast<uint32_t>(v);
+      std::memcpy(dst, &narrowed, sizeof(narrowed));
+    } else {
+      int64_t v;
+      std::memcpy(&v, src, sizeof(v));
+      if (v < std::numeric_limits<int32_t>::min()) {
+        *clamped = true;
+        v = std::numeric_limits<int32_t>::min();
+      } else if (v > std::numeric_limits<int32_t>::max()) {
+        *clamped = true;
+        v = std::numeric_limits<int32_t>::max();
+      }
+      const int32_t narrowed = static_cast<int32_t>(v);
+      std::memcpy(dst, &narrowed, sizeof(narrowed));
+    }
+  }
+
+  std::string new_key =
+      absl::StrCat(value.getRawHandle().getKey().str(), "_x32");
+  return mlir::DenseResourceElementsAttr::get(new_type, new_key,
+                                              std::move(new_blob));
+}
+
+// If `op` is a 64-bit integer constant, narrows its `value` attribute and
+// result type to 32 bits in place and returns true. Values outside the 32-bit
+// range are clamped to that range first, with a warning, so the constant
+// saturates instead of wrapping. Handles both `DenseIntElementsAttr` and
+// `DenseResourceElementsAttr` (the latter is what `lightweight_conversion`
+// produces for large constants). Returns false for any other constant (e.g.
+// f64), leaving it for the cast path.
+//
+// Signed and signless constants clamp to [INT32_MIN, INT32_MAX]; unsigned
+// constants clamp to [0, UINT32_MAX] using unsigned comparisons, so a value
+// with the top bit set is treated as large rather than negative. Narrowing
+// directly (rather than through `tfl.cast`) also keeps ui64 working, since
+// `tfl.cast` does not accept a ui64 input.
+//
+// PyTorch exports routinely use INT64_MAX / INT64_MIN as sentinels (e.g. the
+// fill value of `torch.where(mask, x, iinfo(long).max).min()` or
+// `masked_fill(~m, iinfo(long).min).argmax()`). Truncating those to -1 / 0
+// would silently change the result; saturating keeps their meaning.
+bool NarrowI64Constant(mlir::Operation* op) {
+  if (op->getNumResults() != 1) return false;
+  auto value = op->getAttrOfType<mlir::ElementsAttr>("value");
+  if (!value || !value.getElementType().isInteger(64)) return false;
+  if (value.getType() != op->getResult(0).getType()) return false;
+
+  const bool is_unsigned = value.getElementType().isUnsignedInteger();
+  auto new_type =
+      mlir::cast<mlir::ShapedType>(ConvertX64TypeToX32(value.getType()));
+  bool clamped = false;
+  mlir::ElementsAttr new_value;
+  if (auto dense = mlir::dyn_cast<mlir::DenseIntElementsAttr>(value)) {
+    new_value = NarrowDenseI64(dense, new_type.getElementType(), &clamped);
+  } else if (auto resource =
+                 mlir::dyn_cast<mlir::DenseResourceElementsAttr>(value)) {
+    new_value = NarrowResourceI64(resource, new_type, &clamped);
+  }
+  if (!new_value) return false;
+
+  if (clamped) EmitSaturationWarning(op, is_unsigned);
+  op->setAttr("value", new_value);
+  op->getResult(0).setType(new_type);
+  return true;
+}
+
 }  // namespace
 
 void DowncastX64Pass::runOnOperation() {
@@ -97,8 +251,11 @@ void DowncastX64Pass::runOnOperation() {
       return;
     }
 
-    // Handle constant-like operations by inserting CastOps.
+    // Handle constant-like operations. 64-bit integer constants (dense or
+    // resource-backed) are narrowed in place; anything else gets a CastOp for
+    // the constant folder to resolve.
     if (op->hasTrait<mlir::OpTrait::ConstantLike>()) {
+      if (NarrowI64Constant(op)) return;
       for (mlir::OpResult result : op->getResults()) {
         mlir::Type old_type = result.getType();
         mlir::Type new_type = ConvertX64TypeToX32(old_type);
@@ -120,6 +277,17 @@ void DowncastX64Pass::runOnOperation() {
     for (mlir::Value result : op->getResults()) {
       result.setType(ConvertX64TypeToX32(result.getType()));
     }
+
+    // Update block arguments of any nested regions (e.g. the cond/body of a
+    // region-based `tfl.while`, or a residual `stablehlo.while`/`reduce`), so
+    // they stay consistent with the narrowed operands feeding them.
+    for (mlir::Region& region : op->getRegions()) {
+      for (mlir::Block& block : region) {
+        for (mlir::BlockArgument arg : block.getArguments()) {
+          arg.setType(ConvertX64TypeToX32(arg.getType()));
+        }
+      }
+    }
   });
 }
 

tensorflow/compiler/mlir/lite/transforms/fold_fp16_casts_into_fully_connected_pass.cc

@@ -0,0 +1,130 @@
+/* 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 <memory>
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
+
+namespace mlir {
+namespace TFL {
+namespace {
+
+#define GEN_PASS_DEF_FOLDFP16CASTSINTOFULLYCONNECTEDPASS
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"
+
+// Returns `value`'s type if it is a ranked tensor of exactly `element_type`.
+RankedTensorType GetRankedTensorOf(Value value, Type element_type) {
+  auto type = mlir::dyn_cast<RankedTensorType>(value.getType());
+  if (!type || type.getElementType() != element_type) return nullptr;
+  return type;
+}
+
+// Folds the f16 <-> f32 casts around a float fully_connected:
+//
+//   %x32 = tfl.cast(%x16) : f16 -> f32
+//   %y32 = tfl.fully_connected(%x32, %filter, %bias) : f32
+//   %y16 = tfl.cast(%y32) : f32 -> f16
+//
+// into
+//
+//   %y16 = tfl.fully_connected(%x16, %filter, %bias) : f16
+//
+// The filter and bias keep their types, so the result is a mixed precision
+// fully_connected (f16 activations, f32 or quantized weights). Only runtimes
+// that support that combination can execute the result, which is why the pass
+// is opt-in.
+struct FoldFp16CastsIntoFullyConnectedPattern
+    : public OpRewritePattern<TFL::CastOp> {
+  using OpRewritePattern<TFL::CastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(TFL::CastOp output_cast,
+                                PatternRewriter& rewriter) const override {
+    const Type f16 = rewriter.getF16Type();
+    const Type f32 = rewriter.getF32Type();
+
+    if (!GetRankedTensorOf(output_cast.getOutput(), f16)) return failure();
+    Value fc_output = output_cast.getInput();
+    if (!GetRankedTensorOf(fc_output, f32)) return failure();
+
+    auto fc = fc_output.getDefiningOp<TFL::FullyConnectedOp>();
+    // A shuffled-weights fully_connected has a second result and cannot be
+    // replaced by the single-result op built below.
+    if (!fc || fc.getNumResults() != 1) return failure();
+    if (fc.getWeightsFormat() != "DEFAULT") return failure();
+    // Other users still need the f32 result.
+    if (!fc_output.hasOneUse()) return failure();
+
+    auto input_cast = fc.getInput().getDefiningOp<TFL::CastOp>();
+    if (!input_cast) return failure();
+    Value input = input_cast.getInput();
+    if (!GetRankedTensorOf(input, f16)) return failure();
+    if (!GetRankedTensorOf(input_cast.getOutput(), f32)) return failure();
+
+    auto fc_output_type = mlir::cast<RankedTensorType>(fc_output.getType());
+    auto new_fc = rewriter.create<TFL::FullyConnectedOp>(
+        fc.getLoc(), fc_output_type.clone(f16), input, fc.getFilter(),
+        fc.getBias(), fc.getFusedActivationFunctionAttr(),
+        fc.getWeightsFormatAttr(), fc.getKeepNumDimsAttr(),
+        fc.getAsymmetricQuantizeInputsAttr());
+
+    rewriter.replaceOp(output_cast, new_fc.getOutput());
+    rewriter.eraseOp(fc);
+    // `input_cast` may still feed other ops; it is dead code otherwise and is
+    // erased by the greedy driver.
+    return success();
+  }
+};
+
+struct FoldFp16CastsIntoFullyConnectedPass
+    : public impl::FoldFp16CastsIntoFullyConnectedPassBase<
+          FoldFp16CastsIntoFullyConnectedPass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      FoldFp16CastsIntoFullyConnectedPass)
+
+  void runOnOperation() override {
+    func::FuncOp func = getOperation();
+    MLIRContext* context = func.getContext();
+
+    RewritePatternSet patterns(context);
+    patterns.add<FoldFp16CastsIntoFullyConnectedPattern>(context);
+
+    if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
+      signalPassFailure();
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+CreateFoldFp16CastsIntoFullyConnectedPass() {
+  return std::make_unique<FoldFp16CastsIntoFullyConnectedPass>();
+}
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/transforms/passes.h

@@ -129,6 +129,11 @@ CreateFoldStablehloConstantTransformsPass();
 std::unique_ptr<OperationPass<func::FuncOp>>
 CreateFuseA4W2DRQFullyConnectedPass();
 
+// Folds f16 <-> f32 casts around a float fully_connected into it, producing a
+// fully_connected with f16 activations and unchanged weights.
+std::unique_ptr<OperationPass<func::FuncOp>>
+CreateFoldFp16CastsIntoFullyConnectedPass();
+
 std::unique_ptr<OperationPass<ModuleOp>> CreateLowerQuantAnnotationsPass();
 
 // Creates an instance of the TFLite PropagateQParams pass which propagates

tensorflow/compiler/mlir/lite/transforms/passes.td

@@ -385,6 +385,21 @@ def FuseA4W2DRQFullyConnectedPass : Pass<"tfl-fuse-a4w2-drq-fully-connected", "m
   ];
 }
 
+def FoldFp16CastsIntoFullyConnectedPass : Pass<"tfl-fold-fp16-casts-into-fully-connected", "mlir::func::FuncOp"> {
+  let summary = "Folds f16 <-> f32 casts around a float fully_connected into it.";
+  let description = [{
+      Rewrites `cast(f32 -> f16)(fully_connected(cast(f16 -> f32)(x), w, b))`
+      into `fully_connected(x, w, b)` with an f16 result. The filter and bias
+      keep their types, so the result is a mixed precision fully_connected
+      (f16 activations, f32 or quantized weights) that only runtimes supporting
+      that combination can execute.
+  }];
+  let constructor = "CreateFoldFp16CastsIntoFullyConnectedPass()";
+  let dependentDialects = [
+    "TFL::TensorFlowLiteDialect"
+  ];
+}
+
 def PropagateQParamsPass : Pass<"tfl-propagate-qparams", "mlir::ModuleOp"> {
   let summary = "Propagates Quantization Parameters (scale and zero point) information through the graph.";
   let description = [{

tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc

@@ -22,6 +22,7 @@ limitations under the License.
 
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/log.h"
+#include "absl/log/vlog_is_on.h"
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
 #include "llvm/ADT/StringRef.h"

tensorflow/compiler/mlir/tensorflow/transforms/init_text_file_to_import.cc

@@ -122,7 +122,7 @@ class ConvertInitializeTableFromTextFileV2
     line_nums.resize(lines.size());
     std::iota(line_nums.begin(), line_nums.end(), 0);
 
-    // Create constant ops for keys an values.
+    // Create constant ops for keys and values.
     Value key_constant_tensor = arith::ConstantOp::create(
         rewriter, op.getLoc(),
         DenseStringElementsAttr::get(

tensorflow/compiler/mlir/tensorflow/utils/error_util_test.cc

@@ -73,13 +73,6 @@ TEST_F(StatusScopedDiagnosticHandlerTest,
   EXPECT_THAT(s.message(), HasSubstr("Second diagnostic message reported"));
 }
 
-TEST_F(StatusScopedDiagnosticHandlerTest, VerifyThatWarningsAreIgnored) {
-  // Note: this logic is actually implemented in BaseScopedDiagnosticHandler's
-  // handler() function, but only StatusScopedDiagnosticHandler uses it.
-  StatusScopedDiagnosticHandler handler(&context_);
-  emitWarning(loc_) << "Warning message";
-  TF_EXPECT_OK(handler.ConsumeStatus());
-}
 
 TEST(ErrorUtilTest, StatusScopedDiagnosticHandlerWithFilter) {
   // Filtering logic is based on tensorflow::IsInternalFrameForFilename()

tensorflow/compiler/mlir/tfr/ir/tfr_ops.td

@@ -342,7 +342,7 @@ def TFR_GetElementOp : TFR_Op<"get_element", [Pure]> {
 def TFR_BuildListOp : TFR_Op<"build_list", [Pure]> {
   let description = [{
    The `build_list` operation builds a tensor list from a list of tensors, or
-   an tfr.attr from a list of scalars.
+   a tfr.attr from a list of scalars.
 
     Example:
 

tensorflow/compiler/mlir/tfr/passes/raise_to_tf.cc

@@ -183,7 +183,7 @@ LogicalResult RewriteTFRCallOp::AddDerivedAttrs(
     ArrayRef<Attribute> input_types,
     llvm::StringMap<Attribute>* derived_attrs) const {
   // If there is an attribute associated to the input in the signature, we
-  // store it as an derived attribute.
+  // store it as a derived attribute.
   if (auto tensor_type = mlir::dyn_cast<TFRTensorType>(input_tfr_type)) {
     auto attr_names = tensor_type.getAttrKeys();
     if (attr_names.empty()) return success();
@@ -195,7 +195,7 @@ LogicalResult RewriteTFRCallOp::AddDerivedAttrs(
   }
 
   // If there is an attribute associated to the input in the signature,
-  // we store it as an derived attribute.
+  // we store it as a derived attribute.
   if (auto list_type = mlir::dyn_cast<TFRTensorListType>(input_tfr_type)) {
     auto attr_names = list_type.getAttrKeys();
     if (attr_names.empty()) return success();

tensorflow/compiler/mlir/tools/kernel_gen/transforms/BUILD

@@ -125,6 +125,7 @@ cc_library(
         "//tensorflow/compiler/mlir/tools/kernel_gen/ir:tf_framework_ops",
         "//tensorflow/core:lib",
         "//tensorflow/core/platform:errors",
+        "@com_google_absl//absl/status:statusor",
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:TargetParser",
         "@llvm-project//llvm:TransformUtils",
@@ -182,6 +183,9 @@ cc_library(
         "@xla//xla/stream_executor:device_description",
     ] + if_cuda_is_configured([
         "@tsl//tsl/platform:cuda_root_path",
+        "@xla//xla/stream_executor/cuda:assemble_compilation_provider",
+        "@xla//xla/stream_executor/cuda:compilation_provider",
+        "@xla//xla/stream_executor/cuda:compilation_provider_options",
         "@xla//xla/stream_executor/cuda:cuda_asm_compiler",
         "@xla//xla/service/gpu/llvm_gpu_backend:nvptx_backend",
     ]) + if_rocm_is_configured([

tensorflow/compiler/mlir/tools/kernel_gen/transforms/broadcast_propagation_pass.cc

@@ -128,8 +128,8 @@ DenseIntElementsAttr composeBroadcastDimensionsAttr(OpBuilder &builder,
 }
 
 // Find all the broadcast intents and their dependencies. Start analyzing from
-// the root an collect all broadcast intents that can help broadcast propagation
-// from there.
+// the root and collect all broadcast intents that can help broadcast
+// propagation from there.
 void findBroadcastIntents(
     DynamicBroadcastInDimOp root, Block *parentBlock,
     BroadcastIntent &rootBcastIntent,

tensorflow/compiler/mlir/tools/kernel_gen/transforms/gpu_kernel_to_blob_pass.cc

@@ -16,10 +16,12 @@ limitations under the License.
 #include <algorithm>
 #include <iterator>
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
 
+#include "absl/status/statusor.h"
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
@@ -40,6 +42,9 @@ limitations under the License.
 
 #if GOOGLE_CUDA
 #include "xla/service/gpu/llvm_gpu_backend/nvptx_backend.h"
+#include "xla/stream_executor/cuda/assemble_compilation_provider.h"
+#include "xla/stream_executor/cuda/compilation_provider.h"
+#include "xla/stream_executor/cuda/compilation_provider_options.h"
 #include "xla/stream_executor/cuda/cuda_asm_compiler.h"
 #elif TENSORFLOW_USE_ROCM
 #include "xla/service/gpu/llvm_gpu_backend/amdgpu_backend.h"
@@ -143,6 +148,20 @@ class GpuKernelToBlobPass
     llvmModule->setDataLayout(xla::gpu::nvptx::DataLayout());
     llvmModule->setTargetTriple(llvm::Triple(xla::gpu::nvptx::TargetTriple()));
 
+    std::optional<int> max_ptx_isa_version;
+    absl::StatusOr<std::unique_ptr<tensorflow::se::cuda::CompilationProvider>>
+        compilation_provider =
+            tensorflow::se::cuda::AssembleCompilationProvider(
+                tensorflow::se::cuda::CompilationProviderOptions::
+                    FromDebugOptions(options));
+    if (compilation_provider.ok()) {
+      absl::StatusOr<int> ptx_isa_version =
+          (*compilation_provider)->GetLatestPtxIsaVersion();
+      if (ptx_isa_version.ok()) {
+        max_ptx_isa_version = *ptx_isa_version;
+      }
+    }
+
     // Compile and collect requested cubin and PTX images.
     std::vector<tensorflow::se::CubinOrPTXImage> images;
     auto gpu_asm_opts = xla::gpu::PtxOptsFromDebugOptions(options);
@@ -161,7 +180,7 @@ class GpuKernelToBlobPass
           std::string ptx,
           xla::gpu::nvptx::CompileToPtx(
               llvm_module_copy.get(), stream_executor::GpuComputeCapability(cc),
-              options));
+              options, /*configure_target=*/nullptr, max_ptx_isa_version));
       if (print_ptx_) {
         llvm::dbgs() << "Generated PTX code for module '"
                      << gpu_module.getName() << "' on architecture sm_" << arch

tensorflow/compiler/tests/BUILD

@@ -1213,6 +1213,8 @@ tf_xla_py_strict_test(
     ],
     deps = [
         ":xla_test",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:array_ops_gen",
         "//tensorflow/python/platform:test",
@@ -1390,6 +1392,33 @@ tf_xla_py_strict_test(
     ],
 )
 
+tf_xla_py_strict_test(
+    name = "qr_op_large_test",
+    size = "large",
+    srcs = ["qr_op_large_test.py"],
+    disabled_backends = [
+        # Test is very slow on CPU.
+        "cpu",
+        "cpu_ondemand",
+    ],
+    # Keep shard_count = 2 so testLarge2000x2000 and testLarge17500x128 (if enabled) run in separate shards.
+    shard_count = 2,
+    tags = [
+        "no_pip",  # TODO(b/149738646): fix pip install so these tests run on kokoro pip
+        "nozapfhahn",
+        "optonly",
+        "requires-mem:24g",
+    ],
+    deps = [
+        ":qr_op_test_lib",
+        ":xla_test",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
 tf_xla_py_strict_test(
     name = "unstack_test",
     size = "medium",
@@ -1506,6 +1535,8 @@ tf_xla_py_strict_test(
     ],
     deps = [
         ":xla_test",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/platform:client_testlib",
@@ -1845,6 +1876,7 @@ tf_xla_py_strict_test(
     deps = [
         ":xla_test",
         "//tensorflow/python:extra_py_tests_deps",
+        "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/framework:test_lib",
         "//tensorflow/python/ops:array_ops",
@@ -1869,7 +1901,9 @@ tf_xla_py_strict_test(
     ],
     deps = [
         ":xla_test",
+        "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:array_ops_stack",
         "//tensorflow/python/ops:bitwise_ops",

tensorflow/compiler/tests/categorical_op_test.py

@@ -138,6 +138,31 @@ class CategoricalTest(xla_test.XLATestCase):
       chi2 = self._chi2(probs, freqs)
       self.assertLess(chi2, 1e-3)
 
+  def testHalfPrecisionLargeLogits(self):
+    # Regression test for GitHub issue 128958: half-precision logits were
+    # sampled in their own type, where the spacing between values (32 at 4096
+    # for bfloat16) rounds the Gumbel noise away, so equal logits always gave
+    # the first class.
+    num_samples = 40000
+    probs = [[0.5, 0.5]]
+    for dtype in self.float_types:
+      if dtypes.as_dtype(dtype) not in (dtypes.float16, dtypes.bfloat16):
+        continue
+      logits = np.array([[4096.0, 4096.0]], dtype=dtype)
+      with self.subTest(dtype=dtypes.as_dtype(dtype).name):
+        freqs = self._do_sampling(logits, num_samples)
+        self.assertLess(self._chi2(probs, freqs), 1e-3)
+        if dtypes.as_dtype(dtype) == dtypes.bfloat16:
+          with self.session() as sess:
+            with self.test_scope():
+              seed_t = array_ops.placeholder(dtypes.int32, shape=[2])
+              x = stateless_random_ops.stateless_multinomial(
+                  logits, num_samples, seed_t, output_dtype=dtypes.int32
+              )
+            y = sess.run(x, {seed_t: [0x12345678, 0xABCDEF12]})
+          freqs = [[np.mean(y == k) for k in range(2)]]
+          self.assertLess(self._chi2(probs, freqs), 1e-3)
+
   def testStatelessMultinomialIsInRange(self):
     for dtype in self.float_types.intersection(
         [dtypes.float32, dtypes.bfloat16]):
@@ -190,6 +215,5 @@ class CategoricalTest(xla_test.XLATestCase):
         self.assertEqual(y.shape, (42, 0))
 
 
-
 if __name__ == '__main__':
   googletest.main()

tensorflow/compiler/tests/image_ops_test.py

@@ -1022,6 +1022,103 @@ class NonMaxSuppressionTest(xla_test.XLATestCase):
       self.assertEqual(indices_tf.size, 3)
       self.assertAllClose(indices_tf[:3], [3, 0, 5])
 
+  def _testMixedThresholdDtype(self, nms):
+    # Regression test for GitHub issue 128614: boxes and scores of one
+    # floating-point type with thresholds of another, which
+    # tf.image.non_max_suppression produces for float16 boxes, failed to
+    # compile because XLA doesn't mix floating-point types in one operation.
+    if np.float16 not in self.float_types:
+      self.skipTest("float16 is not supported on this device")
+    boxes_data = [
+        [0, 0, 1, 1],
+        [0, 0.1, 1, 1.1],
+        [0, -0.1, 1, 0.9],
+        [0, 10, 1, 11],
+        [0, 10.1, 1, 11.1],
+        [0, 100, 1, 101],
+    ]
+    scores_data = [0.9, 0.75, 0.6, 0.95, 0.5, 0.3]
+    for dtype, threshold_dtype in (
+        (np.float16, np.float32),
+        (np.float32, np.float16),
+    ):
+      with self.subTest(
+          dtype=dtype.__name__, threshold_dtype=threshold_dtype.__name__
+      ), self.session() as sess:
+        boxes = array_ops.placeholder(dtype, shape=[6, 4])
+        scores = array_ops.placeholder(dtype, shape=[6])
+        iou_threshold = array_ops.placeholder(threshold_dtype, shape=[])
+        score_threshold = array_ops.placeholder(threshold_dtype, shape=[])
+        with self.device_scope():
+          selected_indices, num_valid = nms(
+              boxes, scores, iou_threshold, score_threshold
+          )
+        indices_tf, num_valid_tf = sess.run(
+            [selected_indices, num_valid],
+            feed_dict={
+                boxes: np.array(boxes_data, dtype=dtype),
+                scores: np.array(scores_data, dtype=dtype),
+                iou_threshold: 0.5,
+                score_threshold: 0.4,
+            },
+        )
+        # Box 5 is below the score threshold.
+        self.assertEqual(num_valid_tf, 2)
+        self.assertAllEqual([3, 0], indices_tf[:num_valid_tf])
+
+  def testNMSV3MixedThresholdDtype(self):
+
+    def nms(boxes, scores, iou_threshold, score_threshold):
+      selected_indices = image_ops.non_max_suppression_v3(
+          boxes=boxes,
+          scores=scores,
+          max_output_size=6,
+          iou_threshold=iou_threshold,
+          score_threshold=score_threshold,
+      )
+      return selected_indices, array_ops.size(selected_indices)
+
+    self._testMixedThresholdDtype(nms)
+
+  def testNMSV4MixedThresholdDtype(self):
+
+    def nms(boxes, scores, iou_threshold, score_threshold):
+      return gen_image_ops.non_max_suppression_v4(
+          boxes=boxes,
+          scores=scores,
+          max_output_size=6,
+          iou_threshold=iou_threshold,
+          score_threshold=score_threshold,
+          pad_to_max_output_size=True,
+      )
+
+    self._testMixedThresholdDtype(nms)
+
+  def testNMSV3NonScalarScoreThreshold(self):
+    # Shape inference can't reject a score_threshold of unknown shape, so the
+    # XLA kernel checks that it is a scalar when it compiles.
+    with self.session() as sess:
+      boxes = array_ops.placeholder(np.float32, shape=[6, 4])
+      scores = array_ops.placeholder(np.float32, shape=[6])
+      score_threshold = array_ops.placeholder(np.float32)
+      with self.device_scope():
+        selected_indices = image_ops.non_max_suppression_v3(
+            boxes=boxes,
+            scores=scores,
+            max_output_size=6,
+            iou_threshold=0.5,
+            score_threshold=score_threshold,
+        )
+      with self.assertRaisesOpError("Score Threshold isn't a scalar"):
+        sess.run(
+            selected_indices,
+            feed_dict={
+                boxes: np.zeros([6, 4], np.float32),
+                scores: np.zeros([6], np.float32),
+                score_threshold: [0.4],
+            },
+        )
+
   def testNMSV3EmptyInput(self):
     # Regression test for #117245: with no boxes the suppression loop was
     # built from zero-sized dimensions and segfaulted the compiler. The

tensorflow/compiler/tests/matrix_diag_ops_test.py

@@ -17,6 +17,8 @@
 import numpy as np
 
 from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import gen_array_ops
 from tensorflow.python.platform import googletest
@@ -767,6 +769,118 @@ class MatrixDiagPartTest(xla_test.XLATestCase):
                   "align": align
               }, solution)
 
+  def _diagPartOfSelected(self, op, x, masks, dtype=np.int32, **params):
+    # masks maps axes of x to masks of the entries to keep along them, which
+    # gives those dimensions a size only known at run time.
+    axes = sorted(masks)
+
+    @def_function.function(jit_compile=True)
+    def f(x, *x_masks):
+      for axis, mask in zip(axes, x_masks):
+        x = array_ops.gather(
+            x, array_ops.reshape(array_ops.where(mask), [-1]), axis=axis
+        )
+      return op(x, **params)
+
+    x = np.array(x, dtype=dtype)
+    with self.session() as session:
+      x_ph = array_ops.placeholder(dtypes.as_dtype(dtype), x.shape)
+      mask_phs = [
+          array_ops.placeholder(dtypes.bool, [x.shape[axis]]) for axis in axes
+      ]
+      output = f(x_ph, *mask_phs)
+      feed_dict = {x_ph: x}
+      feed_dict.update({ph: masks[axis] for ph, axis in zip(mask_phs, axes)})
+      return session.run(output, feed_dict)
+
+  def testDynamicDimensions(self):
+    # Regression test for GitHub issue 129053: with a dynamic dimension, the
+    # diagonals had the length they have at its bound. The mask keeps 2 of 5
+    # entries, so the dimensions it selects have size 2 and bound 5.
+    keep = [True, False, True, False, False]
+    x = np.arange(1, 16).reshape(5, 3)  # Keeps [[1, 2, 3], [7, 8, 9]].
+    for dtype in (np.int32, np.float32, np.float64):
+      for op in (array_ops.matrix_diag_part, gen_array_ops.matrix_diag_part):
+        self.assertAllEqual(
+            self._diagPartOfSelected(op, x, {-2: keep}, dtype), [1, 8]
+        )
+      self.assertAllEqual(
+          self._diagPartOfSelected(
+              array_ops.matrix_diag_part,
+              x,
+              {-2: keep},
+              dtype,
+              k=(-1, 1),
+              padding_value=9,
+          ),
+          [[2, 9], [1, 8], [7, 9]],
+      )
+      self.assertAllEqual(
+          self._diagPartOfSelected(
+              array_ops.matrix_diag_part,
+              x,
+              {-2: keep},
+              dtype,
+              k=(-1, 1),
+              padding_value=9,
+              align="LEFT_RIGHT",
+          ),
+          [[2, 9], [1, 8], [9, 7]],
+      )
+      self.assertAllEqual(
+          self._diagPartOfSelected(
+              gen_array_ops.matrix_diag_part_v2,
+              x,
+              {-2: keep},
+              dtype,
+              k=[-1, 1],
+              padding_value=9,
+          ),
+          [[2, 9], [1, 8], [7, 9]],
+      )
+      self.assertAllEqual(
+          self._diagPartOfSelected(
+              array_ops.matrix_diag_part,
+              np.arange(1, 16).reshape(3, 5),
+              {-1: keep},
+              dtype,
+          ),
+          [1, 8],
+      )
+      # With no rows kept, the diagonal is empty.
+      self.assertAllEqual(
+          self._diagPartOfSelected(
+              array_ops.matrix_diag_part, x, {-2: [False] * 5}, dtype
+          ),
+          np.zeros([0]),
+      )
+
+  def testDynamicDimensionsBatch(self):
+    keep = [True, False, True, False, False]
+    x = np.arange(1, 31).reshape(2, 5, 3)
+    self.assertAllEqual(
+        self._diagPartOfSelected(array_ops.matrix_diag_part, x, {-2: keep}),
+        [[1, 8], [16, 23]],
+    )
+    self.assertAllEqual(
+        self._diagPartOfSelected(
+            array_ops.matrix_diag_part,
+            x,
+            {-2: keep},
+            k=(-1, 1),
+            padding_value=9,
+        ),
+        [[[2, 9], [1, 8], [7, 9]], [[17, 24], [16, 23], [22, 9]]],
+    )
+    self.assertAllEqual(
+        self._diagPartOfSelected(
+            array_ops.matrix_diag_part,
+            np.arange(1, 51).reshape(2, 5, 5),
+            {-2: keep, -1: keep},
+        ),
+        [[1, 13], [26, 38]],
+    )
+
 
 if __name__ == "__main__":
   googletest.main()

tensorflow/compiler/tests/qr_op_large_test.py

@@ -0,0 +1,49 @@
+# Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Large tests for tensorflow.ops.linalg_ops.qr."""
+
+import unittest
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.compiler.tests import qr_op_test
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+
+
+@test_util.run_all_without_tensor_float_32(
+    "XLA QR op calls matmul. Also, matmul used for verification. Also with "
+    'TensorFloat-32, mysterious "Unable to launch cuBLAS gemm" error '
+    "occasionally occurs"
+)
+# TODO(b/165435566): Fix "Unable to launch cuBLAS gemm" error
+class QrOpLargeTest(
+    qr_op_test.QrOpTestBase, xla_test.XLATestCase, parameterized.TestCase
+):
+
+  def testLarge2000x2000(self):
+    x_np = self._random_matrix(np.float32, (2000, 2000))
+    self._test(x_np, full_matrices=True)
+
+  @unittest.skip("Test times out on CI")
+  def testLarge17500x128(self):
+    x_np = self._random_matrix(np.float32, (17500, 128))
+    self._test(x_np, full_matrices=True)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/qr_op_test.py

@@ -15,7 +15,6 @@
 """Tests for tensorflow.ops.math_ops.matrix_inverse."""
 
 import itertools
-import unittest
 
 from absl.testing import parameterized
 import numpy as np
@@ -29,12 +28,7 @@ from tensorflow.python.ops import math_ops
 from tensorflow.python.platform import test
 
 
-@test_util.run_all_without_tensor_float_32(
-    "XLA QR op calls matmul. Also, matmul used for verification. Also with "
-    'TensorFloat-32, mysterious "Unable to launch cuBLAS gemm" error '
-    "occasionally occurs")
-# TODO(b/165435566): Fix "Unable to launch cuBLAS gemm" error
-class QrOpTest(xla_test.XLATestCase, parameterized.TestCase):
+class QrOpTestBase:
 
   def AdjustedNorm(self, x):
     """Computes the norm of matrices in 'x', adjusted for dimension and type."""
@@ -111,6 +105,15 @@ class QrOpTest(xla_test.XLATestCase, parameterized.TestCase):
       self.CheckApproximation(x_np, q_tf_val, r_tf_val)
       self.CheckUnitary(q_tf_val)
 
+
+@test_util.run_all_without_tensor_float_32(
+    "XLA QR op calls matmul. Also, matmul used for verification. Also with "
+    'TensorFloat-32, mysterious "Unable to launch cuBLAS gemm" error '
+    "occasionally occurs"
+)
+# TODO(b/165435566): Fix "Unable to launch cuBLAS gemm" error
+class QrOpTest(QrOpTestBase, xla_test.XLATestCase, parameterized.TestCase):
+
   SIZES = [1, 2, 5, 10, 32, 100, 300, 603]
   DTYPES = [np.float32, np.complex64]
   PARAMS = itertools.product(SIZES, SIZES, DTYPES)
@@ -123,15 +126,6 @@ class QrOpTest(xla_test.XLATestCase, parameterized.TestCase):
         x_np = self._random_matrix(dtype, batch_dims + (rows, cols))
         self._test(x_np, full_matrices)
 
-  def testLarge2000x2000(self):
-    x_np = self._random_matrix(np.float32, (2000, 2000))
-    self._test(x_np, full_matrices=True)
-
-  @unittest.skip("Test times out on CI")
-  def testLarge17500x128(self):
-    x_np = self._random_matrix(np.float32, (17500, 128))
-    self._test(x_np, full_matrices=True)
-
   @parameterized.parameters((23, 25), (513, 23))
   def testZeroColumn(self, rows, cols):
     x_np = self._random_matrix(np.complex64, (rows, cols))

tensorflow/compiler/tests/reverse_sequence_op_test.py

@@ -17,6 +17,8 @@
 import numpy as np
 
 from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import dtypes
 from tensorflow.python.ops import array_ops
 from tensorflow.python.platform import test
@@ -84,6 +86,79 @@ class ReverseSequenceTest(xla_test.XLATestCase):
       for seq_dtype in self.all_types & {np.int32, np.int64}:
         self._testBasic(dtype, seq_dtype)
 
+  def _reverseSelected(self, x, seq_lengths, axis, seq_axis, batch_axis):
+    # Gathers the entries of x along axis where mask is true, which gives that
+    # dimension a size only known at run time, and reverses the result.
+    @def_function.function(jit_compile=True)
+    def f(x, seq_lengths, mask):
+      indices = array_ops.reshape(array_ops.where(mask), [-1])
+      x = array_ops.gather(x, indices, axis=axis)
+      if axis == batch_axis:
+        seq_lengths = array_ops.gather(seq_lengths, indices)
+      return array_ops.reverse_sequence(
+          x, seq_lengths, seq_axis=seq_axis, batch_axis=batch_axis
+      )
+
+    x = np.array(x, dtype=np.int32)
+    seq_lengths = np.array(seq_lengths, dtype=np.int32)
+    mask = np.array([True, False, True, True, False])
+    with self.session() as session:
+      with self.device_scope():
+        x_ph = array_ops.placeholder(dtypes.int32, x.shape)
+        seq_lengths_ph = array_ops.placeholder(dtypes.int32, seq_lengths.shape)
+        mask_ph = array_ops.placeholder(dtypes.bool, mask.shape)
+        output = f(x_ph, seq_lengths_ph, mask_ph)
+      result = session.run(
+          output, {x_ph: x, seq_lengths_ph: seq_lengths, mask_ph: mask}
+      )
+    # Eager execution of the compiled function gives the same result, also
+    # when the test itself runs in graph mode.
+    with context.eager_mode(), self.device_scope():
+      self.assertAllEqual(f(x, seq_lengths, mask), result)
+    return result
+
+  def testDynamicDimensions(self):
+    if "TPU" in self.device:
+      self.skipTest(
+          "XLA:TPU fails to compile these dynamic dimensions with a "
+          "bitcast shape size mismatch."
+      )
+    # Regression test for GitHub issue 128953: the output had the bound of a
+    # dynamic input dimension instead of its size. The mask selects 3 of 5
+    # entries, so the selected dimension has size 3 and bound 5. (XLA turns a
+    # gather from 4 or fewer entries into selects, which XLA:GPU doesn't
+    # support with dynamic shapes.)
+    self.assertAllEqual(
+        self._reverseSelected(
+            [[1, 2, 3, 4, 0], [5, 6, 7, 8, 0]],
+            [3, 2],
+            axis=1,
+            seq_axis=1,
+            batch_axis=0,
+        ),
+        [[4, 3, 1], [7, 5, 8]],
+    )
+    self.assertAllEqual(
+        self._reverseSelected(
+            [[1, 5], [2, 6], [3, 7], [4, 8], [0, 0]],
+            [3, 2],
+            axis=0,
+            seq_axis=0,
+            batch_axis=1,
+        ),
+        [[4, 7], [3, 5], [1, 8]],
+    )
+    self.assertAllEqual(
+        self._reverseSelected(
+            [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [0, 0, 0]],
+            [3, 1, 2, 3, 1],
+            axis=0,
+            seq_axis=1,
+            batch_axis=0,
+        ),
+        [[3, 2, 1], [8, 7, 9], [12, 11, 10]],
+    )
+
 
 if __name__ == "__main__":
   test.main()

tensorflow/compiler/tests/ternary_ops_test.py

@@ -19,6 +19,7 @@ import numpy as np
 import scipy.special as sps
 
 from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
@@ -75,6 +76,62 @@ class TernaryOpsTest(xla_test.XLATestCase, parameterized.TestCase):
           dtype(2),
           expected=np.array([1, 3, 5], dtype=dtype),
       )
+    # The distance from start to limit doesn't fit in int64.
+    self._testTernary(
+        math_ops.range,
+        np.int64(-(2**63)),
+        np.int64(2**63 - 1),
+        np.int64(2**62),
+        expected=np.array([-(2**63), -(2**62), 0, 2**62], dtype=np.int64),
+    )
+
+  def _rangeOfMaskedSize(self, args, dtype=np.int32):
+    # n is the number of true entries in mask, a size only known at run time,
+    # and args maps it to the range's start, limit and delta.
+    @def_function.function(jit_compile=True)
+    def f(mask):
+      indices = array_ops.reshape(array_ops.where(mask), [-1])
+      n = array_ops.size(array_ops.gather(math_ops.range(6), indices))
+      return gen_math_ops._range(
+          *[math_ops.cast(a, dtype) for a in args(math_ops.cast(n, dtype))]
+      )
+
+    with self.session() as session:
+      with self.device_scope():
+        mask = array_ops.placeholder(dtypes.bool, [6])
+        output = f(mask)
+      return session.run(
+          output, {mask: [True, False, True, False, True, False]}
+      )
+
+  def testRangeDynamicBounds(self):
+    # Regression test for GitHub issue 128950: with a negative delta, the
+    # output was sized from the wrong bounds of start and limit, and with a
+    # dynamic start, its values began at start's bound. Here n is 3, bound 6.
+    for dtype in (np.int32, np.int64, np.float32):
+      self.assertAllEqual(
+          self._rangeOfMaskedSize(lambda n: (n - 1, -1, -1), dtype), [2, 1, 0]
+      )
+      self.assertAllEqual(
+          self._rangeOfMaskedSize(lambda n: (6, n, -1), dtype), [6, 5, 4]
+      )
+      self.assertAllEqual(
+          self._rangeOfMaskedSize(lambda n: (n, 5, 1), dtype), [3, 4]
+      )
+      self.assertAllEqual(
+          self._rangeOfMaskedSize(lambda n: (n, -n, -2), dtype), [3, 1, -1]
+      )
+      # The CPU kernel rejects runtime values in the wrong order for delta's
+      # sign. XLA can't report that at run time, and returns an empty range.
+      self.assertAllEqual(
+          self._rangeOfMaskedSize(lambda n: (n, 1, 1), dtype), []
+      )
+      self.assertAllEqual(
+          self._rangeOfMaskedSize(lambda n: (n, 4, -1), dtype), []
+      )
+    self.assertAllEqual(
+        self._rangeOfMaskedSize(lambda n: (n, 10, 2), np.uint32), [3, 5, 7, 9]
+    )
 
   def testSelect(self):
     for dtype in self.numeric_types:

tensorflow/compiler/tests/unary_ops_test.py

@@ -20,7 +20,9 @@ import numpy as np
 import six
 
 from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import array_ops_stack  # pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.ops import bitwise_ops
@@ -335,6 +337,52 @@ class UnaryOpsTest(xla_test.XLATestCase):
           ),
       )
 
+  def testQuantizeAndDequantizeV3NumBitsBoundaries(self):
+    for signed_input, invalid_num_bits in (
+        (True, (0, -1, 62, 63)),
+        (False, (0, -1, 63, 64)),
+    ):
+      for num_bits in invalid_num_bits:
+
+        @def_function.function(jit_compile=True)
+        def quantize_and_dequantize_v3(
+            x, num_bits=num_bits, signed_input=signed_input
+        ):
+          return array_ops.quantize_and_dequantize_v3(
+              x, -1.0, 1.0, num_bits=num_bits, signed_input=signed_input
+          )
+
+        with self.session() as sess:
+          with self.device_scope():
+            x = array_ops.placeholder(dtypes.float32)
+            output = quantize_and_dequantize_v3(x)
+          with self.assertRaisesRegex(
+              errors.InvalidArgumentError, "num_bits is out of range"
+          ):
+            sess.run(
+                output, {x: np.array([-1, -0.5, 0, 0.3], dtype=np.float32)}
+            )
+
+    for signed_input, valid_num_bits in (
+        (True, (1, 61)),
+        (False, (1, 62)),
+    ):
+      for num_bits in valid_num_bits:
+
+        @def_function.function(jit_compile=True)
+        def quantize_and_dequantize_v3_valid(
+            x, num_bits=num_bits, signed_input=signed_input
+        ):
+          return array_ops.quantize_and_dequantize_v3(
+              x, -1.0, 1.0, num_bits=num_bits, signed_input=signed_input
+          )
+
+        with self.session() as sess:
+          with self.device_scope():
+            x = array_ops.placeholder(dtypes.float32)
+            output = quantize_and_dequantize_v3_valid(x)
+          sess.run(output, {x: np.array([-1, -0.5, 0, 0.3], dtype=np.float32)})
+
   def testComplexOps(self):
     for dtype in self.complex_types:
 

tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h

@@ -40,7 +40,7 @@ namespace tensorrt {
 class TRTBaseAllocator : public nvinfer1::IGpuAllocator {
   // Base allocator class so we can have a virtual destructor;
  public:
-  // python wrapper seems to be not happy with an pure virtual destructor;
+  // python wrapper seems to be not happy with a pure virtual destructor;
   virtual ~TRTBaseAllocator() = default;
 };
 

tensorflow/compiler/tf2xla/BUILD

@@ -1012,10 +1012,10 @@ cc_library(
     ],
     visibility = [":friends"],
     deps = [
-        "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/common_runtime:dma_helper",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
@@ -1234,6 +1234,8 @@ tf_cc_test(
         "//tensorflow/cc:ops",
         "//tensorflow/cc:resource_variable_ops",
         "//tensorflow/cc:scope",
+        "//tensorflow/compiler/jit:common",
+        "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops",
         "//tensorflow/compiler/tf2xla/kernels:xla_ops",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",

tensorflow/compiler/tf2xla/kernels/BUILD

@@ -549,10 +549,10 @@ tf_kernel_library(
         "check_numerics_op.cc",
     ],
     deps = [
-        "//tensorflow/compiler/tf2xla:mlir_xla_op_kernel",
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/compiler/tf2xla:xla_op_registry",
         "//tensorflow/core:array_ops_op_lib",
+        "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:logging_ops_op_lib",
         "@com_google_absl//absl/log",

tensorflow/compiler/tf2xla/kernels/categorical_op.cc

@@ -101,6 +101,16 @@ class CategoricalOp : public XlaOpKernel {
     }
     xla::PrimitiveType type;
     OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(input_type(0), &type));
+    // Sample in float32 for half-precision logits (the non-XLA kernels sample
+    // in at least float32): the spacing between half-precision values grows
+    // with the logits and can round the Gumbel noise away, which makes equal
+    // logits always pick the first class.
+    xla::XlaOp sampling_logits = logits;
+    if (type == xla::F16 || type == xla::BF16) {
+      type = xla::F32;
+      uniform_shape.set_element_type(type);
+      sampling_logits = xla::ConvertElementType(logits, type);
+    }
     xla::XlaOp log_uniforms = GetLogUniforms(uniform_shape, type, ctx);
 
     if (num_samples_is_dynamic) {
@@ -113,7 +123,7 @@ class CategoricalOp : public XlaOpKernel {
     // https://hips.seas.harvard.edu/blog/2013/04/06/the-gumbel-max-trick-for-discrete-distributions/
     // TODO(b/68769470): Switch to using a cumulative sum approach.
     auto softmax_entries =
-        xla::Sub(logits, log_uniforms,
+        xla::Sub(sampling_logits, log_uniforms,
                  /*broadcast_dimensions=*/{0, class_dimension});
 
     xla::PrimitiveType xla_output_type;
@@ -166,9 +176,6 @@ class StatelessCategoricalOp : public CategoricalOp {
     xla::XlaOp seed = ctx->Input(2);
 
     xla::XlaBuilder* builder = ctx->builder();
-    if (uniform_shape.element_type() == xla::BF16) {
-      uniform_shape.set_element_type(xla::F32);
-    }
     // We want a number in (0, 1) rather than [0, 1) or (0, 1]:
     // * log(-log(0)) is ∞.
     // * log(-log(1)) is -∞.

tensorflow/compiler/tf2xla/kernels/check_numerics_op.cc

@@ -13,12 +13,38 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
+#include "absl/log/log.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/platform/mutex.h"
 
 namespace tensorflow {
 namespace {
 
-REGISTER_XLA_OP(Name("CheckNumerics"), MlirXlaOpKernel);
+class CheckNumericsOp : public XlaOpKernel {
+ public:
+  explicit CheckNumericsOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  ~CheckNumericsOp() override = default;
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    static mutex mu(tensorflow::LINKER_INITIALIZED);
+    static int log_counter = 0;
+
+    mutex_lock l(mu);
+    if (log_counter < 20) {
+      ++log_counter;
+      LOG(WARNING) << "Ignoring CheckNumerics operator " << name();
+    }
+    ctx->SetOutput(0, ctx->Input(0));
+  }
+
+ private:
+  CheckNumericsOp(const CheckNumericsOp&) = delete;
+  void operator=(const CheckNumericsOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("CheckNumerics"), CheckNumericsOp);
 
 }  // anonymous namespace
 }  // namespace tensorflow

tensorflow/compiler/tf2xla/kernels/const_op.cc

@@ -36,9 +36,8 @@ DstT CastTo(SrcT src) {
 }
 
 template <typename DstT,
-          typename std::enable_if<std::is_same<DstT, Eigen::half>::value ||
-                                  std::is_same<DstT, bfloat16>::value>::type* =
-              nullptr>
+          std::enable_if_t<std::is_same_v<DstT, Eigen::half> ||
+                           std::is_same_v<DstT, bfloat16>>* = nullptr>
 DstT CastTo(int32_t src) {
   return absl::bit_cast<DstT>(static_cast<uint16_t>(src));
 }

tensorflow/compiler/tf2xla/kernels/image_ops.cc

@@ -484,6 +484,10 @@ class NonMaxSuppressionOp : public XlaOpKernel {
         context,
         TensorShapeUtils::IsScalar(context->InputShape("iou_threshold")),
         absl::InvalidArgumentError("IOU Threshold isn't a scalar"));
+    OP_REQUIRES(
+        context,
+        TensorShapeUtils::IsScalar(context->InputShape("score_threshold")),
+        absl::InvalidArgumentError("Score Threshold isn't a scalar"));
     OP_REQUIRES_OK(context, context->ConstantInputAsIntScalar(2, &output_size));
     OP_REQUIRES(context, output_size >= 0,
                 absl::InvalidArgumentError(
@@ -491,8 +495,6 @@ class NonMaxSuppressionOp : public XlaOpKernel {
     OP_REQUIRES(context, output_size <= std::numeric_limits<int32_t>::max(),
                 absl::InvalidArgumentError(absl::StrCat(
                     "Need output_size <= kint32Max, got ", output_size)));
-    const xla::XlaOp score_thresh = context->Input("score_threshold");
-    const xla::XlaOp iou_thresh = context->Input("iou_threshold");
     xla::XlaBuilder* const builder = context->builder();
 
     if (num_boxes == 0) {
@@ -521,6 +523,14 @@ class NonMaxSuppressionOp : public XlaOpKernel {
       return;
     }
 
+    // The thresholds' type, T_threshold, can differ from T, the type of boxes
+    // and scores. Compare in T, as the TensorFlow kernels do, since XLA
+    // doesn't mix floating-point types in one operation.
+    const xla::XlaOp score_thresh = xla::ConvertElementType(
+        context->Input("score_threshold"), scores_xla_type);
+    const xla::XlaOp iou_thresh = xla::ConvertElementType(
+        context->Input("iou_threshold"), boxes_xla_type);
+
     // Choose a more convenient layout.
     const xla::XlaOp boxes = xla::Transpose(boxes_input, {1, 0});
     const xla::XlaOp boxes_sorted = xla::GetTupleElement(

tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc

@@ -31,6 +31,7 @@ limitations under the License.
 #include "xla/hlo/builder/lib/constants.h"
 #include "xla/hlo/builder/lib/matrix.h"
 #include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
 #include "tensorflow/core/framework/op_kernel.h"
@@ -416,42 +417,133 @@ class MatrixDiagPartOp : public XlaOpKernel {
 
     // Computes output.
     xla::XlaOp input = context->Input(0);
+
+    // With dynamic dimensions, take the diagonals of the matrices at their
+    // bound sizes, which hold the actual matrices in the top left, and give
+    // the output its actual sizes at the end.
+    xla::Shape input_xla_shape;
+    OP_REQUIRES_VALUE(input_xla_shape, context, context->InputXlaShape(0));
+    std::vector<xla::XlaOp> input_dim_sizes(input_rank);
+    for (int i = 0; i < input_rank; ++i) {
+      if (input_xla_shape.is_dynamic_dimension(i)) {
+        input_dim_sizes[i] = xla::GetDimensionSize(input, i);
+        input = xla::RemoveDynamicDimension(input, i);
+      }
+    }
+    const bool dynamic_matrix =
+        input_xla_shape.is_dynamic_dimension(input_rank - 2) ||
+        input_xla_shape.is_dynamic_dimension(input_rank - 1);
+    const auto actual_size = [&](int dim) {
+      return input_xla_shape.is_dynamic_dimension(dim)
+                 ? input_dim_sizes[dim]
+                 : xla::ConstantR0<int32_t>(
+                       context->builder(),
+                       static_cast<int32_t>(input_shape.dim_size(dim)));
+    };
+    // Only built with dynamic matrix dimensions, so that static inputs build
+    // the same HLO as before.
+    xla::XlaOp actual_num_rows;
+    xla::XlaOp actual_num_cols;
+    // The length of the longest diagonal from lower to upper at the actual
+    // sizes. Diagonals outside the actual matrix, which the CPU kernel
+    // rejects, have length 0.
+    const auto actual_len = [&](int64_t lower, int64_t upper) {
+      return xla::Max(
+          xla::Min(xla::Add(actual_num_rows,
+                            xla::ScalarLike(actual_num_rows,
+                                            std::min(upper, int64_t{0}))),
+                   xla::Sub(actual_num_cols,
+                            xla::ScalarLike(actual_num_cols,
+                                            std::max(lower, int64_t{0})))),
+          xla::ScalarLike(actual_num_rows, 0));
+    };
+    xla::XlaOp actual_max_diag_len;
+    if (dynamic_matrix) {
+      actual_num_rows = actual_size(input_rank - 2);
+      actual_num_cols = actual_size(input_rank - 1);
+      actual_max_diag_len = actual_len(lower_diag_index, upper_diag_index);
+    }
+
+    xla::XlaOp output;
     std::vector<xla::XlaOp> diag_list;
     xla::PaddingConfig padding_config =
         xla::MakeNoPaddingConfig(input_rank - 1);
+    // XLA:GPU can lay out the gather's output with the batch dimensions minor,
+    // and a batched dynamic output laid out that way comes back transposed.
+    const bool use_gather = is_gpu_ && !input_xla_shape.is_dynamic();
     if (num_diags == 1) {
-      context->SetOutput(
-          0, is_gpu_ ? xla::GetMatrixDiagonalViaGather(input, upper_diag_index)
-                     : xla::GetMatrixDiagonal(input, upper_diag_index));
-      return;
-    }
-    for (int diag_index = upper_diag_index; diag_index >= lower_diag_index;
-         --diag_index) {
-      xla::XlaOp single_diag =
-          is_gpu_ ? xla::GetMatrixDiagonalViaGather(input, diag_index)
-                  : xla::GetMatrixDiagonal(input, diag_index);
-      const int64_t diag_len = ComputeDiagLen(diag_index, num_rows, num_cols);
-      const int64_t padding_len = max_diag_len - diag_len;
-      if (padding_len > 0) {
-        if (IsLeftAligned(diag_index, left_align_superdiagonal_,
-                          left_align_subdiagonal_)) {
+      output = use_gather
+                   ? xla::GetMatrixDiagonalViaGather(input, upper_diag_index)
+                   : xla::GetMatrixDiagonal(input, upper_diag_index);
+    } else {
+      auto slice_sizes = input_shape.dim_sizes();
+      slice_sizes.pop_back();
+      slice_sizes.back() = max_diag_len;
+      for (int diag_index = upper_diag_index; diag_index >= lower_diag_index;
+           --diag_index) {
+        xla::XlaOp single_diag =
+            use_gather ? xla::GetMatrixDiagonalViaGather(input, diag_index)
+                       : xla::GetMatrixDiagonal(input, diag_index);
+        const int64_t diag_len = ComputeDiagLen(diag_index, num_rows, num_cols);
+        const int64_t padding_len = max_diag_len - diag_len;
+        const bool left_aligned = IsLeftAligned(
+            diag_index, left_align_superdiagonal_, left_align_subdiagonal_);
+        if (dynamic_matrix) {
+          // Place the diagonal after the padding it has at the actual sizes,
+          // and pad its entries past its actual length, which come from
+          // outside the actual matrix, with padding_value too.
+          const xla::XlaOp len = actual_len(diag_index, diag_index);
+          const xla::XlaOp offset = left_aligned
+                                        ? xla::ScalarLike(len, 0)
+                                        : xla::Sub(actual_max_diag_len, len);
           padding_config.mutable_dimensions(input_rank - 2)
-              ->set_edge_padding_low(0);
+              ->set_edge_padding_low(max_diag_len);
           padding_config.mutable_dimensions(input_rank - 2)
-              ->set_edge_padding_high(padding_len);
-        } else {
-          padding_config.mutable_dimensions(input_rank - 2)
-              ->set_edge_padding_low(padding_len);
-          padding_config.mutable_dimensions(input_rank - 2)
-              ->set_edge_padding_high(0);
+              ->set_edge_padding_high(max_diag_len);
+          single_diag = xla::Pad(single_diag, padding_value, padding_config);
+          std::vector<xla::XlaOp> start_indices(input_rank - 1,
+                                                xla::ScalarLike(len, 0));
+          start_indices.back() =
+              xla::Sub(xla::ScalarLike(len, max_diag_len), offset);
+          single_diag =
+              xla::DynamicSlice(single_diag, start_indices, slice_sizes);
+          const xla::XlaOp in_diagonal =
+              xla::Lt(xla::Iota(context->builder(), xla::S32, max_diag_len),
+                      xla::Add(offset, len));
+          single_diag = xla::Select(
+              xla::BroadcastInDim(in_diagonal, slice_sizes, {input_rank - 2}),
+              single_diag, xla::Broadcast(padding_value, slice_sizes));
+        } else if (padding_len > 0) {
+          if (left_aligned) {
+            padding_config.mutable_dimensions(input_rank - 2)
+                ->set_edge_padding_low(0);
+            padding_config.mutable_dimensions(input_rank - 2)
+                ->set_edge_padding_high(padding_len);
+          } else {
+            padding_config.mutable_dimensions(input_rank - 2)
+                ->set_edge_padding_low(padding_len);
+            padding_config.mutable_dimensions(input_rank - 2)
+                ->set_edge_padding_high(0);
+          }
+          single_diag = xla::Pad(single_diag, padding_value, padding_config);
         }
-        single_diag = xla::Pad(single_diag, padding_value, padding_config);
+        diag_list.emplace_back(single_diag);
       }
-      diag_list.emplace_back(single_diag);
+      auto concat =
+          xla::ConcatInDim(context->builder(), diag_list, input_rank - 2);
+      output = xla::Reshape(concat, output_shape.dim_sizes());
     }
-    auto concat =
-        xla::ConcatInDim(context->builder(), diag_list, input_rank - 2);
-    context->SetOutput(0, xla::Reshape(concat, output_shape.dim_sizes()));
+
+    if (dynamic_matrix) {
+      output = xla::SetDimensionSize(output, actual_max_diag_len,
+                                     output_shape.dims() - 1);
+    }
+    for (int i = 0; i < input_rank - 2; ++i) {
+      if (input_xla_shape.is_dynamic_dimension(i)) {
+        output = xla::SetDimensionSize(output, input_dim_sizes[i], i);
+      }
+    }
+    context->SetOutput(0, output);
   }
 
  private:
@@ -466,10 +558,34 @@ REGISTER_XLA_OP(Name("MatrixDiagPartV2")
                     .CompileTimeConstantInput("k")
                     .CompileTimeConstantInput("padding_value"),
                 MatrixDiagPartOp);
+
+// MatrixDiagPartV3 is lowered through MLIR, whose lowering only takes static
+// shapes, so an input with a dynamic dimension goes through MatrixDiagPartOp.
+class MatrixDiagPartV3Op : public XlaOpKernel {
+ public:
+  explicit MatrixDiagPartV3Op(OpKernelConstruction* context)
+      : XlaOpKernel(context),
+        mlir_kernel_(context),
+        dynamic_shape_kernel_(context) {}
+
+  void Compile(XlaOpKernelContext* context) override {
+    xla::Shape input_xla_shape;
+    OP_REQUIRES_VALUE(input_xla_shape, context, context->InputXlaShape(0));
+    XlaOpKernel& kernel = input_xla_shape.is_dynamic()
+                              ? static_cast<XlaOpKernel&>(dynamic_shape_kernel_)
+                              : static_cast<XlaOpKernel&>(mlir_kernel_);
+    kernel.Compile(context);
+  }
+
+ private:
+  MlirXlaOpKernel mlir_kernel_;
+  MatrixDiagPartOp dynamic_shape_kernel_;
+};
+
 REGISTER_XLA_OP(Name("MatrixDiagPartV3")
                     .CompileTimeConstantInput("k")
                     .CompileTimeConstantInput("padding_value"),
-                MlirXlaOpKernel);
+                MatrixDiagPartV3Op);
 
 class MatrixSetDiagOp : public XlaOpKernel {
  public:

tensorflow/compiler/tf2xla/kernels/pack_op.cc

@@ -64,7 +64,7 @@ class PackOp : public XlaOpKernel {
     std::vector<xla::XlaOp> reshaped_inputs(num);
 
     TensorShape child_shape(shapes[0]);
-    child_shape.InsertDim(axis, 1);
+    OP_REQUIRES_OK(ctx, child_shape.InsertDimWithStatus(axis, 1));
 
     for (int i = 0; i < num; ++i) {
       // Reshape the inputs to have an extra dimension of size 1.

tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc

@@ -112,6 +112,13 @@ class QuantizeAndDequantizeOp : public XlaOpKernel {
       OP_REQUIRES(
           ctx, ctx->num_inputs() == 4,
           errors::Internal("Expected 4 inputs to QuantizeAndDequantize"));
+      int64_t num_bits_val;
+      if (ctx->ConstantInputAsIntScalar(3, &num_bits_val).ok()) {
+        OP_REQUIRES(
+            ctx, num_bits_val > 0 && num_bits_val < (signed_input_ ? 62 : 63),
+            errors::InvalidArgument("num_bits is out of range: ", num_bits_val,
+                                    " with signed_input_ ", signed_input_));
+      }
       num_bits = ctx->Input(3);
     } else {
       num_bits = xla::ConstantR0<int32_t>(b, num_bits_);

tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc

@@ -19,6 +19,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "xla/hlo/builder/lib/constants.h"
 #include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/xla_data.pb.h"
 #include "tensorflow/core/framework/op_kernel.h"
@@ -129,8 +130,21 @@ class ReverseSequenceOp : public XlaOpKernel {
     slice_sizes[batch_dim_] = 1;
     slice_sizes[seq_dim_] = 1;
 
-    context->SetOutput(0,
-                       xla::Gather(input, start_indices, dnums, slice_sizes));
+    xla::XlaOp output = xla::Gather(input, start_indices, dnums, slice_sizes);
+    // The gather indices and slice sizes come from each dimension's static
+    // bound, so the output drops the input's dynamic sizes. Its dimensions are
+    // the input's, in the same order, so give them the input's sizes.
+    xla::Shape input_xla_shape;
+    OP_REQUIRES_VALUE(input_xla_shape, context, context->InputXlaShape(0));
+    if (input_xla_shape.is_dynamic()) {
+      for (int i = 0; i < input_shape.dims(); ++i) {
+        if (input_xla_shape.is_dynamic_dimension(i)) {
+          output =
+              xla::SetDimensionSize(output, xla::GetDimensionSize(input, i), i);
+        }
+      }
+    }
+    context->SetOutput(0, output);
   }
 
  private:

tensorflow/compiler/tf2xla/kernels/sequence_ops.cc

@@ -18,6 +18,7 @@ limitations under the License.
 #include <cmath>
 #include <cstdint>
 #include <cstdlib>
+#include <limits>
 #include <type_traits>
 
 #include "absl/status/statusor.h"
@@ -40,7 +41,8 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
-// The type-specific part of the implementation of Range.
+// The type-specific part of the implementation of Range: checks the arguments
+// and returns delta * [0, 1, ..., size - 1] for the range's size.
 template <typename T>
 absl::StatusOr<xla::XlaOp> CreateRangeTensor(
     const xla::LiteralSlice& start_literal,
@@ -65,13 +67,22 @@ absl::StatusOr<xla::XlaOp> CreateRangeTensor(
     }
   }
   int64_t size;
-  if constexpr (std::is_integral<T>::value) {
-    int64_t start_i = static_cast<int64_t>(start);
-    int64_t limit_i = static_cast<int64_t>(limit);
-    int64_t delta_i = static_cast<int64_t>(delta);
-    size = (limit_i == start_i
-                ? 0
-                : (std::abs(limit_i - start_i) - 1) / std::abs(delta_i) + 1);
+  if constexpr (std::is_integral_v<T>) {
+    // Take |limit - start| and |delta| in uint64_t, which holds them for every
+    // T, as the CPU kernel does, so that a range wider than T can hold doesn't
+    // overflow.
+    const bool count_up = delta > static_cast<T>(0);
+    const uint64_t range =
+        count_up ? static_cast<uint64_t>(limit) - static_cast<uint64_t>(start)
+                 : static_cast<uint64_t>(start) - static_cast<uint64_t>(limit);
+    const uint64_t abs_delta =
+        count_up ? static_cast<uint64_t>(delta) : -static_cast<uint64_t>(delta);
+    const uint64_t size_unsigned = range == 0 ? 0 : (range - 1) / abs_delta + 1;
+    if (size_unsigned > std::numeric_limits<int64_t>::max()) {
+      return errors::InvalidArgument("Requires ((limit - start) / delta) <= ",
+                                     std::numeric_limits<int64_t>::max());
+    }
+    size = static_cast<int64_t>(size_unsigned);
   } else {
     double start_f = static_cast<double>(start);
     double limit_f = static_cast<double>(limit);
@@ -80,10 +91,9 @@ absl::StatusOr<xla::XlaOp> CreateRangeTensor(
         std::ceil(std::abs((limit_f - start_f) / delta_f)));
   }
 
-  return xla::ConstantR0(builder, start) +
-         xla::ConstantR0(builder, delta) *
-             xla::Iota(builder, xla::primitive_util::NativeToPrimitiveType<T>(),
-                       size);
+  return xla::ConstantR0(builder, delta) *
+         xla::Iota(builder, xla::primitive_util::NativeToPrimitiveType<T>(),
+                   size);
 }
 
 class RangeOp : public XlaOpKernel {
@@ -104,11 +114,24 @@ class RangeOp : public XlaOpKernel {
                 errors::InvalidArgument("delta must be a scalar, not shape ",
                                         delta_in_shape.DebugString()));
     xla::Literal start, limit, delta;
-    OP_REQUIRES_OK(ctx, ctx->ConstantInput(
-                            0, &start, xla::ValueInferenceMode::kLowerBound));
-    OP_REQUIRES_OK(ctx, ctx->ConstantInput(
-                            1, &limit, xla::ValueInferenceMode::kUpperBound));
     OP_REQUIRES_OK(ctx, ctx->ConstantInput(2, &delta));
+    // Size the output from the bounds of start and limit that give the longest
+    // range: the lowest start and highest limit when counting up, and the
+    // highest start and lowest limit when counting down.
+    const xla::PrimitiveType delta_type = delta.shape().element_type();
+    const bool count_down =
+        xla::primitive_util::IsFloatingPointType(delta_type)
+            ? delta.GetAsDouble({}).value_or(0) < 0
+            : xla::primitive_util::IsSignedIntegralType(delta_type) &&
+                  delta.GetIntegralAsS64({}).value_or(0) < 0;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInput(
+                            0, &start,
+                            count_down ? xla::ValueInferenceMode::kUpperBound
+                                       : xla::ValueInferenceMode::kLowerBound));
+    OP_REQUIRES_OK(ctx, ctx->ConstantInput(
+                            1, &limit,
+                            count_down ? xla::ValueInferenceMode::kLowerBound
+                                       : xla::ValueInferenceMode::kUpperBound));
 
     DataType type = input_type(0);
     absl::StatusOr<xla::XlaOp> output;
@@ -165,22 +188,50 @@ class RangeOp : public XlaOpKernel {
     bool limit_is_dynamic = false;
     OP_REQUIRES_OK(ctx,
                    ctx->ResolveInputDynamismIntoPred(1, &limit_is_dynamic));
+    // If start is dynamic, its literal is a bound, so add the input instead.
+    output = (start_is_dynamic ? ctx->Input(0)
+                               : xla::ConstantLiteral(ctx->builder(), start)) +
+             output.value();
 
     if (start_is_dynamic || limit_is_dynamic) {
       xla::XlaOp delta = ctx->Input(2);
       xla::XlaOp limit = ctx->Input(1);
       xla::XlaOp start = ctx->Input(0);
+      // The runtime values can be in the wrong order for delta's sign, which
+      // the CPU kernel rejects. Give such a range size 0 rather than a size
+      // past the output's static size.
+      const xla::XlaOp is_empty =
+          count_down ? xla::Le(start, limit) : xla::Ge(start, limit);
+      xla::XlaOp dynamic_size;
       if (DataTypeIsInteger(type)) {
-        auto dynamic_size = (xla::Abs(limit - start) + xla::Abs(delta) -
-                             xla::One(ctx->builder(), ctx->input_xla_type(0))) /
-                            xla::Abs(delta);
-        dynamic_size = xla::ConvertElementType(dynamic_size, xla::S32);
-        output = xla::SetDimensionSize(output.value(), dynamic_size, 0);
+        // Take |limit - start| and |delta| in the unsigned type of the same
+        // width, which holds them. XLA's Abs doesn't take unsigned types, whose
+        // delta is positive anyway.
+        const xla::PrimitiveType xla_type = ctx->input_xla_type(0);
+        const xla::PrimitiveType unsigned_type =
+            xla::primitive_util::UnsignedIntegralTypeForBitWidth(
+                xla::primitive_util::BitWidth(xla_type));
+        const auto as_unsigned = [&](xla::XlaOp x) {
+          return xla::BitcastConvertType(x, unsigned_type);
+        };
+        const xla::XlaOp range = count_down
+                                     ? as_unsigned(start) - as_unsigned(limit)
+                                     : as_unsigned(limit) - as_unsigned(start);
+        const xla::XlaOp abs_delta =
+            as_unsigned(xla::primitive_util::IsSignedIntegralType(xla_type)
+                            ? xla::Abs(delta)
+                            : delta);
+        const xla::XlaOp one = xla::One(ctx->builder(), unsigned_type);
+        dynamic_size =
+            xla::Select(is_empty, xla::Zero(ctx->builder(), unsigned_type),
+                        (range - one) / abs_delta + one);
       } else {
-        auto dynamic_size = (xla::Ceil(xla::Abs((limit - start) / delta)));
-        dynamic_size = xla::ConvertElementType(dynamic_size, xla::S32);
-        output = xla::SetDimensionSize(output.value(), dynamic_size, 0);
+        dynamic_size = xla::Select(
+            is_empty, xla::Zero(ctx->builder(), ctx->input_xla_type(0)),
+            xla::Ceil(xla::Abs((limit - start) / delta)));
       }
+      dynamic_size = xla::ConvertElementType(dynamic_size, xla::S32);
+      output = xla::SetDimensionSize(output.value(), dynamic_size, 0);
     }
 
     ctx->SetOutput(0, output.value());

tensorflow/compiler/tf2xla/kernels/shape_op.cc

@@ -289,6 +289,12 @@ class ExpandDimsOp : public XlaOpKernel {
                                         " for tensor with ", input_shape.dims(),
                                         " dimensions."));
 
+    OP_REQUIRES(ctx, input_shape.dims() < TensorShape::MaxDimensions(),
+                errors::InvalidArgument(
+                    "Cannot expand a tensor of rank ", input_shape.dims(),
+                    ": expanding would exceed the maximum supported rank of ",
+                    TensorShape::MaxDimensions(), "."));
+
     auto existing_dims = input_shape.dim_sizes();
     // Safe - # elements in tensor dims bounded.
     const int existing_dims_size = static_cast<int>(existing_dims.size());

tensorflow/compiler/tf2xla/kernels/unary_ops.cc

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

tensorflow/compiler/tf2xla/resource_util.h

@@ -38,7 +38,7 @@ class ResourceUsageAnalysis {
     std::string node_name_;
     std::string op_;
 
-    NodeInfo() {}
+    NodeInfo() = default;
 
     NodeInfo(const std::optional<std::string>& function_name,
              std::string node_name, std::string op)