Branch 172647355

tensorflow/c/c_api.cc

@@ -1799,6 +1799,17 @@ void TF_GraphToGraphDef(TF_Graph* graph, TF_Buffer* output_graph_def,
   status->status = MessageToBuffer(def, output_graph_def);
 }
 
+void TF_GraphGetOpDef(TF_Graph* graph, const char* op_name,
+                      TF_Buffer* output_op_def, TF_Status* status) {
+  const OpDef* op_def;
+  {
+    mutex_lock l(graph->mu);
+    status->status = graph->graph.op_registry()->LookUpOpDef(op_name, &op_def);
+    if (!status->status.ok()) return;
+  }
+  status->status = MessageToBuffer(*op_def, output_op_def);
+}
+
 TF_ImportGraphDefOptions* TF_NewImportGraphDefOptions() {
   return new TF_ImportGraphDefOptions;
 }

tensorflow/c/c_api.h

@@ -864,6 +864,13 @@ TF_CAPI_EXPORT extern void TF_GraphToGraphDef(TF_Graph* graph,
                                               TF_Buffer* output_graph_def,
                                               TF_Status* status);
 
+// Returns the serialized OpDef proto with name `op_name`, or a bad status if no
+// such op exists. This can return OpDefs of functions copied into the graph.
+TF_CAPI_EXPORT extern void TF_GraphGetOpDef(TF_Graph* graph,
+                                            const char* op_name,
+                                            TF_Buffer* output_op_def,
+                                            TF_Status* status);
+
 // TF_ImportGraphDefOptions holds options that can be passed to
 // TF_GraphImportGraphDef.
 typedef struct TF_ImportGraphDefOptions TF_ImportGraphDefOptions;

tensorflow/c/c_api_function_test.cc

@@ -1465,5 +1465,26 @@ TEST_F(CApiFunctionTest, AppendHash) {
   ASSERT_EQ(string("func_name_base_qaJ8jA8UmGY"), fdef.signature().name());
 }
 
+TEST_F(CApiFunctionTest, GetOpDef) {
+  DefineFunction(func_name_, &func_);
+  TF_GraphCopyFunction(host_graph_, func_, nullptr, s_);
+  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);
+
+  // Test we can retrieve function OpDef from graph
+  TF_Buffer* buffer = TF_NewBuffer();
+  TF_GraphGetOpDef(host_graph_, func_name_, buffer, s_);
+  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);
+
+  // Sanity check returned OpDef
+  string data(static_cast<const char*>(buffer->data), buffer->length);
+  OpDef op_def;
+  op_def.ParseFromString(data);
+  EXPECT_EQ(op_def.name(), func_name_);
+  EXPECT_EQ(op_def.input_arg_size(), 1);
+  EXPECT_EQ(op_def.output_arg_size(), 1);
+
+  TF_DeleteBuffer(buffer);
+}
+
 }  // namespace
 }  // namespace tensorflow

tensorflow/c/c_api_test.cc

@@ -37,6 +37,7 @@ limitations under the License.
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/graph/tensor_id.h"
 #include "tensorflow/core/lib/core/error_codes.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
 #include "tensorflow/core/lib/io/path.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
