[pull] master from tensorflow:master

tensorflow/compiler/mlir/tf2xla/tests/legalize-tf.mlir

@@ -2708,6 +2708,17 @@ func.func @neg_dynamic(%arg0: tensor<?xf32>) -> tensor<?xf32> {
 
 // -----
 
+// CHECK-LABEL: @random_shuffle_float_vector
+func.func @random_shuffle_float_vector(%arg0: tensor<2xf32>) -> tensor<2xf32> {
+  // CHECK: "mhlo.sort"
+  // CHECK: ^bb0([[KEY_LHS:%.*]]: tensor<i32>, [[KEY_RHS:%.*]]: tensor<i32>, {{.*}}: tensor<f32>, {{.*}}: tensor<f32>):
+  // CHECK: mhlo.compare  LT, [[KEY_LHS]], [[KEY_RHS]], TOTALORDER : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %0 = "tf.RandomShuffle"(%arg0) {seed = 1 : i64, seed2 = 2 : i64} : (tensor<2xf32>) -> tensor<2xf32>
+  func.return %0 : tensor<2xf32>
+}
+
+// -----
+
 // CHECK-LABEL: @sigmoid
 func.func @sigmoid(%arg0: tensor<2xf32>) -> tensor<2xf32> {
   // CHECK: mhlo.logistic
@@ -2798,4 +2809,3 @@ func.func @func_xla_sharding_consistent(%arg0: tensor<4x8xi32>) -> (tensor<4x8xi
   %1 = "tf.A"(%0) : (tensor<4x8xi32>) -> (tensor<4x8xi32>)
   func.return %1 : tensor<4x8xi32>
 }
-

tensorflow/python/eager/pywrap_gradient_exclusions.cc

@@ -50,7 +50,7 @@ auto OpGradientInfoInit(const T &a) {
 
 absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedInputIndices(
     const tensorflow::string &op_name) {
-  static std::array<OpIndexInfo, 366> a = {{
+  static std::array<OpIndexInfo, 365> a = {{
       {"Acosh"},
       {"AllToAll", 1, {0}},
       {"ApproximateEqual"},
@@ -292,7 +292,6 @@ absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedInputIndices(
       {"SdcaFprint"},
       {"SegmentSum", 1, {0}},
       {"Select", 1, {2}},
-      {"Selu"},
       {"SerializeTensor"},
       {"SetSize"},
       {"Shape"},
@@ -429,7 +428,7 @@ absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedInputIndices(
 
 absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedOutputIndices(
     const tensorflow::string &op_name) {
-  static std::array<OpIndexInfo, 489> a = {{
+  static std::array<OpIndexInfo, 490> a = {{
       {"Abs"},
       {"AccumulateNV2"},
       {"Acos"},
@@ -760,6 +759,7 @@ absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedOutputIndices(
       {"SegmentMean"},
       {"SegmentSum"},
       {"Select"},
+      {"Selu"},
       {"SeluGrad"},
       {"SerializeTensor"},
       {"SetSize"},

tensorflow/python/ops/BUILD

@@ -2208,6 +2208,7 @@ py_library(
         ":array_ops_stack",
         ":math_ops",
         ":nn_ops_gen",
+        "//tensorflow/python/framework:constant_op",
         "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/framework:ops",
     ],

tensorflow/python/ops/nn_grad.py

@@ -18,6 +18,7 @@
 import itertools
 import operator
 
+from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
@@ -466,7 +467,18 @@ def _EluGrad(op: ops.Operation, grad):
 
 @ops.RegisterGradient("Selu")
 def _SeluGrad(op: ops.Operation, grad):
-  return gen_nn_ops.selu_grad(grad, op.outputs[0])
+  x = op.inputs[0]
+  scale = constant_op.constant(1.0507009873554804934193349852946, dtype=x.dtype)
+  scale_alpha = constant_op.constant(
+      1.7580993408473768599402175208123, dtype=x.dtype
+  )
+  # Reconstructing the negative-branch derivative from the SELU output loses
+  # precision when the output rounds to -scale_alpha. Compute it from x.
+  derivative = array_ops.where_v2(
+      x < 0.0, scale_alpha * math_ops.exp(math_ops.minimum(x, 0.0)), scale
+  )
+  derivative = array_ops.where_v2(math_ops.is_nan(x), x, derivative)
+  return grad * derivative
 
 
 @ops.RegisterGradient("Softplus")

tensorflow/python/ops/nn_grad_test.py

@@ -237,6 +237,57 @@ def testEluGradGradWRTinputs(self):
 
 class SeluGradOpTest(test.TestCase):
 
+  @test_util.run_in_graph_and_eager_modes
+  def testSeluGradPreservesSmallNegativeGradients(self):
+    scale_alpha = 1.7580993408473768599402175208123
+    test_cases = (
+        (dtypes.float32, -20.0, 1e-6),
+        (dtypes.float64, -700.0, 1e-14),
+    )
+    for dtype, value, rtol in test_cases:
+      with self.subTest(dtype=dtype.name):
+        inputs = constant_op.constant(value, dtype=dtype)
+        with backprop.GradientTape() as tape:
+          tape.watch(inputs)
+          selu = gen_nn_ops.selu(inputs)
+
+        selu_grad = tape.gradient(selu, inputs)
+        np_dtype = dtype.as_numpy_dtype
+        expected = np_dtype(scale_alpha) * np.exp(np_dtype(value))
+        self.assertAllClose(
+            expected, self.evaluate(selu_grad), rtol=rtol, atol=0
+        )
+
+  @test_util.run_in_graph_and_eager_modes
+  def testSeluGradEdgeCases(self):
+    scale = 1.0507009873554804934193349852946
+    scale_alpha = 1.7580993408473768599402175208123
+    for dtype in (dtypes.float32, dtypes.float64):
+      with self.subTest(dtype=dtype.name):
+        np_dtype = dtype.as_numpy_dtype
+        values = np.array(
+            [-np.inf, -2.0, -0.0, 0.0, 2.0, np.inf, np.nan], dtype=np_dtype
+        )
+        inputs = constant_op.constant(values, dtype=dtype)
+        with backprop.GradientTape() as tape:
+          tape.watch(inputs)
+          selu = gen_nn_ops.selu(inputs)
+
+        selu_grad = self.evaluate(tape.gradient(selu, inputs))
+        expected = np.array(
+            [
+                0.0,
+                np_dtype(scale_alpha) * np.exp(np_dtype(-2.0)),
+                scale,
+                scale,
+                scale,
+                scale,
+                np.nan,
+            ],
+            dtype=np_dtype,
+        )
+        self.assertAllClose(expected, selu_grad)
+
   @test_util.run_deprecated_v1
   def testSeluGradGradWRTgrad_ys(self):
     inputs = constant_op.constant(

tensorflow/python/ops/numpy_ops/np_array_ops.py

@@ -100,7 +100,7 @@ def ones_like(a, dtype=None):
 def eye(N, M=None, k=0, dtype=float):  # pylint: disable=invalid-name,missing-docstring
   if dtype:
     dtype = np_utils.result_type(dtype)
-  if not M:
+  if M is None:
     M = N
   # Making sure N, M and k are `int`
   N = int(N)

tensorflow/python/ops/numpy_ops/np_array_ops_test.py

@@ -230,6 +230,13 @@ def testEye(self):
                 np_array_ops.eye(n, m, k, dtype=dtype),
                 np.eye(n, m, k, dtype=dtype))
 
+    # Test M=0 and N=0 zero-dimension edge cases
+    for n in (0, 1, 3):
+      for m in (0, 1, 3):
+        self.match(np_array_ops.eye(n, m), np.eye(n, m))
+        for k in range(-n - 1, m + 2):
+          self.match(np_array_ops.eye(n, m, k), np.eye(n, m, k))
+
   def testIdentity(self):
     n_max = 3
 

tensorflow/python/training/training_ops_test.py

@@ -603,6 +603,43 @@ def testSparseApplyOpsRejectLowerRankGrad(self):
       with self.assertRaises(errors.InvalidArgumentError):
         self.evaluate(apply_op())
 
+  @test_util.run_v2_only
+  def testResourceSparseApplyAdagradDAInvalidGradRank(self):
+    # A scalar `grad` with a higher-rank `var` used to hit a fatal CHECK
+    # instead of raising InvalidArgumentError (see GitHub issue #94130).
+    # ResourceSparseApplyAdagradDA only has a CPU kernel, so pin the whole
+    # test to CPU rather than relying on default placement.
+    with ops.device("/cpu:0"):
+      var = variables.Variable([[0.0, 0.0]] * 10, dtype=dtypes.float32)
+      gradient_accumulator = variables.Variable(
+          [[0.0, 0.0]] * 10, dtype=dtypes.float32
+      )
+      gradient_squared_accumulator = variables.Variable(
+          [[0.0, 0.0]] * 10, dtype=dtypes.float32
+      )
+      self.evaluate(variables.global_variables_initializer())
+
+      grad = constant_op.constant(0.0, dtype=dtypes.float32)  # wrong rank
+      indices = constant_op.constant([0, 0], dtype=dtypes.int32)
+
+      with self.assertRaisesRegex(
+          errors.InvalidArgumentError,
+          "grad must have the same number of dimensions as var",
+      ):
+        self.evaluate(
+            gen_training_ops.resource_sparse_apply_adagrad_da(
+                var.handle,
+                gradient_accumulator.handle,
+                gradient_squared_accumulator.handle,
+                grad,
+                indices,
+                constant_op.constant(0.0, dtype=dtypes.float32),
+                constant_op.constant(0.0, dtype=dtypes.float32),
+                constant_op.constant(0.0, dtype=dtypes.float32),
+                constant_op.constant(1, dtype=dtypes.int64),
+            )
+        )
+
 
 if __name__ == '__main__':
   googletest.main()