[pull] master from tensorflow:master

tensorflow/compiler/mlir/tensorflow/utils/translate_utils.cc

@@ -120,9 +120,8 @@ namespace {
 // Sets type list attribute with the given `name` to the given `types`. If the
 // attribute already exists with a different value, returns an error.
 template <typename ContainerT,
-          typename = typename std::enable_if<
-              std::is_same<mlir::Type, decltype(*std::declval<ContainerT>()
-                                                     .begin())>::value>::type>
+          typename = std::enable_if_t<std::is_same_v<
+              mlir::Type, decltype(*std::declval<ContainerT>().begin())>>>
 absl::Status SetTypeAttribute(absl::string_view name, ContainerT types,
                               AttrValueMap* values) {
   AttrValue value;
@@ -143,9 +142,9 @@ absl::Status SetTypeAttribute(absl::string_view name, ContainerT types,
 // Sets shape list attribute with the given `name` to the given `shapes`. If the
 // attribute already exists then this will just retain the set value.
 template <typename ContainerT,
-          typename = typename std::enable_if<std::is_same<
-              std::optional<llvm::ArrayRef<int64_t>>,
-              decltype(*std::declval<ContainerT>().begin())>::value>::type>
+          typename = std::enable_if_t<
+              std::is_same_v<std::optional<llvm::ArrayRef<int64_t>>,
+                             decltype(*std::declval<ContainerT>().begin())>>>
 void SetShapeAttribute(absl::string_view name, ContainerT shapes,
                        AttrValueMap* values) {
   AttrValue value;

tensorflow/core/common_runtime/gpu/BUILD

@@ -327,7 +327,7 @@ tf_cc_test(
 
 tf_cuda_cc_test(
     name = "gpu_bfc_allocator_test",
-    size = "small",
+    size = "medium",
     srcs = [
         "gpu_bfc_allocator_test.cc",
     ],

tensorflow/core/grappler/optimizers/BUILD

@@ -320,7 +320,7 @@ cc_library(
 
 tf_cuda_cc_test(
     name = "arithmetic_optimizer_test",
-    size = "small",
+    size = "medium",
     srcs = ["arithmetic_optimizer_test.cc"],
     deps = [
         ":arithmetic_optimizer",

tensorflow/python/_pywrap_tensorflow.def

@@ -92,6 +92,7 @@ EXPORTS
  ??0MetaGraphDef@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
  ??0MetaOptimizer@grappler@tensorflow@@QEAA@PEAVDeviceBase@2@AEBVConfigProto@2@@Z
  ??0ModelAnalyzer@grappler@tensorflow@@QEAA@AEBUGrapplerItem@12@@Z
+ ??0ModelProto_Node@model@data@tensorflow@@IEAA@PEAVArena@protobuf@google@@AEBV0123@@Z
  ??0Name@register_kernel@tensorflow@@QEAA@PEBD@Z
  ??0NamedDevice@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
  ??0NamedDevice@tensorflow@@IEAA@PEAVArena@protobuf@google@@AEBV01@@Z
@@ -163,6 +164,7 @@ EXPORTS
  ??1MeshProto@dtensor@tensorflow@@UEAA@XZ
  ??1MessageDifferencer@util@protobuf@google@@QEAA@XZ
  ??1MetaGraphDef@tensorflow@@UEAA@XZ
+ ??1ModelProto_Node@model@data@tensorflow@@UEAA@XZ
  ??1NamedDevice@tensorflow@@UEAA@XZ
  ??1NodeDef@tensorflow@@UEAA@XZ
  ??1OpDef@tensorflow@@UEAA@XZ

tensorflow/python/kernel_tests/linalg/matrix_exponential_op_test.py

@@ -19,11 +19,15 @@
 import numpy as np
 
 from tensorflow.python.client import session
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import forwardprop
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
 from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import random_ops
 from tensorflow.python.ops import variables
 from tensorflow.python.ops.linalg import linalg_impl
@@ -162,6 +166,75 @@ def testConcurrentExecutesWithoutError(self):
       expm = self.evaluate([expm1, expm2])
       self.assertAllEqual(expm[0], expm[1])
 
+  @test_util.run_v2_only
+  def testSecondDerivativeNestedForwardAccumulator(self):
+    # Regression test for Issue #127224.
+    def target(t):
+      one = constant_op.constant(1.0, dtype=t.dtype)
+      a = array_ops_stack.stack([
+          array_ops_stack.stack([t, one]),
+          array_ops_stack.stack([-one, t]),
+      ])
+      return math_ops.reduce_sum(linalg_impl.matrix_exponential(a))
+
+    def forward_ad(fn):
+      def jvp(x):
+        with forwardprop.ForwardAccumulator(x, array_ops.ones_like(x)) as acc:
+          y = fn(x)
+        return acc.jvp(y)
+
+      return jvp
+
+    for dtype in [ops.dtypes.float32, ops.dtypes.float64]:
+      x = constant_op.constant(0.7, dtype=dtype)
+      expected = 2.0 * np.exp(0.7) * np.cos(1.0)
+
+      if dtype == ops.dtypes.float64:
+        atol_1st = 1e-12
+        rtol_1st = 1e-12
+        # Nested ForwardAccumulator (forward-over-forward) diverges from the
+        # analytical value by ~9% here, independent of dtype: matrix_exponential
+        # chains matmuls three levels deep (matrix_4 = matrix_2 * matrix_2,
+        # matrix_6 = matrix_4 * matrix_2, ...), and TF's forward-over-forward
+        # differentiation is inaccurate past two chained matmul levels
+        # regardless of precision, so this isn't something this PR's
+        # coefficient normalization can improve further.
+        atol_2nd = 0.25
+        rtol_2nd = 1e-1
+        atol_hess = 1e-12
+        rtol_hess = 1e-12
+      else:
+        atol_1st = 1e-3
+        rtol_1st = 1e-3
+        atol_2nd = 0.25
+        rtol_2nd = 1e-1
+        atol_hess = 1e-4
+        rtol_hess = 1e-4
+
+      # 1st order forward-mode JVP
+      actual_1st = forward_ad(target)(x)
+      self.assertAllClose(
+          expected, self.evaluate(actual_1st), rtol=rtol_1st, atol=atol_1st
+      )
+
+      # 2nd order nested forward-mode JVP
+      actual_2nd = forward_ad(forward_ad(target))(x)
+      self.assertAllClose(
+          expected, self.evaluate(actual_2nd), rtol=rtol_2nd, atol=atol_2nd
+      )
+
+      # 2nd order reverse-mode Hessian
+      with backprop.GradientTape() as t2:
+        t2.watch(x)
+        with backprop.GradientTape() as t1:
+          t1.watch(x)
+          y = target(x)
+        g1 = t1.gradient(y, x)
+      g2 = t2.gradient(g1, x)
+      self.assertAllClose(
+          expected, self.evaluate(g2), rtol=rtol_hess, atol=atol_hess
+      )
+
 
 class MatrixExponentialBenchmark(test.Benchmark):
 

tensorflow/python/ops/linalg/linalg_impl.py

@@ -134,59 +134,67 @@ def adjoint(matrix, name=None):
 # https://eigen.tuxfamily.org/dox/unsupported/MatrixExponential_8h_source.html
 def _matrix_exp_pade3(matrix):
   """3rd-order Pade approximant for matrix exponential."""
-  b = [120.0, 60.0, 12.0]
-  b = [constant_op.constant(x, matrix.dtype) for x in b]
+  b = [120.0, 60.0, 12.0, 1.0]
+  c = [constant_op.constant(x / b[0], matrix.dtype) for x in b]
   ident = linalg_ops.eye(
       array_ops.shape(matrix)[-2],
       batch_shape=array_ops.shape(matrix)[:-2],
       dtype=matrix.dtype)
   matrix_2 = math_ops.matmul(matrix, matrix)
-  tmp = matrix_2 + b[1] * ident
+  tmp = c[3] * matrix_2 + c[1] * ident
   matrix_u = math_ops.matmul(matrix, tmp)
-  matrix_v = b[2] * matrix_2 + b[0] * ident
+  matrix_v = c[2] * matrix_2 + ident
   return matrix_u, matrix_v
 
 
 def _matrix_exp_pade5(matrix):
   """5th-order Pade approximant for matrix exponential."""
-  b = [30240.0, 15120.0, 3360.0, 420.0, 30.0]
-  b = [constant_op.constant(x, matrix.dtype) for x in b]
+  b = [30240.0, 15120.0, 3360.0, 420.0, 30.0, 1.0]
+  c = [constant_op.constant(x / b[0], matrix.dtype) for x in b]
   ident = linalg_ops.eye(
       array_ops.shape(matrix)[-2],
       batch_shape=array_ops.shape(matrix)[:-2],
       dtype=matrix.dtype)
   matrix_2 = math_ops.matmul(matrix, matrix)
   matrix_4 = math_ops.matmul(matrix_2, matrix_2)
-  tmp = matrix_4 + b[3] * matrix_2 + b[1] * ident
+  tmp = c[5] * matrix_4 + c[3] * matrix_2 + c[1] * ident
   matrix_u = math_ops.matmul(matrix, tmp)
-  matrix_v = b[4] * matrix_4 + b[2] * matrix_2 + b[0] * ident
+  matrix_v = c[4] * matrix_4 + c[2] * matrix_2 + ident
   return matrix_u, matrix_v
 
 
 def _matrix_exp_pade7(matrix):
   """7th-order Pade approximant for matrix exponential."""
-  b = [17297280.0, 8648640.0, 1995840.0, 277200.0, 25200.0, 1512.0, 56.0]
-  b = [constant_op.constant(x, matrix.dtype) for x in b]
+  b = [17297280.0, 8648640.0, 1995840.0, 277200.0, 25200.0, 1512.0, 56.0, 1.0]
+  c = [constant_op.constant(x / b[0], matrix.dtype) for x in b]
   ident = linalg_ops.eye(
       array_ops.shape(matrix)[-2],
       batch_shape=array_ops.shape(matrix)[:-2],
       dtype=matrix.dtype)
   matrix_2 = math_ops.matmul(matrix, matrix)
   matrix_4 = math_ops.matmul(matrix_2, matrix_2)
   matrix_6 = math_ops.matmul(matrix_4, matrix_2)
-  tmp = matrix_6 + b[5] * matrix_4 + b[3] * matrix_2 + b[1] * ident
+  tmp = c[7] * matrix_6 + c[5] * matrix_4 + c[3] * matrix_2 + c[1] * ident
   matrix_u = math_ops.matmul(matrix, tmp)
-  matrix_v = b[6] * matrix_6 + b[4] * matrix_4 + b[2] * matrix_2 + b[0] * ident
+  matrix_v = c[6] * matrix_6 + c[4] * matrix_4 + c[2] * matrix_2 + ident
   return matrix_u, matrix_v
 
 
 def _matrix_exp_pade9(matrix):
   """9th-order Pade approximant for matrix exponential."""
   b = [
-      17643225600.0, 8821612800.0, 2075673600.0, 302702400.0, 30270240.0,
-      2162160.0, 110880.0, 3960.0, 90.0
+      17643225600.0,
+      8821612800.0,
+      2075673600.0,
+      302702400.0,
+      30270240.0,
+      2162160.0,
+      110880.0,
+      3960.0,
+      90.0,
+      1.0,
   ]
-  b = [constant_op.constant(x, matrix.dtype) for x in b]
+  c = [constant_op.constant(x / b[0], matrix.dtype) for x in b]
   ident = linalg_ops.eye(
       array_ops.shape(matrix)[-2],
       batch_shape=array_ops.shape(matrix)[:-2],
@@ -196,23 +204,42 @@ def _matrix_exp_pade9(matrix):
   matrix_6 = math_ops.matmul(matrix_4, matrix_2)
   matrix_8 = math_ops.matmul(matrix_6, matrix_2)
   tmp = (
-      matrix_8 + b[7] * matrix_6 + b[5] * matrix_4 + b[3] * matrix_2 +
-      b[1] * ident)
+      c[9] * matrix_8
+      + c[7] * matrix_6
+      + c[5] * matrix_4
+      + c[3] * matrix_2
+      + c[1] * ident
+  )
   matrix_u = math_ops.matmul(matrix, tmp)
   matrix_v = (
-      b[8] * matrix_8 + b[6] * matrix_6 + b[4] * matrix_4 + b[2] * matrix_2 +
-      b[0] * ident)
+      c[8] * matrix_8
+      + c[6] * matrix_6
+      + c[4] * matrix_4
+      + c[2] * matrix_2
+      + ident
+  )
   return matrix_u, matrix_v
 
 
 def _matrix_exp_pade13(matrix):
   """13th-order Pade approximant for matrix exponential."""
   b = [
-      64764752532480000.0, 32382376266240000.0, 7771770303897600.0,
-      1187353796428800.0, 129060195264000.0, 10559470521600.0, 670442572800.0,
-      33522128640.0, 1323241920.0, 40840800.0, 960960.0, 16380.0, 182.0
+      64764752532480000.0,
+      32382376266240000.0,
+      7771770303897600.0,
+      1187353796428800.0,
+      129060195264000.0,
+      10559470521600.0,
+      670442572800.0,
+      33522128640.0,
+      1323241920.0,
+      40840800.0,
+      960960.0,
+      16380.0,
+      182.0,
+      1.0,
   ]
-  b = [constant_op.constant(x, matrix.dtype) for x in b]
+  c = [constant_op.constant(x / b[0], matrix.dtype) for x in b]
   ident = linalg_ops.eye(
       array_ops.shape(matrix)[-2],
       batch_shape=array_ops.shape(matrix)[:-2],
@@ -221,13 +248,23 @@ def _matrix_exp_pade13(matrix):
   matrix_4 = math_ops.matmul(matrix_2, matrix_2)
   matrix_6 = math_ops.matmul(matrix_4, matrix_2)
   tmp_u = (
-      math_ops.matmul(matrix_6, matrix_6 + b[11] * matrix_4 + b[9] * matrix_2) +
-      b[7] * matrix_6 + b[5] * matrix_4 + b[3] * matrix_2 + b[1] * ident)
+      math_ops.matmul(
+          matrix_6, c[13] * matrix_6 + c[11] * matrix_4 + c[9] * matrix_2
+      )
+      + c[7] * matrix_6
+      + c[5] * matrix_4
+      + c[3] * matrix_2
+      + c[1] * ident
+  )
   matrix_u = math_ops.matmul(matrix, tmp_u)
-  tmp_v = b[12] * matrix_6 + b[10] * matrix_4 + b[8] * matrix_2
+  tmp_v = c[12] * matrix_6 + c[10] * matrix_4 + c[8] * matrix_2
   matrix_v = (
-      math_ops.matmul(matrix_6, tmp_v) + b[6] * matrix_6 + b[4] * matrix_4 +
-      b[2] * matrix_2 + b[0] * ident)
+      math_ops.matmul(matrix_6, tmp_v)
+      + c[6] * matrix_6
+      + c[4] * matrix_4
+      + c[2] * matrix_2
+      + ident
+  )
   return matrix_u, matrix_v