Branch 200727545

SECURITY.md

@@ -245,4 +245,4 @@ v//Fw6ZeY+HmRDFdirjD7wXtIuER4vqCryIqR6Xe9X8oJXz9L/Jhslc=
 ### Known Vulnerabilities
 
 For a list of known vulnerabilities and security advisories for TensorFlow,
-(https://github.com/tensorflow/tensorflow/blob/master/tensorflow/security/index.md)[click here].
+[click here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/security/index.md).

tensorflow/compiler/jit/encapsulate_subgraphs_pass.h

@@ -28,6 +28,9 @@ limitations under the License.
 namespace tensorflow {
 
 // A rewriting function to apply to each subgraph during encapsulation.
+// 'arg_source_tensors' are the tensors corresponding to the arguments in the
+// original source graph (*not* 'graph').
+//
 // 'graph' is the subgraph. The rewriting may renumber the inputs and outputs;
 // 'input_permutation' is a mapping from old argument numbers to new argument
 // numbers, whereas 'output_permutation' is the same for outputs. Both
@@ -37,6 +40,7 @@ namespace tensorflow {
 // The rewrite may also change the NodeDef's operator name, and that
 // name will be used as the name of the generated function.
 typedef std::function<Status(
+    const std::vector<OutputTensor>& arg_source_tensors,
     std::unique_ptr<Graph>* graph, std::vector<int>* input_permutation,
     std::vector<int>* output_permutation, NodeDef* node_def)>
     RewriteSubgraphFn;

tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc

@@ -757,7 +757,8 @@ TEST(EncapsulateSubgraphsWithGuaranteeConstOpTest, Simple) {
   TF_ASSERT_OK(EncapsulateSubgraphsInFunctions(
       "_encapsulate", "_outside", graph_before,
       /*rewrite_subgraph_fn=*/
-      [&guaranteed_consts](std::unique_ptr<Graph>* graph_ptr,
+      [&guaranteed_consts](const std::vector<OutputTensor>& arg_source_tensors,
+                           std::unique_ptr<Graph>* graph_ptr,
                            std::vector<int>* input_permutation,
                            std::vector<int>* output_permutation,
                            NodeDef* call_def) {
@@ -801,7 +802,8 @@ TEST(EncapsulateSubgraphsWithGuaranteeConstOpTest, Add) {
   TF_ASSERT_OK(EncapsulateSubgraphsInFunctions(
       "_encapsulate", "_outside", graph_before,
       /*rewrite_subgraph_fn=*/
-      [&guaranteed_consts](std::unique_ptr<Graph>* graph_ptr,
+      [&guaranteed_consts](const std::vector<OutputTensor>& arg_source_tensors,
+                           std::unique_ptr<Graph>* graph_ptr,
                            std::vector<int>* input_permutation,
                            std::vector<int>* output_permutation,
                            NodeDef* call_def) {

tensorflow/compiler/tests/BUILD

@@ -51,6 +51,15 @@ py_library(
     ],
 )
 
+py_test(
+    name = "xla_test_test",
+    size = "small",
+    srcs = ["xla_test_test.py"],
+    deps = [
+        ":xla_test",
+    ],
+)
+
 tf_xla_py_test(
     name = "adagrad_test",
     size = "small",

tensorflow/compiler/tests/unary_ops_test.py

@@ -201,6 +201,16 @@ def testFloatOps(self):
           expected=np.array([1.54308063, 3.76219569, 10.067662, 27.30823284],
                             dtype=dtype))
 
+      # Disable float16 testing for now
+      if dtype != np.float16:
+        x = np.arange(-10, 10, 1).astype(dtype)
+        with self.test_session() as session:
+          erf_x = session.run(math_ops.erf(x))
+          erfc_x = session.run(math_ops.erfc(x))
+
+        self._assertOpOutputMatchesExpected(math_ops.erf, x, expected=erf_x)
+        self._assertOpOutputMatchesExpected(math_ops.erfc, x, expected=erfc_x)
+
       self._assertOpOutputMatchesExpected(
           math_ops.exp,
           np.array([[-1, 1]], dtype=dtype),

tensorflow/compiler/tests/xla_test.py

@@ -49,6 +49,32 @@
                     'Value to set the TF_XLA_FLAGS environment variable to')
 
 
+def parse_disabled_manifest(manifest_content):
+  comments_re = re.compile('#.*$')
+  disabled_tests = []
+  disabled_method_types = []
+  for l in manifest_content.splitlines():
+    stripped = comments_re.sub('', l).strip()
+    if not stripped:
+      continue
+    entry = stripped.split(' ')
+    if len(entry) == 1:
+      disabled_tests.append(entry[0])
+    elif len(entry) == 2:
+      disabled_method_types.append((entry[0], entry[1].strip().split(',')))
+    else:
+      raise ValueError('Bad entry in manifest file.')
+
+  disabled_regex = '|'.join(disabled_tests)
+  method_types_filter = dict()
+  for method, types in disabled_method_types:
+    method_types_filter[method] = set([
+        dtypes.as_dtype(types_pb2.DataType.Value(name)).as_numpy_dtype
+        for name in types
+    ])
+  return disabled_regex, method_types_filter
+
+
 class XLATestCase(test.TestCase):
   """XLA test cases are parameterized test cases."""
 
@@ -85,38 +111,21 @@ def __init__(self, method_name='runTest'):
 
     # Parse the manifest file, if any, into a regex identifying tests to
     # disable
-    self.disabled_regex = None
-    self._method_types_filter = dict()
     # TODO(xpan): Make it text proto if it doesn't scale.
     # Each line of the manifest file specifies an entry. The entry can be
     # 1) TestNameRegex  // E.g. CumprodTest.* Or
     # 2) TestName TypeName  // E.g. AdamOptimizerTest.testSharing DT_BFLOAT16
     # The 1) disables the entire test. While 2) only filter some numeric types
     # so that they are not used in those tests.
+    self.disabled_regex = None
+    self._method_types_filter = {}
 
     if FLAGS.disabled_manifest is not None:
-      comments_re = re.compile('#.*$')
-      manifest_file = open(FLAGS.disabled_manifest, 'r')
-      disabled_tests = []
-      disabled_method_types = []
-      for l in manifest_file.read().splitlines():
-        if not l:
-          continue
-        entry = comments_re.sub('', l).strip().split(' ')
-        if len(entry) == 1:
-          disabled_tests.append(entry[0])
-        elif len(entry) == 2:
-          disabled_method_types.append(
-              (entry[0], entry[1].strip().split(',')))
-        else:
-          raise ValueError('Bad entry in manifest file.')
-
-      self.disabled_regex = re.compile('|'.join(disabled_tests))
-      for method, types in disabled_method_types:
-        self._method_types_filter[method] = set([
-            dtypes.as_dtype(types_pb2.DataType.Value(name)).as_numpy_dtype
-            for name in types])
-      manifest_file.close()
+      with open(FLAGS.disabled_manifest, 'r') as manifest_file:
+        disabled_regex, self._method_types_filter = (
+            parse_disabled_manifest(manifest_file.read()))
+        if disabled_regex:
+          self.disabled_regex = re.compile(disabled_regex)
 
     if FLAGS.tf_xla_flags is not None:
       os.environ['TF_XLA_FLAGS'] = FLAGS.tf_xla_flags

tensorflow/compiler/tests/xla_test_test.py

@@ -0,0 +1,44 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for the XLATestCase test fixture base class."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.platform import test
+
+
+class XlaTestCaseTestCase(test.TestCase):
+
+  def testManifestEmptyLineDoesNotCatchAll(self):
+    manifest = """
+testCaseOne
+"""
+    disabled_regex, _ = xla_test.parse_disabled_manifest(manifest)
+    self.assertEqual(disabled_regex, "testCaseOne")
+
+  def testManifestWholeLineCommentDoesNotCatchAll(self):
+    manifest = """# I am a comment
+testCaseOne
+testCaseTwo
+"""
+    disabled_regex, _ = xla_test.parse_disabled_manifest(manifest)
+    self.assertEqual(disabled_regex, "testCaseOne|testCaseTwo")
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tf2xla/kernels/unary_ops.cc

@@ -16,9 +16,11 @@ limitations under the License.
 // Native XLA implementations of simple unary Ops
 
 #include "tensorflow/compiler/tf2xla/kernels/cwise_ops.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "tensorflow/compiler/xla/client/client_library.h"
+#include "tensorflow/compiler/xla/client/lib/arithmetic.h"
 #include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/core/framework/kernel_def_builder.h"
 
@@ -185,5 +187,49 @@ XLAJIT_MAKE_UNARY(Imag, b->Imag(x));
 
 #undef XLAJIT_MAKE_UNARY
 
+// Erf/Erfc.  For x in (-1, 1), the erf approximation is used; erfc polynomial
+// is used outside of this range.
+class ErfOp : public XlaOpKernel {
+ public:
+  explicit ErfOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaBuilder* b = ctx->builder();
+    xla::PrimitiveType primitive_type;
+    xla::XlaOp one = XlaHelpers::One(b, input_type(0));
+    xla::XlaOp x = ctx->Input(0);
+    xla::XlaOp abs_x = b->Abs(x);
+
+    OP_REQUIRES_OK(ctx,
+                   DataTypeToPrimitiveType(input_type(0), &primitive_type));
+
+    auto y = b->Select(b->Gt(abs_x, one),
+                       b->Sub(one, ComputeErfc(b, x, primitive_type)),
+                       ComputeErf(b, x, primitive_type));
+    ctx->SetOutput(0, y);
+  }
+};
+REGISTER_XLA_OP(Name("Erf"), ErfOp);
+
+class ErfcOp : public XlaOpKernel {
+ public:
+  explicit ErfcOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp one = XlaHelpers::One(b, input_type(0));
+    xla::XlaOp x = ctx->Input(0);
+    xla::XlaOp abs_x = b->Abs(x);
+
+    xla::PrimitiveType primitive_type;
+    OP_REQUIRES_OK(ctx,
+                   DataTypeToPrimitiveType(input_type(0), &primitive_type));
+
+    auto y = b->Select(b->Lt(abs_x, one),
+                       b->Sub(one, ComputeErf(b, x, primitive_type)),
+                       ComputeErfc(b, x, primitive_type));
+    ctx->SetOutput(0, y);
+  }
+};
+REGISTER_XLA_OP(Name("Erfc"), ErfcOp);
+
 }  // namespace
 }  // namespace tensorflow

tensorflow/compiler/tf2xla/lib/batch_dot.cc

@@ -71,8 +71,8 @@ xla::StatusOr<xla::XlaOp> BatchDot(xla::XlaBuilder* builder, xla::XlaOp x,
   }
 
   // Check for zero lhs/rhs dim size.
-  if (xla::ShapeUtil::HasZeroElements(x_shape) ||
-      xla::ShapeUtil::HasZeroElements(y_shape)) {
+  if (xla::ShapeUtil::IsZeroElementArray(x_shape) ||
+      xla::ShapeUtil::IsZeroElementArray(y_shape)) {
     std::vector<int64> dimensions(batch_dimension_numbers.size());
     for (int i = 0; i < batch_dimension_numbers.size(); ++i) {
       dimensions[i] = x_shape.dimensions(batch_dimension_numbers[i]);

tensorflow/compiler/xla/BUILD

@@ -309,7 +309,6 @@ cc_library(
         ":types",
         ":util",
         ":xla_data_proto",
-        "//tensorflow/core:framework",
         "//tensorflow/core:lib",
     ],
 )

tensorflow/compiler/xla/client/lib/arithmetic.cc

@@ -121,4 +121,88 @@ StatusOr<XlaOp> Any(const XlaOp& predicates, XlaBuilder* builder) {
   return builder->Reduce(predicates, f, logical_or, all_dimensions);
 }
 
+namespace {
+xla::XlaOp FloatLiteral(xla::XlaBuilder* b, PrimitiveType data_type,
+                        float value) {
+  return b->ConvertElementType(b->ConstantR0(value), data_type);
+}
+
+// Polynomials for computing erf/erfc.  Originally from cephes.
+// Note we use float for compatibility across devices, at the cost of some
+// precision for 64 bit computations.
+//
+// Coefficients are in descending order.
+std::array<float, 9> kErfcPCoefficient = {
+    2.46196981473530512524E-10, 5.64189564831068821977E-1,
+    7.46321056442269912687E0,   4.86371970985681366614E1,
+    1.96520832956077098242E2,   5.26445194995477358631E2,
+    9.34528527171957607540E2,   1.02755188689515710272E3,
+    5.57535335369399327526E2};
+std::array<float, 9> kErfcQCoefficient = {
+    1.00000000000000000000E0, 1.32281951154744992508E1,
+    8.67072140885989742329E1, 3.54937778887819891062E2,
+    9.75708501743205489753E2, 1.82390916687909736289E3,
+    2.24633760818710981792E3, 1.65666309194161350182E3,
+    5.57535340817727675546E2};
+std::array<float, 6> kErfcRCoefficient = {
+    5.64189583547755073984E-1, 1.27536670759978104416E0,
+    5.01905042251180477414E0,  6.16021097993053585195E0,
+    7.40974269950448939160E0,  2.97886665372100240670E0};
+std::array<float, 7> kErfcSCoefficient = {
+    1.00000000000000000000E0, 2.26052863220117276590E0,
+    9.39603524938001434673E0, 1.20489539808096656605E1,
+    1.70814450747565897222E1, 9.60896809063285878198E0,
+    3.36907645100081516050E0};
+std::array<float, 5> kErfTCoefficient = {
+    9.60497373987051638749E0, 9.00260197203842689217E1,
+    2.23200534594684319226E3, 7.00332514112805075473E3,
+    5.55923013010394962768E4};
+std::array<float, 6> kErfUCoefficient = {
+    1.00000000000000000000E0, 3.35617141647503099647E1,
+    5.21357949780152679795E2, 4.59432382970980127987E3,
+    2.26290000613890934246E4, 4.92673942608635921086E4};
+}  // namespace
+
+// Evaluate the polynomial given coefficients and `x`.
+// N.B. Coefficients should be supplied in decreasing order.
+xla::XlaOp EvaluatePolynomial(xla::XlaBuilder* b, const xla::XlaOp& x,
+                              tensorflow::gtl::ArraySlice<float> coefficients,
+                              PrimitiveType data_type) {
+  xla::XlaOp poly = FloatLiteral(b, data_type, 0.0);
+  for (float c : coefficients) {
+    poly = b->Add(b->Mul(poly, x), FloatLiteral(b, data_type, c));
+  }
+  return poly;
+}
+
+// Compute an approximation of the error function complement (1 - erf(x)).
+xla::XlaOp ComputeErfc(xla::XlaBuilder* b, const xla::XlaOp& x,
+                       PrimitiveType data_type) {
+  xla::XlaOp zero = FloatLiteral(b, data_type, 0.0);
+  xla::XlaOp two = FloatLiteral(b, data_type, 2.0);
+  xla::XlaOp eight = FloatLiteral(b, data_type, 8.0);
+
+  xla::XlaOp abs_x = b->Abs(x);
+  xla::XlaOp z = b->Exp(b->Mul(b->Neg(x), x));
+
+  xla::XlaOp pp = EvaluatePolynomial(b, abs_x, kErfcPCoefficient, data_type);
+  xla::XlaOp pq = EvaluatePolynomial(b, abs_x, kErfcQCoefficient, data_type);
+  xla::XlaOp pr = EvaluatePolynomial(b, abs_x, kErfcRCoefficient, data_type);
+  xla::XlaOp ps = EvaluatePolynomial(b, abs_x, kErfcSCoefficient, data_type);
+
+  xla::XlaOp y = b->Select(b->Lt(abs_x, eight), b->Div(b->Mul(z, pp), pq),
+                           b->Div(b->Mul(z, pr), ps));
+
+  return b->Select(b->Lt(x, zero), b->Sub(two, y), y);
+}
+
+// Compute a polynomial approximation of the error function.
+xla::XlaOp ComputeErf(xla::XlaBuilder* b, const xla::XlaOp& x,
+                      PrimitiveType data_type) {
+  xla::XlaOp z = b->Mul(x, x);
+  xla::XlaOp pt = EvaluatePolynomial(b, z, kErfTCoefficient, data_type);
+  xla::XlaOp pu = EvaluatePolynomial(b, z, kErfUCoefficient, data_type);
+  return b->Div(b->Mul(x, pt), pu);
+}
+
 }  // namespace xla

tensorflow/compiler/xla/client/lib/arithmetic.h

@@ -55,6 +55,20 @@ XlaComputation CreateScalarOrComputation(XlaBuilder* builder);
 // Note: if predicates is zero-sized, Any() vacuously returns false.
 StatusOr<XlaOp> Any(const XlaOp& predicates, XlaBuilder* builder);
 
+// Evaluate the polynomial given coefficients and `x`.
+// N.B. Coefficients should be supplied in decreasing order.
+xla::XlaOp EvaluatePolynomial(xla::XlaBuilder* b, const xla::XlaOp& x,
+                              tensorflow::gtl::ArraySlice<double> coefficients,
+                              PrimitiveType data_type);
+
+// Compute an approximation of the error function complement (1 - erf(x)).
+xla::XlaOp ComputeErfc(xla::XlaBuilder* b, const xla::XlaOp& x,
+                       PrimitiveType data_type);
+
+// Compute an approximation of the error function.
+xla::XlaOp ComputeErf(xla::XlaBuilder* b, const xla::XlaOp& x,
+                      PrimitiveType data_type);
+
 }  // namespace xla
 
 #endif  // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_ARITHMETIC_H_