[pull] master from tensorflow:master

MODULE.bazel

@@ -133,6 +133,7 @@ use_repo(
     "nvshmem",
     "onednn",
     "onednn_async",
+    "onednn_async_aarch64",
     "pthreadpool",
     "raft",
     "riegeli",

tensorflow/core/kernels/BUILD

@@ -3821,6 +3821,15 @@ tf_kernel_library(
     deps = MATH_DEPS + ["//tensorflow/core:framework_internal"],
 )
 
+tf_kernel_library(
+    name = "integer_pow_gpu",
+    prefix = "integer_pow_gpu",
+    deps = MATH_DEPS + [
+        "//tensorflow/core:framework_internal",
+        "@com_google_absl//absl/status",
+    ],
+)
+
 tf_kernel_library(
     name = "cross_op",
     prefix = "cross_op",
@@ -3842,6 +3851,7 @@ tf_kernel_library(
     ),
     prefix = "cwise_op",
     deps = MATH_DEPS + [
+        ":integer_pow_gpu",
         "//tensorflow/core/kernels/mlir_generated:cwise_op",
         "@com_google_absl//absl/base:prefetch",
         "@com_google_absl//absl/status",

tensorflow/core/kernels/cwise_op_pow.cc

@@ -15,6 +15,10 @@ limitations under the License.
 
 #include "tensorflow/core/kernels/cwise_ops_common.h"
 
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#include "tensorflow/core/kernels/integer_pow_gpu.h"
+#endif
+
 namespace tensorflow {
 REGISTER6(BinaryOp, CPU, "Pow", functor::pow, float, Eigen::half, bfloat16,
           double, complex64, complex128);
@@ -24,7 +28,10 @@ REGISTER4(BinaryOp, CPU, "Pow", functor::safe_pow, int8_t, int16_t, int32_t,
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 #if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 REGISTER3(BinaryOp, GPU, "Pow", functor::pow, float, Eigen::half, double);
-REGISTER(BinaryOp, GPU, "Pow", functor::safe_pow_ignore_error, int64);
+using Int64GpuPowOp = IntegerPowGpuOp<
+    int64_t, BinaryOp<GPUDevice, functor::safe_pow_ignore_error<int64_t>>>;
+REGISTER_KERNEL_BUILDER(
+    Name("Pow").Device(DEVICE_GPU).TypeConstraint<int64_t>("T"), Int64GpuPowOp);
 #endif
 REGISTER(BinaryOp, GPU, "Pow", functor::pow, bfloat16);
 #endif

tensorflow/core/kernels/cwise_ops.h

@@ -1221,8 +1221,9 @@ struct safe_pow : base<T, Eigen::internal::safe_scalar_binary_pow_op<T, T>> {
   static constexpr bool has_errors = true;
 };
 
-// Version of safe_pow for integers which returns 0 if RHS is negative and LHS
-// is not 1 or -1. For use on GPUs, where we cannot raise an error.
+// Version of safe_pow for GPU computation. The integer Pow GPU kernel checks
+// negative exponents separately and reports an error after the device work
+// completes; this functor must still produce a defined result on the device.
 template <typename T>
 struct safe_pow_ignore_error_op {
   static_assert(std::is_integral<T>::value, "Integer type expected");

tensorflow/core/kernels/integer_pow_gpu.cu.cc

@@ -0,0 +1,58 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#define EIGEN_USE_GPU
+#include "tensorflow/core/kernels/integer_pow_gpu.h"
+
+#include <algorithm>
+#include <cstdint>
+
+#include "tensorflow/core/util/gpu_kernel_helper.h"
+
+namespace tensorflow {
+namespace {
+
+template <typename T>
+__global__ void CheckNegativeExponentKernel(const T* exponents, int64_t size,
+                                            int32_t* found_negative) {
+  for (int64_t i : GpuGridRangeX<int64_t>(size)) {
+    if (exponents[i] < 0) {
+      atomicExch(found_negative, 1);
+    }
+  }
+}
+
+}  // namespace
+
+template <typename T>
+absl::Status CheckNegativeExponentGpu<T>::operator()(
+    const Eigen::GpuDevice& device, const T* exponents, int64_t size,
+    int32_t* found_negative) {
+  constexpr int kThreads = 256;
+  const int blocks = static_cast<int>(
+      std::min<int64_t>((size + kThreads - 1) / kThreads, 1024));
+  return GpuLaunchKernel(CheckNegativeExponentKernel<T>, blocks, kThreads, 0,
+                         device.stream(), exponents, size, found_negative);
+}
+
+template struct CheckNegativeExponentGpu<int8_t>;
+template struct CheckNegativeExponentGpu<int16_t>;
+template struct CheckNegativeExponentGpu<int64_t>;
+
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM

tensorflow/core/kernels/integer_pow_gpu.h

@@ -0,0 +1,119 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_KERNELS_INTEGER_POW_GPU_H_
+#define TENSORFLOW_CORE_KERNELS_INTEGER_POW_GPU_H_
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#include <cstdint>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor_types.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/util/bcast.h"
+
+namespace tensorflow {
+
+// Instantiated in integer_pow_gpu.cu.cc for the signed GPU Pow types.
+template <typename T>
+struct CheckNegativeExponentGpu {
+  absl::Status operator()(const Eigen::GpuDevice& device, const T* exponents,
+                          int64_t size, int32_t* found_negative);
+};
+
+// Applies the same exponent validation to the Eigen and MLIR GPU kernels.
+// The check runs before Pow can reuse an exponent buffer for its output. Only
+// a single flag is copied back to the host, without blocking the caller.
+template <typename T, typename Kernel>
+class IntegerPowGpuOp : public AsyncOpKernel {
+ public:
+  explicit IntegerPowGpuOp(OpKernelConstruction* context)
+      : AsyncOpKernel(context), kernel_(context) {}
+
+  void ComputeAsync(OpKernelContext* context, DoneCallback done) override {
+    const BCast bcast(BCast::FromShape(context->input(0).shape()),
+                      BCast::FromShape(context->input(1).shape()));
+    if (!bcast.IsValid() ||
+        BCast::ToShape(bcast.output_shape()).num_elements() == 0) {
+      kernel_.Compute(context);
+      done();
+      return;
+    }
+
+    auto* stream = context->op_device_context()->stream();
+    OP_REQUIRES_ASYNC(context, stream != nullptr,
+                      absl::InternalError("No GPU stream available."), done);
+
+    Tensor found_negative;
+    OP_REQUIRES_OK_ASYNC(
+        context,
+        context->allocate_temp(DT_INT32, TensorShape({}), &found_negative),
+        done);
+    stream_executor::DeviceAddressBase device_flag(
+        found_negative.flat<int32_t>().data(), sizeof(int32_t));
+    OP_REQUIRES_OK_ASYNC(
+        context, stream->Memset32(&device_flag, 0, sizeof(int32_t)), done);
+    const auto& exponents = context->input(1).flat<T>();
+    OP_REQUIRES_OK_ASYNC(
+        context,
+        CheckNegativeExponentGpu<T>()(context->eigen_device<Eigen::GpuDevice>(),
+                                      exponents.data(), exponents.size(),
+                                      found_negative.flat<int32_t>().data()),
+        done);
+
+    AllocatorAttributes attributes;
+    attributes.set_on_host(true);
+    attributes.set_gpu_compatible(true);
+    Tensor host_flag;
+    OP_REQUIRES_OK_ASYNC(context,
+                         context->allocate_temp(DT_INT32, TensorShape({}),
+                                                &host_flag, attributes),
+                         done);
+    OP_REQUIRES_ASYNC(
+        context,
+        stream
+            ->Memcpy(host_flag.flat<int32_t>().data(), device_flag,
+                     sizeof(int32_t))
+            .ok(),
+        absl::InternalError("GPU memcpy from device to host failed"), done);
+
+    kernel_.Compute(context);
+
+    auto check = [found_negative, host_flag, context, done]() {
+      // Keep the device flag alive until the queued GPU work has completed.
+      (void)found_negative;
+      if (context->status().ok() && host_flag.scalar<int32_t>()() != 0) {
+        context->SetStatus(absl::InvalidArgumentError(
+            "Integers to negative integer powers are not allowed"));
+      }
+      done();
+    };
+    context->device()
+        ->tensorflow_accelerator_device_info()
+        ->event_mgr->ThenExecute(stream, std::move(check));
+  }
+
+ private:
+  Kernel kernel_;
+};
+
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#endif  // TENSORFLOW_CORE_KERNELS_INTEGER_POW_GPU_H_

tensorflow/core/kernels/mlir_generated/BUILD

@@ -234,6 +234,7 @@ tf_kernel_library(
     tags = ["manual"],
     deps = if_mlir_generated_gpu_kernels_enabled([
         ":base_gpu_op",
+        "//tensorflow/core/kernels:integer_pow_gpu",
         ":gpu_atan2_kernels",
         ":gpu_bitwise_and_kernels",
         ":gpu_bitwise_or_kernels",

tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc

@@ -1304,6 +1304,13 @@ absl::InlinedVector<T, 10> PowInput() {
   return test::InputAsVector<T, double>({-2, -1, -1, 1, 1, 3});
 }
 
+template <typename T>
+absl::InlinedVector<T, 10> IntegerPowExponents() {
+  // These generated tests expect successful results. Negative exponents are
+  // covered separately by the GPU Pow error tests.
+  return test::InputAsVector<T>({0, 1, 2, 1, 2, 3});
+}
+
 template <>
 Eigen::half baseline_pow(Eigen::half lhs, Eigen::half rhs) {
   return static_cast<Eigen::half>(
@@ -1327,18 +1334,18 @@ GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
 GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
     Pow,
     /*test_name=*/Int64, int64_t, int64_t, PowInput<int64_t>(),
-    PowInput<int64_t>(), baseline_pow,
+    IntegerPowExponents<int64_t>(), baseline_pow,
     test::OpsTestConfig().ExpectStrictlyEqual())
 
 /// Test the JIT-compiled kernels.
 #if defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
 GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
     Pow, /*test_name=*/Int8, int8_t, int8_t, PowInput<int8_t>(),
-    PowInput<int8_t>(), baseline_pow,
+    IntegerPowExponents<int8_t>(), baseline_pow,
     test::OpsTestConfig().ExpectStrictlyEqual())
 GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES(
     Pow, /*test_name=*/Int16, int16_t, int16_t, PowInput<int16_t>(),
-    PowInput<int16_t>(), baseline_pow,
+    IntegerPowExponents<int16_t>(), baseline_pow,
     test::OpsTestConfig().ExpectStrictlyEqual())
 #endif
 

tensorflow/core/kernels/mlir_generated/gpu_op_pow.cc

@@ -14,17 +14,37 @@ limitations under the License.
 ==============================================================================*/
 #include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
 #include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/kernels/integer_pow_gpu.h"
 #include "tensorflow/core/kernels/mlir_generated/base_gpu_op.h"
 
 namespace tensorflow {
 
 GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Pow, DT_HALF);
 GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Pow, DT_FLOAT);
 GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Pow, DT_DOUBLE);
-GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Pow, DT_INT64);
+GENERATE_BINARY_GPU_KERNEL(Pow, DT_INT64);
+using Int64GpuPowOp =
+    IntegerPowGpuOp<int64_t, MLIR_OP(Pow, GPU, DT_INT64, DT_INT64)>;
+REGISTER_KERNEL_BUILDER(
+    Name("Pow").Device(DEVICE_GPU).TypeConstraint<int64_t>("T"), Int64GpuPowOp);
 
 // These kernels are JIT-compiled.
-GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Pow, DT_INT8);
-GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Pow, DT_INT16);
+GENERATE_BINARY_GPU_KERNEL(Pow, DT_INT8);
+using Int8GpuPowOp =
+    IntegerPowGpuOp<int8_t, MLIR_OP(Pow, GPU, DT_INT8, DT_INT8)>;
+REGISTER_KERNEL_BUILDER(Name("Pow")
+                            .Device(DEVICE_GPU)
+                            .TypeConstraint<int8_t>("T")
+                            .Label(kJitKernelLabel),
+                        Int8GpuPowOp);
+
+GENERATE_BINARY_GPU_KERNEL(Pow, DT_INT16);
+using Int16GpuPowOp =
+    IntegerPowGpuOp<int16_t, MLIR_OP(Pow, GPU, DT_INT16, DT_INT16)>;
+REGISTER_KERNEL_BUILDER(Name("Pow")
+                            .Device(DEVICE_GPU)
+                            .TypeConstraint<int16_t>("T")
+                            .Label(kJitKernelLabel),
+                        Int16GpuPowOp);
 
 }  // namespace tensorflow

tensorflow/lite/kernels/BUILD

@@ -2527,10 +2527,17 @@ cc_test(
     deps = [
         ":test_main",
         ":test_util",
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels/internal:reference_base",
+        "//tensorflow/lite/kernels/internal:tensor",
+        "//tensorflow/lite/kernels/internal:types",
         "//tensorflow/lite/schema:schema_fbs",
         "//tensorflow/lite/types:half",
+        "@com_google_benchmark//:benchmark",
         "@com_google_googletest//:gtest",
         "@eigen_archive//:eigen3",
+        "@ruy//ruy/profiler:instrumentation",
     ],
 )