[pull] master from tensorflow:master

tensorflow/compiler/mlir/tensorflow/analysis/resource_value_typed_analyzer.cc

@@ -148,10 +148,10 @@ LogicalResult ResourceAnalyzer::AnalyzeRegion(Region& region) {
       return;
     }
     if (auto if_op = dyn_cast<TF::IfRegionOp>(op)) {
-      PropagatePotentiallyWrittenUpFromCallee(if_op.getThenBranch(),
-                                              if_op.getODSOperands(1));
-      PropagatePotentiallyWrittenUpFromCallee(if_op.getElseBranch(),
-                                              if_op.getODSOperands(1));
+      PropagatePotentiallyWrittenUpFromCallee(
+          if_op.getThenBranch(), if_op->getOperands().drop_front());
+      PropagatePotentiallyWrittenUpFromCallee(
+          if_op.getElseBranch(), if_op->getOperands().drop_front());
       return;
     }
     if (auto while_op = dyn_cast<TF::WhileOp>(op)) {

tensorflow/compiler/tests/bucketize_op_test.py

@@ -34,6 +34,15 @@ def testInt(self):
       self.assertAllEqual(expected_out,
                           sess.run(op, {p: [-5, 0, 2, 3, 5, 8, 10, 11, 12]}))
 
+  def testIntWithFloatBoundaries(self):
+    for dtype in [dtypes.int32, dtypes.int64]:
+      with self.session() as sess:
+        p = array_ops.placeholder(dtype)
+        with self.test_scope():
+          op = math_ops._bucketize(p, boundaries=[0.1, 1.1])
+        expected_out = [0, 1, 2, 2]
+        self.assertAllEqual(expected_out, sess.run(op, {p: [0, 1, 2, 3]}))
+
   def testFloat(self):
     with self.session() as sess:
       p = array_ops.placeholder(dtypes.float32)

tensorflow/compiler/tf2xla/kernels/in_topk_op.cc

@@ -28,8 +28,6 @@ limitations under the License.
 #include "tensorflow/core/framework/op_requires.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.pb.h"
-#include "tensorflow/core/platform/errors.h"
-#include "tensorflow/core/platform/types.h"
 
 namespace tensorflow {
 namespace {
@@ -79,6 +77,17 @@ class InTopKOp : public XlaOpKernel {
     xla::XlaOp iota_r2 = xla::Broadcast(iota_r1, {batch_size});
 
     xla::XlaOp eq_r2 = xla::Eq(targets_r1, iota_r2, {0});
+    // The masked reduction below treats an out-of-range target as zero, which
+    // can otherwise produce a false hit. Check bounds per target, rather than
+    // reducing the class-match matrix again.
+    xla::XlaOp targets_s64 = xla::ConvertElementType(targets_r1, xla::S64);
+    // This folds to a constant for static dimensions, while using the actual
+    // class count instead of the static bound for a dynamic dimension.
+    xla::XlaOp num_classes_s64 = xla::ConvertElementType(
+        xla::GetDimensionSize(predictions_r2, 1), xla::S64);
+    xla::XlaOp valid_target_r1 =
+        xla::And(xla::Ge(targets_s64, xla::ConstantR0<int64_t>(xla_builder, 0)),
+                 xla::Lt(targets_s64, num_classes_s64));
     xla::XlaOp zero_r0_f32 = xla::Zero(xla_builder, xla::F32);
     xla::XlaOp zero_r2_f32 = xla::ZerosLike(predictions_r2);
     xla::XlaOp select_r2 = xla::Select(eq_r2, predictions_r2, zero_r2_f32);
@@ -99,9 +108,10 @@ class InTopKOp : public XlaOpKernel {
         one_hot_r2, zero_r0,
         xla::CreateScalarAddComputation(xla::S32, xla_builder), {1});
 
-    xla::XlaOp result =
+    xla::XlaOp result = xla::And(
         xla::And(xla::Lt(num_gt_r1, xla::ConstantR0<int32_t>(xla_builder, k)),
-                 xla::IsFinite(targets_values_r1));
+                 xla::IsFinite(targets_values_r1)),
+        valid_target_r1);
 
     context->SetOutput(0, result);
   }

tensorflow/core/kernels/bucketize_op_gpu.cu.cc

@@ -17,6 +17,8 @@ limitations under the License.
 
 #define EIGEN_USE_GPU
 
+#include <type_traits>
+
 #include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
@@ -55,15 +57,21 @@ __global__ void BucketizeCustomKernel(
     boundaries = shared_mem_boundaries;
   }
 
+  // For double input, compare as double. For float and integer inputs, compare
+  // as float to match CPU and XLA behavior without truncating fractional
+  // boundaries.
+  using CompType = std::conditional_t<std::is_same_v<T, double>, double, float>;
+
   GPU_1D_KERNEL_LOOP(i, size_in) {
     T value = in[i];
+    CompType value_cast = static_cast<CompType>(value);
     int32_t bucket = 0;
     int32_t count = size_boundaries;
     while (count > 0) {
       int32_t l = bucket;
       int32_t step = count / 2;
       l += step;
-      if (!(value < static_cast<T>(boundaries[l]))) {
+      if (!(value_cast < static_cast<CompType>(boundaries[l]))) {
         bucket = ++l;
         count -= step + 1;
       } else {

tensorflow/core/kernels/topk_op.cc

@@ -20,6 +20,10 @@ limitations under the License.
 #include "tensorflow/core/kernels/topk_op.h"
 
 #include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <limits>
+#include <memory>
 #include <numeric>
 #include <vector>
 
@@ -30,6 +34,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/lib/gtl/top_n.h"
+#include "tensorflow/core/platform/macros.h"  // NOLINT(misc-include-cleaner)
 #include "tensorflow/core/util/work_sharder.h"
 
 namespace tensorflow {
@@ -132,53 +137,76 @@ class TopK : public OpKernel {
 
 namespace functor {
 
+namespace {
+
+// Returns +1 for +NaN, -1 for -NaN, and 0 for any other value. Ordering by
+// value and then by this rank gives +NaN > +Inf > ... > -Inf > -NaN, which
+// matches the GPU radix-sort key order and the XLA total order for NaNs.
+// +0.0 and -0.0 compare equal and are ordered by index.
+template <typename T>
+EIGEN_ALWAYS_INLINE int NanRank(const T v) {  // NOLINT(misc-include-cleaner)
+  // NOLINTNEXTLINE(misc-include-cleaner)
+  if (TF_PREDICT_TRUE(!Eigen::numext::isnan(v))) return 0;
+  // Eigen::numext::signbit returns a Scalar bit mask (all ones, itself a NaN
+  // for floating types), not a bool, so test the sign with std::signbit. The
+  // conversion to double preserves the sign bit of a NaN for every float type.
+  return std::signbit(static_cast<double>(v)) ? -1 : 1;
+}
+
+}  // namespace
+
 template <typename T, typename Tidx>
 struct TopKFunctor<CPUDevice, T, Tidx> {
   static EIGEN_ALWAYS_INLINE absl::Status Compute(
       OpKernelContext* context, bool sorted, int k,
       const typename TTypes<T, 2>::ConstTensor& input, const int64_t num_rows,
       const int64_t num_cols, typename TTypes<T, 2>::Tensor values,
       typename TTypes<Tidx, 2>::Tensor indices) {
-    const CPUDevice& d = context->eigen_device<CPUDevice>();
-
-    // Special case for k == 1.
-    if (k == 1) {
-      typename Eigen::IndexList<Eigen::type2index<1>> reduce_on_cols;
-      typename Eigen::IndexList<int, Eigen::type2index<1>> rows_by_one;
-      rows_by_one.set(0, num_rows);
-
-      values.device(d) =
-          input.maximum(/*dims=*/reduce_on_cols).eval().reshape(rows_by_one);
-      // Get the indices of the maximum values.
-      for (int r = 0; r < num_rows; ++r) {
-        indices(r, 0) = Tidx(0);
-        for (int c = 0; c < num_cols; ++c) {
-          if (values(r, 0) == input(r, c)) {
-            indices(r, 0) = static_cast<Tidx>(c);
-            break;
-          }
-        }
-        values(r, 0) = input(r, indices(r, 0));
-      }
-
-      return absl::OkStatus();
-    }
-
     auto SortIndices = [&](int64_t start_batch, int64_t limit_batch) {
-      for (int32_t b = start_batch; b < limit_batch; ++b) {
+      for (int64_t b = start_batch; b < limit_batch; ++b) {
         const T* input_data = &input(b, 0);
-        const auto stable_comp = [input_data](const int32_t a,
-                                              const int32_t b) {
-          if (input_data[b] < input_data[a]) {
+        // Both comparators order values descending with the relational
+        // comparisons first (the common path). Equal values (never NaN, and
+        // including +0/-0) tie; only a NaN operand falls back to the NaN rank,
+        // so NaNs are ordered +NaN first and -NaN last.
+        const auto stable_comp = [input_data](const Tidx a, const Tidx b) {
+          const T val_a = input_data[a];
+          const T val_b = input_data[b];
+          if (val_b < val_a) {
             return true;
-          } else if (input_data[b] > input_data[a]) {
+          }
+          if (val_a < val_b) {
             return false;
-          } else {
+          }
+          if (val_a == val_b) {
             return a < b;
           }
+          if constexpr (!Eigen::NumTraits<T>::IsInteger) {
+            const int rank_a = NanRank(val_a);
+            const int rank_b = NanRank(val_b);
+            if (rank_a != rank_b) {
+              return rank_b < rank_a;
+            }
+          }
+          return a < b;
         };
-        const auto comp = [input_data](const int32_t a, const int32_t b) {
-          return input_data[b] < input_data[a];
+        const auto comp = [input_data](const Tidx a, const Tidx b) {
+          const T val_a = input_data[a];
+          const T val_b = input_data[b];
+          if (val_b < val_a) {
+            return true;
+          }
+          if (val_a < val_b) {
+            return false;
+          }
+          if (val_a == val_b) {
+            return false;
+          }
+          if constexpr (!Eigen::NumTraits<T>::IsInteger) {
+            return NanRank(val_b) < NanRank(val_a);
+          } else {
+            return false;
+          }
         };
         // TODO(ebrevdo): For large k < num_cols, instead of using
         // TopN, it may be faster to create a temporary vector of
@@ -194,19 +222,50 @@ struct TopKFunctor<CPUDevice, T, Tidx> {
           // indices that started out sorted.  First, do a std::sort, which
           // is notably faster than std::stable_sort.
           std::sort(begin, end, comp);
-          // Then, for runs of adjacent elements that were equal, sort the
-          // indices in those runs in increasing order.
+          // Then, for runs of adjacent elements that are equivalent under
+          // `comp` (equal values, or NaNs of the same sign), sort the indices
+          // in those runs in increasing order.
           for (auto* run_begin = begin; run_begin != end;) {
             auto* run_end = run_begin + 1;
-            if (run_end == end) break;
-            if (input_data[*run_begin] == input_data[*run_end]) {
-              while (++run_end != end) {
-                if (input_data[*run_begin] != input_data[*run_end]) break;
+            while (run_end != end && !comp(*run_begin, *run_end)) ++run_end;
+            if (run_end - run_begin > 1) std::sort(run_begin, run_end);
+            run_begin = run_end;
+          }
+        } else if (k == 1) {
+          // A single scan instead of an Eigen `maximum` reduction: the NaN
+          // propagation of the reduction is unspecified, and locating its
+          // result by equality can never find a NaN. Returns the first +NaN
+          // if any, else the first maximum non-NaN value, else the first -NaN.
+          Tidx best_c = 0;
+          T best_val = input_data[0];
+          Tidx first_c = 1;
+          if constexpr (!Eigen::NumTraits<T>::IsInteger) {
+            // A leading +NaN is the maximum: skip the scan.
+            if (NanRank(best_val) > 0) first_c = static_cast<Tidx>(num_cols);
+          }
+          for (Tidx c = first_c; c < num_cols; ++c) {
+            const T val = input_data[c];
+            // '<=' is false only if val > best_val or either operand is NaN.
+            if (val <= best_val) continue;
+            if (best_val < val) {
+              best_val = val;
+              best_c = c;
+              continue;
+            }
+            if constexpr (!Eigen::NumTraits<T>::IsInteger) {
+              const int rank = NanRank(val);
+              if (rank > 0) {  // +NaN: the global maximum, stop.
+                best_c = c;
+                break;
+              }
+              if (rank == 0) {  // best_val is -NaN; any non-NaN beats it.
+                best_val = val;
+                best_c = c;
               }
-              std::sort(run_begin, run_end);
+              // rank < 0: a -NaN never beats the current best.
             }
-            run_begin = run_end;
           }
+          indices(b, 0) = best_c;
         } else {
           // Use the TopN heap object to sort.
           gtl::TopN<Tidx, decltype(stable_comp)> filter(k, stable_comp);
@@ -215,7 +274,7 @@ struct TopKFunctor<CPUDevice, T, Tidx> {
             filter.push(c);
           }
 
-          int32_t i = 0;
+          Tidx i = 0;
           if (sorted) {
             std::unique_ptr<std::vector<Tidx>> top_k(filter.Extract());
             for (auto top_k_it = top_k->begin(); top_k_it != top_k->end();

tensorflow/python/eager/BUILD

@@ -1151,7 +1151,7 @@ cuda_py_strict_test(
     srcs = ["remote_test.py"],
     exec_properties = select({
         # copybara:uncomment_begin(google-only)
-        # "//tools/cpp:msan_build": {"mem": "30g"},
+        # "//tools/cpp:msan_build": {"mem": "28g"},
         # copybara:uncomment_end
         "//conditions:default": None,
     }),

tensorflow/python/kernel_tests/math_ops/BUILD

@@ -360,7 +360,11 @@ cuda_py_strict_test(
     size = "small",
     srcs = ["in_topk_op_test.py"],
     deps = [
+        "//tensorflow/compiler/jit:xla_cpu_jit",
+        "//tensorflow/compiler/jit:xla_gpu_jit",
+        "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/ops:nn_ops",
         "//tensorflow/python/platform:client_testlib",
         "//third_party/py/numpy",