[pull] master from tensorflow:master

tensorflow/compiler/tests/BUILD

@@ -594,6 +594,7 @@ tf_xla_py_strict_test(
     ],
     deps = [
         ":xla_test",
+        "//tensorflow/python/framework:errors",
         "//tensorflow/python/framework:tensor_shape",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:linalg_ops",
@@ -1371,6 +1372,7 @@ tf_xla_py_strict_test(
     ],
     deps = [
         ":xla_test",
+        "//tensorflow/python/framework:errors",
         "//tensorflow/python/framework:test_lib",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:linalg_ops",

tensorflow/compiler/tests/qr_op_test.py

@@ -21,6 +21,7 @@
 import numpy as np
 
 from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import linalg_ops
@@ -143,6 +144,18 @@ def testRepeatedColumn(self, rows, cols):
     x_np[:, 1] = x_np[:, 2]
     self._test(x_np, full_matrices=True, full_rank=False)
 
+  def testVectorInputRaisesError(self):
+    # Regression test for GitHub issue 110798. The graph-level shape check
+    # only runs for inputs of known rank, so a rank-1 input reaching the
+    # compiler through an unknown-rank placeholder must be rejected at
+    # compile time instead of reaching a fatal check inside QrExplicit.
+    with self.session() as sess:
+      x_tf = array_ops.placeholder(np.float32)
+      with self.device_scope():
+        q_tf, r_tf = linalg_ops.qr(x_tf, full_matrices=True)
+      with self.assertRaisesRegex(errors.InvalidArgumentError, "rank >= 2"):
+        sess.run([q_tf, r_tf], feed_dict={x_tf: np.zeros([8], np.float32)})
+
 
 if __name__ == "__main__":
   test.main()

tensorflow/compiler/tests/svd_op_test.py

@@ -20,6 +20,7 @@
 import numpy as np
 
 from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import tensor_shape
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import gen_linalg_ops
@@ -86,6 +87,18 @@ def testSvd(self, n, dtype):
       self._testSvdCorrectness(dtype, batch_dims + (2 * n, n))
       self._testSvdCorrectness(dtype, batch_dims + (n, 2 * n))
 
+  def testVectorInputRaisesError(self):
+    # Like GitHub issue 110798 for QR: the graph-level shape check only runs
+    # for inputs of known rank, so a rank-1 input reaching the compiler
+    # through an unknown-rank placeholder must be rejected at compile time
+    # instead of reading a negative trailing dimension.
+    with self.session() as sess:
+      x_tf = array_ops.placeholder(np.float32)
+      with self.test_scope():
+        s, u, v = linalg_ops.svd(x_tf, full_matrices=True)
+      with self.assertRaisesRegex(errors.InvalidArgumentError, "rank >= 2"):
+        sess.run([s, u, v], feed_dict={x_tf: np.zeros([8], np.float32)})
+
 
 if __name__ == "__main__":
   test.main()

tensorflow/compiler/tf2xla/kernels/BUILD

@@ -850,6 +850,8 @@ tf_kernel_library(
         "//tensorflow/compiler/tf2xla/ops:xla_ops",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
         "@xla//xla:shape_util",
         "@xla//xla:xla_data_proto_cc",
         "@xla//xla/hlo/builder/lib:constants",
@@ -1563,6 +1565,8 @@ tf_kernel_library(
         "//tensorflow/compiler/tf2xla:xla_resource",
         "//tensorflow/compiler/tf2xla/ops:xla_ops",
         "//tensorflow/core:framework",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
         "@xla//xla/hlo/builder:xla_builder",
         "@xla//xla/hlo/builder/lib:qr",
     ],

tensorflow/compiler/tf2xla/kernels/qr_op.cc

@@ -13,12 +13,15 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "xla/hlo/builder/lib/qr.h"
 #include "xla/hlo/builder/xla_builder.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 
 namespace tensorflow {
 namespace {
@@ -29,6 +32,14 @@ class QROp : public XlaOpKernel {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("full_matrices", &full_matrices_));
   }
   void Compile(XlaOpKernelContext* ctx) override {
+    // The rank is only known at compile time when the graph-level shape
+    // inference saw an unknown rank, so it must be validated here before
+    // QrExplicit reads the trailing two dimensions.
+    const TensorShape input_shape = ctx->InputShape(0);
+    OP_REQUIRES(ctx, input_shape.dims() >= 2,
+                absl::InvalidArgumentError(
+                    absl::StrCat("Input must have rank >= 2, got shape ",
+                                 input_shape.DebugString())));
     xla::XlaOp q, r;
     xla::QrExplicit(ctx->Input(0), full_matrices_, q, r);
     ctx->SetOutput(0, q);

tensorflow/compiler/tf2xla/kernels/xla_svd_op.cc

@@ -17,6 +17,8 @@ limitations under the License.
 #include <cstdint>
 #include <string>
 
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "xla/hlo/builder/lib/constants.h"
@@ -51,6 +53,11 @@ class XlaSvdOp : public XlaOpKernel {
     }
   }
   void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape input_shape = ctx->InputShape(0);
+    OP_REQUIRES(ctx, input_shape.dims() >= 2,
+                absl::InvalidArgumentError(
+                    absl::StrCat("Input must have rank >= 2, got shape ",
+                                 input_shape.DebugString())));
     auto result = xla::SVD(ctx->Input(0), max_iter_, epsilon_,
                            precision_config_.operand_precision(0));
     ctx->SetOutput(0, result.d);
@@ -71,7 +78,14 @@ class SvdOp : public XlaOpKernel {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("full_matrices", &full_matrices_));
   }
   void Compile(XlaOpKernelContext* ctx) override {
+    // The rank is only known at compile time when the graph-level shape
+    // inference saw an unknown rank, so it must be validated here before
+    // the trailing two dimensions are read.
     const TensorShape input_shape = ctx->InputShape("input");
+    OP_REQUIRES(ctx, input_shape.dims() >= 2,
+                absl::InvalidArgumentError(
+                    absl::StrCat("Input must have rank >= 2, got shape ",
+                                 input_shape.DebugString())));
     int m = input_shape.dim_size(input_shape.dims() - 2);
     int n = input_shape.dim_size(input_shape.dims() - 1);
     // This is based on heuristics that approx log(n) sweep updates are needed.

tensorflow/lite/kernels/transpose.cc

@@ -50,26 +50,12 @@ TfLiteStatus ResizeOutputTensor(TfLiteContext* context,
   // Ensure validity of the permutations tensor as a 1D tensor.
   TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->perm), 1);
   TF_LITE_ENSURE_EQ(context, op_context->perm->dims->data[0], dims);
-  // `perm` must be a permutation of [0, dims), not merely a set of in-range
-  // values: the output shape and the element offsets are both derived from it,
-  // and a repeated entry makes them disagree with the input extent.
-  // `dims` is bounded by kTransposeMaxDimensions, which Prepare() enforces
-  // before this function is reachable, so a 64-bit mask is sufficient and
-  // avoids a heap allocation in the kernel.
-  static_assert(kTransposeMaxDimensions <= 64,
-                "Permutation bitmask assumes at most 64 dimensions.");
-  uint64_t seen = 0;
   for (int idx = 0; idx < dims; ++idx) {
     TF_LITE_ENSURE_MSG(context,
                        (perm_data[idx] >= -dims && perm_data[idx] < dims),
                        "Transpose op permutations array is out of bounds.");
     new_perm_data[idx] = perm_data[idx];
     if (new_perm_data[idx] < 0) new_perm_data[idx] += dims;
-    const uint64_t bit = uint64_t{1} << new_perm_data[idx];
-    TF_LITE_ENSURE_MSG(
-        context, (seen & bit) == 0,
-        "Transpose op permutations array must not contain duplicate values.");
-    seen |= bit;
   }
 
   // Determine size of output tensor.

tensorflow/lite/kernels/transpose_test.cc

@@ -193,23 +193,6 @@ TEST(TransposeTest, TestPermOutOfBounds) {
   EXPECT_DEATH(TransposeOpConstModel({1, 3, 3, 1}, {4}, {0, 1, 2, 4}),
                "Transpose op permutations array is out of bounds.");
 }
-
-// `perm` must be a permutation of [0, dims). A repeated entry passes the range
-// check but makes the derived output shape disagree with the input extent, so
-// the kernel reads outside the input tensor.
-TEST(TransposeTest, TestPermDuplicateValues) {
-  EXPECT_DEATH(
-      TransposeOpConstModel({1, 3, 3, 1}, {4}, {0, 1, 2, 2}),
-      "Transpose op permutations array must not contain duplicate values.");
-}
-
-// Duplicates must also be rejected after negative entries are normalised:
-// on a rank-2 input {0, -2} normalises to {0, 0}.
-TEST(TransposeTest, TestPermDuplicateValuesAfterNegativeNormalization) {
-  EXPECT_DEATH(
-      TransposeOpConstModel({2, 3}, {2}, {0, -2}),
-      "Transpose op permutations array must not contain duplicate values.");
-}
 #endif
 
 TEST(TransposeTest, TestInt41DInputConstTensor) {

tensorflow/python/ops/numpy_ops/np_math_ops.py

@@ -529,25 +529,41 @@ def f(a, b):
 @tf_export.tf_export('experimental.numpy.logaddexp', v1=[])
 @np_utils.np_doc('logaddexp')
 def logaddexp(x1, x2):
-  amax = maximum(x1, x2)
-  delta = x1 - x2
-  return np_array_ops.where(
-      isnan(delta),
-      x1 + x2,  # NaNs or infinities of the same sign.
-      amax + log1p(exp(-abs(delta))),
-  )
+
+  def f(x1, x2):
+    if not np.issubdtype(x1.dtype.as_numpy_dtype, np.inexact):
+      float_dtype = np_utils.result_type(float)
+      x1 = math_ops.cast(x1, float_dtype)
+      x2 = math_ops.cast(x2, float_dtype)
+    amax = maximum(x1, x2)
+    delta = x1 - x2
+    return np_array_ops.where(
+        isnan(delta),
+        x1 + x2,  # NaNs or infinities of the same sign.
+        amax + log1p(exp(-abs(delta))),
+    )
+
+  return _bin_op(f, x1, x2)
 
 
 @tf_export.tf_export('experimental.numpy.logaddexp2', v1=[])
 @np_utils.np_doc('logaddexp2')
 def logaddexp2(x1, x2):
-  amax = maximum(x1, x2)
-  delta = x1 - x2
-  return np_array_ops.where(
-      isnan(delta),
-      x1 + x2,  # NaNs or infinities of the same sign.
-      amax + log1p(exp2(-abs(delta))) / np.log(2),
-  )
+
+  def f(x1, x2):
+    if not np.issubdtype(x1.dtype.as_numpy_dtype, np.inexact):
+      float_dtype = np_utils.result_type(float)
+      x1 = math_ops.cast(x1, float_dtype)
+      x2 = math_ops.cast(x2, float_dtype)
+    amax = maximum(x1, x2)
+    delta = x1 - x2
+    return np_array_ops.where(
+        isnan(delta),
+        x1 + x2,  # NaNs or infinities of the same sign.
+        amax + log1p(exp2(-abs(delta))) / math_ops.cast(np.log(2), x1.dtype),
+    )
+
+  return _bin_op(f, x1, x2)
 
 
 @tf_export.tf_export('experimental.numpy.polyval', v1=[])

tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py

@@ -126,6 +126,28 @@ def testAverageAcceptsIntegerInputs(self):
           err_msg='average({})'.format(arg),
       )
 
+  def testLogaddexpAcceptsIntegerInputs(self):
+    # `logaddexp` and `logaddexp2` promote integer arguments to a float dtype,
+    # which must not depend on `enable_numpy_methods_on_tensor()`.
+    x1 = np.array([1, 2, 3], dtype=np.int32)
+    for x2 in (np.int32(1), np.array([4, 5, 6], dtype=np.int32)):
+      actual = np_math_ops.logaddexp(x1, x2)
+      np.testing.assert_allclose(
+          np.asarray(actual),
+          np.logaddexp(x1, x2),
+          rtol=1e-6,
+          atol=1e-6,
+          err_msg='logaddexp({}, {})'.format(x1, x2),
+      )
+      actual2 = np_math_ops.logaddexp2(x1, x2)
+      np.testing.assert_allclose(
+          np.asarray(actual2),
+          np.logaddexp2(x1, x2),
+          rtol=1e-6,
+          atol=1e-6,
+          err_msg='logaddexp2({}, {})'.format(x1, x2),
+      )
+
   def testFloatInputsAreUnchanged(self):
     arg = np.array([1.5, 2.5, 3.5], dtype=np.float32)
     for name, tf_fun, np_fun in _PROMOTING_UNARY_OPS:

tensorflow/python/ops/numpy_ops/np_math_ops_test.py

@@ -207,6 +207,30 @@ def testHypotSpecialCases(self):
     expected = np.hypot(x, y)
     np.testing.assert_equal(actual.tolist(), expected.tolist())
 
+  def testLogaddexp(self):
+    self._testBinaryOp(np_math_ops.logaddexp, np.logaddexp, 'logaddexp')
+    self._testBinaryOp(np_math_ops.logaddexp2, np.logaddexp2, 'logaddexp2')
+
+  def testLogaddexpNonFloatInputs(self):
+    int_args = [
+        ([1, 2, 3], [4, 5, 6]),
+        (np.array([1, 2], dtype=np.int32), np.array([3, 4], dtype=np.int32)),
+        (np.array([1, 2], dtype=np.int64), np.array([3, 4], dtype=np.int64)),
+        (1, 2),
+        ([1.0, 2.0], [3.0, 4.0]),
+    ]
+    for x1, x2 in int_args:
+      self.match(
+          np_math_ops.logaddexp(x1, x2),
+          np.logaddexp(np.asarray(x1), np.asarray(x2)),
+          msg='logaddexp({}, {})'.format(x1, x2),
+      )
+      self.match(
+          np_math_ops.logaddexp2(x1, x2),
+          np.logaddexp2(np.asarray(x1), np.asarray(x2)),
+          msg='logaddexp2({}, {})'.format(x1, x2),
+      )
+
   def match(self, actual, expected, msg='', check_dtype=True):
     self.assertIsInstance(actual, np_arrays.ndarray)
     if check_dtype:

third_party/xla/xla/service/BUILD

@@ -2311,6 +2311,7 @@ xla_cc_test(
     srcs = ["hlo_schedule_test.cc"],
     deps = [
         ":buffer_value",
+        ":hlo_module_config",
         "//xla:literal_util",
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
@@ -2325,6 +2326,8 @@ xla_cc_test(
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_benchmark//:benchmark",
         "@com_google_googletest//:gtest",
     ],
 )

third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc

@@ -683,27 +683,6 @@ e {
                              ::testing::HasSubstr("Unsupported broadcast")));
 }
 
-TEST_F(TritonDotAnalysisTest, OutputBroadcastIsNotAccepted) {
-  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                       ParseAndReturnVerifiedModule(R"(
-HloModule t
-
-ENTRY e {
-  p0 = f16[2,35] parameter(0)
-  p0c = bf16[2,35] convert(p0)
-  p1 = bf16[35,2] parameter(1)
-  dot = bf16[2,2] dot(p0c, p1),
-    lhs_contracting_dims={1}, rhs_contracting_dims={0}
-  ROOT bc = bf16[2,2,100] broadcast(dot), dimensions={0,1}
-})"));
-  EXPECT_TRUE(GemmFusion(se::GpuComputeCapability{se::CudaComputeCapability{
-                             se::CudaComputeCapability::kAmpere, 0}})
-                  .Run(module.get())
-                  .value());
-  EXPECT_EQ(module->entry_computation()->root_instruction()->opcode(),
-            HloOpcode::kBroadcast);
-}
-
 TEST_F(TritonDotAnalysisTest, DegenerateSplitFragmentIsHandled) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(