@@ -50,6 +51,11 @@ Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst);
 
 namespace {
 
+static void ExpectHasSubstr(StringPiece s, StringPiece expected) {
+  EXPECT_TRUE(StringPiece(s).contains(expected))
+      << "'" << s << "' does not contain '" << expected << "'";
+}
+
 TEST(CAPI, Version) { EXPECT_STRNE("", TF_Version()); }
 
 TEST(CAPI, Status) {
@@ -837,6 +843,31 @@ TEST(CAPI, ShapeInferenceError) {
   TF_DeleteStatus(status);
 }
 
+TEST(CAPI, GetOpDef) {
+  TF_Status* status = TF_NewStatus();
+  TF_Graph* graph = TF_NewGraph();
+  TF_Buffer* buffer = TF_NewBuffer();
+
+  TF_GraphGetOpDef(graph, "Add", buffer, status);
+  ASSERT_EQ(TF_OK, TF_GetCode(status));
+  const OpDef* expected_op_def;
+  TF_ASSERT_OK(OpRegistry::Global()->LookUpOpDef("Add", &expected_op_def));
+  string expected_serialized;
+  expected_op_def->SerializeToString(&expected_serialized);
+  string actual_string(reinterpret_cast<const char*>(buffer->data),
+                       buffer->length);
+  EXPECT_EQ(expected_serialized, actual_string);
+
+  TF_GraphGetOpDef(graph, "MyFakeOp", buffer, status);
+  EXPECT_EQ(TF_NOT_FOUND, TF_GetCode(status));
+  ExpectHasSubstr(TF_Message(status),
+                  "Op type not registered 'MyFakeOp' in binary");
+
+  TF_DeleteBuffer(buffer);
+  TF_DeleteGraph(graph);
+  TF_DeleteStatus(status);
+}
+
 void StringVectorToArrays(const std::vector<string>& v,
                           std::unique_ptr<const void* []>* ptrs,
                           std::unique_ptr<size_t[]>* lens) {

tensorflow/compiler/xla/BUILD

@@ -171,6 +171,7 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         ":status",
+        ":status_macros",
         ":types",
         ":xla_data_proto",
         "//tensorflow/core:lib",

tensorflow/compiler/xla/service/BUILD

@@ -579,12 +579,14 @@ cc_library(
         ":shaped_buffer",
         ":versioned_computation_handle",
         "//tensorflow/compiler/xla:executable_run_options",
+        "//tensorflow/compiler/xla:status",
         "//tensorflow/compiler/xla:status_macros",
         "//tensorflow/compiler/xla:statusor",
         "//tensorflow/compiler/xla:util",
         "//tensorflow/compiler/xla:xla_data_proto",
         "//tensorflow/compiler/xla/legacy_flags:debug_options_flags",
         "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
         "//tensorflow/core:stream_executor_no_cuda",
         "//tensorflow/stream_executor",
     ],

tensorflow/compiler/xla/service/executable.cc

@@ -17,7 +17,9 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h"
 #include "tensorflow/compiler/xla/service/hlo_graph_dumper.h"
+#include "tensorflow/compiler/xla/status.h"
 #include "tensorflow/compiler/xla/status_macros.h"
+#include "tensorflow/core/lib/hash/hash.h"
 #include "tensorflow/core/lib/io/path.h"
 #include "tensorflow/core/lib/strings/stringprintf.h"
 #include "tensorflow/core/platform/env.h"
@@ -82,7 +84,11 @@ Status Executable::DumpSessionModule() {
   }
   filename = SanitizeFileName(std::move(filename));
   string file_path = tensorflow::io::JoinPath(directory_path, filename);
-  return tensorflow::WriteBinaryProto(env, file_path, session_module);
+  string result;
+  TF_RET_CHECK(
+      tensorflow::SerializeToStringDeterministic(session_module, &result));
+  return tensorflow::WriteStringToFile(tensorflow::Env::Default(), file_path,
+                                       result);
 }
 
 }  // namespace xla

tensorflow/compiler/xla/service/hlo_module.cc

@@ -154,8 +154,8 @@ string HloModule::ToString() const {
   std::ostringstream s;
   s << "HloModule " << name() << ":\n\n";
   s << "ENTRY " << entry_computation()->ToString() << "\n\n";
-  for (const std::unique_ptr<HloComputation>& computation : computations_) {
-    if (computation.get() != entry_computation()) {
+  for (const HloComputation* computation : MakeNonfusionComputations()) {
+    if (computation != entry_computation()) {
       s << computation->ToString() << "\n\n";
     }
   }

tensorflow/compiler/xla/util.h

@@ -24,6 +24,7 @@ limitations under the License.
 #include <vector>
 
 #include "tensorflow/compiler/xla/status.h"
+#include "tensorflow/compiler/xla/status_macros.h"
 #include "tensorflow/compiler/xla/types.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
 #include "tensorflow/core/lib/core/status.h"

tensorflow/contrib/data/python/kernel_tests/BUILD

@@ -11,6 +11,9 @@ py_test(
     size = "small",
     srcs = ["batch_dataset_op_test.py"],
     srcs_version = "PY2AND3",
+    tags = [
+        "manual",  # b/67958604
+    ],
     deps = [
         "//tensorflow/contrib/data/python/ops:dataset_ops",
         "//tensorflow/contrib/data/python/ops:transformation_ops",
@@ -352,6 +355,9 @@ py_test(
     size = "small",
     srcs = ["sloppy_transformation_dataset_op_test.py"],
     srcs_version = "PY2AND3",
+    tags = [
+        "manual",  # b/67958761
+    ],
     deps = [
         "//tensorflow/contrib/data/python/ops:dataset_ops",
         "//tensorflow/contrib/data/python/ops:transformation_ops",

tensorflow/contrib/distributions/python/kernel_tests/poisson_lognormal_test.py

@@ -18,6 +18,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import numpy as np
+
 from tensorflow.contrib.distributions.python.ops import poisson_lognormal
 from tensorflow.contrib.distributions.python.ops import test_util
 from tensorflow.python.platform import test
@@ -32,7 +34,8 @@ def testSampleProbConsistent(self):
       pln = poisson_lognormal.PoissonLogNormalQuadratureCompound(
           loc=-2.,
           scale=1.1,
-          quadrature_polynomial_degree=10,
+          quadrature_grid_and_probs=(
+              np.polynomial.hermite.hermgauss(deg=10)),
           validate_args=True)
       self.run_test_sample_consistent_log_prob(
           sess, pln, rtol=0.1)
@@ -42,7 +45,8 @@ def testMeanVariance(self):
       pln = poisson_lognormal.PoissonLogNormalQuadratureCompound(
           loc=0.,
           scale=1.,
-          quadrature_polynomial_degree=10,
+          quadrature_grid_and_probs=(
+              np.polynomial.hermite.hermgauss(deg=10)),
           validate_args=True)
       self.run_test_sample_consistent_mean_variance(
           sess, pln, rtol=0.02)
@@ -52,7 +56,8 @@ def testSampleProbConsistentBroadcastScalar(self):
       pln = poisson_lognormal.PoissonLogNormalQuadratureCompound(
           loc=[0., -0.5],
           scale=1.,
-          quadrature_polynomial_degree=10,
+          quadrature_grid_and_probs=(
+              np.polynomial.hermite.hermgauss(deg=10)),
           validate_args=True)
       self.run_test_sample_consistent_log_prob(
           sess, pln, rtol=0.1, atol=0.01)
@@ -62,7 +67,8 @@ def testMeanVarianceBroadcastScalar(self):
       pln = poisson_lognormal.PoissonLogNormalQuadratureCompound(
           loc=[0., -0.5],
           scale=1.,
-          quadrature_polynomial_degree=10,
+          quadrature_grid_and_probs=(
+              np.polynomial.hermite.hermgauss(deg=10)),
           validate_args=True)
       self.run_test_sample_consistent_mean_variance(
           sess, pln, rtol=0.1, atol=0.01)
@@ -72,7 +78,8 @@ def testSampleProbConsistentBroadcastBoth(self):
       pln = poisson_lognormal.PoissonLogNormalQuadratureCompound(
           loc=[[0.], [-0.5]],
           scale=[[1., 0.9]],
-          quadrature_polynomial_degree=10,
+          quadrature_grid_and_probs=(
+              np.polynomial.hermite.hermgauss(deg=10)),
           validate_args=True)
       self.run_test_sample_consistent_log_prob(
           sess, pln, rtol=0.1, atol=0.08)
@@ -82,7 +89,8 @@ def testMeanVarianceBroadcastBoth(self):
       pln = poisson_lognormal.PoissonLogNormalQuadratureCompound(
           loc=[[0.], [-0.5]],
           scale=[[1., 0.9]],
-          quadrature_polynomial_degree=10,
+          quadrature_grid_and_probs=(
+              np.polynomial.hermite.hermgauss(deg=10)),
           validate_args=True)
       self.run_test_sample_consistent_mean_variance(
           sess, pln, rtol=0.1, atol=0.01)

tensorflow/contrib/distributions/python/ops/poisson_lognormal.py

@@ -93,7 +93,7 @@ class PoissonLogNormalQuadratureCompound(distribution_lib.Distribution):
           : d=0, ..., deg-1 }
   ```
 
-  where, [`grid, w = numpy.polynomial.hermite.hermgauss(deg)`](
+  where, [e.g., `grid, w = numpy.polynomial.hermite.hermgauss(deg)`](
   https://docs.scipy.org/doc/numpy-1.10.0/reference/generated/numpy.polynomial.hermite.hermgauss.html)
   and `prob = w / sqrt(pi)`.
 
@@ -106,14 +106,15 @@ class PoissonLogNormalQuadratureCompound(distribution_lib.Distribution):
   pln = ds.PoissonLogNormalQuadratureCompound(
       loc=[0., -0.5],
       scale=1.,
-      quadrature_polynomial_degree=10,
+      quadrature_grid_and_probs=(
+        np.polynomial.hermite.hermgauss(deg=10)),
       validate_args=True)
   """
 
   def __init__(self,
                loc,
                scale,
-               quadrature_polynomial_degree=8,
+               quadrature_grid_and_probs=None,
                validate_args=False,
                allow_nan_stats=True,
                name="PoissonLogNormalQuadratureCompound"):
@@ -124,8 +125,9 @@ def __init__(self,
         the LogNormal prior.
       scale: `float`-like (batch of) scalar `Tensor`; the scale parameter of
         the LogNormal prior.
-      quadrature_polynomial_degree: Python `int`-like scalar.
-        Default value: 8.
+      quadrature_grid_and_probs: Python pair of `list`-like objects representing
+        the sample points and the corresponding (possibly normalized) weight.
+        When `None`, defaults to: `np.polynomial.hermite.hermgauss(deg=8)`.
       validate_args: Python `bool`, default `False`. When `True` distribution
         parameters are checked for validity despite possibly degrading runtime
         performance. When `False` invalid inputs may silently render incorrect
@@ -138,6 +140,8 @@ def __init__(self,
 
     Raises:
       TypeError: if `loc.dtype != scale[0].dtype`.
+      ValueError: if `quadrature_grid_and_probs is not None` and
+        `len(quadrature_grid_and_probs[0]) != len(quadrature_grid_and_probs[1])`
     """
     parameters = locals()
     with ops.name_scope(name, values=[loc, scale]):
@@ -153,18 +157,21 @@ def __init__(self,
             "loc.dtype(\"{}\") does not match scale.dtype(\"{}\")".format(
                 loc.dtype.name, scale.dtype.name))
 
-      self._degree = quadrature_polynomial_degree
-
-      grid, prob = np.polynomial.hermite.hermgauss(
-          deg=quadrature_polynomial_degree)
-
-      # It should be that `sum(prob) == sqrt(pi)`, but self-normalization is
-      # more numerically stable.
-      prob = prob.astype(dtype.as_numpy_dtype)
-      prob /= np.linalg.norm(prob, ord=1)
+      if quadrature_grid_and_probs is None:
+        grid, probs = np.polynomial.hermite.hermgauss(deg=8)
+      else:
+        grid, probs = tuple(quadrature_grid_and_probs)
+        if len(grid) != len(probs):
+          raise ValueError("`quadrature_grid_and_probs` must be a `tuple` of "
+                           "same-length list-like objects")
+      grid = grid.astype(dtype.as_numpy_dtype)
+      probs = probs.astype(dtype.as_numpy_dtype)
+      probs /= np.linalg.norm(probs, ord=1)
+      self._quadrature_grid = grid
+      self._quadrature_probs = probs
 
       self._mixture_distribution = categorical_lib.Categorical(
-          logits=np.log(prob),
+          logits=np.log(probs),
           validate_args=validate_args,
           allow_nan_stats=allow_nan_stats)
 
@@ -210,9 +217,14 @@ def scale(self):
     return self._scale
 
   @property
-  def quadrature_polynomial_degree(self):
-    """Polynomial largest exponent used for Gauss-Hermite quadrature."""
-    return self._degree
+  def quadrature_grid(self):
+    """Quadrature grid points."""
+    return self._quadrature_grid
+
+  @property
+  def quadrature_probs(self):
+    """Quadrature normalized weights."""
+    return self._quadrature_probs
 
   def _batch_shape_tensor(self):
     return array_ops.broadcast_dynamic_shape(
@@ -242,10 +254,10 @@ def _sample_n(self, n, seed=None):
                 [batch_size])),
         seed=distribution_util.gen_new_seed(
             seed, "poisson_lognormal_quadrature_compound"))
-    # Stride `quadrature_polynomial_degree` for `batch_size` number of times.
+    # Stride `quadrature_degree` for `batch_size` number of times.
     offset = math_ops.range(start=0,
-                            limit=batch_size * self._degree,
-                            delta=self._degree,
+                            limit=batch_size * len(self.quadrature_probs),
+                            delta=len(self.quadrature_probs),
                             dtype=ids.dtype)
     ids += offset
     rate = array_ops.gather(