[pull] master from tensorflow:master

RELEASE.md

@@ -18,7 +18,11 @@
 
 ### Bug Fixes and Other Changes
 
-*   <BUG FIXES GO HERE>
+*   `tf.linalg.band_part`: `num_lower` and `num_upper` values that exceed the
+    number of rows or columns of the input are now clamped to the matrix
+    dimensions instead of raising an `InvalidArgumentError`. This matches the
+    documented definition of the op and the behavior of the XLA lowering used by
+    `tf.function(jit_compile=True)`.
 
 ## Thanks to our Contributors
 

tensorflow/c/experimental/gradients/tape/tape_operation.cc

@@ -109,10 +109,10 @@ absl::Status TapeOperation::SetAttrShape(const char* attr_name,
                                          const int64_t* dims,
                                          const int num_dims) {
   if (num_dims > TensorShape::MaxDimensions()) {
-    return errors::InvalidArgument("Value specified for `", attr_name, "` has ",
-                                   num_dims,
-                                   " dimensions which is over the limit of ",
-                                   TensorShape::MaxDimensions(), ".");
+    return absl::InvalidArgumentError(
+        absl::StrCat("Value specified for `", attr_name, "` has ", num_dims,
+                     " dimensions which is over the limit of ",
+                     TensorShape::MaxDimensions(), "."));
   }
   TensorShapeProto proto;
   if (num_dims < 0) {
@@ -128,19 +128,19 @@ absl::Status TapeOperation::SetAttrShape(const char* attr_name,
 }
 absl::Status TapeOperation::SetAttrFunction(const char* attr_name,
                                             const AbstractOperation* value) {
-  return tensorflow::errors::Unimplemented(
+  return absl::UnimplementedError(
       "SetAttrFunction has not been implemented yet.");
 }
 absl::Status TapeOperation::SetAttrFunctionName(const char* attr_name,
                                                 const char* value,
                                                 size_t length) {
-  return tensorflow::errors::Unimplemented(
+  return absl::UnimplementedError(
       "SetAttrFunctionName has not been implemented "
       "yet.");
 }
 absl::Status TapeOperation::SetAttrTensor(const char* attr_name,
                                           AbstractTensorInterface* tensor) {
-  return tensorflow::errors::Unimplemented(
+  return absl::UnimplementedError(
       "SetAttrTensor has not been implemented yet.");
 }
 absl::Status TapeOperation::SetAttrStringList(const char* attr_name,
@@ -196,7 +196,7 @@ absl::Status TapeOperation::SetAttrShapeList(const char* attr_name,
     const auto num_dims_i = num_dims[i];
 
     if (num_dims_i > TensorShape::MaxDimensions()) {
-      return errors::InvalidArgument(
+      return absl::InvalidArgumentError(
           strings::StrCat("Value specified for `", attr_name, "` has ",
                           num_dims_i, " dimensions which is over the limit of ",
                           TensorShape::MaxDimensions(), "."));
@@ -217,7 +217,7 @@ absl::Status TapeOperation::SetAttrShapeList(const char* attr_name,
 }
 absl::Status TapeOperation::SetAttrFunctionList(
     const char* attr_name, absl::Span<const AbstractOperation*> values) {
-  return tensorflow::errors::Unimplemented(
+  return absl::UnimplementedError(
       "SetAttrFunctionList has not been "
       "implemented yet.");
 }

tensorflow/compiler/jit/kernels/BUILD

@@ -48,7 +48,6 @@ XLA_OPS_DEPS = [
     "@xla//xla/client:client_library",
     "@xla//xla/client:local_client",
     "@xla//xla/service:compiler",
-    "@xla//xla/service/gpu:gpu_executable_run_options",
     "//tensorflow/core:core_cpu_internal",
     "//tensorflow/core:framework",
     "//tensorflow/core:lib",

tensorflow/compiler/mlir/lite/core/c/BUILD

@@ -49,6 +49,7 @@ exports_files(
     ],
     visibility = [
         # copybara:uncomment_begin(google-only)
+        # "//third_party/odml/litert/litert/swift:__subpackages__",
         # "//tensorflow/lite:__subpackages__",
         # copybara:uncomment_end_and_comment_begin
         "//visibility:public",
@@ -63,7 +64,14 @@ cc_library(
     ],
     compatible_with = get_compatible_with_portable(),
     copts = tflite_copts(),
-    visibility = ["//visibility:public"],
+    visibility = [
+        # copybara:uncomment_begin(google-only)
+        # "//third_party/odml/litert/litert/swift:__subpackages__",
+        # "//visibility:public",
+        # copybara:uncomment_end_and_comment_begin
+        "//visibility:public",
+        # copybara:comment_end
+    ],
 )
 
 # LINT.IfChange(common)

tensorflow/compiler/mlir/quantization/tensorflow/debugging/dump_tensor_op.cc

@@ -63,10 +63,10 @@ REGISTER_OP("DumpTensor")
 class DumpTensorOp : public OpKernel {
  public:
   explicit DumpTensorOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
-    string log_dir_path;
-    string file_name;
-    string func_name;
-    string node_name;
+    std::string log_dir_path;
+    std::string file_name;
+    std::string func_name;
+    std::string node_name;
     OP_REQUIRES_OK(ctx, ctx->GetAttr("log_dir_path", &log_dir_path));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("enabled", &enabled_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("file_name", &file_name));
@@ -92,7 +92,7 @@ class DumpTensorOp : public OpKernel {
     quant_unit_proto.set_func_name(func_name);
     quant_unit_proto.set_node_name(node_name);
 
-    string quant_unit_path = io::JoinPath(log_dir_path, "quant_unit.pb");
+    std::string quant_unit_path = io::JoinPath(log_dir_path, "quant_unit.pb");
     OP_REQUIRES_OK(
         ctx, SaveSerializedProtoToFile(quant_unit_proto.SerializeAsString(),
                                        quant_unit_path, ctx->env()));

tensorflow/core/common_runtime/BUILD

@@ -3115,11 +3115,11 @@ tf_cc_test(
 
 tf_cuda_cc_test(
     name = "direct_session_test",
-    size = "medium",
+    size = "large",
     srcs = ["direct_session_test.cc"],
     args = [] + if_cuda(["--heap_check="]),  # The GPU tracer leaks memory
     features = ["-layering_check"],
-    shard_count = 4,
+    shard_count = 8,
     deps = [
         ":core_cpu",
         ":core_cpu_internal",

tensorflow/core/kernels/barrier_ops.cc

@@ -587,6 +587,11 @@ class TakeManyOp : public BarrierOpKernel {
         ctx, TensorShapeUtils::IsScalar(Tnum_elements->shape()),
         absl::InvalidArgumentError("num_elements must be a scalar."), callback);
     const int32_t num_elements = Tnum_elements->scalar<int32_t>()();
+    OP_REQUIRES_ASYNC(
+        ctx, num_elements >= 0,
+        absl::InvalidArgumentError(absl::StrCat("BarrierTakeMany requested ",
+                                                num_elements, " < 0 elements")),
+        callback);
 
     DataTypeVector expected_inputs = {DT_STRING_REF, DT_INT32};
     // The first output is the insertion index, the second output is the key.

tensorflow/core/kernels/linalg/matrix_band_part_op.cc

@@ -74,24 +74,25 @@ class MatrixBandPartOp : public OpKernel {
         return tensor.scalar<int64_t>()();
       }
     };
-    const int64_t num_lower = as_int64_scalar(num_lower_in);
-    OP_REQUIRES(
-        context, num_lower <= input_reshaped.dimension(1),
-        errors::InvalidArgument(
-            "num_lower must be negative or less or equal to number of rows (",
-            input_reshaped.dimension(1), ") got: ", num_lower));
+    // Band limits that are out of range are clamped rather than rejected. The
+    // op is defined by
+    //   in_band(m, n) = (num_lower < 0 || (m - n) <= num_lower) &&
+    //                   (num_upper < 0 || (n - m) <= num_upper),
+    // which is well defined for any limit, and this is what the XLA lowering
+    // already computes. Clamping to [-1, dimension] additionally keeps the
+    // limits exactly representable in the `int` parameters of
+    // MatrixBandPartFunctor below; every negative limit is equivalent to -1,
+    // i.e. "keep the entire triangle".
+    const int64_t num_lower = std::clamp<int64_t>(
+        as_int64_scalar(num_lower_in), -1, input_reshaped.dimension(1));
 
     const Tensor& num_upper_in = context->input(2);
     OP_REQUIRES(context, TensorShapeUtils::IsScalar(num_upper_in.shape()),
                 absl::InvalidArgumentError(
                     absl::StrCat("num_upper must be scalar, got shape ",
                                  num_upper_in.shape().DebugString())));
-    const int64_t num_upper = as_int64_scalar(num_upper_in);
-    OP_REQUIRES(context, num_upper <= input_reshaped.dimension(2),
-                errors::InvalidArgument("num_upper must be negative or less or "
-                                        "equal to number of columns (",
-                                        input_reshaped.dimension(2),
-                                        ") got: ", num_upper));
+    const int64_t num_upper = std::clamp<int64_t>(
+        as_int64_scalar(num_upper_in), -1, input_reshaped.dimension(2));
 
     if (input.NumElements() == 0 ||
         ((num_lower < 0 || num_lower == input_reshaped.dimension(1)) &&

tensorflow/core/kernels/list_kernels.h

@@ -1012,8 +1012,24 @@ class TensorListScatter : public OpKernel {
           highest_index = i;
         }
       }
-      output_list.tensors().resize(std::max(highest_index + 1, num_elements),
-                                   Tensor(DT_INVALID));
+      // TensorListLength returns int32, so highest_index + 1 must be
+      // representable as int32 before the list is resized. Same guard as
+      // TensorListScatterIntoExistingList (#list_kernels.h).
+      OP_REQUIRES(
+          c, highest_index < std::numeric_limits<int32_t>::max(),
+          absl::InvalidArgumentError(absl::StrCat(
+              "TensorListScatter: index ", highest_index,
+              " would produce a list length that is not representable as "
+              "int32")));
+      // Compute the required size as size_t: highest_index == INT32_MAX would
+      // overflow the signed `highest_index + 1` (undefined behavior).
+      const size_t needed_from_indices =
+          highest_index < 0 ? 0 : static_cast<size_t>(highest_index) + 1;
+      const size_t needed =
+          num_elements > 0
+              ? std::max(needed_from_indices, static_cast<size_t>(num_elements))
+              : needed_from_indices;
+      output_list.tensors().resize(needed, Tensor(DT_INVALID));
     }
 
     OP_REQUIRES_OK(c,

tensorflow/core/kernels/lrn_op.cc

@@ -352,6 +352,12 @@ class LRNOp : public OpKernel {
                    context->allocate_output(
                        0, TensorShape({batch, rows, cols, depth}), &output));
 
+    // NOMUTANTS(LCR): Early return is a no-op guard on CPU but prevents fatal
+    // cuDNN crashes on GPU.
+    if (in.NumElements() == 0) {
+      return;
+    }
+
     LaunchLRN<Device, T> launcher(depth_radius_, bias_, alpha_, beta_);
     launcher.launch(context, this, in, output);
   }
@@ -696,6 +702,12 @@ class LRNGradOp : public OpKernel {
                    context->allocate_output(
                        0, TensorShape({batch, rows, cols, depth}), &output));
 
+    // NOMUTANTS(LCR): Early return is a no-op guard on CPU but prevents fatal
+    // cuDNN crashes on GPU.
+    if (in_grads.NumElements() == 0) {
+      return;
+    }
+
     LaunchLRNGrad<Device, T> launcher(depth_radius_, bias_, alpha_, beta_);
     launcher.launch(context, this, in_grads, in_image, out_image, output);
   }

tensorflow/core/kernels/lrn_op_test.cc

@@ -159,6 +159,44 @@ TEST_F(LRNFloatTest, Depth16) {
   EXPECT_TRUE(Compare());
 }
 
+TEST_F(LRNFloatTest, ZeroSizeInput) {
+  TF_ASSERT_OK(NodeDefBuilder("lrn_op", "LRN")
+                   .Input(FakeInput())
+                   .Attr("depth_radius", 2)
+                   .Attr("bias", 1.0f)
+                   .Attr("alpha", 1.0f)
+                   .Attr("beta", 0.5f)
+                   .Finalize(node_def()));
+  TF_ASSERT_OK(InitOp());
+  AddInput<float>(TensorShape({0, 0, 0, 0}),
+                  [](int i) -> float { return 0.0f; });
+  TF_ASSERT_OK(RunOpKernel());
+  EXPECT_EQ(GetOutput(0)->NumElements(), 0);
+  EXPECT_EQ(GetOutput(0)->shape(), TensorShape({0, 0, 0, 0}));
+}
+
+TEST_F(LRNFloatTest, ZeroSizeInputGrad) {
+  TF_ASSERT_OK(NodeDefBuilder("lrn_grad_op", "LRNGrad")
+                   .Input(FakeInput())
+                   .Input(FakeInput())
+                   .Input(FakeInput())
+                   .Attr("depth_radius", 2)
+                   .Attr("bias", 1.0f)
+                   .Attr("alpha", 1.0f)
+                   .Attr("beta", 0.5f)
+                   .Finalize(node_def()));
+  TF_ASSERT_OK(InitOp());
+  AddInput<float>(TensorShape({0, 0, 0, 0}),
+                  [](int i) -> float { return 0.0f; });
+  AddInput<float>(TensorShape({0, 0, 0, 0}),
+                  [](int i) -> float { return 0.0f; });
+  AddInput<float>(TensorShape({0, 0, 0, 0}),
+                  [](int i) -> float { return 0.0f; });
+  TF_ASSERT_OK(RunOpKernel());
+  EXPECT_EQ(GetOutput(0)->NumElements(), 0);
+  EXPECT_EQ(GetOutput(0)->shape(), TensorShape({0, 0, 0, 0}));
+}
+
 static double RndGaussian(random::SimplePhilox* rnd) {
   // Box-Muller transformation.
   // See, for example, http://www.taygeta.com/random/gaussian.html

tensorflow/core/kernels/spectrogram_op.cc

@@ -34,11 +34,12 @@ class AudioSpectrogramOp : public OpKernel {
     OP_REQUIRES_OK(context, context->GetAttr("stride", &stride_));
     OP_REQUIRES_OK(context,
                    context->GetAttr("magnitude_squared", &magnitude_squared_));
-    OP_REQUIRES(
-        context, window_size_ > 1,
-        errors::InvalidArgument("window_size must be > 1, got ", window_size_));
+    OP_REQUIRES(context, window_size_ > 1,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "window_size must be > 1, got ", window_size_)));
     OP_REQUIRES(context, stride_ > 0,
-                errors::InvalidArgument("stride must be > 0, got ", stride_));
+                absl::InvalidArgumentError(
+                    absl::StrCat("stride must be > 0, got ", stride_)));
   }
 
   void Compute(OpKernelContext* context) override {

tensorflow/lite/core/async/c/BUILD

@@ -37,7 +37,12 @@ exports_files(
         "types.h",
     ],
     visibility = [
-        "//tensorflow/lite:__subpackages__",
+        # copybara:uncomment_begin(google-only)
+        # "//third_party/odml/litert/litert/swift:__subpackages__",
+        # "//tensorflow/lite:__subpackages__",
+        # copybara:uncomment_end_and_comment_begin
+        "//visibility:public",
+        # copybara:comment_end
     ],
 )
 

tensorflow/lite/core/c/BUILD

@@ -50,7 +50,12 @@ exports_files(
         "operator.h",
     ],
     visibility = [
-        "//tensorflow/lite:__subpackages__",
+        # copybara:uncomment_begin(google-only)
+        # "//third_party/odml/litert/litert/swift:__subpackages__",
+        # "//tensorflow/lite:__subpackages__",
+        # copybara:uncomment_end_and_comment_begin
+        "//visibility:public",
+        # copybara:comment_end
     ],
 )
 

tensorflow/lite/core/c/special_rules.bzl

@@ -16,15 +16,15 @@
 
 def c_api_visibility_allowlist():
     """Returns a list of packages that can depend on c_api."""
-    return []
+    return ["//litert/swift:__subpackages__"]
 
 def c_api_experimental_visibility_allowlist():
     """Returns a list of packages that can depend on third_party/tensorflow/lite/core/c/c_api_experimental.h."""
-    return []
+    return ["//litert/swift:__subpackages__"]
 
 def common_header_visibility_allowlist():
     """Returns a list of packages that can depend on tensorflow/lite/core/c/common.h."""
-    return []
+    return ["//litert/swift:__subpackages__"]
 
 def experimental_acceleration_api_allowlist():
     """Returns a list of packages that can depend on tensorflow/lite/core/experimental/acceleration/ subpackages.."""

tensorflow/lite/profiling/telemetry/c/BUILD

@@ -33,7 +33,12 @@ exports_files(
         "telemetry_setting.h",
     ],
     visibility = [
-        "//tensorflow/lite:__subpackages__",
+        # copybara:uncomment_begin(google-only)
+        # "//third_party/odml/litert/litert/swift:__subpackages__",
+        # "//tensorflow/lite:__subpackages__",
+        # copybara:uncomment_end_and_comment_begin
+        "//visibility:public",
+        # copybara:comment_end
     ],
 )