From 758f25e8168bf1ff76c63a5b54dfd50ff54e4e27 Mon Sep 17 00:00:00 2001
From: Sunitha Kambhampati <skambha@us.ibm.com>
Date: Tue, 17 Apr 2018 15:47:33 -0700
Subject: [PATCH 001/395] Fix calculation of the histogram buckets and writing
 to the tensor and add a unit test

---
 .../tensorboard/db/summary_db_writer.cc       | 21 +++++---
 .../tensorboard/db/summary_db_writer_test.cc  | 49 +++++++++++++++++++
 2 files changed, 62 insertions(+), 8 deletions(-)

diff --git a/tensorflow/contrib/tensorboard/db/summary_db_writer.cc b/tensorflow/contrib/tensorboard/db/summary_db_writer.cc
index 6590d6f7df..046a2d3884 100644
--- a/tensorflow/contrib/tensorboard/db/summary_db_writer.cc
+++ b/tensorflow/contrib/tensorboard/db/summary_db_writer.cc
@@ -1182,14 +1182,19 @@ class SummaryDbWriter : public SummaryWriterInterface {
     // See tensorboard/plugins/histogram/summary.py and data_compat.py
     Tensor t{DT_DOUBLE, {k, 3}};
     auto data = t.flat<double>();
-    for (int i = 0; i < k; ++i) {
-      double left_edge = ((i - 1 >= 0) ? histo.bucket_limit(i - 1)
-                                       : std::numeric_limits<double>::min());
-      double right_edge = ((i + 1 < k) ? histo.bucket_limit(i + 1)
-                                       : std::numeric_limits<double>::max());
-      data(i + 0) = left_edge;
-      data(i + 1) = right_edge;
-      data(i + 2) = histo.bucket(i);
+    for (int i = 0, j = 0; i < k; ++i) {
+      // From summary.proto
+      // Parallel arrays encoding the bucket boundaries and the bucket values.
+      // bucket(i) is the count for the bucket i.  The range for
+      // a bucket is:
+      //   i == 0:  -DBL_MAX .. bucket_limit(0)
+      //   i != 0:  bucket_limit(i-1) .. bucket_limit(i)
+      double left_edge = (i == 0) ? std::numeric_limits<double>::min()
+                                  : histo.bucket_limit(i - 1);
+
+      data(j++) = left_edge;
+      data(j++) = histo.bucket_limit(i);
+      data(j++) = histo.bucket(i);
     }
     int64 tag_id;
     PatchPluginName(s->mutable_metadata(), kHistogramPluginName);
diff --git a/tensorflow/contrib/tensorboard/db/summary_db_writer_test.cc b/tensorflow/contrib/tensorboard/db/summary_db_writer_test.cc
index 29b8063218..cb51325d15 100644
--- a/tensorflow/contrib/tensorboard/db/summary_db_writer_test.cc
+++ b/tensorflow/contrib/tensorboard/db/summary_db_writer_test.cc
@@ -100,6 +100,55 @@ class SummaryDbWriterTest : public ::testing::Test {
   SummaryWriterInterface* writer_ = nullptr;
 };
 
+TEST_F(SummaryDbWriterTest, WriteHistogram_VerifyTensorValues) {
+  TF_ASSERT_OK(CreateSummaryDbWriter(db_, "histtest", "test1", "user1", &env_,
+                                     &writer_));
+  int step = 0;
+  std::unique_ptr<Event> e{new Event};
+  e->set_step(step);
+  e->set_wall_time(123);
+  Summary::Value* s = e->mutable_summary()->add_value();
+  s->set_tag("normal/myhisto");
+
+  double dummy_value = 10.123;
+  HistogramProto* proto = s->mutable_histo();
+  proto->Clear();
+  proto->set_min(dummy_value);
+  proto->set_max(dummy_value);
+  proto->set_num(dummy_value);
+  proto->set_sum(dummy_value);
+  proto->set_sum_squares(dummy_value);
+
+  int size = 3;
+  double bucket_limits[] = {-30.5, -10.5, -5.5};
+  double bucket[] = {-10, 10, 20};
+  for (int i = 0; i < size; i++) {
+    proto->add_bucket_limit(bucket_limits[i]);
+    proto->add_bucket(bucket[i]);
+  }
+  TF_ASSERT_OK(writer_->WriteEvent(std::move(e)));
+  TF_ASSERT_OK(writer_->Flush());
+  writer_->Unref();
+  writer_ = nullptr;
+
+  // Verify the data
+  string result = QueryString("SELECT data FROM Tensors");
+  const double* val = reinterpret_cast<const double*>(result.data());
+  double histarray[] = {std::numeric_limits<double>::min(),
+                        -30.5,
+                        -10,
+                        -30.5,
+                        -10.5,
+                        10,
+                        -10.5,
+                        -5.5,
+                        20};
+  int histarray_size = 9;
+  for (int i = 0; i < histarray_size; i++) {
+    EXPECT_EQ(histarray[i], val[i]);
+  }
+}
+
 TEST_F(SummaryDbWriterTest, NothingWritten_NoRowsCreated) {
   TF_ASSERT_OK(CreateSummaryDbWriter(db_, "mad-science", "train", "jart", &env_,
                                      &writer_));
-- 
GitLab


From f0df6701d01954073e912f24f7c983de4f091a1e Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Fri, 20 Apr 2018 14:01:02 +0300
Subject: [PATCH 002/395] [tf.data] Check in a strictly faster rejection
 resampling transformation.

This transformation is faster because it rejects fewer data. This
is done by occasionally sampling from the original data distribution
in an efficient way.

Tested:
bazel test :resample_test
---
 .../data/python/kernel_tests/resample_test.py | 128 +++++++--
 .../contrib/data/python/ops/resampling.py     | 267 ++++++++++++++----
 2 files changed, 327 insertions(+), 68 deletions(-)

diff --git a/tensorflow/contrib/data/python/kernel_tests/resample_test.py b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
index 5f47dcb339..9e1273eba1 100644
--- a/tensorflow/contrib/data/python/kernel_tests/resample_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
@@ -18,6 +18,8 @@ from __future__ import division
 from __future__ import print_function
 
 import numpy as np
+import time
+from absl.testing import parameterized
 
 from tensorflow.contrib.data.python.ops import resampling
 from tensorflow.python.data.ops import dataset_ops
@@ -30,47 +32,70 @@ from tensorflow.python.platform import test
 from tensorflow.python.util import compat
 
 
-class ResampleTest(test.TestCase):
+def _time_resampling(
+    test_obj, data_np, target_dist, init_dist, use_v2, num_to_sample):
+  dataset = dataset_ops.Dataset.from_tensor_slices(data_np).repeat()
 
-  def testInitialKnownDistribution(self):
-    self._testDistribution(initial_known=True)
+  # Reshape distribution via rejection sampling.
+  apply_fn = (resampling.rejection_resample_v2 if use_v2 else
+              resampling.rejection_resample)
+  dataset = dataset.apply(
+      apply_fn(
+          class_func=lambda x: x,
+          target_dist=target_dist,
+          initial_dist=init_dist,
+          seed=142))
 
-  def testInitialNotKnownDistribution(self):
-    self._testDistribution(initial_known=False)
+  get_next = dataset.make_one_shot_iterator().get_next()
 
-  def _testDistribution(self, initial_known):
+  with test_obj.test_session() as sess:
+    start_time = time.time()
+    for _ in xrange(num_to_sample):
+      sess.run(get_next)
+    end_time = time.time()
+
+  return end_time - start_time
+
+
+class ResampleTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      ('InitialnDistributionKnown', True, False),
+      ('InitialDistributionUnknown', False, False),
+      ('InitialDistributionKnownV2', True, True),
+      ('InitialDistributionUnknownV2', False, True))
+  def testDistribution(self, initial_known, use_v2):
     classes = np.random.randint(5, size=(20000,))  # Uniformly sampled
     target_dist = [0.9, 0.05, 0.05, 0.0, 0.0]
     initial_dist = [0.2] * 5 if initial_known else None
-    iterator = (dataset_ops.Dataset.from_tensor_slices(classes).shuffle(
-        200, seed=21).map(lambda c: (c, string_ops.as_string(c))).apply(
-            resampling.rejection_resample(
-                target_dist=target_dist,
-                initial_dist=initial_dist,
-                class_func=lambda c, _: c,
-                seed=27)).make_one_shot_iterator())
-    get_next = iterator.get_next()
+    dataset = dataset_ops.Dataset.from_tensor_slices(classes).shuffle(
+        200, seed=21).map(lambda c: (c, string_ops.as_string(c))).repeat()
+    apply_fn = (resampling.rejection_resample_v2 if use_v2 else
+                resampling.rejection_resample)
+    get_next = dataset.apply(
+        apply_fn(
+            target_dist=target_dist,
+            initial_dist=initial_dist,
+            class_func=lambda c, _: c,
+            seed=27)).make_one_shot_iterator().get_next()
 
     with self.test_session() as sess:
       returned = []
-      with self.assertRaises(errors.OutOfRangeError):
-        while True:
-          returned.append(sess.run(get_next))
+      while len(returned) < 4000:
+        returned.append(sess.run(get_next))
 
     returned_classes, returned_classes_and_data = zip(*returned)
     _, returned_data = zip(*returned_classes_and_data)
     self.assertAllEqual([compat.as_bytes(str(c))
                          for c in returned_classes], returned_data)
     total_returned = len(returned_classes)
-    # Subsampling rejects a large percentage of the initial data in
-    # this case.
-    self.assertGreater(total_returned, 20000 * 0.2)
     class_counts = np.array([
         len([True for v in returned_classes if v == c])
         for c in range(5)])
     returned_dist = class_counts / total_returned
     self.assertAllClose(target_dist, returned_dist, atol=1e-2)
 
+
   def testRandomClasses(self):
     init_dist = [0.25, 0.25, 0.25, 0.25]
     target_dist = [0.0, 0.0, 0.0, 1.0]
@@ -109,5 +134,68 @@ class ResampleTest(test.TestCase):
 
     self.assertAllClose(target_dist, bincount, atol=1e-2)
 
+  @parameterized.named_parameters(
+        ('InitialnDistributionKnown', True, False),
+        ('InitialDistributionUnknown', False, False),
+        ('InitialDistributionKnownV2', True, True),
+        ('InitialDistributionUnknownV2', False, True))
+  def _testNewResampleIsFaster(self, target_dist, num_to_sample):
+    init_dist = [0.25, 0.25, 0.25, 0.25]
+    num_classes = len(init_dist)
+    num_samples = 1000
+    data_np = np.random.choice(num_classes, num_samples, p=init_dist)
+
+    fast_time = _time_resampling(self, data_np, target_dist, init_dist,
+                                 use_v2=True, num_to_sample=num_to_sample)
+    slow_time = _time_resampling(self, data_np, target_dist, init_dist,
+                                 use_v2=False, num_to_sample=num_to_sample)
+
+    self.assertLess(fast_time, slow_time)
+
+
+  def testNewResampleIsFasterSmallSkewManySamples(self):
+    self._testNewResampleIsFaster([0.1, 0.1, 0.1, 0.7], 1000)
+
+  def testNewResampleIsFasterBigSkewManySamples(self):
+    self._testNewResampleIsFaster([0.01, 0.01, 0.01, 0.97], 1000)
+
+  def testNewResampleIsFasterSmallSkewFewSamples(self):
+    self._testNewResampleIsFaster([0.1, 0.1, 0.1, 0.7], 100)
+
+  def testNewResampleIsFasterBigSkewFewSamples(self):
+    self._testNewResampleIsFaster([0.01, 0.01, 0.01, 0.97], 100)
+
+
+class MapDatasetBenchmark(test.Benchmark):
+
+  def benchmarkResamplePerformance(self):
+    init_dist = [0.25, 0.25, 0.25, 0.25]
+    target_dist = [0.0, 0.0, 0.0, 1.0]
+    num_classes = len(init_dist)
+    # We don't need many samples to test a dirac-delta target distribution
+    num_samples = 1000
+    data_np = np.random.choice(num_classes, num_samples, p=init_dist)
+
+    resample_time = _time_resampling(
+        self, data_np, target_dist, init_dist, use_v2=False, num_to_sample=1000)
+
+    self.report_benchmark(
+        iters=1000, wall_time=resample_time, name="benchmark_resample")
+
+  def benchmarkResampleAndBatchPerformance(self):
+    init_dist = [0.25, 0.25, 0.25, 0.25]
+    target_dist = [0.0, 0.0, 0.0, 1.0]
+    num_classes = len(init_dist)
+    # We don't need many samples to test a dirac-delta target distribution
+    num_samples = 1000
+    data_np = np.random.choice(num_classes, num_samples, p=init_dist)
+
+    resample_time = _time_resampling(
+        self, data_np, target_dist, init_dist, use_v2=True, num_to_sample=1000)
+
+    self.report_benchmark(
+        iters=1000, wall_time=resample_time, name="benchmark_resample_v2")
+
+
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index b465397437..94e28b9a2d 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -20,6 +20,7 @@ from __future__ import print_function
 import numpy as np
 
 from tensorflow.contrib.data.python.ops import batching
+from tensorflow.contrib.data.python.ops import interleave_ops
 from tensorflow.contrib.data.python.ops import scan_ops
 from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.framework import dtypes
@@ -50,14 +51,15 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
     A `Dataset` transformation function, which can be passed to
     @{tf.data.Dataset.apply}.
   """
-
   def _apply_fn(dataset):
     """Function from `Dataset` to `Dataset` that applies the transformation."""
-    dist_estimation_batch_size = 32
     target_dist_t = ops.convert_to_tensor(target_dist, name="target_dist")
     class_values_ds = dataset.map(class_func)
+
+    # Get initial distribution.
     if initial_dist is not None:
-      initial_dist_t = ops.convert_to_tensor(initial_dist, name="initial_dist")
+      initial_dist_t = ops.convert_to_tensor(
+          initial_dist, name="initial_dist")
       acceptance_dist = _calculate_acceptance_probs(initial_dist_t,
                                                     target_dist_t)
       initial_dist_ds = dataset_ops.Dataset.from_tensors(
@@ -65,55 +67,181 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
       acceptance_dist_ds = dataset_ops.Dataset.from_tensors(
           acceptance_dist).repeat()
     else:
-      num_classes = (target_dist_t.shape[0].value or
-                     array_ops.shape(target_dist_t)[0])
-      smoothing_constant = 10
-      initial_examples_per_class_seen = array_ops.fill(
-          [num_classes], np.int64(smoothing_constant))
-
-      def update_estimate_and_tile(num_examples_per_class_seen, c):
-        updated_examples_per_class_seen, dist = _estimate_data_distribution(
-            c, num_examples_per_class_seen)
-        tiled_dist = array_ops.tile(
-            array_ops.expand_dims(dist, 0), [dist_estimation_batch_size, 1])
-        return updated_examples_per_class_seen, tiled_dist
-
-      initial_dist_ds = (class_values_ds.batch(dist_estimation_batch_size)
-                         .apply(scan_ops.scan(initial_examples_per_class_seen,
-                                              update_estimate_and_tile))
-                         .apply(batching.unbatch()))
+      initial_dist_ds = _estimate_initial_dist_ds(
+          target_dist_t, class_values_ds)
       acceptance_dist_ds = initial_dist_ds.map(
           lambda initial: _calculate_acceptance_probs(initial, target_dist_t))
+    return _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds,
+                      class_values_ds, seed)
+
+  return _apply_fn
+
 
-    def maybe_warn_on_large_rejection(accept_dist, initial_dist):
-      proportion_rejected = math_ops.reduce_sum(
-          (1 - accept_dist) * initial_dist)
-      return control_flow_ops.cond(
-          math_ops.less(proportion_rejected, .5),
-          lambda: accept_dist,
-          lambda: logging_ops.Print(  # pylint: disable=g-long-lambda
-              accept_dist, [proportion_rejected, initial_dist, accept_dist],
-              message="Proportion of examples rejected by sampler is high: ",
-              summarize=100,
-              first_n=10))
-
-    acceptance_dist_ds = (dataset_ops.Dataset.zip((acceptance_dist_ds,
-                                                   initial_dist_ds))
-                          .map(maybe_warn_on_large_rejection))
-
-    def _gather_and_copy(class_val, acceptance_prob, data):
-      return (class_val, array_ops.gather(acceptance_prob, class_val), data)
-    current_probabilities_and_class_and_data_ds = dataset_ops.Dataset.zip(
-        (class_values_ds, acceptance_dist_ds, dataset)).map(_gather_and_copy)
-    filtered_ds = (
-        current_probabilities_and_class_and_data_ds
-        .filter(lambda _1, p, _2: random_ops.random_uniform([], seed=seed) < p))
-    return filtered_ds.map(lambda class_value, _, data: (class_value, data))
+def rejection_resample_v2(class_func, target_dist, initial_dist=None,
+                          seed=None):
+  """A transformation that resamples a dataset to achieve a target distribution.
 
+  This differs from v1 in that it will also sample from the original dataset
+  with some probability, so it makes strictly fewer data rejections. This
+  transformation is faster than the original.
+
+  **NOTE** Resampling is performed via rejection sampling; some fraction
+  of the input values will be dropped.
+
+  Args:
+    class_func: A function mapping an element of the input dataset to a scalar
+      `tf.int32` tensor. Values should be in `[0, num_classes)`.
+    target_dist: A floating point type tensor, shaped `[num_classes]`.
+    initial_dist: (Optional.)  A floating point type tensor, shaped
+      `[num_classes]`.  If not provided, the true class distribution is
+      estimated live in a streaming fashion.
+    seed: (Optional.) Python integer seed for the resampler.
+
+  Returns:
+    A `Dataset` transformation function, which can be passed to
+    @{tf.data.Dataset.apply}.
+  """
+  def _apply_fn(dataset):
+    """Function from `Dataset` to `Dataset` that applies the transformation."""
+    target_dist_t = ops.convert_to_tensor(target_dist, name="target_dist")
+    class_values_ds = dataset.map(class_func)
+
+    # Get initial distribution.
+    if initial_dist is not None:
+      initial_dist_t = ops.convert_to_tensor(
+          initial_dist, name="initial_dist")
+      acceptance_dist, prob_of_original = (
+          _calculate_acceptance_probs_with_mixing(initial_dist_t,
+                                                  target_dist_t))
+      initial_dist_ds = dataset_ops.Dataset.from_tensors(
+          initial_dist_t).repeat()
+      acceptance_dist_ds = dataset_ops.Dataset.from_tensors(
+          acceptance_dist).repeat()
+      prob_of_original_ds = dataset_ops.Dataset.from_tensors(
+          prob_of_original).repeat()
+    else:
+      initial_dist_ds = _estimate_initial_dist_ds(
+          target_dist_t, class_values_ds)
+      acceptance_and_original_prob_ds = initial_dist_ds.map(
+          lambda initial: _calculate_acceptance_probs_with_mixing(
+              initial, target_dist_t))
+      acceptance_dist_ds = acceptance_and_original_prob_ds.map(
+          lambda accept_prob, _: accept_prob)
+      prob_of_original_ds = acceptance_and_original_prob_ds.map(
+          lambda _, prob_original: prob_original)
+    filtered_ds = _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds,
+                             class_values_ds, seed)
+    # Prefetch filtered dataset for speed.
+    filtered_ds = filtered_ds.prefetch(3)
+
+    return interleave_ops.sample_from_datasets(
+        [dataset_ops.Dataset.zip((class_values_ds, dataset)), filtered_ds],
+        weights=prob_of_original_ds.map(lambda prob: [(prob, 1.0 - prob)]),
+        seed=seed)
 
   return _apply_fn
 
 
+def _random_interleave_datasets(ds1, ds1_classes, ds2, prob_of_ds1, seed=None):
+  """Randomly interleave datasets.
+
+  We carefully combine `ds1` and 'ds2' so that we don't needlessly compute the
+  filtering.
+
+  Args:
+    ds1: A dataset to interleave.
+    ds1_classes: Dataset of class values associated with ds1.
+    ds2: Another dataset to interleave.
+    prob_of_ds1: A dataset of probabilities. Each probability represents the
+      likelihood of drawing from `ds1`.
+    seed: (Optional.) Python integer seed for the resampler.
+
+  Returns:
+    A single dataset, combined from `ds1` and `ds2`.
+  """
+  num_filtered_to_prefetch = 3
+  ds2 = ds2.prefetch(num_filtered_to_prefetch)
+  filtered_iterator = ds2.make_one_shot_iterator()
+  combined_ds = dataset_ops.Dataset.zip(
+        (ds1_classes, ds1, prob_of_ds1)).map(
+        lambda ds1_class, original_data, prob_of_original:
+        control_flow_ops.cond(
+              random_ops.random_uniform([], seed=seed) < prob_of_original,
+              lambda: (ds1_class, original_data),
+              filtered_iterator.get_next))
+  return combined_ds
+
+
+def _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds, class_values_ds,
+               seed):
+  """Filters a dataset based on per-class acceptance probabilities.
+
+  Args:
+    dataset: The dataset to be filtered.
+    acceptance_dist_ds: A dataset of acceptance probabilities.
+    initial_dist_ds: A dataset of the initial probability distribution, given or
+        estimated.
+    class_values_ds: A dataset of the corresponding classes.
+    seed: (Optional.) Python integer seed for the resampler.
+
+  Returns:
+    A dataset of (class value, data) after filtering.
+  """
+  def maybe_warn_on_large_rejection(accept_dist, initial_dist):
+    proportion_rejected = math_ops.reduce_sum((1 - accept_dist) * initial_dist)
+    return control_flow_ops.cond(
+        math_ops.less(proportion_rejected, .5),
+        lambda: accept_dist,
+        lambda: logging_ops.Print(  # pylint: disable=g-long-lambda
+            accept_dist, [proportion_rejected, initial_dist, accept_dist],
+            message="Proportion of examples rejected by sampler is high: ",
+            summarize=100,
+            first_n=10))
+
+  acceptance_dist_ds = (dataset_ops.Dataset.zip((acceptance_dist_ds,
+                                                 initial_dist_ds))
+                        .map(maybe_warn_on_large_rejection))
+
+  def _gather_and_copy(class_val, acceptance_prob, data):
+    return class_val, array_ops.gather(acceptance_prob, class_val), data
+
+  current_probabilities_and_class_and_data_ds = dataset_ops.Dataset.zip(
+      (class_values_ds, acceptance_dist_ds, dataset)).map(_gather_and_copy)
+  filtered_ds = (
+      current_probabilities_and_class_and_data_ds
+      .filter(lambda _1, p, _2: random_ops.random_uniform([], seed=seed) < p))
+  return filtered_ds.map(lambda class_value, _, data: (class_value, data))
+
+
+def _estimate_initial_dist_ds(
+    target_dist_t, class_values_ds, dist_estimation_batch_size=32,
+    smoothing_constant=10):
+  num_classes = (target_dist_t.shape[0].value or
+                 array_ops.shape(target_dist_t)[0])
+  initial_examples_per_class_seen = array_ops.fill(
+      [num_classes], np.int64(smoothing_constant))
+
+  def update_estimate_and_tile(num_examples_per_class_seen, c):
+    updated_examples_per_class_seen, dist = _estimate_data_distribution(
+        c, num_examples_per_class_seen)
+    tiled_dist = array_ops.tile(
+        array_ops.expand_dims(dist, 0), [dist_estimation_batch_size, 1])
+    return updated_examples_per_class_seen, tiled_dist
+
+  initial_dist_ds = (class_values_ds.batch(dist_estimation_batch_size)
+                     .apply(scan_ops.scan(initial_examples_per_class_seen,
+                                          update_estimate_and_tile))
+                     .apply(batching.unbatch()))
+
+  return initial_dist_ds
+
+
+def _get_target_to_initial_ratio(initial_probs, target_probs):
+  # Add tiny to initial_probs to avoid divide by zero.
+  denom = (initial_probs + np.finfo(initial_probs.dtype.as_numpy_dtype).tiny)
+  return target_probs / denom
+
+
 def _calculate_acceptance_probs(initial_probs, target_probs):
   """Calculate the per-class acceptance rates.
 
@@ -152,13 +280,10 @@ def _calculate_acceptance_probs(initial_probs, target_probs):
   0 <= t_i <= 1, sum_i(t_i) = 1
   ```
 
-
   A solution for a_i in terms of the other variables is the following:
     ```a_i = (t_i / p_i) / max_i[t_i / p_i]```
   """
-  # Add tiny to initial_probs to avoid divide by zero.
-  denom = (initial_probs + np.finfo(initial_probs.dtype.as_numpy_dtype).tiny)
-  ratio_l = target_probs / denom
+  ratio_l = _get_target_to_initial_ratio(initial_probs, target_probs)
 
   # Calculate list of acceptance probabilities.
   max_ratio = math_ops.reduce_max(ratio_l)
@@ -188,3 +313,49 @@ def _estimate_data_distribution(c, num_examples_per_class_seen):
       math_ops.reduce_sum(num_examples_per_class_seen))
   dist = math_ops.cast(init_prob_estimate, dtypes.float32)
   return num_examples_per_class_seen, dist
+
+
+def _calculate_acceptance_probs_with_mixing(initial_probs, target_probs):
+  """Calculates the acceptance probabilities and mixing ratio.
+
+  In this case, we assume that we can *either* sample from the original data
+  distribution with probability `m`, or sample from a reshaped distribution
+  that comes from rejection sampling on the original distribution. This
+  rejection sampling is done on a per-class basis, with `a_i` representing the
+  probability of accepting data from class `i`.
+
+  If we try to minimize the amount of data rejected, we get the following:
+
+  M_max = max_i [ t_i / p_i ]
+  M_min = min_i [ t_i / p_i ]
+
+  The desired probability of accepting data if it comes from class `i`:
+
+  a_i = (t_i/p_i - m) / (M_max - m)
+
+  The desired probability of pulling a data element from the original dataset,
+  rather than the filtered one:
+
+  m = M_min
+
+  See the docstring for `_calculate_acceptance_probs` for more details.
+
+  Args:
+    initial_probs: A Tensor of the initial probability distribution, given or
+      estimated.
+    target_probs: A Tensor of the corresponding classes.
+
+  Returns:
+    (A 1D Tensor with the per-class acceptance probabilities, the desired
+    probability of pull from the original distribution.)
+  """
+  ratio_l = _get_target_to_initial_ratio(initial_probs, target_probs)
+  max_ratio = math_ops.reduce_max(ratio_l)
+  min_ratio = math_ops.reduce_min(ratio_l)
+
+  # Target prob to sample from original distribution.
+  m = min_ratio
+
+  # TODO(joelshor): Simplify fraction, if possible.
+  a_i = (ratio_l - m) / (max_ratio - m)
+  return a_i, m
-- 
GitLab


From b1067116c6a2351f4c597a9391b21ad0f513565b Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Fri, 20 Apr 2018 14:27:30 +0300
Subject: [PATCH 003/395] [tf.data] Clean up resampler and update BUILD files.

---
 .../contrib/data/python/kernel_tests/BUILD    |  6 ++-
 .../data/python/kernel_tests/resample_test.py | 32 +++++----------
 tensorflow/contrib/data/python/ops/BUILD      |  2 +
 .../contrib/data/python/ops/resampling.py     | 40 ++++---------------
 4 files changed, 23 insertions(+), 57 deletions(-)

diff --git a/tensorflow/contrib/data/python/kernel_tests/BUILD b/tensorflow/contrib/data/python/kernel_tests/BUILD
index b15b9663f4..a6b46b37e7 100644
--- a/tensorflow/contrib/data/python/kernel_tests/BUILD
+++ b/tensorflow/contrib/data/python/kernel_tests/BUILD
@@ -308,13 +308,17 @@ py_test(
     srcs_version = "PY2AND3",
     tags = ["noasan"],
     deps = [
+        "//third_party/py/absl/testing:parameterized",
+        "//third_party/py/numpy",
         "//tensorflow/contrib/data/python/ops:resampling",
         "//tensorflow/python:client_testlib",
+        "//tensorflow/python:dtypes",
         "//tensorflow/python:errors",
+        "//tensorflow/python:math_ops",
+        "//tensorflow/python:random_ops",
         "//tensorflow/python:string_ops",
         "//tensorflow/python:util",
         "//tensorflow/python/data/ops:dataset_ops",
-        "//third_party/py/numpy",
     ],
 )
 
diff --git a/tensorflow/contrib/data/python/kernel_tests/resample_test.py b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
index 9e1273eba1..97c4b68cb6 100644
--- a/tensorflow/contrib/data/python/kernel_tests/resample_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
@@ -60,10 +60,10 @@ def _time_resampling(
 class ResampleTest(test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(
-      ('InitialnDistributionKnown', True, False),
-      ('InitialDistributionUnknown', False, False),
-      ('InitialDistributionKnownV2', True, True),
-      ('InitialDistributionUnknownV2', False, True))
+      ("InitialnDistributionKnown", True, False),
+      ("InitialDistributionUnknown", False, False),
+      ("InitialDistributionKnownV2", True, True),
+      ("InitialDistributionUnknownV2", False, True))
   def testDistribution(self, initial_known, use_v2):
     classes = np.random.randint(5, size=(20000,))  # Uniformly sampled
     target_dist = [0.9, 0.05, 0.05, 0.0, 0.0]
@@ -95,7 +95,6 @@ class ResampleTest(test.TestCase, parameterized.TestCase):
     returned_dist = class_counts / total_returned
     self.assertAllClose(target_dist, returned_dist, atol=1e-2)
 
-
   def testRandomClasses(self):
     init_dist = [0.25, 0.25, 0.25, 0.25]
     target_dist = [0.0, 0.0, 0.0, 1.0]
@@ -135,11 +134,11 @@ class ResampleTest(test.TestCase, parameterized.TestCase):
     self.assertAllClose(target_dist, bincount, atol=1e-2)
 
   @parameterized.named_parameters(
-        ('InitialnDistributionKnown', True, False),
-        ('InitialDistributionUnknown', False, False),
-        ('InitialDistributionKnownV2', True, True),
-        ('InitialDistributionUnknownV2', False, True))
-  def _testNewResampleIsFaster(self, target_dist, num_to_sample):
+      ("SmallSkewManySamples", [0.1, 0.1, 0.1, 0.7], 1000),
+      ("BigSkewManySamples", [0.01, 0.01, 0.01, 0.97], 1000),
+      ("SmallSkewFewSamples", [0.1, 0.1, 0.1, 0.7], 100),
+      ("BigSkewFewSamples", [0.01, 0.01, 0.01, 0.97], 100))
+  def testNewResampleIsFaster(self, target_dist, num_to_sample):
     init_dist = [0.25, 0.25, 0.25, 0.25]
     num_classes = len(init_dist)
     num_samples = 1000
@@ -153,19 +152,6 @@ class ResampleTest(test.TestCase, parameterized.TestCase):
     self.assertLess(fast_time, slow_time)
 
 
-  def testNewResampleIsFasterSmallSkewManySamples(self):
-    self._testNewResampleIsFaster([0.1, 0.1, 0.1, 0.7], 1000)
-
-  def testNewResampleIsFasterBigSkewManySamples(self):
-    self._testNewResampleIsFaster([0.01, 0.01, 0.01, 0.97], 1000)
-
-  def testNewResampleIsFasterSmallSkewFewSamples(self):
-    self._testNewResampleIsFaster([0.1, 0.1, 0.1, 0.7], 100)
-
-  def testNewResampleIsFasterBigSkewFewSamples(self):
-    self._testNewResampleIsFaster([0.01, 0.01, 0.01, 0.97], 100)
-
-
 class MapDatasetBenchmark(test.Benchmark):
 
   def benchmarkResamplePerformance(self):
diff --git a/tensorflow/contrib/data/python/ops/BUILD b/tensorflow/contrib/data/python/ops/BUILD
index e00f2304cc..8cb4fa7f14 100644
--- a/tensorflow/contrib/data/python/ops/BUILD
+++ b/tensorflow/contrib/data/python/ops/BUILD
@@ -193,7 +193,9 @@ py_library(
     srcs_version = "PY2AND3",
     deps = [
         ":batching",
+        ":interleave_ops",
         ":scan_ops",
+        "//third_party/py/numpy",
         "//tensorflow/python:array_ops",
         "//tensorflow/python:control_flow_ops",
         "//tensorflow/python:dtypes",
diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index 94e28b9a2d..16d851bf96 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -82,8 +82,12 @@ def rejection_resample_v2(class_func, target_dist, initial_dist=None,
   """A transformation that resamples a dataset to achieve a target distribution.
 
   This differs from v1 in that it will also sample from the original dataset
-  with some probability, so it makes strictly fewer data rejections. This
-  transformation is faster than the original.
+  with some probability, so it makes strictly fewer data rejections. Due to an
+  implementation detail it must initialize a separate dataset initializer, so
+  the dataset becomes stateful after this transformation is applied
+  (`make_one_shot_iterator` won't work; users must use
+  `make_initializable_iterator`). This transformation is faster than the
+  original, except for overhead.
 
   **NOTE** Resampling is performed via rejection sampling; some fraction
   of the input values will be dropped.
@@ -142,36 +146,6 @@ def rejection_resample_v2(class_func, target_dist, initial_dist=None,
   return _apply_fn
 
 
-def _random_interleave_datasets(ds1, ds1_classes, ds2, prob_of_ds1, seed=None):
-  """Randomly interleave datasets.
-
-  We carefully combine `ds1` and 'ds2' so that we don't needlessly compute the
-  filtering.
-
-  Args:
-    ds1: A dataset to interleave.
-    ds1_classes: Dataset of class values associated with ds1.
-    ds2: Another dataset to interleave.
-    prob_of_ds1: A dataset of probabilities. Each probability represents the
-      likelihood of drawing from `ds1`.
-    seed: (Optional.) Python integer seed for the resampler.
-
-  Returns:
-    A single dataset, combined from `ds1` and `ds2`.
-  """
-  num_filtered_to_prefetch = 3
-  ds2 = ds2.prefetch(num_filtered_to_prefetch)
-  filtered_iterator = ds2.make_one_shot_iterator()
-  combined_ds = dataset_ops.Dataset.zip(
-        (ds1_classes, ds1, prob_of_ds1)).map(
-        lambda ds1_class, original_data, prob_of_original:
-        control_flow_ops.cond(
-              random_ops.random_uniform([], seed=seed) < prob_of_original,
-              lambda: (ds1_class, original_data),
-              filtered_iterator.get_next))
-  return combined_ds
-
-
 def _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds, class_values_ds,
                seed):
   """Filters a dataset based on per-class acceptance probabilities.
@@ -358,4 +332,4 @@ def _calculate_acceptance_probs_with_mixing(initial_probs, target_probs):
 
   # TODO(joelshor): Simplify fraction, if possible.
   a_i = (ratio_l - m) / (max_ratio - m)
-  return a_i, m
+  return a_i, m
\ No newline at end of file
-- 
GitLab


From 0cba8b7c66bead25ed2e6e1c6bf5a23d6cbe9557 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Fri, 20 Apr 2018 14:44:47 +0300
Subject: [PATCH 004/395] [tf.data] Fix `absl` build rule.

---
 tensorflow/contrib/data/python/kernel_tests/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/data/python/kernel_tests/BUILD b/tensorflow/contrib/data/python/kernel_tests/BUILD
index a6b46b37e7..f90b17e79e 100644
--- a/tensorflow/contrib/data/python/kernel_tests/BUILD
+++ b/tensorflow/contrib/data/python/kernel_tests/BUILD
@@ -308,7 +308,6 @@ py_test(
     srcs_version = "PY2AND3",
     tags = ["noasan"],
     deps = [
-        "//third_party/py/absl/testing:parameterized",
         "//third_party/py/numpy",
         "//tensorflow/contrib/data/python/ops:resampling",
         "//tensorflow/python:client_testlib",
@@ -319,6 +318,7 @@ py_test(
         "//tensorflow/python:string_ops",
         "//tensorflow/python:util",
         "//tensorflow/python/data/ops:dataset_ops",
+        "@absl_py//absl/testing:parameterized",
     ],
 )
 
-- 
GitLab


From 8cc506f8f6c3e9071069ede1cd5c91a9f3da7c11 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Fri, 20 Apr 2018 15:00:02 +0300
Subject: [PATCH 005/395] [tf.data] Reorder BUILD rule deps and add `xrange`
 from `six`.

---
 tensorflow/contrib/data/python/kernel_tests/BUILD            | 2 +-
 tensorflow/contrib/data/python/kernel_tests/resample_test.py | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/contrib/data/python/kernel_tests/BUILD b/tensorflow/contrib/data/python/kernel_tests/BUILD
index f90b17e79e..92c6967933 100644
--- a/tensorflow/contrib/data/python/kernel_tests/BUILD
+++ b/tensorflow/contrib/data/python/kernel_tests/BUILD
@@ -308,7 +308,6 @@ py_test(
     srcs_version = "PY2AND3",
     tags = ["noasan"],
     deps = [
-        "//third_party/py/numpy",
         "//tensorflow/contrib/data/python/ops:resampling",
         "//tensorflow/python:client_testlib",
         "//tensorflow/python:dtypes",
@@ -318,6 +317,7 @@ py_test(
         "//tensorflow/python:string_ops",
         "//tensorflow/python:util",
         "//tensorflow/python/data/ops:dataset_ops",
+        "//third_party/py/numpy",
         "@absl_py//absl/testing:parameterized",
     ],
 )
diff --git a/tensorflow/contrib/data/python/kernel_tests/resample_test.py b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
index 97c4b68cb6..7f007fede8 100644
--- a/tensorflow/contrib/data/python/kernel_tests/resample_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
@@ -18,6 +18,7 @@ from __future__ import division
 from __future__ import print_function
 
 import numpy as np
+from six.moves import xrange  # pylint: disable=redefined-builtin
 import time
 from absl.testing import parameterized
 
-- 
GitLab


From a10708db0d587831cafcb2e7dbdcbbcf11aede95 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Fri, 20 Apr 2018 15:09:50 +0300
Subject: [PATCH 006/395] [tf.data] Second reorder BUILD rule deps.

---
 tensorflow/contrib/data/python/ops/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/data/python/ops/BUILD b/tensorflow/contrib/data/python/ops/BUILD
index 8cb4fa7f14..d9a5502508 100644
--- a/tensorflow/contrib/data/python/ops/BUILD
+++ b/tensorflow/contrib/data/python/ops/BUILD
@@ -195,7 +195,6 @@ py_library(
         ":batching",
         ":interleave_ops",
         ":scan_ops",
-        "//third_party/py/numpy",
         "//tensorflow/python:array_ops",
         "//tensorflow/python:control_flow_ops",
         "//tensorflow/python:dtypes",
@@ -204,6 +203,7 @@ py_library(
         "//tensorflow/python:math_ops",
         "//tensorflow/python:random_ops",
         "//tensorflow/python/data/ops:dataset_ops",
+        "//third_party/py/numpy",
     ],
 )
 
-- 
GitLab


From fc6510b506731bf2ffc2520e30fba73b79e5b687 Mon Sep 17 00:00:00 2001
From: Chris Ying <chrisying@google.com>
Date: Tue, 17 Apr 2018 15:28:12 -0700
Subject: [PATCH 007/395] Fix CheckpointSaverHook to properly save every
 save_checkpoints_steps for TPU workloads.

PiperOrigin-RevId: 193266515
(cherry picked from commit 5aba07dce5b9e924183efcd05cd82f2fbb70edc8)
---
 .../contrib/tpu/python/tpu/tpu_estimator.py   |  9 ++
 .../training/basic_session_run_hooks.py       | 10 +-
 .../training/basic_session_run_hooks_test.py  | 93 +++++++++++++++++++
 ...sorflow.train.-checkpoint-saver-hook.pbtxt |  2 +-
 4 files changed, 111 insertions(+), 3 deletions(-)

diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
index 1332108d04..c8c4cc6c68 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
@@ -2054,6 +2054,14 @@ class TPUEstimator(estimator_lib.Estimator):
                   },
                   every_n_secs=30)
           ] + input_hooks
+          chief_hooks = [
+              training.CheckpointSaverHook(
+                  self.model_dir,
+                  save_secs=self._config.save_checkpoints_secs,
+                  save_steps=self._config.save_checkpoints_steps,
+                  steps_per_run=self._config.tpu_config.iterations_per_loop,
+                  scaffold=scaffold)
+          ]
           summary.scalar(model_fn_lib.LOSS_METRIC_KEY, loss)
           with ops.control_dependencies([loss]):
             update_ops = _sync_variables_ops()
@@ -2067,6 +2075,7 @@ class TPUEstimator(estimator_lib.Estimator):
           return model_fn_lib.EstimatorSpec(
               mode,
               loss=loss,
+              training_chief_hooks=chief_hooks,
               training_hooks=hooks,
               train_op=train_op,
               scaffold=scaffold)
diff --git a/tensorflow/python/training/basic_session_run_hooks.py b/tensorflow/python/training/basic_session_run_hooks.py
index 094a9e886b..3651291bdf 100644
--- a/tensorflow/python/training/basic_session_run_hooks.py
+++ b/tensorflow/python/training/basic_session_run_hooks.py
@@ -391,7 +391,8 @@ class CheckpointSaverHook(session_run_hook.SessionRunHook):
                saver=None,
                checkpoint_basename="model.ckpt",
                scaffold=None,
-               listeners=None):
+               listeners=None,
+               steps_per_run=1):
     """Initializes a `CheckpointSaverHook`.
 
     Args:
@@ -404,6 +405,9 @@ class CheckpointSaverHook(session_run_hook.SessionRunHook):
       listeners: List of `CheckpointSaverListener` subclass instances.
         Used for callbacks that run immediately before or after this hook saves
         the checkpoint.
+      steps_per_run: `int`, number of steps that occur between each invocation
+        of the hook. Primarily used for TPU workloads which run multiple steps
+        in a while loop in a single Session.run.
 
     Raises:
       ValueError: One of `save_steps` or `save_secs` should be set.
@@ -419,6 +423,7 @@ class CheckpointSaverHook(session_run_hook.SessionRunHook):
     self._timer = SecondOrStepTimer(every_secs=save_secs,
                                     every_steps=save_steps)
     self._listeners = listeners or []
+    self._steps_per_run = steps_per_run
 
   def begin(self):
     self._summary_writer = SummaryWriterCache.get(self._checkpoint_dir)
@@ -450,7 +455,8 @@ class CheckpointSaverHook(session_run_hook.SessionRunHook):
 
   def after_run(self, run_context, run_values):
     stale_global_step = run_values.results
-    if self._timer.should_trigger_for_step(stale_global_step+1):
+    if self._timer.should_trigger_for_step(
+        stale_global_step + self._steps_per_run):
       # get the real value after train op.
       global_step = run_context.session.run(self._global_step_tensor)
       if self._timer.should_trigger_for_step(global_step):
diff --git a/tensorflow/python/training/basic_session_run_hooks_test.py b/tensorflow/python/training/basic_session_run_hooks_test.py
index f39a5261a9..25962f6bf7 100644
--- a/tensorflow/python/training/basic_session_run_hooks_test.py
+++ b/tensorflow/python/training/basic_session_run_hooks_test.py
@@ -719,6 +719,99 @@ class CheckpointSaverHookTest(test.TestCase):
     fake_summary_writer.FakeSummaryWriter.uninstall()
 
 
+class CheckpointSaverHookMultiStepTest(test.TestCase):
+
+  def setUp(self):
+    self.model_dir = tempfile.mkdtemp()
+    self.graph = ops.Graph()
+    self.steps_per_run = 5
+    with self.graph.as_default():
+      self.scaffold = monitored_session.Scaffold()
+      self.global_step = variables.get_or_create_global_step()
+      self.train_op = training_util._increment_global_step(self.steps_per_run)
+
+  def tearDown(self):
+    shutil.rmtree(self.model_dir, ignore_errors=True)
+
+  def test_save_steps_saves_in_first_step(self):
+    with self.graph.as_default():
+      hook = basic_session_run_hooks.CheckpointSaverHook(
+          self.model_dir,
+          save_steps=2*self.steps_per_run,
+          scaffold=self.scaffold,
+          steps_per_run=self.steps_per_run)
+      hook.begin()
+      self.scaffold.finalize()
+      with session_lib.Session() as sess:
+        sess.run(self.scaffold.init_op)
+        mon_sess = monitored_session._HookedSession(sess, [hook])
+        mon_sess.run(self.train_op)
+        self.assertEqual(5,
+                         checkpoint_utils.load_variable(self.model_dir,
+                                                        self.global_step.name))
+
+  def test_save_steps_saves_periodically(self):
+    with self.graph.as_default():
+      hook = basic_session_run_hooks.CheckpointSaverHook(
+          self.model_dir,
+          save_steps=2*self.steps_per_run,
+          scaffold=self.scaffold,
+          steps_per_run=self.steps_per_run)
+      hook.begin()
+      self.scaffold.finalize()
+      with session_lib.Session() as sess:
+        sess.run(self.scaffold.init_op)
+        mon_sess = monitored_session._HookedSession(sess, [hook])
+        mon_sess.run(self.train_op)
+        # Saved (step=5)
+        self.assertEqual(5,
+                         checkpoint_utils.load_variable(self.model_dir,
+                                                        self.global_step.name))
+
+        mon_sess.run(self.train_op)
+        # Not saved (step=10)
+        self.assertEqual(5,
+                         checkpoint_utils.load_variable(self.model_dir,
+                                                        self.global_step.name))
+
+        mon_sess.run(self.train_op)
+        # Saved (step=15)
+        self.assertEqual(15,
+                         checkpoint_utils.load_variable(self.model_dir,
+                                                        self.global_step.name))
+
+        mon_sess.run(self.train_op)
+        # Not saved (step=20)
+        self.assertEqual(15,
+                         checkpoint_utils.load_variable(self.model_dir,
+                                                        self.global_step.name))
+
+        mon_sess.run(self.train_op)
+        # Saved (step=25)
+        self.assertEqual(25,
+                         checkpoint_utils.load_variable(self.model_dir,
+                                                        self.global_step.name))
+
+  def test_save_steps_saves_at_end(self):
+    with self.graph.as_default():
+      hook = basic_session_run_hooks.CheckpointSaverHook(
+          self.model_dir,
+          save_steps=2*self.steps_per_run,
+          scaffold=self.scaffold,
+          steps_per_run=self.steps_per_run)
+      hook.begin()
+      self.scaffold.finalize()
+      with session_lib.Session() as sess:
+        sess.run(self.scaffold.init_op)
+        mon_sess = monitored_session._HookedSession(sess, [hook])
+        mon_sess.run(self.train_op)
+        mon_sess.run(self.train_op)
+        hook.end(sess)
+        self.assertEqual(10,
+                         checkpoint_utils.load_variable(self.model_dir,
+                                                        self.global_step.name))
+
+
 class ResourceCheckpointSaverHookTest(test.TestCase):
 
   def setUp(self):
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-checkpoint-saver-hook.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-checkpoint-saver-hook.pbtxt
index c3037baa8c..327799729c 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-checkpoint-saver-hook.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-checkpoint-saver-hook.pbtxt
@@ -5,7 +5,7 @@ tf_class {
   is_instance: "<type \'object\'>"
   member_method {
     name: "__init__"
-    argspec: "args=[\'self\', \'checkpoint_dir\', \'save_secs\', \'save_steps\', \'saver\', \'checkpoint_basename\', \'scaffold\', \'listeners\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'model.ckpt\', \'None\', \'None\'], "
+    argspec: "args=[\'self\', \'checkpoint_dir\', \'save_secs\', \'save_steps\', \'saver\', \'checkpoint_basename\', \'scaffold\', \'listeners\', \'steps_per_run\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'model.ckpt\', \'None\', \'None\', \'1\'], "
   }
   member_method {
     name: "after_create_session"
-- 
GitLab


From e1cc34d34b3a811da7c7a2d7cc6c60398c50fdfb Mon Sep 17 00:00:00 2001
From: Chris Ying <chrisying@google.com>
Date: Tue, 17 Apr 2018 20:31:30 -0700
Subject: [PATCH 008/395] Disable CheckpointSaverHook when both
 save_checkpoints_secs and save_checkpoints_steps are None

PiperOrigin-RevId: 193299688
(cherry picked from commit 41e2cd187b31e9e6d88bc042e21e73f7be0ed729)
---
 .../contrib/tpu/python/tpu/tpu_estimator.py    | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
index c8c4cc6c68..8df631b475 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
@@ -2054,14 +2054,16 @@ class TPUEstimator(estimator_lib.Estimator):
                   },
                   every_n_secs=30)
           ] + input_hooks
-          chief_hooks = [
-              training.CheckpointSaverHook(
-                  self.model_dir,
-                  save_secs=self._config.save_checkpoints_secs,
-                  save_steps=self._config.save_checkpoints_steps,
-                  steps_per_run=self._config.tpu_config.iterations_per_loop,
-                  scaffold=scaffold)
-          ]
+          chief_hooks = []
+          if (self._config.save_checkpoints_secs or
+              self._config.save_checkpoints_steps):
+            chief_hooks.append(
+                training.CheckpointSaverHook(
+                    self.model_dir,
+                    save_secs=self._config.save_checkpoints_secs,
+                    save_steps=self._config.save_checkpoints_steps,
+                    steps_per_run=self._config.tpu_config.iterations_per_loop,
+                    scaffold=scaffold))
           summary.scalar(model_fn_lib.LOSS_METRIC_KEY, loss)
           with ops.control_dependencies([loss]):
             update_ops = _sync_variables_ops()
-- 
GitLab


From a722cdf7a62a3ee82ca6ee1b3d33f3d03dba49ee Mon Sep 17 00:00:00 2001
From: Francois Chollet <fchollet@google.com>
Date: Wed, 18 Apr 2018 15:04:21 -0700
Subject: [PATCH 009/395] Fix loss computation bug in Model training/eval
 methods with eager execution enabled. Fixes #18642.

PiperOrigin-RevId: 193423288
---
 .../_impl/keras/engine/training_eager.py      |  2 +-
 .../_impl/keras/engine/training_eager_test.py | 25 +++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/keras/_impl/keras/engine/training_eager.py b/tensorflow/python/keras/_impl/keras/engine/training_eager.py
index 4cdb5f108a..695669d9ee 100644
--- a/tensorflow/python/keras/_impl/keras/engine/training_eager.py
+++ b/tensorflow/python/keras/_impl/keras/engine/training_eager.py
@@ -150,7 +150,7 @@ def _model_loss(model, inputs, targets, sample_weights=None, training=False):
       weighted_masked_fn = training_utils.weighted_masked_objective(loss_fn)
       with backend.name_scope(model.output_names[i] + '_loss'):
         output_loss = weighted_masked_fn(
-            outs[i], targets[i], weights, mask=mask)
+            targets[i], outs[i], weights, mask=mask)
       loss_metrics.append(backend.mean(output_loss))
 
       loss_weight = model.loss_weights_list[i]
diff --git a/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py b/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py
index 6cdb6b0753..ed0f91ee1e 100644
--- a/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py
+++ b/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py
@@ -21,6 +21,7 @@ from __future__ import print_function
 import numpy as np
 
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util as tf_test_util
 from tensorflow.python.keras._impl import keras
 from tensorflow.python.keras._impl.keras import testing_utils
 from tensorflow.python.platform import test
@@ -625,6 +626,30 @@ class LossWeightingTest(test.TestCase):
       model.fit(x_np, [y_np, y_np], epochs=1, sample_weight={'1': bad_w_np})
 
 
+class CorrectnessTest(test.TestCase):
+
+  @tf_test_util.run_in_graph_and_eager_modes()
+  def test_loss_correctness(self):
+    # Test that training loss is the same in eager and graph
+    # (by comparing it to a reference value in a deterministic case)
+    model = keras.Sequential()
+    model.add(keras.layers.Dense(3,
+                                 activation='relu',
+                                 input_dim=4,
+                                 kernel_initializer='ones'))
+    model.add(keras.layers.Dense(2,
+                                 activation='softmax',
+                                 kernel_initializer='ones'))
+    model.compile(loss='sparse_categorical_crossentropy',
+                  optimizer=RMSPropOptimizer(learning_rate=0.001))
+    x = np.ones((100, 4))
+    np.random.seed(123)
+    y = np.random.randint(0, 1, size=(100, 1))
+    history = model.fit(x, y, epochs=1, batch_size=10)
+    self.assertEqual(
+        np.around(history.history['loss'][-1], decimals=4), 0.6173)
+
+
 if __name__ == '__main__':
   ops.enable_eager_execution()
   test.main()
-- 
GitLab


From 955c1edb2f92871597aaf74f5684da4d22843064 Mon Sep 17 00:00:00 2001
From: zhangyaobit <yaozhang@google.com>
Date: Mon, 23 Apr 2018 13:46:26 -0700
Subject: [PATCH 010/395] Update layout_optimizer.cc

Place data format op on CPU:0.
---
 tensorflow/core/grappler/optimizers/layout_optimizer.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/core/grappler/optimizers/layout_optimizer.cc b/tensorflow/core/grappler/optimizers/layout_optimizer.cc
index 561226f945..8fb30d116d 100644
--- a/tensorflow/core/grappler/optimizers/layout_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/layout_optimizer.cc
@@ -919,6 +919,7 @@ class NodeProcessor : public GraphProcessor {
         ParseNodeName(input_name, &port);
         if (IsHostMemory(*input, port)) {
           parsed_name.type = "CPU";
+          parsed_name.id = 0;
           device = DeviceNameUtils::ParsedNameToString(parsed_name);
         }
       }
-- 
GitLab


From aaf1e32d53e1b473e9d1700afba71662e28150ff Mon Sep 17 00:00:00 2001
From: zhangyaobit <yaozhang@google.com>
Date: Mon, 23 Apr 2018 13:49:22 -0700
Subject: [PATCH 011/395] Update layout_optimizer_test.cc

Place data format op on CPU:0.
---
 tensorflow/core/grappler/optimizers/layout_optimizer_test.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/grappler/optimizers/layout_optimizer_test.cc b/tensorflow/core/grappler/optimizers/layout_optimizer_test.cc
index 260347b0e8..b913f2b004 100644
--- a/tensorflow/core/grappler/optimizers/layout_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/layout_optimizer_test.cc
@@ -36,7 +36,7 @@ class LayoutOptimizerTest : public ::testing::Test {
     DeviceProperties device_properties;
     device_properties.set_type("GPU");
     device_properties.mutable_environment()->insert({"architecture", "6"});
-    virtual_cluster_.reset(new VirtualCluster({{"/GPU:0", device_properties}}));
+    virtual_cluster_.reset(new VirtualCluster({{"/GPU:1", device_properties}}));
   }
 
   Output SimpleConv2D(tensorflow::Scope* s, int input_size, int filter_size,
-- 
GitLab


From 9ad432781fce95a397d7d4a8ce506932160b83f1 Mon Sep 17 00:00:00 2001
From: Amit Patankar <amitpatankar@google.com>
Date: Mon, 23 Apr 2018 14:00:28 -0700
Subject: [PATCH 012/395] Update install_linux.md

---
 tensorflow/docs_src/install/install_linux.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/docs_src/install/install_linux.md b/tensorflow/docs_src/install/install_linux.md
index f19f827e25..63b8eb30e9 100644
--- a/tensorflow/docs_src/install/install_linux.md
+++ b/tensorflow/docs_src/install/install_linux.md
@@ -48,7 +48,7 @@ must be installed on your system:
     Toolkit.
   * The libcupti-dev library, which is the NVIDIA CUDA Profile Tools Interface.
     This library provides advanced profiling support. To install this library,
-    issue the following command for CUDA Toolkit >= 8.0:
+    issue the following command for CUDA Toolkit >= 9.0:
 
     <pre>
     $ <b>sudo apt-get install cuda-command-line-tools</b>
-- 
GitLab


From dd9ee4a2f13c2219ebd7c6f8754b8dd32188e2a5 Mon Sep 17 00:00:00 2001
From: Amit Patankar <amitpatankar@google.com>
Date: Tue, 24 Apr 2018 10:59:10 -0700
Subject: [PATCH 013/395] Update README.md

---
 tensorflow/tools/docker/README.md | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensorflow/tools/docker/README.md b/tensorflow/tools/docker/README.md
index f46c56e11a..525f2995ce 100644
--- a/tensorflow/tools/docker/README.md
+++ b/tensorflow/tools/docker/README.md
@@ -16,12 +16,12 @@ quick links here:
 
 We currently maintain two Docker container images:
 
-* `gcr.io/tensorflow/tensorflow` - TensorFlow with all dependencies - CPU only!
+* `tensorflow/tensorflow` - TensorFlow with all dependencies - CPU only!
 
-* `gcr.io/tensorflow/tensorflow:latest-gpu` - TensorFlow with all dependencies
+* `tensorflow/tensorflow:latest-gpu` - TensorFlow with all dependencies
   and support for NVidia CUDA
 
-Note: We also publish the same containers into
+Note: We store all our containers on 
 [Docker Hub](https://hub.docker.com/r/tensorflow/tensorflow/tags/).
 
 
@@ -29,12 +29,12 @@ Note: We also publish the same containers into
 
 Run non-GPU container using
 
-    $ docker run -it -p 8888:8888 gcr.io/tensorflow/tensorflow
+    $ docker run -it -p 8888:8888 tensorflow/tensorflow
 
 For GPU support install NVidia drivers (ideally latest) and
 [nvidia-docker](https://github.com/NVIDIA/nvidia-docker). Run using
 
-    $ nvidia-docker run -it -p 8888:8888 gcr.io/tensorflow/tensorflow:latest-gpu
+    $ nvidia-docker run -it -p 8888:8888 tensorflow/tensorflow:latest-gpu
 
 
 Note: If you would have a problem running nvidia-docker you may try the old method
@@ -44,7 +44,7 @@ it there and try using nvidia-docker as described above.
     $ # The old, not recommended way to run docker with gpu support:
     $ export CUDA_SO=$(\ls /usr/lib/x86_64-linux-gnu/libcuda.* | xargs -I{} echo '-v {}:{}')
     $ export DEVICES=$(\ls /dev/nvidia* | xargs -I{} echo '--device {}:{}')
-    $ docker run -it -p 8888:8888 $CUDA_SO $DEVICES gcr.io/tensorflow/tensorflow:latest-gpu
+    $ docker run -it -p 8888:8888 $CUDA_SO $DEVICES tensorflow/tensorflow:latest-gpu
 
 
 ## More containers
-- 
GitLab


From 61c463020618ef6441392db770bdb0ec23375c73 Mon Sep 17 00:00:00 2001
From: Nick Felt <nfelt@users.noreply.github.com>
Date: Tue, 24 Apr 2018 14:51:20 -0700
Subject: [PATCH 014/395] Update tensorboard dep to 1.8.x

---
 tensorflow/tools/pip_package/setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/tools/pip_package/setup.py b/tensorflow/tools/pip_package/setup.py
index 6da3223d33..bcf6c1e515 100644
--- a/tensorflow/tools/pip_package/setup.py
+++ b/tensorflow/tools/pip_package/setup.py
@@ -38,7 +38,7 @@ REQUIRED_PACKAGES = [
     'numpy >= 1.13.3',
     'six >= 1.10.0',
     'protobuf >= 3.4.0',
-    'tensorboard >= 1.7.0, < 1.8.0',
+    'tensorboard >= 1.8.0, < 1.9.0',
     'termcolor >= 1.1.0',
 ]
 
-- 
GitLab


From a8654769c1faf6327b715edae614eb48775394a1 Mon Sep 17 00:00:00 2001
From: anj-s <32556631+anj-s@users.noreply.github.com>
Date: Tue, 24 Apr 2018 16:28:41 -0700
Subject: [PATCH 015/395] 1.8r Cherrypick request-cherrypicks_30740: Fix for
 dropped metrics in evaluate function for Keras models. (#18799)

---
 .../keras/_impl/keras/engine/training.py      | 29 ++-------
 .../_impl/keras/engine/training_eager.py      | 39 ++++--------
 .../_impl/keras/engine/training_eager_test.py | 11 ++--
 .../keras/_impl/keras/engine/training_test.py | 26 ++++++++
 .../_impl/keras/engine/training_utils.py      | 62 +++++++++++++++++++
 5 files changed, 109 insertions(+), 58 deletions(-)

diff --git a/tensorflow/python/keras/_impl/keras/engine/training.py b/tensorflow/python/keras/_impl/keras/engine/training.py
index 71de657da8..2b72e0e33d 100644
--- a/tensorflow/python/keras/_impl/keras/engine/training.py
+++ b/tensorflow/python/keras/_impl/keras/engine/training.py
@@ -276,6 +276,8 @@ class Model(Network):
           self.metrics_names.append(self.output_names[i] + '_loss')
       self.nested_metrics = training_utils.collect_metrics(metrics,
                                                            self.output_names)
+      with K.name_scope('metrics'):
+        training_utils.populate_metric_names(self)
       self._feed_sample_weight_modes = []
       for i in range(len(self.outputs)):
         self._feed_sample_weight_modes.append(None)
@@ -462,7 +464,6 @@ class Model(Network):
         output_weighted_metrics = nested_weighted_metrics[i]
 
         def handle_metrics(metrics, weights=None):
-          metric_name_prefix = 'weighted_' if weights is not None else ''
 
           for metric in metrics:
             if metric in ('accuracy', 'acc', 'crossentropy', 'ce'):
@@ -489,39 +490,19 @@ class Model(Network):
                   metric_fn = metrics_module.categorical_accuracy
                 elif metric in ('crossentropy', 'ce'):
                   metric_fn = metrics_module.categorical_crossentropy
-              if metric in ('accuracy', 'acc'):
-                suffix = 'acc'
-              elif metric in ('crossentropy', 'ce'):
-                suffix = 'ce'
               weighted_metric_fn = training_utils.weighted_masked_objective(
                   metric_fn)
-              metric_name = metric_name_prefix + suffix
             else:
               metric_fn = metrics_module.get(metric)
               weighted_metric_fn = training_utils.weighted_masked_objective(
                   metric_fn)
-              # Get metric name as string
-              if hasattr(metric_fn, 'name'):
-                metric_name = metric_fn.name
-              else:
-                metric_name = metric_fn.__name__
-              metric_name = metric_name_prefix + metric_name
-
+            metric_name = training_utils.get_base_metric_name(
+                metric, weighted=weights is not None)
             with K.name_scope(metric_name):
               metric_result = weighted_metric_fn(
                   y_true, y_pred, weights=weights, mask=masks[i])
 
-            # Append to self.metrics_names, self.metric_tensors,
-            # self.stateful_metric_names
-            if len(self.output_names) > 1:
-              metric_name = '%s_%s' % (self.output_names[i], metric_name)
-            # Dedupe name
-            j = 1
-            base_metric_name = metric_name
-            while metric_name in self.metrics_names:
-              metric_name = '%s_%d' % (base_metric_name, j)
-              j += 1
-            self.metrics_names.append(metric_name)
+            training_utils.add_metric_name(self, metric_name, i)
             self.metrics_tensors.append(metric_result)
 
             # Keep track of state updates created by
diff --git a/tensorflow/python/keras/_impl/keras/engine/training_eager.py b/tensorflow/python/keras/_impl/keras/engine/training_eager.py
index 695669d9ee..ad239d6151 100644
--- a/tensorflow/python/keras/_impl/keras/engine/training_eager.py
+++ b/tensorflow/python/keras/_impl/keras/engine/training_eager.py
@@ -100,7 +100,7 @@ def _eager_metrics_fn(model, outputs, targets):
         metric_names.append(metric_name)
         metric_results.append(backend.mean(metric_result))
 
-  return metric_names, metric_results
+  return metric_results
 
 
 def _model_loss(model, inputs, targets, sample_weights=None, training=False):
@@ -151,7 +151,12 @@ def _model_loss(model, inputs, targets, sample_weights=None, training=False):
       with backend.name_scope(model.output_names[i] + '_loss'):
         output_loss = weighted_masked_fn(
             targets[i], outs[i], weights, mask=mask)
-      loss_metrics.append(backend.mean(output_loss))
+      # If the number of outputs is 1 then we don't append the loss metric
+      # associated with each model output. When there are multiple outputs
+      # associated with a model, each output's loss is calculated and returned
+      # as part of the loss_metrics.
+      if len(model.outputs) > 1:
+        loss_metrics.append(backend.mean(output_loss))
 
       loss_weight = model.loss_weights_list[i]
       if total_loss is None:
@@ -274,7 +279,7 @@ def train_on_batch(model, inputs, targets, sample_weights=None):
       model, inputs, targets, sample_weights=sample_weights, training=True)
   if not isinstance(outs, list):
     outs = [outs]
-  _, metrics_results = _eager_metrics_fn(
+  metrics_results = _eager_metrics_fn(
       model, outs, targets)
   if not isinstance(loss, list):
     loss = [loss]
@@ -304,7 +309,7 @@ def test_on_batch(model, inputs, targets, sample_weights=None):
       model, inputs, targets, sample_weights=sample_weights, training=False)
   if not isinstance(outs, list):
     outs = [outs]
-  _, metrics_results = _eager_metrics_fn(
+  metrics_results = _eager_metrics_fn(
       model, outs, targets)
   if not isinstance(loss, list):
     loss = [loss]
@@ -498,34 +503,12 @@ def fit_loop(
         for l, o in zip(out_labels, outs):
           batch_logs[l] = o
         # Required for Eager mode
-        metrics_names, metrics_results = _eager_metrics_fn(
-            model, outs, targets_batch)
+        metrics_results = _eager_metrics_fn(model, outs, targets_batch)
         batch_logs['loss'] = tensor_util.constant_value(backend.mean(loss))
 
-        # TODO(anjalisridhar): Move this to compile to avoid duplicate code.
-        # In graph mode we set the metric names in compile. However in
-        # Eager mode we calculate the metrics for each batch in fit_loop.
-        # We could calculate the metric names and functions in compile.
-        # This would avoid setting the callback parameters separately.
-        # We need to do this for the first iteration alone
-        for m in metrics_names:
-          if m not in callback_metrics:
-            callback_metrics.append(m)
-
-        callbacks.set_params({
-            'batch_size': batch_size,
-            'epochs': epochs,
-            'steps': steps_per_epoch,
-            'samples': num_train_samples,
-            'verbose': verbose,
-            'do_validation': do_validation,
-            'metrics': callback_metrics or [],
-        })
-
         for k, v in zip(model.metrics_names,
                         [backend.mean(loss)] + loss_metrics + metrics_results):
           batch_logs[k] = tensor_util.constant_value(v)
-
         callbacks.on_batch_end(batch_index, batch_logs)
         if callback_model.stop_training:
           break
@@ -611,7 +594,7 @@ def test_loop(model, inputs, targets,
           targets_batch,
           sample_weights=sample_weights_batch,
           training=False)
-      _, metrics_results = _eager_metrics_fn(model, loss_outs, targets_batch)
+      metrics_results = _eager_metrics_fn(model, loss_outs, targets_batch)
       batch_outs = []
       for _, v in zip(model.metrics_names,
                       [backend.mean(loss)] + loss_metrics + metrics_results):
diff --git a/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py b/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py
index ed0f91ee1e..c45e07e08b 100644
--- a/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py
+++ b/tensorflow/python/keras/_impl/keras/engine/training_eager_test.py
@@ -212,7 +212,7 @@ class TrainingTest(test.TestCase):
     optimizer = RMSPropOptimizer(learning_rate=0.001)
     loss = 'mse'
     loss_weights = [1., 0.5]
-    metrics = ['mae']
+    metrics = ['acc', 'mae']
     model.compile(
         optimizer,
         loss,
@@ -231,20 +231,20 @@ class TrainingTest(test.TestCase):
         [input_a_np, input_b_np], [output_d_np, output_e_np],
         batch_size=5,
         verbose=0)
-    self.assertEqual(len(out), 5)
+    self.assertEqual(len(out), 7)
     out = model.evaluate(
         [input_a_np, input_b_np], [output_d_np, output_e_np],
         batch_size=5,
         verbose=1)
-    self.assertEqual(len(out), 5)
+    self.assertEqual(len(out), 7)
     out = model.evaluate(
         [input_a_np, input_b_np], [output_d_np, output_e_np],
         batch_size=5,
         verbose=2)
-    self.assertEqual(len(out), 5)
+    self.assertEqual(len(out), 7)
     out = model.test_on_batch([input_a_np, input_b_np],
                               [output_d_np, output_e_np])
-    self.assertEqual(len(out), 5)
+    self.assertEqual(len(out), 7)
 
     # Test evaluate with dictionary inputs
     model.evaluate(
@@ -625,7 +625,6 @@ class LossWeightingTest(test.TestCase):
       bad_w_np = np.random.random((10, 2, 2))
       model.fit(x_np, [y_np, y_np], epochs=1, sample_weight={'1': bad_w_np})
 
-
 class CorrectnessTest(test.TestCase):
 
   @tf_test_util.run_in_graph_and_eager_modes()
diff --git a/tensorflow/python/keras/_impl/keras/engine/training_test.py b/tensorflow/python/keras/_impl/keras/engine/training_test.py
index 08fd26dd18..47d80704cf 100644
--- a/tensorflow/python/keras/_impl/keras/engine/training_test.py
+++ b/tensorflow/python/keras/_impl/keras/engine/training_test.py
@@ -23,11 +23,14 @@ import unittest
 
 import numpy as np
 
+from tensorflow.python.framework import test_util as tf_test_util
 from tensorflow.python.keras._impl import keras
 from tensorflow.python.keras._impl.keras import testing_utils
 from tensorflow.python.keras._impl.keras.engine.training_utils import weighted_masked_objective
 from tensorflow.python.keras._impl.keras.utils.generic_utils import slice_arrays
 from tensorflow.python.platform import test
+from tensorflow.python.training.rmsprop import RMSPropOptimizer
+
 
 try:
   import scipy.sparse as scipy_sparse  # pylint: disable=g-import-not-at-top
@@ -1667,6 +1670,29 @@ class TestTrainingWithDataTensors(test.TestCase):
       model.train_on_batch([input_a_np, input_b_np],
                            [output_a_np, output_b_np])
 
+  @tf_test_util.run_in_graph_and_eager_modes()
+  def test_metric_names_are_identical_in_graph_and_eager(self):
+    a = keras.layers.Input(shape=(3,), name='input_a')
+    b = keras.layers.Input(shape=(3,), name='input_b')
+
+    dense = keras.layers.Dense(4, name='dense')
+    c = dense(a)
+    d = dense(b)
+    e = keras.layers.Dropout(0.5, name='dropout')(c)
+
+    model = keras.models.Model([a, b], [d, e])
+
+    optimizer = RMSPropOptimizer(learning_rate=0.001)
+    loss = 'mse'
+    loss_weights = [1., 0.5]
+    metrics = ['mae', 'acc']
+    model.compile(optimizer, loss, metrics=metrics, loss_weights=loss_weights)
+    reference_metric_names = ['loss', 'dense_loss', 'dropout_loss',
+                              'dense_mean_absolute_error',
+                              'dense_acc',
+                              'dropout_mean_absolute_error',
+                              'dropout_acc']
+    self.assertEqual(reference_metric_names, model.metrics_names)
 
 if __name__ == '__main__':
   # Bazel sets these environment variables to very long paths.
diff --git a/tensorflow/python/keras/_impl/keras/engine/training_utils.py b/tensorflow/python/keras/_impl/keras/engine/training_utils.py
index a3fc8ef2a0..34c0738f26 100644
--- a/tensorflow/python/keras/_impl/keras/engine/training_utils.py
+++ b/tensorflow/python/keras/_impl/keras/engine/training_utils.py
@@ -26,6 +26,7 @@ from tensorflow.python.eager import context
 from tensorflow.python.framework import tensor_util
 from tensorflow.python.keras._impl.keras import backend as K
 from tensorflow.python.keras._impl.keras import losses
+from tensorflow.python.keras._impl.keras import metrics as metrics_module
 from tensorflow.python.ops import math_ops
 
 
@@ -553,3 +554,64 @@ def standardize_weights(y,
 def has_symbolic_tensors(ls):
   return (any(tensor_util.is_tensor(v) for v in ls)
           and not context.executing_eagerly())
+
+
+def populate_metric_names(model):
+  for i in range(len(model.outputs)):
+    metrics = model.nested_metrics[i]
+    for metric in metrics:
+      base_metric_name = get_base_metric_name(metric)
+      add_metric_name(model, base_metric_name, i)
+
+
+def get_base_metric_name(metric, weighted=False):
+  """Returns the metric name given the metric function.
+
+  Arguments:
+      metric: Metric function name or reference.
+      weighted: Boolean indicating if the metric for which we are adding
+          names is weighted.
+
+  Returns:
+      a metric name.
+  """
+  metric_name_prefix = 'weighted_' if weighted else ''
+  if metric in ('accuracy', 'acc', 'crossentropy', 'ce'):
+    if metric in ('accuracy', 'acc'):
+      suffix = 'acc'
+    elif metric in ('crossentropy', 'ce'):
+      suffix = 'ce'
+    metric_name = metric_name_prefix + suffix
+  else:
+    metric_fn = metrics_module.get(metric)
+    # Get metric name as string
+    if hasattr(metric_fn, 'name'):
+      metric_name = metric_fn.name
+    else:
+      metric_name = metric_fn.__name__
+    metric_name = metric_name_prefix + metric_name
+
+  return metric_name
+
+
+def add_metric_name(model, metric_name, index):
+  """Makes the metric name unique and adds it to the model's metric name list.
+
+    If there are multiple outputs for which the metrics are calculated, the
+    metric names have to be made unique by appending an integer.
+
+  Arguments:
+    model: Model to which we are adding metric names.
+    metric_name: Metric name that corresponds to the metric specified by the
+        user. For example: 'acc'
+    index: The index of the model output for which the metric name is being
+        added.
+  """
+  if len(model.output_names) > 1:
+    metric_name = '%s_%s' % (model.output_names[index], metric_name)
+  j = 1
+  base_metric_name = metric_name
+  while metric_name in model.metrics_names:
+    metric_name = '%s_%d' % (base_metric_name, j)
+    j += 1
+  model.metrics_names.append(metric_name)
-- 
GitLab


From d1d5fc27ad8d84f1468ce459ba8fab208b174c6f Mon Sep 17 00:00:00 2001
From: Francois Chollet <>
Date: Tue, 24 Apr 2018 17:00:40 -0700
Subject: [PATCH 016/395] Fix critical metrics computation bug with Model in
 Eager mode.

---
 tensorflow/python/keras/_impl/keras/engine/training_eager.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/keras/_impl/keras/engine/training_eager.py b/tensorflow/python/keras/_impl/keras/engine/training_eager.py
index 4cdb5f108a..924f74e5b6 100644
--- a/tensorflow/python/keras/_impl/keras/engine/training_eager.py
+++ b/tensorflow/python/keras/_impl/keras/engine/training_eager.py
@@ -96,7 +96,7 @@ def _eager_metrics_fn(model, outputs, targets):
           model.metrics_names.append(metric_name)
 
       with backend.name_scope(metric_name):
-        metric_result = metric_fn(outputs[i], targets[i])
+        metric_result = metric_fn(targets[i], outputs[i])
         metric_names.append(metric_name)
         metric_results.append(backend.mean(metric_result))
 
-- 
GitLab


From 7316e5af78c583d75b7e39d022a22248c9d11ab9 Mon Sep 17 00:00:00 2001
From: Jiri Simsa <jiri.simsa@gmail.com>
Date: Wed, 25 Apr 2018 10:25:57 -0700
Subject: [PATCH 017/395] Updating release notes.

---
 RELEASE.md | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/RELEASE.md b/RELEASE.md
index 2717c75740..55923a2c9b 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -6,7 +6,7 @@
 * Added Gradient Boosted Trees as pre-made Estimators: BoostedTreesClassifier, BoostedTreesRegressor.
 * Add 3rd generation pipeline config for Cloud TPUs which improves performance and usability.
 * `tf.contrib.bayesflow` is moving out to it's own repo.
-* Added `tf.contrib.{proto,rpc}` to allow generic proto parsing and RPC communication.
+* Added `tf.contrib.{proto,rpc}` to allow generic proto parsing and RPC communication<sup>[1](#rpc-issue)</sup>.
 
 ## Bug Fixes and Other Changes
 * `tf.data`:
@@ -49,13 +49,14 @@
   * Fix non-uniformity of orthogonal matrices.
   * Fix bug where multi-image Estimator eval summaries were not displayed correctly.
 
+<a name="rpc-issue"><sup>1</sup></a> The cancellation logic of the RPC op contains a concurrency error. A fix has been submitted to master and will be part of the next release.
+
 ## Thanks to our Contributors
 
 This release contains contributions from many people at Google, as well as:
 
 4d55397500, Aghasy, Alan Du, Alan Lee, Alan Yee, Alex Wiltschko, Animesh Karnewar, Ankit Gupta, Anton Matosov, Aris L, Ben Barsdell, Brent Yi, Brett Koonce, Carl Thomé, cbockman, Chikanaga Tomoyuki, Chris Tava, CéDric Deltheil, Dahan Gong, Dalmo Cirne, Daniel Erenrich, David Norman, DavidNorman, Edd Wilder-James, Fanjin Zeng, Felix Abecassis, fo40225, George Sterpu, Giovanni Terlingen, Gor Baghdasaryan, Guillaume Klein, Hanchen Li, Ilya Polenov, Jakub Kolodziejczyk, Jason Sadler, Jayaram Bobba, Jerry Liu, jinghuangintel, Jiongyan Zhang (张炯衍), Joel Shor, Jong Wook Kim, Julian Eisenschlos, Karl Lessard, Krish Ravindranath, Loo Rong Jie, Lukas Geiger, Luke Iwanski, Mahmoud Abuzaina, ManHyuk, Marvin Richter, Maximilian Mitchell, Mohammad Ashraf Bhuiyan, msofka, Mustafa Kasap, Nathan Burnham, Nathan Luehr, Naveen Marri, ngc92, nio1814, Oleg Zabluda, Ou Changkun, Panos Ipeirotis, Paul Van Eck, Peter Lee, Piotr Czapla, qjivy, Rholais Lii, Rodrigo Formigone, Russell Klopfer, ryantimjohn, Sang Han, SebastiáN RamíRez, shengfuintel, Siby Jose Plathottam, Silver Chan, Stanislaw Antol, Taehoon Lee, Tarang Chugh, Ted Chang, Thomas Bastiani, Xian Xu, Xiaoming (Jason) Cui, Yan Facai (颜发才), yaox12, Yashal Shakti Kanungo, Yong Tang, Yuan (Terry) Tang, Yuxin Wu, Ziyue(Louis) Lu
 
-
 # Release 1.7.0
 
 ## Major Features And Improvements
-- 
GitLab


From 36d2e178c6d7790dd78cece70056d429aea6b917 Mon Sep 17 00:00:00 2001
From: Yifei Feng <yifeif@google.com>
Date: Wed, 25 Apr 2018 11:08:42 -0700
Subject: [PATCH 018/395] Update version string to 1.8.0.

---
 tensorflow/core/public/version.h              |  2 +-
 tensorflow/docs_src/install/install_c.md      |  2 +-
 tensorflow/docs_src/install/install_go.md     |  2 +-
 tensorflow/docs_src/install/install_java.md   | 22 +++++++++----------
 tensorflow/docs_src/install/install_linux.md  | 22 +++++++++----------
 tensorflow/docs_src/install/install_mac.md    | 10 ++++-----
 .../docs_src/install/install_sources.md       |  4 ++--
 tensorflow/tools/pip_package/setup.py         |  2 +-
 8 files changed, 33 insertions(+), 33 deletions(-)

diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index ba69efb289..522a9d84fd 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -24,7 +24,7 @@ limitations under the License.
 
 // TF_VERSION_SUFFIX is non-empty for pre-releases (e.g. "-alpha", "-alpha.1",
 // "-beta", "-rc", "-rc.1")
-#define TF_VERSION_SUFFIX "-rc1"
+#define TF_VERSION_SUFFIX ""
 
 #define TF_STR_HELPER(x) #x
 #define TF_STR(x) TF_STR_HELPER(x)
diff --git a/tensorflow/docs_src/install/install_c.md b/tensorflow/docs_src/install/install_c.md
index 8c165aad52..1abd840ab3 100644
--- a/tensorflow/docs_src/install/install_c.md
+++ b/tensorflow/docs_src/install/install_c.md
@@ -38,7 +38,7 @@ enable TensorFlow for C:
          OS="linux" # Change to "darwin" for macOS
          TARGET_DIRECTORY="/usr/local"
          curl -L \
-           "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-${OS}-x86_64-1.8.0-rc1.tar.gz" |
+           "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-${OS}-x86_64-1.8.0.tar.gz" |
            sudo tar -C $TARGET_DIRECTORY -xz
 
      The `tar` command extracts the TensorFlow C library into the `lib`
diff --git a/tensorflow/docs_src/install/install_go.md b/tensorflow/docs_src/install/install_go.md
index 26cbcc9a9b..52a2a3f8a6 100644
--- a/tensorflow/docs_src/install/install_go.md
+++ b/tensorflow/docs_src/install/install_go.md
@@ -38,7 +38,7 @@ steps to install this library and enable TensorFlow for Go:
          TF_TYPE="cpu" # Change to "gpu" for GPU support
          TARGET_DIRECTORY='/usr/local'
          curl -L \
-           "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-$(go env GOOS)-x86_64-1.8.0-rc1.tar.gz" |
+           "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-$(go env GOOS)-x86_64-1.8.0.tar.gz" |
          sudo tar -C $TARGET_DIRECTORY -xz
 
      The `tar` command extracts the TensorFlow C library into the `lib`
diff --git a/tensorflow/docs_src/install/install_java.md b/tensorflow/docs_src/install/install_java.md
index 1b0bbdba7b..700ae01236 100644
--- a/tensorflow/docs_src/install/install_java.md
+++ b/tensorflow/docs_src/install/install_java.md
@@ -36,7 +36,7 @@ following to the project's `pom.xml` to use the TensorFlow Java APIs:
 <dependency>
   <groupId>org.tensorflow</groupId>
   <artifactId>tensorflow</artifactId>
-  <version>1.8.0-rc1</version>
+  <version>1.8.0</version>
 </dependency>
 ```
 
@@ -65,7 +65,7 @@ As an example, these steps will create a Maven project that uses TensorFlow:
                <dependency>
                  <groupId>org.tensorflow</groupId>
                  <artifactId>tensorflow</artifactId>
-                 <version>1.8.0-rc1</version>
+                 <version>1.8.0</version>
                </dependency>
              </dependencies>
          </project>
@@ -123,12 +123,12 @@ instead:
 <dependency>
   <groupId>org.tensorflow</groupId>
   <artifactId>libtensorflow</artifactId>
-  <version>1.8.0-rc1</version>
+  <version>1.8.0</version>
 </dependency>
 <dependency>
   <groupId>org.tensorflow</groupId>
   <artifactId>libtensorflow_jni_gpu</artifactId>
-  <version>1.8.0-rc1</version>
+  <version>1.8.0</version>
 </dependency>
 ```
 
@@ -147,7 +147,7 @@ refer to the simpler instructions above instead.
 Take the following steps to install TensorFlow for Java on Linux or macOS:
 
   1. Download
-     [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.8.0-rc1.jar),
+     [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.8.0.jar),
      which is the TensorFlow Java Archive (JAR).
 
   2. Decide whether you will run TensorFlow for Java on CPU(s) only or with
@@ -166,7 +166,7 @@ Take the following steps to install TensorFlow for Java on Linux or macOS:
          OS=$(uname -s | tr '[:upper:]' '[:lower:]')
          mkdir -p ./jni
          curl -L \
-           "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-${TF_TYPE}-${OS}-x86_64-1.8.0-rc1.tar.gz" |
+           "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-${TF_TYPE}-${OS}-x86_64-1.8.0.tar.gz" |
            tar -xz -C ./jni
 
 ### Install on Windows
@@ -174,10 +174,10 @@ Take the following steps to install TensorFlow for Java on Linux or macOS:
 Take the following steps to install TensorFlow for Java on Windows:
 
   1. Download
-     [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.8.0-rc1.jar),
+     [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.8.0.jar),
      which is the TensorFlow Java Archive (JAR).
   2. Download the following Java Native Interface (JNI) file appropriate for
-     [TensorFlow for Java on Windows](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-cpu-windows-x86_64-1.8.0-rc1.zip).
+     [TensorFlow for Java on Windows](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-cpu-windows-x86_64-1.8.0.zip).
   3. Extract this .zip file.
 
 
@@ -225,7 +225,7 @@ must be part of your `classpath`. For example, you can include the
 downloaded `.jar` in your `classpath` by using the `-cp` compilation flag
 as follows:
 
-<pre><b>javac -cp libtensorflow-1.8.0-rc1.jar HelloTF.java</b></pre>
+<pre><b>javac -cp libtensorflow-1.8.0.jar HelloTF.java</b></pre>
 
 
 ### Running
@@ -239,11 +239,11 @@ two files are available to the JVM:
 For example, the following command line executes the `HelloTF` program on Linux
 and macOS X:
 
-<pre><b>java -cp libtensorflow-1.8.0-rc1.jar:. -Djava.library.path=./jni HelloTF</b></pre>
+<pre><b>java -cp libtensorflow-1.8.0.jar:. -Djava.library.path=./jni HelloTF</b></pre>
 
 And the following command line executes the `HelloTF` program on Windows:
 
-<pre><b>java -cp libtensorflow-1.8.0-rc1.jar;. -Djava.library.path=jni HelloTF</b></pre>
+<pre><b>java -cp libtensorflow-1.8.0.jar;. -Djava.library.path=jni HelloTF</b></pre>
 
 If the program prints <tt>Hello from <i>version</i></tt>, you've successfully
 installed TensorFlow for Java and are ready to use the API.  If the program
diff --git a/tensorflow/docs_src/install/install_linux.md b/tensorflow/docs_src/install/install_linux.md
index 63b8eb30e9..42d218c4bc 100644
--- a/tensorflow/docs_src/install/install_linux.md
+++ b/tensorflow/docs_src/install/install_linux.md
@@ -194,7 +194,7 @@ Take the following steps to install TensorFlow with Virtualenv:
      Virtualenv environment:
 
      <pre>(tensorflow)$ <b>pip3 install --upgrade \
-     https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc1-cp34-cp34m-linux_x86_64.whl</b></pre>
+     https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0-cp34-cp34m-linux_x86_64.whl</b></pre>
 
 If you encounter installation problems, see
 [Common Installation Problems](#common_installation_problems).
@@ -299,7 +299,7 @@ take the following steps:
 
      <pre>
      $ <b>sudo pip3 install --upgrade \
-     https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc1-cp34-cp34m-linux_x86_64.whl</b>
+     https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0-cp34-cp34m-linux_x86_64.whl</b>
      </pre>
 
      If this step fails, see
@@ -485,7 +485,7 @@ Take the following steps to install TensorFlow in an Anaconda environment:
 
      <pre>
      (tensorflow)$ <b>pip install --ignore-installed --upgrade \
-     https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc1-cp34-cp34m-linux_x86_64.whl</b></pre>
+     https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0-cp34-cp34m-linux_x86_64.whl</b></pre>
 
 <a name="ValidateYourInstallation"></a>
 ## Validate your installation
@@ -659,14 +659,14 @@ This section documents the relevant values for Linux installations.
 CPU only:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc1-cp27-none-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0-cp27-none-linux_x86_64.whl
 </pre>
 
 
 GPU support:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0rc1-cp27-none-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0-cp27-none-linux_x86_64.whl
 </pre>
 
 Note that GPU support requires the NVIDIA hardware and software described in
@@ -678,14 +678,14 @@ Note that GPU support requires the NVIDIA hardware and software described in
 CPU only:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc1-cp34-cp34m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0-cp34-cp34m-linux_x86_64.whl
 </pre>
 
 
 GPU support:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0rc1-cp34-cp34m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0-cp34-cp34m-linux_x86_64.whl
 </pre>
 
 Note that GPU support requires the NVIDIA hardware and software described in
@@ -697,14 +697,14 @@ Note that GPU support requires the NVIDIA hardware and software described in
 CPU only:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc1-cp35-cp35m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0-cp35-cp35m-linux_x86_64.whl
 </pre>
 
 
 GPU support:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0rc1-cp35-cp35m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0-cp35-cp35m-linux_x86_64.whl
 </pre>
 
 
@@ -716,14 +716,14 @@ Note that GPU support requires the NVIDIA hardware and software described in
 CPU only:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc1-cp36-cp36m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0-cp36-cp36m-linux_x86_64.whl
 </pre>
 
 
 GPU support:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0rc1-cp36-cp36m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0-cp36-cp36m-linux_x86_64.whl
 </pre>
 
 
diff --git a/tensorflow/docs_src/install/install_mac.md b/tensorflow/docs_src/install/install_mac.md
index ff6c2f5e44..c79075b09d 100644
--- a/tensorflow/docs_src/install/install_mac.md
+++ b/tensorflow/docs_src/install/install_mac.md
@@ -119,7 +119,7 @@ Take the following steps to install TensorFlow with Virtualenv:
      TensorFlow in the active Virtualenv is as follows:
 
      <pre> $ <b>pip3 install --upgrade \
-     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc1-py3-none-any.whl</b></pre>
+     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0-py3-none-any.whl</b></pre>
 
 If you encounter installation problems, see
 [Common Installation Problems](#common-installation-problems).
@@ -242,7 +242,7 @@ take the following steps:
      issue the following command:
 
      <pre> $ <b>sudo pip3 install --upgrade \
-     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc1-py3-none-any.whl</b> </pre>
+     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0-py3-none-any.whl</b> </pre>
 
      If the preceding command fails, see
      [installation problems](#common-installation-problems).
@@ -350,7 +350,7 @@ Take the following steps to install TensorFlow in an Anaconda environment:
      TensorFlow for Python 2.7:
 
      <pre> (<i>targetDirectory</i>)$ <b>pip install --ignore-installed --upgrade \
-     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc1-py2-none-any.whl</b></pre>
+     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0-py2-none-any.whl</b></pre>
 
 
 <a name="ValidateYourInstallation"></a>
@@ -524,7 +524,7 @@ The value you specify depends on your Python version.
 
 
 <pre>
-https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc1-py2-none-any.whl
+https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0-py2-none-any.whl
 </pre>
 
 
@@ -532,5 +532,5 @@ https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc1-py2-none-a
 
 
 <pre>
-https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc1-py3-none-any.whl
+https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0-py3-none-any.whl
 </pre>
diff --git a/tensorflow/docs_src/install/install_sources.md b/tensorflow/docs_src/install/install_sources.md
index d48a6ee550..3d937367fc 100644
--- a/tensorflow/docs_src/install/install_sources.md
+++ b/tensorflow/docs_src/install/install_sources.md
@@ -350,10 +350,10 @@ Invoke `pip install` to install that pip package.
 The filename of the `.whl` file depends on your platform.
 For example, the following command will install the pip package
 
-for TensorFlow 1.8.0rc1 on Linux:
+for TensorFlow 1.8.0 on Linux:
 
 <pre>
-$ <b>sudo pip install /tmp/tensorflow_pkg/tensorflow-1.8.0rc1-py2-none-any.whl</b>
+$ <b>sudo pip install /tmp/tensorflow_pkg/tensorflow-1.8.0-py2-none-any.whl</b>
 </pre>
 
 ## Validate your installation
diff --git a/tensorflow/tools/pip_package/setup.py b/tensorflow/tools/pip_package/setup.py
index bcf6c1e515..3ec5ea9af5 100644
--- a/tensorflow/tools/pip_package/setup.py
+++ b/tensorflow/tools/pip_package/setup.py
@@ -29,7 +29,7 @@ from setuptools.dist import Distribution
 # This version string is semver compatible, but incompatible with pip.
 # For pip, we will remove all '-' characters from this string, and use the
 # result for pip.
-_VERSION = '1.8.0-rc1'
+_VERSION = '1.8.0'
 
 REQUIRED_PACKAGES = [
     'absl-py >= 0.1.6',
-- 
GitLab


From f20740204f970e40e6238da5ad6507887f9bd95f Mon Sep 17 00:00:00 2001
From: Sami Kama <skama@nvidia.com>
Date: Thu, 26 Apr 2018 11:12:23 -0700
Subject: [PATCH 019/395] Introducing TRTOptimizationPass Use TF allocator for
 allocating TensorRT memory Fix an issue in build_pip_package.sh

---
 tensorflow/contrib/tensorrt/BUILD             |  10 +-
 .../contrib/tensorrt/convert/convert_graph.cc | 105 +++++--
 .../contrib/tensorrt/convert/convert_graph.h  |  12 +
 .../contrib/tensorrt/convert/convert_nodes.cc |  25 +-
 .../contrib/tensorrt/convert/convert_nodes.h  |  15 +-
 .../tensorrt/convert/trt_optimization_pass.cc | 236 +++++++++++++++
 .../tensorrt/convert/trt_optimization_pass.h  |  70 +++++
 .../contrib/tensorrt/kernels/trt_engine_op.cc |  51 +++-
 .../contrib/tensorrt/kernels/trt_engine_op.h  |   5 +-
 .../tensorrt/resources/trt_allocator.cc       |  57 ++++
 .../tensorrt/resources/trt_allocator.h        |  65 ++++
 .../tensorrt/resources/trt_resources.h        |   3 +
 .../contrib/tensorrt/segment/segment.cc       | 283 ++++++++++++++----
 tensorflow/contrib/tensorrt/segment/segment.h |  99 +++++-
 .../contrib/tensorrt/segment/segment_test.cc  |   6 +-
 .../contrib/tensorrt/test/test_tftrt.py       |  60 +++-
 .../tools/pip_package/build_pip_package.sh    |   2 +-
 17 files changed, 988 insertions(+), 116 deletions(-)
 create mode 100644 tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
 create mode 100644 tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
 create mode 100644 tensorflow/contrib/tensorrt/resources/trt_allocator.cc
 create mode 100644 tensorflow/contrib/tensorrt/resources/trt_allocator.h

diff --git a/tensorflow/contrib/tensorrt/BUILD b/tensorflow/contrib/tensorrt/BUILD
index f80b4f1b11..f7328ff228 100644
--- a/tensorflow/contrib/tensorrt/BUILD
+++ b/tensorflow/contrib/tensorrt/BUILD
@@ -88,6 +88,7 @@ cc_library(
         ":trt_logging",
         ":trt_resources",
         "//tensorflow/core:gpu_headers_lib",
+        "//tensorflow/core:gpu_runtime",
         "//tensorflow/core:lib_proto_parsing",
         "//tensorflow/core:stream_executor_headers_lib",
     ] + if_tensorrt([
@@ -194,10 +195,12 @@ tf_py_wrap_cc(
 tf_cuda_library(
     name = "trt_resources",
     srcs = [
+        "resources/trt_allocator.cc",
         "resources/trt_int8_calibrator.cc",
         "resources/trt_resource_manager.cc",
     ],
     hdrs = [
+        "resources/trt_allocator.h",
         "resources/trt_int8_calibrator.h",
         "resources/trt_resource_manager.h",
         "resources/trt_resources.h",
@@ -206,6 +209,7 @@ tf_cuda_library(
         ":trt_logging",
         "//tensorflow/core:framework_headers_lib",
         "//tensorflow/core:framework_lite",
+        "//tensorflow/core:core_cpu_lib",
         "//tensorflow/core:lib_proto_parsing",
     ] + if_tensorrt([
         "@local_config_tensorrt//:nv_infer",
@@ -218,10 +222,12 @@ tf_cuda_library(
     srcs = [
         "convert/convert_graph.cc",
         "convert/convert_nodes.cc",
+        "convert/trt_optimization_pass.cc",
     ],
     hdrs = [
         "convert/convert_graph.h",
         "convert/convert_nodes.h",
+        "convert/trt_optimization_pass.h",
     ],
     deps = [
         ":segment",
@@ -230,6 +236,7 @@ tf_cuda_library(
         "//tensorflow/core/grappler:grappler_item",
         "//tensorflow/core/grappler:utils",
         "//tensorflow/core:framework",
+        "//tensorflow/core:gpu_runtime",
         "//tensorflow/core:framework_lite",
         "//tensorflow/core:graph",
         "//tensorflow/core:lib",
@@ -238,8 +245,7 @@ tf_cuda_library(
         "//tensorflow/core/grappler:devices",
         "//tensorflow/core/grappler/clusters:virtual_cluster",
         "//tensorflow/core/grappler/costs:graph_properties",
-        "//tensorflow/core/grappler/optimizers:constant_folding",
-        "//tensorflow/core/grappler/optimizers:layout_optimizer",
+        "//tensorflow/core/grappler/optimizers:meta_optimizer",
     ] + if_tensorrt([
         "@local_config_tensorrt//:nv_infer",
     ]) + tf_custom_op_library_additional_deps(),
diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.cc b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
index b412b296e0..785c33c4c4 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_graph.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
@@ -24,15 +24,20 @@ limitations under the License.
 
 #include "tensorflow/contrib/tensorrt/convert/convert_nodes.h"
 #include "tensorflow/contrib/tensorrt/segment/segment.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h"
+#include "tensorflow/core/common_runtime/gpu/process_state.h"
 #include "tensorflow/core/graph/algorithm.h"
 #include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/graph/graph_constructor.h"
+#include "tensorflow/core/grappler/clusters/utils.h"
 #include "tensorflow/core/grappler/clusters/virtual_cluster.h"
 #include "tensorflow/core/grappler/costs/graph_properties.h"
 #include "tensorflow/core/grappler/devices.h"
 #include "tensorflow/core/grappler/grappler_item.h"
 #include "tensorflow/core/grappler/optimizers/constant_folding.h"
 #include "tensorflow/core/grappler/optimizers/layout_optimizer.h"
+#include "tensorflow/core/grappler/optimizers/meta_optimizer.h"
 #include "tensorflow/core/grappler/utils.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status.h"
@@ -115,8 +120,8 @@ std::pair<string, int> ParseTensorName(string name, int default_idx = 0) {
   int idx = default_idx;
   size_t sep = name.find_last_of(':');
   if (sep != string::npos) {
-    name = name.substr(0, sep);
     idx = std::stoi(name.substr(sep + 1));
+    name = name.substr(0, sep);
   }
   return std::make_pair(name, idx);
 }
@@ -141,7 +146,8 @@ struct ConvertGraphParams {
       size_t max_supported_batch_size, size_t max_consumed_workspace_size_bytes,
       const tensorflow::grappler::GraphProperties& current_graph_properties,
       std::unordered_map<string, std::pair<int, string>>* output_edges,
-      int engine_precision_mode)
+      int engine_precision_mode, const string& device_name,
+      std::shared_ptr<nvinfer1::IGpuAllocator> allocator, int cuda_device_id)
       : graph(inp_graph),
         output_names(output_node_names),
         subgraph_node_ids(subgraph_node_id_numbers),
@@ -149,7 +155,10 @@ struct ConvertGraphParams {
         max_workspace_size_bytes(max_consumed_workspace_size_bytes),
         graph_properties(current_graph_properties),
         output_edge_map(output_edges),
-        precision_mode(engine_precision_mode) {}
+        precision_mode(engine_precision_mode),
+        device_name_(device_name),
+        allocator_(allocator),
+        cuda_device_id_(cuda_device_id) {}
   tensorflow::Graph& graph;
   const std::vector<string>& output_names;
   const std::set<int>& subgraph_node_ids;
@@ -158,6 +167,9 @@ struct ConvertGraphParams {
   const tensorflow::grappler::GraphProperties& graph_properties;
   std::unordered_map<string, std::pair<int, string>>* output_edge_map;
   int precision_mode;
+  string device_name_;
+  std::shared_ptr<nvinfer1::IGpuAllocator> allocator_;
+  int cuda_device_id_;
   std::vector<std::pair<int, int>> subgraph_inputs;
   std::vector<std::pair<int, int>> subgraph_outputs;
   tensorflow::EdgeSet subgraph_incoming_edges;
@@ -200,7 +212,8 @@ tensorflow::Status GetCalibNode(ConvertGraphParams* params) {
                    params->subgraph_inputs, params->subgraph_outputs,
                    params->max_batch_size, params->max_workspace_size_bytes,
                    params->graph_properties, params->output_edge_map,
-                   &trt_node_def, params->precision_mode);
+                   &trt_node_def, params->precision_mode, params->device_name_,
+                   params->allocator_, params->cuda_device_id_);
   TF_RETURN_IF_ERROR(InjectCalibrationNode(s));
   tensorflow::Status status;
   tensorflow::Node* trt_node = params->graph.AddNode(trt_node_def, &status);
@@ -214,7 +227,7 @@ tensorflow::Status GetCalibNode(ConvertGraphParams* params) {
     auto src_output = in_edge->src_output();
     auto dst_node = in_edge->dst();
     auto dst_input = in_edge->dst_input();
-    VLOG(1) << " update edge " << trt_node->name() << ":" << src_output
+    VLOG(0) << " update edge " << trt_node->name() << ":" << src_output
             << " -> " << dst_node->name() << ":" << dst_input;
     TF_RETURN_IF_ERROR(
         params->graph.UpdateEdge(trt_node, src_output, dst_node, dst_input));
@@ -230,7 +243,8 @@ tensorflow::Status ConvertSubGraphToTensorRT(ConvertGraphParams* params) {
                    params->subgraph_inputs, params->subgraph_outputs,
                    params->max_batch_size, params->max_workspace_size_bytes,
                    params->graph_properties, params->output_edge_map,
-                   &trt_node_def, params->precision_mode);
+                   &trt_node_def, params->precision_mode, params->device_name_,
+                   params->allocator_, params->cuda_device_id_);
   TF_RETURN_IF_ERROR(ConvertSubGraphToTensorRTNodeDef(s));
   tensorflow::Status status;
   tensorflow::Node* trt_node = params->graph.AddNode(trt_node_def, &status);
@@ -348,27 +362,41 @@ tensorflow::Status ConvertGraphDefToTensorRT(
   int num_gpus = tensorflow::grappler::GetNumAvailableGPUs();
   VLOG(2) << "cpu_cores: " << num_cpu_cores;
   VLOG(2) << "gpus: " << num_gpus;
-
-  TF_RETURN_IF_ERROR(optimizer.Optimize(cluster, item, &gdef));
-
+  tensorflow::RewriterConfig rw_cfg;
+  tensorflow::grappler::MetaOptimizer meta_opt(nullptr, rw_cfg);
+  // TF_RETURN_IF_ERROR(optimizer.Optimize(cluster, item, &gdef));
+  TF_RETURN_IF_ERROR(meta_opt.Optimize(cluster, item, &gdef));
   // constant folding
   item.graph = gdef;
-  tensorflow::grappler::ConstantFolding fold(nullptr);
-  TF_RETURN_IF_ERROR(fold.Optimize(nullptr, item, &gdef));
+  // tensorflow::grappler::ConstantFolding fold(nullptr);
+  // TF_RETURN_IF_ERROR(fold.Optimize(nullptr, item, &gdef));
 
   // AJ refactoring shape inference through grappler/GraphProperties.
   tensorflow::grappler::GraphProperties static_graph_properties(item);
-  TF_RETURN_IF_ERROR(static_graph_properties.InferStatically(false));
+  TF_RETURN_IF_ERROR(static_graph_properties.InferStatically(true));
   // Build full graph
+
+  return ConvertAfterShapes(gdef, output_names, max_batch_size,
+                            max_workspace_size_bytes, new_graph_def,
+                            precision_mode, minimum_segment_size,
+                            static_graph_properties, nullptr);
+}
+
+tensorflow::Status ConvertAfterShapes(
+    const tensorflow::GraphDef& gdef, const std::vector<string>& output_names,
+    size_t max_batch_size, size_t max_workspace_size_bytes,
+    tensorflow::GraphDef* new_graph_def, int precision_mode,
+    int minimum_segment_size,
+    const tensorflow::grappler::GraphProperties& graph_properties,
+    const tensorflow::grappler::Cluster* cluster) {
+  // Segment the graph into subgraphs that can be converted to TensorRT
+  tensorflow::tensorrt::segment::SegmentOptions segment_options;
   tensorflow::FunctionLibraryDefinition flib(tensorflow::OpRegistry::Global(),
                                              gdef.library());
   tensorflow::Graph graph(flib);
   TF_RETURN_IF_ERROR(tensorflow::ConvertGraphDefToGraph(
       tensorflow::GraphConstructorOptions(), gdef, &graph));
 
-  // Segment the graph into subgraphs that can be converted to TensorRT
-  tensorflow::tensorrt::segment::SegmentOptions segment_options;
-
   // TODO(ben,jie,sami): exclude output nodes (DISCUSS IT)
   for (auto node : output_names) {
     segment_options.exclude_node_list.insert(node);
@@ -378,7 +406,7 @@ tensorflow::Status ConvertGraphDefToTensorRT(
   segment_options.minimum_segment_size = minimum_segment_size;
   tensorflow::tensorrt::segment::SegmentNodesVector segments;
   TF_RETURN_IF_ERROR(tensorrt::segment::SegmentGraph(
-      gdef, IsTensorRTCandidate, segment_options, &segments));
+      &graph, IsTensorRTCandidate, segment_options, &segments));
   if (segments.size() > 1) {
     VLOG(0) << "MULTIPLE tensorrt candidate conversion: " << segments.size();
   }
@@ -388,9 +416,17 @@ tensorflow::Status ConvertGraphDefToTensorRT(
   int count = 0;
   float total_num_nodes_in_segments = 0.;
   for (auto s : segments) {
-    total_num_nodes_in_segments += s.size();
+    total_num_nodes_in_segments += s.first.size();
+  }
+  std::map<string, tensorflow::Device*> name_to_device_map;
+  if (cluster) {
+    for (const auto dm : cluster->GetDeviceSet()->devices()) {
+      name_to_device_map[dm->name()] = dm;
+    }
   }
-  for (const std::set<string>& subgraph_node_names : segments) {
+  for (const auto& segment_nodes_and_device : segments) {
+    const std::set<string>& subgraph_node_names =
+        segment_nodes_and_device.first;
     std::set<int> subgraph_node_ids;
     size_t max_mem_per_engine =
         max_workspace_size_bytes *
@@ -400,10 +436,37 @@ tensorflow::Status ConvertGraphDefToTensorRT(
       oss << " " << node_name;
       subgraph_node_ids.insert(node_map.at(node_name)->id());
     }
-    VLOG(2) << "Subgraph nodes" << oss.str();
+    VLOG(1) << "Subgraph nodes at device " << segment_nodes_and_device.second
+            << " : " << oss.str();
+    auto target_device =
+        name_to_device_map.find(segment_nodes_and_device.second);
+    std::shared_ptr<nvinfer1::IGpuAllocator> allocator(0);
+
+    int cuda_device_id = 0;
+    if (target_device != name_to_device_map.end()) {
+      tensorflow::TfGpuId tf_gpu_id(target_device->second->parsed_name().id);
+      CudaGpuId cuda_gpu_id;
+      Status s = GpuIdManager::TfToCudaGpuId(tf_gpu_id, &cuda_gpu_id);
+      if (!s.ok()) {
+        LOG(ERROR)
+            << "Cuda device identification failed, using device 0. Error= " << s;
+      } else {
+        cuda_device_id = cuda_gpu_id.value();
+      }
+      tensorflow::GPUOptions gpuoptions;
+      auto pm = tensorflow::ProcessState::singleton();
+      // this should be instantiated by now
+      auto dev_allocator = pm->GetGPUAllocator(gpuoptions, tf_gpu_id, 1);
+      VLOG(1) << "Got an allocator for device tf_device=" << tf_gpu_id.value()
+              << " cuda device= " << cuda_device_id << " at " << dev_allocator;
+      allocator = std::make_shared<TRTDeviceAllocator>(dev_allocator);
+    } else {  // device unknown or not available
+      allocator = std::make_shared<TRTCudaAllocator>();
+    }
     ConvertGraphParams p(graph, output_names, subgraph_node_ids, max_batch_size,
-                         max_mem_per_engine, static_graph_properties,
-                         &output_edge_map, precision_mode);
+                         max_mem_per_engine, graph_properties, &output_edge_map,
+                         precision_mode, segment_nodes_and_device.second,
+                         allocator, cuda_device_id);
     if (precision_mode == INT8MODE) {
       tensorflow::Status status = GetCalibNode(&p);
       if (status != tensorflow::Status::OK()) {
diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.h b/tensorflow/contrib/tensorrt/convert/convert_graph.h
index e01e4a5328..23a83b5094 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_graph.h
+++ b/tensorflow/contrib/tensorrt/convert/convert_graph.h
@@ -17,7 +17,11 @@ limitations under the License.
 
 #include <vector>
 
+#include "tensorflow/contrib/tensorrt/segment/segment.h"
 #include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/grappler/clusters/cluster.h"
+#include "tensorflow/core/grappler/costs/graph_properties.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/types.h"
 
@@ -43,6 +47,14 @@ tensorflow::Status ConvertGraphDefToTensorRT(
     size_t max_workspace_size_bytes, tensorflow::GraphDef* new_graph_def,
     int precision_mode, int minimum_segment_size);
 
+// Method to call from optimization pass
+tensorflow::Status ConvertAfterShapes(
+    const tensorflow::GraphDef& graph, const std::vector<string>& output_names,
+    size_t max_batch_size, size_t max_workspace_size_bytes,
+    tensorflow::GraphDef* new_graph_def, int precision_mode,
+    int minimum_segment_size,
+    const tensorflow::grappler::GraphProperties& graph_properties,
+    const tensorflow::grappler::Cluster* cluster);
 }  // namespace convert
 }  // namespace tensorrt
 }  // namespace tensorflow
diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
index b81ae9dc3e..b37c535736 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
@@ -346,10 +346,11 @@ void ReorderCKtoKC(const TRT_ShapedWeights& iweights,
       break;
     }
     case tensorflow::DataType::DT_HALF: {
-      Reorder2({k, c}, static_cast<Eigen::half const*>(iweights.GetValues()),
-               istrides, static_cast<Eigen::half*>(
-                             const_cast<void*>(oweights->GetValues())),
-               ostrides);
+      Reorder2(
+          {k, c}, static_cast<Eigen::half const*>(iweights.GetValues()),
+          istrides,
+          static_cast<Eigen::half*>(const_cast<void*>(oweights->GetValues())),
+          ostrides);
       break;
     }
     default:
@@ -2246,8 +2247,12 @@ tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) {
   auto op_res = new tensorflow::tensorrt::TRTCalibrationResource();
   TF_CHECK_OK(op_rmgr->Create(calib_op_name, calib_op_name, op_res));
   op_res->logger_ = new tensorflow::tensorrt::Logger();
+  cudaSetDevice(s.cuda_device_id_);
   op_res->builder_ = nvinfer1::createInferBuilder(*(op_res->logger_));
-
+  op_res->allocator_=s.allocator_;
+#if NV_TENSORRT_MAJOR >4 
+  op_res->builder_->setGpuAllocator(s.allocator_.get());
+#endif
   if (!op_res->builder_) {
     return tensorflow::errors::Internal(
         "failed to create TensorRT builder object");
@@ -2476,13 +2481,15 @@ tensorflow::Status ConvertSubGraphToTensorRTNodeDef(
   // Topological order is needed to build TRT network
 
   tensorflow::tensorrt::Logger trt_logger;
-
+cudaSetDevice(s.cuda_device_id_);
   auto trt_builder = infer_object(nvinfer1::createInferBuilder(trt_logger));
   if (!trt_builder) {
     return tensorflow::errors::Internal(
         "Failed to create TensorRT builder object");
   }
-
+#if NV_TENSORRT_MAJOR >3 
+  trt_builder->setGpuAllocator(s.allocator_.get());
+#endif
   auto trt_network = infer_object(trt_builder->createNetwork());
   if (!trt_network) {
     return tensorflow::errors::Internal(
@@ -2707,9 +2714,11 @@ tensorflow::Status ConvertSubGraphToTensorRTNodeDef(
                     .Attr("input_nodes", input_names)
                     .Attr("output_nodes", output_names)
                     .Attr("OutT", output_dtypes)
+                    .Device(s.device_name_)
                     .Finalize(s.trt_node);
 
-  VLOG(0) << status.ToString() << " finished op building";
+  VLOG(0) << status.ToString() << " finished op building for " << engine_name
+          << " on device " << s.device_name_ ;
 
   return tensorflow::Status::OK();
 }
diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.h b/tensorflow/contrib/tensorrt/convert/convert_nodes.h
index 954a1e72f8..ecccaf36e3 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_nodes.h
+++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.h
@@ -22,6 +22,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/grappler/costs/graph_properties.h"
@@ -29,7 +30,7 @@ limitations under the License.
 
 #if GOOGLE_CUDA
 #if GOOGLE_TENSORRT
-
+#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h"
 namespace tensorflow {
 namespace tensorrt {
 namespace convert {
@@ -48,7 +49,9 @@ struct SubGraphParams {
       const tensorflow::grappler::GraphProperties& current_graph_properties,
       std::unordered_map<string, std::pair<int, string>>* output_edges,
       tensorflow::NodeDef* constructed_trt_node,
-      int engine_precision_mode = FP32MODE)
+      int engine_precision_mode = FP32MODE, const string& device_name = "",
+      std::shared_ptr<nvinfer1::IGpuAllocator> allocator = 0,
+      int cuda_device_id = 0)
       : graph(inp_graph),
         subgraph_node_ids(subgraph_node_id_numbers),
         input_inds(input_indices),
@@ -58,7 +61,10 @@ struct SubGraphParams {
         graph_properties(current_graph_properties),
         output_edge_map(output_edges),
         trt_node(constructed_trt_node),
-        precision_mode(engine_precision_mode) {}
+        precision_mode(engine_precision_mode),
+        device_name_(device_name),
+        allocator_(allocator),
+        cuda_device_id_(cuda_device_id) {}
 
   tensorflow::Graph& graph;
   const std::set<int>& subgraph_node_ids;
@@ -70,6 +76,9 @@ struct SubGraphParams {
   std::unordered_map<string, std::pair<int, string>>* output_edge_map;
   tensorflow::NodeDef* trt_node;
   const int precision_mode;
+  const string device_name_;
+  std::shared_ptr<nvinfer1::IGpuAllocator> allocator_;
+  const int cuda_device_id_;
 };
 
 // TODO(sami): Replace references with const reference or pointers
diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
new file mode 100644
index 0000000000..880ffe1b3a
--- /dev/null
+++ b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
@@ -0,0 +1,236 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+1;4804;0c
+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.
+==============================================================================*/
+
+#include "tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h"
+#include "tensorflow/contrib/tensorrt/convert/convert_graph.h"
+#include "tensorflow/core/grappler/clusters/cluster.h"
+#include "tensorflow/core/grappler/grappler_item.h"
+#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/public/session_options.h"
+
+using tensorflow::str_util::Uppercase;
+using tensorflow::strings::StrAppend;
+using tensorflow::strings::StrCat;
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+tensorflow::Status TRTOptimizationPass::Init(
+    const tensorflow::RewriterConfig_CustomGraphOptimizer* config) {
+  VLOG(1) << "Called INIT for " << m_name_ << " with config = " << config;
+  if (config == nullptr) {
+    maximum_workspace_size_ = 2 << 30;
+    return tensorflow::Status::OK();
+  }
+  const auto params = config->parameter_map();
+  if (params.count("minimum_segment_size")) {
+    minimum_segment_size_ = params.at("minimum_segment_size").i();
+  }
+  if (params.count("max_batch_size")) {
+    maximum_batch_size_ = params.at("max_batch_size").i();
+  }
+  if (params.count("max_workspace_size_bytes"))
+    maximum_workspace_size_ = params.at("max_workspace_size_bytes").i();
+  if (params.count("precision_mode")) {
+    string pm = Uppercase(params.at("precision_mode").s());
+    if (pm == "FP32") {
+      precision_mode_ = 0;
+    } else if (pm == "FP16") {
+      precision_mode_ = 1;
+    } else if (pm == "INT8") {
+      precision_mode_ = 2;
+    } else {
+      LOG(ERROR) << "Unknown precision mode '" << pm << "'";
+      return tensorflow::errors::InvalidArgument(
+          "Unknown precision mode argument" + pm +
+          " Valid values are FP32, FP16, INT8");
+    }
+  }
+  return tensorflow::Status::OK();
+};
+
+tensorflow::Status TRTOptimizationPass::Optimize(
+    tensorflow::grappler::Cluster* cluster,
+    const tensorflow::grappler::GrapplerItem& item, GraphDef* optimized_graph) {
+  VLOG(1) << "Called TRTOptimization Pass " << m_name_;
+  VLOG(1) << "Cluster = " << cluster;
+  string offset("  ");
+  string offset2 = StrCat(offset, offset);
+  string offset3 = StrCat(offset2, offset);
+  string offset4 = StrCat(offset2, offset2);
+  if (cluster) {
+    VLOG(1) << offset << "type             = " << cluster->type();
+    VLOG(1) << offset << "num warmup steps = " << cluster->NumWarmupSteps();
+    const auto devNames = cluster->GetDeviceNames();
+    if (devNames.size()) {
+      VLOG(1) << offset << " Device names:";
+      for (const auto s : devNames) {
+        VLOG(1) << offset2 << s;
+      }
+    }
+    std::unordered_map<string, uint64> peak_mem;
+    auto status = cluster->GetPeakMemoryUsage(&peak_mem);
+    if (status == tensorflow::Status::OK()) {
+      VLOG(1) << offset << "Peak Memory Usage :";
+      for (auto s : peak_mem) {
+        VLOG(1) << offset2 << s.first << " = " << s.second;
+      }
+    }
+
+    const auto dev_props = cluster->GetDevices();
+    if (dev_props.size()) {
+      VLOG(1) << offset << "Device properties:";
+      for (auto k : dev_props) {
+        VLOG(1) << offset2 << k.first;
+        const auto& dt = k.second;
+        VLOG(1) << offset3 << "type          = " << dt.type();
+        VLOG(1) << offset3 << "vendor        = " << dt.vendor();
+        VLOG(1) << offset3 << "model         = " << dt.model();
+        VLOG(1) << offset3 << "frequency     = " << dt.frequency();
+        VLOG(1) << offset3 << "num cores     = " << dt.num_cores();
+        VLOG(1) << offset3 << "num registers = " << dt.num_registers();
+        VLOG(1) << offset3 << "L1 cache size = " << dt.l1_cache_size();
+        VLOG(1) << offset3 << "L2 cache size = " << dt.l2_cache_size();
+        VLOG(1) << offset3 << "L3 cache size = " << dt.l3_cache_size();
+        VLOG(1) << offset3 << "SHMem per SMP = "
+                << dt.shared_memory_size_per_multiprocessor();
+        VLOG(1) << offset3 << "memory size   = " << dt.memory_size();
+        VLOG(1) << offset3 << "bandwidth     = " << dt.bandwidth();
+        if (dt.environment_size()) {
+          VLOG(1) << offset3 << "environment   :";
+          for (const auto e : dt.environment()) {
+            VLOG(1) << offset4 << e.first << " = " << e.second;
+          }
+        }
+      }
+    }
+  }
+  VLOG(1) << "item: " << item.id;
+  int max_dim = -1;
+  if (item.feed.size()) {
+    VLOG(1) << offset << "Feeds  :";
+    for (const auto& f : item.feed) {
+      const auto& shape = f.second.shape();
+      if (shape.dims() > 0) {
+        if (shape.dim_size(0) > max_dim) max_dim = shape.dim_size(0);
+      }
+      VLOG(1) << offset2 << f.first << " = shaped "
+              << f.second.shape().DebugString();
+    }
+  } else {
+    VLOG(1) << offset << "No Feeds";
+  }
+  if (maximum_batch_size_ < 0) {  // automatic batch size from input
+    if (max_dim > 0) {
+      maximum_batch_size_ = max_dim;
+      VLOG(1) << "Setting maximum batch size to " << max_dim;
+    } else {
+      maximum_batch_size_ = 128;
+      LOG(WARNING) << "Maximum batch size is not set"
+                      " and can't be deduced from inputs setting it to"
+                   << maximum_batch_size_
+                   << ". Suggest configuring it from configuration parameters";
+    }
+  } else {
+    if (max_dim > maximum_batch_size_) {
+      LOG(WARNING) << "Configured batch size " << maximum_batch_size_
+                   << " is less than input batch size " << max_dim
+                   << " adjusting maximum batch size to match input batch size";
+    }
+  }
+  if (item.fetch.size()) {
+    VLOG(1) << offset << "Fetches  :";
+    for (const auto& f : item.fetch) {
+      VLOG(1) << offset2 << f;
+    }
+  } else {
+    VLOG(1) << offset << "No Fetches";
+  }
+
+  if (item.init_ops.size()) {
+    VLOG(1) << offset << "init ops  :";
+    for (const auto& f : item.init_ops) {
+      VLOG(1) << offset2 << f;
+    }
+  } else {
+    VLOG(1) << offset << "No init ops";
+  }
+  VLOG(1) << "Save Op = " << item.save_op;
+  VLOG(1) << "Restore Op = " << item.restore_op;
+  VLOG(1) << "save_restore_loc_tensor = " << item.save_restore_loc_tensor;
+  if (item.keep_ops.size()) {
+    VLOG(1) << offset << "keep ops  :";
+    for (const auto& f : item.keep_ops) {
+      VLOG(1) << offset2 << f;
+    }
+  } else {
+    VLOG(1) << offset << "No keep ops";
+  }
+  VLOG(1) << item.graph.DebugString();
+  tensorflow::grappler::GraphProperties static_graph_properties(item);
+  TF_RETURN_IF_ERROR(static_graph_properties.InferStatically(true));
+  for (const auto dev : cluster->GetDeviceSet()->devices()) {
+    const auto& pname = dev->parsed_name();
+    VLOG(1) << "Device name= " << dev->name()
+            << " parsedname job= " << pname.job << " id= " << pname.id
+            << " has_id: " << pname.has_id << " has_job: " << pname.has_job<< 
+            "has_type: "<<pname.has_type<<" type ="<<pname.type;
+  }
+  auto status = tensorflow::tensorrt::convert::ConvertAfterShapes(
+      item.graph, item.fetch, maximum_batch_size_, maximum_workspace_size_,
+      optimized_graph, precision_mode_, minimum_segment_size_,
+      static_graph_properties, cluster);
+  VLOG(2) << optimized_graph->DebugString();
+  return status;
+}
+
+void TRTOptimizationPass::Feedback(
+    tensorflow::grappler::Cluster* cluster,
+    const tensorflow::grappler::GrapplerItem& item,
+    const GraphDef& optimized_graph, double result) {}
+
+using tensorflow::grappler::CustomGraphOptimizerRegistrar;
+namespace {
+
+class samiReg : public CustomGraphOptimizerRegistrar {
+ public:
+  samiReg(const tensorflow::grappler::CustomGraphOptimizerRegistry::Creator& cr,
+          const string& name)
+      : CustomGraphOptimizerRegistrar(cr, name) {
+    VLOG(1) << "Constructing a CustomOptimizationPass registration object for "
+            << name;
+  }
+};
+// static CustomGraphOptimizerRegistrar TRTOptimizationPass_Registrar([]() {
+static samiReg TRTOptimizationPass_Registrar(
+    []() {
+      VLOG(1)
+          << "Instantiating CustomOptimizationPass object TensorRTOptimizer";
+      return new tensorflow::tensorrt::convert::TRTOptimizationPass(
+          "TensorRTOptimizer");
+    },
+    ("TensorRTOptimizer"));
+}  // namespace
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif
+#endif
diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
new file mode 100644
index 0000000000..5b1462f573
--- /dev/null
+++ b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
@@ -0,0 +1,70 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CONTRIB_TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_
+#define TENSORFLOW_CONTRIB_TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_
+
+#include <set>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h"
+#include "tensorflow/core/platform/logging.h"
+
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer {
+ public:
+  TRTOptimizationPass(string optName = "TRTOptimizationPass")
+      : m_name_(optName),
+        minimum_segment_size_(3),
+        precision_mode_(0),
+        maximum_batch_size_(-1),
+        maximum_workspace_size_(-1) {
+    VLOG(1) << "Constructing " << m_name_;
+  };
+  // tensorflow::Status Run(const tensorflow::GraphOptimizationPassOptions
+  // &options) override;
+  string name() const override { return m_name_; };
+  tensorflow::Status Init(const tensorflow::RewriterConfig_CustomGraphOptimizer*
+                              config = nullptr) override;
+
+  tensorflow::Status Optimize(tensorflow::grappler::Cluster* cluster,
+                              const tensorflow::grappler::GrapplerItem& item,
+                              GraphDef* optimized_graph) override;
+  void Feedback(tensorflow::grappler::Cluster* cluster,
+                const tensorflow::grappler::GrapplerItem& item,
+                const GraphDef& optimized_graph, double result) override;
+
+ private:
+  string m_name_;
+  int minimum_segment_size_;
+  int precision_mode_;
+  int maximum_batch_size_;
+  int64_t maximum_workspace_size_;
+};
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif
+#endif
+#endif
\ No newline at end of file
diff --git a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
index b32371b642..9c59fd973b 100644
--- a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
+++ b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
@@ -18,6 +18,9 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/stream_executor.h"
 #include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/common_runtime/gpu/process_state.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h"
 
 #if GOOGLE_CUDA
 #if GOOGLE_TENSORRT
@@ -33,9 +36,8 @@ namespace tensorrt {
 
 TRTEngineOp::TRTEngineOp(OpKernelConstruction* context) : OpKernel(context) {
   // read serialized_engine
-  string serialized_engine;
   OP_REQUIRES_OK(context,
-                 context->GetAttr("serialized_engine", &serialized_engine));
+                 context->GetAttr("serialized_engine", &serialized_engine_));
 
   // register input output node name in trt_sub_graph
   OP_REQUIRES_OK(context, context->GetAttr("input_nodes", &input_nodes_));
@@ -46,25 +48,43 @@ TRTEngineOp::TRTEngineOp(OpKernelConstruction* context) : OpKernel(context) {
   // from resourcemanager
   // TODO(jie): cudaSetDevice make sure trt engine is allocated on the same
   // gpu where the input/output is also located.
-  int gpu_id = context->device()->tensorflow_gpu_device_info()->gpu_id;
-  cudaSetDevice(gpu_id);
-  int device;
-  cudaGetDevice(&device);
-  if (gpu_id != device) LOG(FATAL) << "set device failed!";
+  // int gpu_id = context->device()->tensorflow_gpu_device_info()->gpu_id;
+  // cudaSetDevice(gpu_id);
+  // int device;
+  // cudaGetDevice(&device);
+  // if (gpu_id != device) LOG(FATAL) << "set device failed!";
 
   // TODO(samikama) runtime should be taken from a resourcemanager as well.
   // Only engine should be in the op and context and runtime should be taken
   // from resourcemanager
 
-  IRuntime* infer = nvinfer1::createInferRuntime(logger);
-  trt_engine_ptr_.reset(infer->deserializeCudaEngine(
-      serialized_engine.c_str(), serialized_engine.size(), nullptr));
-  trt_execution_context_ptr_.reset(trt_engine_ptr_->createExecutionContext());
+  // IRuntime* infer = nvinfer1::createInferRuntime(logger);
+  // trt_engine_ptr_.reset(infer->deserializeCudaEngine(
+  //     serialized_engine.c_str(), serialized_engine.size(), nullptr));
+  // trt_execution_context_ptr_.reset(trt_engine_ptr_->createExecutionContext());
   // Runtime is safe to delete after engine creation
-  infer->destroy();
+  // infer->destroy();
 }
 
 void TRTEngineOp::Compute(OpKernelContext* context) {
+  if(!trt_execution_context_ptr_){
+    tensorflow::TfGpuId tf_gpu_id(context->device()->tensorflow_gpu_device_info()->gpu_id);
+    tensorflow::GPUOptions gpuoptions;
+    auto pm = tensorflow::ProcessState::singleton();
+    auto dev_allocator = pm->GetGPUAllocator(gpuoptions, tf_gpu_id, 1);
+    IRuntime* infer = nvinfer1::createInferRuntime(logger);
+    if(!dev_allocator){
+      LOG(FATAL)<<"Can't find device allocator for gpu device"<<tf_gpu_id;
+    }
+    allocator_ = std::make_shared<TRTDeviceAllocator>(dev_allocator);
+    infer->setGpuAllocator(allocator_.get());
+    trt_engine_ptr_.reset(infer->deserializeCudaEngine(
+        serialized_engine_.c_str(), serialized_engine_.size(), nullptr));
+    trt_execution_context_ptr_.reset(trt_engine_ptr_->createExecutionContext());
+    // Runtime is safe to delete after engine creation
+    infer->destroy();
+    serialized_engine_.clear();
+  }
   int num_binding = context->num_inputs() + context->num_outputs();
   std::vector<void*> buffers(num_binding);
 
@@ -147,7 +167,12 @@ void TRTEngineOp::Compute(OpKernelContext* context) {
   VLOG(2) << "enqueue returns: " << ret;
   // sync should be done by TF.
 }
-
+TRTEngineOp::~TRTEngineOp(){
+  // Order matters!
+  trt_execution_context_ptr_.reset();
+  trt_engine_ptr_.reset();
+  allocator_.reset();
+}
 REGISTER_KERNEL_BUILDER(Name("TRTEngineOp").Device(DEVICE_GPU), TRTEngineOp);
 
 }  // namespace tensorrt
diff --git a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h
index 0964b4b18a..791bb6f583 100644
--- a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h
+++ b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h
@@ -26,6 +26,7 @@ limitations under the License.
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorrt/include/NvInfer.h"
+#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h"
 
 namespace tensorflow {
 namespace tensorrt {
@@ -36,7 +37,7 @@ class TRTEngineOp : public OpKernel {
   explicit TRTEngineOp(OpKernelConstruction* context);
 
   void Compute(OpKernelContext* context) override;
-
+  ~TRTEngineOp();
  private:
   template <typename T>
   struct Destroyer {
@@ -51,6 +52,8 @@ class TRTEngineOp : public OpKernel {
 
   std::vector<string> input_nodes_;
   std::vector<string> output_nodes_;
+  std::shared_ptr<nvinfer1::IGpuAllocator> allocator_;
+  string serialized_engine_;
 };
 
 }  // namespace tensorrt
diff --git a/tensorflow/contrib/tensorrt/resources/trt_allocator.cc b/tensorflow/contrib/tensorrt/resources/trt_allocator.cc
new file mode 100644
index 0000000000..4705f6d20f
--- /dev/null
+++ b/tensorflow/contrib/tensorrt/resources/trt_allocator.cc
@@ -0,0 +1,57 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h"
+
+#include "tensorflow/core/platform/logging.h"
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
+#if NV_TENSORRT_MAJOR > 2
+#include "cuda/include/cuda_runtime_api.h"
+
+namespace tensorflow {
+namespace tensorrt {
+void* TRTCudaAllocator::allocate(uint64_t size, uint64_t alignment,
+                                 uint32_t flags) {
+  assert((alignment & (alignment - 1)) == 0);  // zero or a power of 2.
+  void* memory;
+  cudaMalloc(&memory, size);
+  return memory;
+}
+void TRTCudaAllocator::free(void* memory) { cudaFree(memory); }
+
+void* TRTDeviceAllocator::allocate(uint64_t size, uint64_t alignment,
+                                   uint32_t flags) {
+  assert((alignment & (alignment - 1)) == 0);  // zero or a power of 2.
+  void* mem = allocator_->AllocateRaw(alignment, size);
+  VLOG(2) << "Allocated " << size << " bytes with alignment " << alignment
+          << " @ " << mem;
+  return mem;
+}
+
+TRTDeviceAllocator::TRTDeviceAllocator(tensorflow::Allocator* allocator)
+    : allocator_(allocator) {
+  VLOG(1) << "Using " << allocator->Name() << " allocator from TensorFlow";
+};
+void TRTDeviceAllocator::free(void* memory) {
+  VLOG(2) << "Deallocating " << memory;
+  allocator_->DeallocateRaw(memory);
+}
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif
+#endif
+#endif
\ No newline at end of file
diff --git a/tensorflow/contrib/tensorrt/resources/trt_allocator.h b/tensorflow/contrib/tensorrt/resources/trt_allocator.h
new file mode 100644
index 0000000000..8bdb0519ba
--- /dev/null
+++ b/tensorflow/contrib/tensorrt/resources/trt_allocator.h
@@ -0,0 +1,65 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_ALLOCATOR_H_
+#define TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_ALLOCATOR_H_
+
+#include <list>
+#include <sstream>
+#include <string>
+#include <thread>
+#include <vector>
+#include "tensorflow/contrib/tensorrt/log/trt_logger.h"
+#include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
+#include "tensorrt/include/NvInfer.h"
+#if NV_TENSORRT_MAJOR == 3
+// define interface here temporarily until TRT 4.0 is released
+namespace nvinfer1 {
+class IGpuAllocator {
+  virtual void* allocate(uint64_t size, uint64_t alignment, uint32_t flags) = 0;
+  virtual void free(void* memory) = 0;
+};
+}  // namespace nvinfer1
+#endif
+namespace tensorflow {
+namespace tensorrt {
+class TRTCudaAllocator : public nvinfer1::IGpuAllocator {
+ public:
+  TRTCudaAllocator() {}
+  virtual ~TRTCudaAllocator(){};
+  void* allocate(uint64_t size, uint64_t alignment, uint32_t flags) override;
+  void free(void* memory) override;
+};
+class TRTDeviceAllocator : public nvinfer1::IGpuAllocator {
+ public:
+  TRTDeviceAllocator(tensorflow::Allocator* allocator);
+  virtual ~TRTDeviceAllocator(){};
+  void* allocate(uint64_t size, uint64_t alignment, uint32_t flags) override;
+  void free(void* memory) override;
+
+ private:
+  tensorflow::Allocator* allocator_;
+};
+class AllocatorFactory {};
+
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+#endif
+#endif
+#endif
\ No newline at end of file
diff --git a/tensorflow/contrib/tensorrt/resources/trt_resources.h b/tensorflow/contrib/tensorrt/resources/trt_resources.h
index 3c85968ae7..166ca9c3de 100644
--- a/tensorflow/contrib/tensorrt/resources/trt_resources.h
+++ b/tensorflow/contrib/tensorrt/resources/trt_resources.h
@@ -28,6 +28,7 @@ limitations under the License.
 #if GOOGLE_TENSORRT
 #include "tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h"
 #include "tensorrt/include/NvInfer.h"
+#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h"
 
 namespace tensorflow {
 namespace tensorrt {
@@ -47,6 +48,7 @@ class TRTCalibrationResource : public tensorflow::ResourceBase {
         << " Network    = " << std::hex << network_ << std::dec << std::endl
         << " Engine     = " << std::hex << engine_ << std::dec << std::endl
         << " Logger     = " << std::hex << logger_ << std::dec << std::endl
+        << " Allocator  = " << std::hex << allocator_.get()<< std::dec << std::endl
         << " Thread     = " << std::hex << thr_ << std::dec << std::endl;
     return oss.str();
   }
@@ -57,6 +59,7 @@ class TRTCalibrationResource : public tensorflow::ResourceBase {
   nvinfer1::IBuilder* builder_;
   nvinfer1::INetworkDefinition* network_;
   nvinfer1::ICudaEngine* engine_;
+  std::shared_ptr<nvinfer1::IGpuAllocator> allocator_;
   tensorflow::tensorrt::Logger* logger_;
   // TODO(sami): Use threadpool threads!
   std::thread* thr_;
diff --git a/tensorflow/contrib/tensorrt/segment/segment.cc b/tensorflow/contrib/tensorrt/segment/segment.cc
index 8fc4697c51..8f335f2bf1 100644
--- a/tensorflow/contrib/tensorrt/segment/segment.cc
+++ b/tensorflow/contrib/tensorrt/segment/segment.cc
@@ -25,18 +25,58 @@ limitations under the License.
 #include "tensorflow/core/graph/graph_constructor.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/types.h"
 
 namespace tensorflow {
 namespace tensorrt {
 namespace segment {
-
+using ::tensorflow::strings::StrAppend;
 namespace {
 
-bool CanContractEdge(const tensorflow::Edge* edge,
-                     const tensorflow::Graph& graph) {
-  const tensorflow::Node* src = edge->src();
-  const tensorflow::Node* dst = edge->dst();
+bool check_cycles(const Graph* g, const Node* src,
+                  const std::vector<Node*>& start) {
+  struct Work {
+    Node* node;
+    bool leave;  // Are we entering or leaving n?
+  };
+
+  std::vector<Work> stack(start.size());
+  for (int i = 0; i < start.size(); ++i) {
+    stack[i] = Work{start[i], false};
+  }
+
+  std::vector<bool> visited(g->num_node_ids(), false);
+  while (!stack.empty()) {
+    Work w = stack.back();
+    stack.pop_back();
+
+    auto n = w.node;
+    if (w.leave) {
+      if (n == src) {
+        return true;
+      }
+      continue;
+    }
+
+    if (visited[n->id()]) continue;
+    visited[n->id()] = true;
+    // Arrange to call leave(n) when all done with descendants.
+    stack.push_back(Work{n, true});
+
+    auto nodes = n->in_nodes();
+    for (const auto node : nodes) {
+      if (!visited[node->id()]) {
+        stack.push_back(Work{node, false});
+      }
+    }
+  }
+  return false;
+}
+
+bool CanContractEdge(const Edge* edge, const Graph* graph) {
+  const auto src = edge->src();
+  const auto dst = edge->dst();
 
   // Can't contract edge if doing so would cause a cycle in the
   // graph. So, if there is a directed path from 'src' to 'dst', other
@@ -48,46 +88,131 @@ bool CanContractEdge(const tensorflow::Edge* edge,
   //   1. Get all nodes incoming to 'dst', excluding 'src'
   //   2. Reverse DFS from those nodes
   //   3. If reverse DFS reaches 'src' then we have a cycle
-  std::vector<tensorflow::Node*> dfs_start_nodes;
-  for (tensorflow::Node* node : dst->in_nodes()) {
+  std::vector<Node*> dfs_start_nodes;
+  for (Node* node : dst->in_nodes()) {
     if (node != src) {
       dfs_start_nodes.push_back(node);
     }
   }
 
-  bool is_cycle = false;
-  if (!dfs_start_nodes.empty()) {
-    tensorflow::ReverseDFSFrom(graph, dfs_start_nodes, {},
-                               [&is_cycle, src](tensorflow::Node* node) {
-                                 if (node == src) {
-                                   is_cycle = true;
-                                 }
-                               });
-  }
+  bool is_cycle = check_cycles(graph, src, dfs_start_nodes);
+  // if (!dfs_start_nodes.empty()) {
+  //   tensorflow::ReverseDFSFrom(graph, dfs_start_nodes, {},
+  //                              [&is_cycle, src](tensorflow::Node* node) {
+  //                                if (node == src) {
+  //                                  is_cycle = true;
+  //                                }
+  //                              });
+  // }
 
   return !is_cycle;
 }
+}  // namespace
+Node::Node(const tensorflow::Node* node, const int id) : node_(node), id_(id) {
+  if (node_) {
+    in_edges_.reserve(node_->in_edges().size());
+    out_edges_.reserve(node_->out_edges().size());
+  }
+}
+
+Graph::Graph(const tensorflow::Graph* g) : g_(g) {
+  int n_nodes = g_->num_node_ids();
+  nodes_.resize(n_nodes, nullptr);
+  nodes_[g->kSourceId] = new Node(g->source_node(), g->kSourceId);
+  nodes_[g->kSinkId] = new Node(g->sink_node(), g->kSinkId);
+  int n_edges = g->num_edge_ids();
+  edges_.resize(n_edges, nullptr);
+  for (int i = 2; i < n_nodes; i++) {
+    const auto n = g->FindNodeId(i);
+    if (n) {
+      nodes_[i] = new Node(n, i);
+    } else {
+      node_ids_.insert(i);
+    }
+  }
+  for (int i = 0; i < n_edges; i++) {
+    const auto e = g->FindEdgeId(i);
+    if (e) {
+      const auto tfsrc = e->src();
+      const auto tfdst = e->dst();
+      bool is_control = e->IsControlEdge();
+      auto src = nodes_[tfsrc->id()];
+      auto dst = nodes_[tfdst->id()];
+      auto edge =
+          new Edge(i, src, e->src_output(), dst, e->dst_input(), is_control);
+      edges_[i]=edge;
+      src->out_edges_.push_back(edge);
+      dst->in_edges_.push_back(edge);
+    } else {
+      edge_ids_.insert(i);
+    }
+  }
+}
+
+void Graph::AddEdge(Node* src, int out_port, Node* dst, int in_port) {
+  int i = edges_.size();
+  if (edge_ids_.size()) {
+    auto it = edge_ids_.begin();
+    i = *it;
+    edge_ids_.erase(it);
+  } else {
+    edges_.push_back(0);
+  }
+  bool is_control = (out_port == tensorflow::Graph::kControlSlot);
+  is_control |= (in_port == tensorflow::Graph::kControlSlot);
+  auto edge = new Edge(i, src, out_port, dst, in_port, is_control);
+  edges_[i] = edge;
+  src->out_edges_.push_back(edge);
+  dst->in_edges_.push_back(edge);
+}
+
+void Graph::AddControlEdge(Node* src, Node* dst) {
+  AddEdge(src, tensorflow::Graph::kControlSlot, dst,
+          tensorflow::Graph::kControlSlot);
+}
 
-void ContractEdge(tensorflow::Edge* edge, tensorflow::Graph* graph,
-                  std::vector<const tensorflow::Edge*>* remove_edges) {
+void Graph::RemoveEdge(const Edge* edge) {
+  auto src = edge->src();
+  auto dst = edge->dst();
+  for (auto it = src->out_edges_.begin(); it != src->out_edges_.end(); ++it) {
+    if (*it == edge) {
+      src->out_edges_.erase(it);
+      break;
+    }
+  }
+  for (auto it = dst->in_edges_.begin(); it != dst->in_edges_.end(); ++it) {
+    if (*it == edge) {
+      dst->in_edges_.erase(it);
+      break;
+    }
+  }
+}
+
+Graph::~Graph() {
+  for (auto x : nodes_) delete x;
+  for (auto x : edges_) delete x;
+}
+
+void ContractEdge(Edge* edge, Graph* graph,
+                  std::vector<const Edge*>* remove_edges) {
   // Transfer all inputs and outputs of 'dst' to 'src' except edges
   // connecting the two.
-  tensorflow::Node* src = edge->src();
-  tensorflow::Node* dst = edge->dst();
+  auto src = edge->src();
+  auto dst = edge->dst();
 
   // We can use '0' for input/output index because we don't need them
   // to be accurate for the way we are using the graph.
-  std::vector<const tensorflow::Edge*> in_edges(dst->in_edges().begin(),
-                                                dst->in_edges().end());
-  for (const tensorflow::Edge* in_edge : in_edges) {
+  std::vector<const Edge*> in_edges(dst->in_edges().begin(),
+                                    dst->in_edges().end());
+  for (const Edge* in_edge : in_edges) {
     if (in_edge->IsControlEdge()) {
       if (in_edge->src() != src) {
-        tensorflow::Edge* e = const_cast<tensorflow::Edge*>(in_edge);
+        Edge* e = const_cast<Edge*>(in_edge);
         graph->AddControlEdge(e->src(), src);
       }
     } else {
       if (in_edge->src() != src) {
-        tensorflow::Edge* e = const_cast<tensorflow::Edge*>(in_edge);
+        Edge* e = const_cast<Edge*>(in_edge);
         if (e->src() == graph->source_node()) {
           graph->AddEdge(e->src(), e->src_output(), src,
                          tensorflow::Graph::kControlSlot);
@@ -98,14 +223,14 @@ void ContractEdge(tensorflow::Edge* edge, tensorflow::Graph* graph,
     }
   }
 
-  std::vector<const tensorflow::Edge*> out_edges(dst->out_edges().begin(),
-                                                 dst->out_edges().end());
-  for (const tensorflow::Edge* out_edge : out_edges) {
+  std::vector<const Edge*> out_edges(dst->out_edges().begin(),
+                                     dst->out_edges().end());
+  for (const Edge* out_edge : out_edges) {
     if (out_edge->IsControlEdge()) {
-      tensorflow::Edge* e = const_cast<tensorflow::Edge*>(out_edge);
+      Edge* e = const_cast<Edge*>(out_edge);
       graph->AddControlEdge(src, e->dst());
     } else {
-      tensorflow::Edge* e = const_cast<tensorflow::Edge*>(out_edge);
+      Edge* e = const_cast<Edge*>(out_edge);
       if (e->dst() == graph->sink_node()) {
         VLOG(1) << " edge to sink node " << src->name() << " -> "
                 << e->dst()->name();
@@ -128,8 +253,6 @@ void ContractEdge(tensorflow::Edge* edge, tensorflow::Graph* graph,
   }
 }
 
-}  // namespace
-
 tensorflow::Status SegmentGraph(
     const tensorflow::GraphDef& gdef,
     const std::function<bool(const tensorflow::Node*)>& candidate_fn,
@@ -140,17 +263,23 @@ tensorflow::Status SegmentGraph(
   tensorflow::Graph graph(flib);
   TF_RETURN_IF_ERROR(tensorflow::ConvertGraphDefToGraph(
       tensorflow::GraphConstructorOptions(), gdef, &graph));
+  return SegmentGraph(&graph, candidate_fn, options, segments);
+}
 
+tensorflow::Status SegmentGraph(
+    tensorflow::Graph* tf_graph,
+    const std::function<bool(const tensorflow::Node*)>& candidate_fn,
+    const SegmentOptions& options, SegmentNodesVector* segments) {
   // tensorflow::DumpGraph("Pre-Segment", &graph);
-
+  Graph* graph= new Graph(tf_graph);
   // Use a union-find to collect the nodes that belong to the same
-  // segment. A node value of nullptr indicates that the node is not a
-  // candidate for TRT.
-  std::vector<UnionFind<tensorflow::Node*>> node_segments;
-  for (int i = 0; i < graph.num_node_ids(); ++i) {
-    tensorflow::Node* node = graph.FindNodeId(i);
+  // segment. A node value of nullptr indicates that tusing
+  // ::tensorflow::strings::StrAppendhe node is not a candidate for TRT.
+  std::vector<UnionFind<Node*>> node_segments;
+  for (int i = 0; i < graph->num_node_ids(); ++i) {
+    Node* node = graph->FindNodeId(i);
     if (options.exclude_node_list.count(node->name()) != 0 ||
-        !candidate_fn(node)) {
+        !candidate_fn(node->tf_node())) {
       node = nullptr;
     }
     node_segments.emplace_back(node);
@@ -164,10 +293,16 @@ tensorflow::Status SegmentGraph(
   // a measure of how beneficial it is to include a given node in a
   // TRT subgraph then we can revisit this algorithm to take advantage
   // of that information.
-  std::vector<tensorflow::Node*> order;
-  tensorflow::GetPostOrder(graph, &order);
-
-  for (const tensorflow::Node* node : order) {
+  std::vector<tensorflow::Node*> tforder;
+  tensorflow::GetPostOrder(*tf_graph, &tforder);
+  // use postorder implementation from tensorflow and construct mirror in
+  // internal format
+  std::vector<Node*> order;
+  order.reserve(tforder.size());
+  for (const auto tfnode : tforder) {
+    order.push_back(graph->FindNodeId(tfnode->id()));
+  }
+  for (const Node* node : order) {
     // All output nodes of 'node' have been visited...
     VLOG(2) << "Trying node " << node->name() << " id=" << node->id();
 
@@ -181,8 +316,8 @@ tensorflow::Status SegmentGraph(
     // nodes. Iterate since combining two nodes may unblock other
     // combining.
     while (true) {
-      std::set<const tensorflow::Edge*> contract_edges;
-      for (const tensorflow::Edge* out_edge : node->out_edges()) {
+      std::set<const Edge*> contract_edges;
+      for (const Edge* out_edge : node->out_edges()) {
         VLOG(2) << "... out node " << out_edge->dst()->name() << " ( "
                 << out_edge->dst()->id() << " <- " << node->id() << " )";
         if (out_edge->IsControlEdge()) {
@@ -210,9 +345,9 @@ tensorflow::Status SegmentGraph(
       // Contract edges and collect the adjacent nodes into the same
       // segment/subgraph.
       while (!contract_edges.empty()) {
-        const tensorflow::Edge* contract_edge = *contract_edges.begin();
-        const tensorflow::Node* src = contract_edge->src();
-        const tensorflow::Node* dst = contract_edge->dst();
+        const Edge* contract_edge = *contract_edges.begin();
+        const Node* src = contract_edge->src();
+        const Node* dst = contract_edge->dst();
 
         VLOG(2) << "Merge " << src->name() << " <- " << dst->name() << " ("
                 << src->id() << " <- " << dst->id();
@@ -221,13 +356,13 @@ tensorflow::Status SegmentGraph(
         // Contracting the edge leaves disconnected graph edges.
         // Remove these from the graph and from 'contract_edges' so we
         // don't visit them again.
-        tensorflow::Edge* e = const_cast<tensorflow::Edge*>(contract_edge);
-        std::vector<const tensorflow::Edge*> remove_edges;
-        ContractEdge(e, &graph, &remove_edges);
+        Edge* e = const_cast<Edge*>(contract_edge);
+        std::vector<const Edge*> remove_edges;
+        ContractEdge(e, graph, &remove_edges);
 
-        for (const tensorflow::Edge* r : remove_edges) {
+        for (const Edge* r : remove_edges) {
           contract_edges.erase(r);
-          graph.RemoveEdge(r);
+          graph->RemoveEdge(r);
         }
       }
     }
@@ -236,9 +371,22 @@ tensorflow::Status SegmentGraph(
   // Collect the segments/subgraphs. Each subgraph is represented by a
   // set of the names of the nodes in that subgraph.
   std::unordered_map<string, std::set<string>> sg_map;
+  std::unordered_map<string, std::set<string>> device_maps;
   for (auto& u : node_segments) {
     if ((u.Value() != nullptr) && (u.ParentValue() != nullptr)) {
       sg_map[u.ParentValue()->name()].insert(u.Value()->name());
+      auto tf_node = u.Value()->tf_node();
+      if (tf_node->has_assigned_device_name()) {
+        device_maps[u.ParentValue()->name()].insert(
+            tf_node->assigned_device_name());
+      } else if (tf_node->requested_device().size() > 0) {
+        device_maps[u.ParentValue()->name()].insert(
+            tf_node->requested_device());
+      } else {
+        VLOG(1) << "Node " << tf_node->name()
+                << " has no device assigned requested device is: "
+                << tf_node->requested_device();
+      }
     }
   }
 
@@ -260,10 +408,33 @@ tensorflow::Status SegmentGraph(
               << segment_node_names.size() << " nodes, dropping";
       continue;
     }
-
-    segments->emplace_back(segment_node_names);
+    const auto& dev_itr = device_maps.find(itr.first);
+    if (dev_itr == device_maps.end() || dev_itr->second.size() == 0) {
+      VLOG(1) << "No device assigned to segment " << segments->size();
+      segments->emplace_back(std::make_pair(segment_node_names, string()));
+    } else if (dev_itr->second.size() > 1) {
+      string s("Segment ");
+      StrAppend(&s, segments->size(), " has multiple devices attached: ");
+      for (const auto& dev : dev_itr->second) {
+        StrAppend(&s, dev, ", ");
+      }
+      LOG(WARNING) << s << " choosing " << *(dev_itr->second.begin());
+      segments->emplace_back(
+          std::make_pair(segment_node_names, *(dev_itr->second.begin())));
+    } else {
+      segments->emplace_back(
+          std::make_pair(segment_node_names, *(dev_itr->second.begin())));
+    }
   }
-
+  for (const auto& d : device_maps) {
+    string s("Segment ");
+    StrAppend(&s, ": '", d.first, "' ");
+    for (const auto& dd : d.second) {
+      StrAppend(&s, dd, ", ");
+    }
+    VLOG(1) << "Devices " << s;
+  }
+  delete graph;
   return tensorflow::Status::OK();
 }
 
diff --git a/tensorflow/contrib/tensorrt/segment/segment.h b/tensorflow/contrib/tensorrt/segment/segment.h
index 7e8685f44a..659fea1859 100644
--- a/tensorflow/contrib/tensorrt/segment/segment.h
+++ b/tensorflow/contrib/tensorrt/segment/segment.h
@@ -29,25 +29,116 @@ namespace tensorflow {
 namespace tensorrt {
 namespace segment {
 
-using SegmentNodesVector = std::vector<std::set<string>>;
+using SegmentNodesVector = std::vector<std::pair<std::set<string>, string>>;
+class Node;
+class Graph;
+class Edge {
+ public:
+  Edge(int id, Node* src, int src_port, Node* dst, int dst_port,
+       bool is_control = false)
+      : id_(id),
+        src_(src),
+        src_port_(src_port),
+        dst_(dst),
+        dst_port_(dst_port),
+        control_(is_control){};
+  Node* src() const { return src_; }
+  Node* dst() const { return dst_; }
+  int src_output() const { return src_port_; }
+  int dst_input() const { return dst_port_; }
+  int id() const { return id_; }
+  bool IsControlEdge() const { return control_; }
+  ~Edge() {}
 
+ private:
+  int id_;
+  Node* src_;
+  int src_port_;
+  Node* dst_;
+  int dst_port_;
+  bool control_;
+};
+class Node {
+  friend class Graph;
+
+ public:
+  Node(const tensorflow::Node* node, const int id);
+  const std::vector<Edge*>& in_edges() const { return in_edges_; };
+  const std::vector<Edge*>& out_edges() const { return out_edges_; };
+  std::vector<Node*> in_nodes() const {
+    std::vector<Node*> res;
+    res.reserve(in_edges_.size());
+    for (const auto e : in_edges_) {
+      if (e) res.push_back(e->src());
+    }
+    return res;
+  }
+  const string& name() const { return node_->name(); }
+  const tensorflow::Node* tf_node() const { return node_; }
+  int id() const { return id_; }
+
+ private:
+  const tensorflow::Node* node_;
+  std::vector<Edge*> in_edges_;
+  std::vector<Edge*> out_edges_;
+  int id_;
+};
+
+class Graph {
+ public:
+  Graph(const tensorflow::Graph* g);
+  void AddControlEdge(Node* src, Node* dst);
+  void AddEdge(Node* src, int out_port, Node* dst, int in_port);
+  void RemoveEdge(const Edge*);
+  Node* FindNodeId(int node_id) {
+    if (node_id < 0 || node_id > (int)nodes_.size()) return nullptr;
+    return nodes_[node_id];
+  }
+  ~Graph();
+  int num_node_ids() const { return nodes_.size(); }
+  const Node* source_node() const {
+    return nodes_[tensorflow::Graph::kSourceId];
+  }
+  const Node* sink_node() const { return nodes_[tensorflow::Graph::kSinkId]; }
+
+ private:
+  const tensorflow::Graph* g_;
+  std::vector<Node*> nodes_;
+  std::vector<Edge*> edges_;
+  std::set<int> edge_ids_;
+  std::set<int> node_ids_;
+};
 struct SegmentOptions {
   // Segment must contain at least this many nodes.
   int minimum_segment_size = 2;
   std::set<string> exclude_node_list;
 };
 
+// // Get the subgraphs of a graph that can be handled by TensorRT.
+// //
+// // @param gdef The GraphDef describing the network
+// // @param candidate_fn A function that returns true for a NodeDef if
+// // that node can be handled by TensorRT.
+// // @param segments Returns the TensorRT segments/subgraphs. Each entry
+// // in the vector describes a subgraph by giving a set of the names of
+// // all the NodeDefs in that subgraph.
+// // @return the status.
+tensorflow::Status SegmentGraph(
+    const tensorflow::GraphDef& gdef,
+    const std::function<bool(const tensorflow::Node*)>& candidate_fn,
+    const SegmentOptions& options, SegmentNodesVector* segments);
+
 // Get the subgraphs of a graph that can be handled by TensorRT.
 //
-// @param gdef The GraphDef describing the network
-// @param candidate_fn A function that returns true for a NodeDef if
+// @param graph tensorflow::Graph of the network
+// @param candidate_fn A function that returns true for a Node* if
 // that node can be handled by TensorRT.
 // @param segments Returns the TensorRT segments/subgraphs. Each entry
 // in the vector describes a subgraph by giving a set of the names of
 // all the NodeDefs in that subgraph.
 // @return the status.
 tensorflow::Status SegmentGraph(
-    const tensorflow::GraphDef& gdef,
+    tensorflow::Graph* graph,
     const std::function<bool(const tensorflow::Node*)>& candidate_fn,
     const SegmentOptions& options, SegmentNodesVector* segments);
 
diff --git a/tensorflow/contrib/tensorrt/segment/segment_test.cc b/tensorflow/contrib/tensorrt/segment/segment_test.cc
index 7ddabec268..7fe824b12f 100644
--- a/tensorflow/contrib/tensorrt/segment/segment_test.cc
+++ b/tensorflow/contrib/tensorrt/segment/segment_test.cc
@@ -35,7 +35,7 @@ class SegmentTest : public ::testing::Test {
   TF_Operation* Add(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
                     TF_Status* s, const char* name);
 
-  std::function<bool(const Node*)> MakeCandidateFn(
+  std::function<bool(const tensorflow::Node*)> MakeCandidateFn(
       const std::set<string>& node_names);
 
  protected:
@@ -60,9 +60,9 @@ bool SegmentTest::GetGraphDef(TF_Graph* graph,
   return ret;
 }
 
-std::function<bool(const Node*)> SegmentTest::MakeCandidateFn(
+std::function<bool(const tensorflow::Node*)> SegmentTest::MakeCandidateFn(
     const std::set<string>& node_names) {
-  return [node_names](const Node* node) -> bool {
+  return [node_names](const tensorflow::Node* node) -> bool {
     return node_names.find(node->name()) != node_names.end();
   };
 }
diff --git a/tensorflow/contrib/tensorrt/test/test_tftrt.py b/tensorflow/contrib/tensorrt/test/test_tftrt.py
index ad01bedd8f..aaaed0c30f 100644
--- a/tensorflow/contrib/tensorrt/test/test_tftrt.py
+++ b/tensorflow/contrib/tensorrt/test/test_tftrt.py
@@ -18,7 +18,9 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import argparse
 import numpy as np
+
 # normally we should do import tensorflow as tf and then
 # tf.placeholder, tf.constant, tf.nn.conv2d etc but
 # it looks like internal builds don't like it so
@@ -26,6 +28,7 @@ import numpy as np
 
 from tensorflow.contrib import tensorrt as trt
 from tensorflow.core.protobuf import config_pb2 as cpb2
+from tensorflow.core.protobuf import rewriter_config_pb2 as rwpb2
 from tensorflow.python.client import session as csess
 from tensorflow.python.framework import constant_op as cop
 from tensorflow.python.framework import dtypes as dtypes
@@ -59,9 +62,12 @@ def get_simple_graph_def():
   return g.as_graph_def()
 
 
-def run_graph(gdef, dumm_inp):
+def execute_graph(gdef, dumm_inp):
   """Run given graphdef once."""
+  print("executing")
   gpu_options = cpb2.GPUOptions(per_process_gpu_memory_fraction=0.50)
+  #graph_options = cpb2.GraphOptions(rewrite_options=opt_config)
+  sessconfig = cpb2.ConfigProto(gpu_options=gpu_options)
   ops.reset_default_graph()
   g = ops.Graph()
   with g.as_default():
@@ -69,15 +75,18 @@ def run_graph(gdef, dumm_inp):
         graph_def=gdef, return_elements=["input", "output"])
     inp = inp.outputs[0]
     out = out.outputs[0]
+  # with csess.Session(
+  #     config=cpb2.ConfigProto(gpu_options=gpu_options), graph=g) as sess:
+  #   val = sess.run(out, {inp: dumm_inp})
   with csess.Session(
-      config=cpb2.ConfigProto(gpu_options=gpu_options), graph=g) as sess:
+      config=sessconfig, graph=g) as sess:
     val = sess.run(out, {inp: dumm_inp})
   return val
 
 
 # Use real data that is representative of the inference dataset
 # for calibration. For this test script it is random data.
-def run_calibration(gdef, dumm_inp):
+def execute_calibration(gdef, dumm_inp):
   """Run given calibration graph multiple times."""
   gpu_options = cpb2.GPUOptions(per_process_gpu_memory_fraction=0.50)
   ops.reset_default_graph()
@@ -96,7 +105,9 @@ def run_calibration(gdef, dumm_inp):
   return val
 
 
-if "__main__" in __name__:
+def user(run_graph=execute_graph, run_calibration=execute_calibration):
+  """ Example function that converts a graph to TFTRT graph """
+
   inp_dims = (100, 24, 24, 2)
   dummy_input = np.random.random_sample(inp_dims)
   orig_graph = get_simple_graph_def()  # use a frozen graph for inference
@@ -137,3 +148,44 @@ if "__main__" in __name__:
   assert np.allclose(o1, o4)
   assert np.allclose(o1, o5)
   print("Pass")
+
+def auto():
+  """ Run the conversion as an optimization pass"""
+  inp_dims = (100, 24, 24, 2)
+  dummy_input = np.random.random_sample(inp_dims)
+  orig_graph = get_simple_graph_def()
+  opt_config = rwpb2.RewriterConfig()
+  opt_config.optimizers.extend(["constfold", "layout"])
+  custom_op = opt_config.custom_optimizers.add()
+  custom_op.name = "TensorRTOptimizer"
+  custom_op.parameter_map["minimum_segment_size"].i = 3
+  custom_op.parameter_map["precision_mode"].s = "FP32"
+  custom_op.parameter_map["max_batch_size"].i = inp_dims[0]
+  custom_op.parameter_map["max_workspace_size_bytes"].i = 1 << 25
+  print(custom_op)
+  gpu_options = cpb2.GPUOptions(per_process_gpu_memory_fraction=0.50)
+  graph_options = cpb2.GraphOptions(rewrite_options=opt_config)
+  sessconfig = cpb2.ConfigProto(gpu_options=gpu_options,
+                                graph_options=graph_options)
+  print(sessconfig)
+  g = ops.Graph()
+  ops.reset_default_graph()
+  with g.as_default():
+    inp, out = importer.import_graph_def(
+        graph_def=orig_graph, return_elements=["input", "output"])
+    inp = inp.outputs[0]
+    out = out.outputs[0]
+    with csess.Session(config=sessconfig, graph=g) as sess:
+      val = sess.run(out, {inp: dummy_input})
+  print(val.shape)
+
+if "__main__" in __name__:
+  P = argparse.ArgumentParser(prog="tftrt_test",
+    description="Example utilization of TensorFlow-TensorRT integration")
+  P.add_argument("--automatic", "-a", action="store_true",
+                 help="Do TRT conversion automatically", default=False)
+  flags, unparsed = P.parse_known_args()
+  if flags.automatic:
+    auto()
+  else:
+    user()
diff --git a/tensorflow/tools/pip_package/build_pip_package.sh b/tensorflow/tools/pip_package/build_pip_package.sh
index 8f0cf8c3d1..3af79ee170 100755
--- a/tensorflow/tools/pip_package/build_pip_package.sh
+++ b/tensorflow/tools/pip_package/build_pip_package.sh
@@ -24,7 +24,7 @@ function real_path() {
 function cp_external() {
   local src_dir=$1
   local dest_dir=$2
-  for f in `find "$src_dir" -maxdepth 1 -mindepth 1 ! -name '*local_config_cuda*' ! -name '*org_tensorflow*'`; do
+  for f in `find "$src_dir" -maxdepth 1 -mindepth 1 ! -name '*local_config_cuda*' ! -name '*local_config_tensorrt*' ! -name '*org_tensorflow*'`; do
     cp -R "$f" "$dest_dir"
   done
   mkdir -p "${dest_dir}/local_config_cuda/cuda/cuda/"
-- 
GitLab


From 81a34fb835e8389dd2523335c5d186405294f95e Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Fri, 27 Apr 2018 02:21:44 +0300
Subject: [PATCH 020/395] [tf.data] Just replace old resample with new.

Also, add an optimization / bug fix that shortcircuits combining the two datasets if one should always be sampled from.

Tested:

bazel test :resample_test
---
 .../data/python/kernel_tests/resample_test.py |  85 ++++-----
 .../contrib/data/python/ops/resampling.py     | 178 +++++++-----------
 2 files changed, 107 insertions(+), 156 deletions(-)

diff --git a/tensorflow/contrib/data/python/kernel_tests/resample_test.py b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
index 7f007fede8..fc84301b17 100644
--- a/tensorflow/contrib/data/python/kernel_tests/resample_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
@@ -34,14 +34,12 @@ from tensorflow.python.util import compat
 
 
 def _time_resampling(
-    test_obj, data_np, target_dist, init_dist, use_v2, num_to_sample):
+    test_obj, data_np, target_dist, init_dist, num_to_sample):
   dataset = dataset_ops.Dataset.from_tensor_slices(data_np).repeat()
 
   # Reshape distribution via rejection sampling.
-  apply_fn = (resampling.rejection_resample_v2 if use_v2 else
-              resampling.rejection_resample)
   dataset = dataset.apply(
-      apply_fn(
+      resampling.rejection_resample(
           class_func=lambda x: x,
           target_dist=target_dist,
           initial_dist=init_dist,
@@ -61,20 +59,17 @@ def _time_resampling(
 class ResampleTest(test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(
-      ("InitialnDistributionKnown", True, False),
-      ("InitialDistributionUnknown", False, False),
-      ("InitialDistributionKnownV2", True, True),
-      ("InitialDistributionUnknownV2", False, True))
-  def testDistribution(self, initial_known, use_v2):
+      ("InitialnDistributionKnown", True),
+      ("InitialDistributionUnknown", False))
+  def testDistribution(self, initial_known):
     classes = np.random.randint(5, size=(20000,))  # Uniformly sampled
     target_dist = [0.9, 0.05, 0.05, 0.0, 0.0]
     initial_dist = [0.2] * 5 if initial_known else None
     dataset = dataset_ops.Dataset.from_tensor_slices(classes).shuffle(
         200, seed=21).map(lambda c: (c, string_ops.as_string(c))).repeat()
-    apply_fn = (resampling.rejection_resample_v2 if use_v2 else
-                resampling.rejection_resample)
+
     get_next = dataset.apply(
-        apply_fn(
+        resampling.rejection_resample(
             target_dist=target_dist,
             initial_dist=initial_dist,
             class_func=lambda c, _: c,
@@ -96,11 +91,39 @@ class ResampleTest(test.TestCase, parameterized.TestCase):
     returned_dist = class_counts / total_returned
     self.assertAllClose(target_dist, returned_dist, atol=1e-2)
 
+  @parameterized.named_parameters(
+      ("OnlyInitial", True),
+      ("NotInitial", False))
+  def testEdgeCasesSampleFromInitialDataset(self, only_initial_dist):
+    init_dist = [0.5, 0.5]
+    target_dist = [0.5, 0.5] if only_initial_dist else [0.0, 1.0]
+    num_classes = len(init_dist)
+    # We don't need many samples to test that this works.
+    num_samples = 100
+    data_np = np.random.choice(num_classes, num_samples, p=init_dist)
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(data_np)
+
+    # Reshape distribution.
+    dataset = dataset.apply(
+        resampling.rejection_resample(
+            class_func=lambda x: x,
+            target_dist=target_dist,
+            initial_dist=init_dist))
+
+    get_next = dataset.make_one_shot_iterator().get_next()
+
+    with self.test_session() as sess:
+      returned = []
+      with self.assertRaises(errors.OutOfRangeError):
+        while True:
+          returned.append(sess.run(get_next))
+
   def testRandomClasses(self):
     init_dist = [0.25, 0.25, 0.25, 0.25]
     target_dist = [0.0, 0.0, 0.0, 1.0]
     num_classes = len(init_dist)
-    # We don't need many samples to test a dirac-delta target distribution
+    # We don't need many samples to test a dirac-delta target distribution.
     num_samples = 100
     data_np = np.random.choice(num_classes, num_samples, p=init_dist)
 
@@ -134,26 +157,8 @@ class ResampleTest(test.TestCase, parameterized.TestCase):
 
     self.assertAllClose(target_dist, bincount, atol=1e-2)
 
-  @parameterized.named_parameters(
-      ("SmallSkewManySamples", [0.1, 0.1, 0.1, 0.7], 1000),
-      ("BigSkewManySamples", [0.01, 0.01, 0.01, 0.97], 1000),
-      ("SmallSkewFewSamples", [0.1, 0.1, 0.1, 0.7], 100),
-      ("BigSkewFewSamples", [0.01, 0.01, 0.01, 0.97], 100))
-  def testNewResampleIsFaster(self, target_dist, num_to_sample):
-    init_dist = [0.25, 0.25, 0.25, 0.25]
-    num_classes = len(init_dist)
-    num_samples = 1000
-    data_np = np.random.choice(num_classes, num_samples, p=init_dist)
-
-    fast_time = _time_resampling(self, data_np, target_dist, init_dist,
-                                 use_v2=True, num_to_sample=num_to_sample)
-    slow_time = _time_resampling(self, data_np, target_dist, init_dist,
-                                 use_v2=False, num_to_sample=num_to_sample)
-
-    self.assertLess(fast_time, slow_time)
-
 
-class MapDatasetBenchmark(test.Benchmark):
+class ResampleDatasetBenchmark(test.Benchmark):
 
   def benchmarkResamplePerformance(self):
     init_dist = [0.25, 0.25, 0.25, 0.25]
@@ -164,25 +169,11 @@ class MapDatasetBenchmark(test.Benchmark):
     data_np = np.random.choice(num_classes, num_samples, p=init_dist)
 
     resample_time = _time_resampling(
-        self, data_np, target_dist, init_dist, use_v2=False, num_to_sample=1000)
+        self, data_np, target_dist, init_dist, num_to_sample=1000)
 
     self.report_benchmark(
         iters=1000, wall_time=resample_time, name="benchmark_resample")
 
-  def benchmarkResampleAndBatchPerformance(self):
-    init_dist = [0.25, 0.25, 0.25, 0.25]
-    target_dist = [0.0, 0.0, 0.0, 1.0]
-    num_classes = len(init_dist)
-    # We don't need many samples to test a dirac-delta target distribution
-    num_samples = 1000
-    data_np = np.random.choice(num_classes, num_samples, p=init_dist)
-
-    resample_time = _time_resampling(
-        self, data_np, target_dist, init_dist, use_v2=True, num_to_sample=1000)
-
-    self.report_benchmark(
-        iters=1000, wall_time=resample_time, name="benchmark_resample_v2")
-
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index 16d851bf96..66eaf9b69a 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -58,62 +58,7 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
 
     # Get initial distribution.
     if initial_dist is not None:
-      initial_dist_t = ops.convert_to_tensor(
-          initial_dist, name="initial_dist")
-      acceptance_dist = _calculate_acceptance_probs(initial_dist_t,
-                                                    target_dist_t)
-      initial_dist_ds = dataset_ops.Dataset.from_tensors(
-          initial_dist_t).repeat()
-      acceptance_dist_ds = dataset_ops.Dataset.from_tensors(
-          acceptance_dist).repeat()
-    else:
-      initial_dist_ds = _estimate_initial_dist_ds(
-          target_dist_t, class_values_ds)
-      acceptance_dist_ds = initial_dist_ds.map(
-          lambda initial: _calculate_acceptance_probs(initial, target_dist_t))
-    return _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds,
-                      class_values_ds, seed)
-
-  return _apply_fn
-
-
-def rejection_resample_v2(class_func, target_dist, initial_dist=None,
-                          seed=None):
-  """A transformation that resamples a dataset to achieve a target distribution.
-
-  This differs from v1 in that it will also sample from the original dataset
-  with some probability, so it makes strictly fewer data rejections. Due to an
-  implementation detail it must initialize a separate dataset initializer, so
-  the dataset becomes stateful after this transformation is applied
-  (`make_one_shot_iterator` won't work; users must use
-  `make_initializable_iterator`). This transformation is faster than the
-  original, except for overhead.
-
-  **NOTE** Resampling is performed via rejection sampling; some fraction
-  of the input values will be dropped.
-
-  Args:
-    class_func: A function mapping an element of the input dataset to a scalar
-      `tf.int32` tensor. Values should be in `[0, num_classes)`.
-    target_dist: A floating point type tensor, shaped `[num_classes]`.
-    initial_dist: (Optional.)  A floating point type tensor, shaped
-      `[num_classes]`.  If not provided, the true class distribution is
-      estimated live in a streaming fashion.
-    seed: (Optional.) Python integer seed for the resampler.
-
-  Returns:
-    A `Dataset` transformation function, which can be passed to
-    @{tf.data.Dataset.apply}.
-  """
-  def _apply_fn(dataset):
-    """Function from `Dataset` to `Dataset` that applies the transformation."""
-    target_dist_t = ops.convert_to_tensor(target_dist, name="target_dist")
-    class_values_ds = dataset.map(class_func)
-
-    # Get initial distribution.
-    if initial_dist is not None:
-      initial_dist_t = ops.convert_to_tensor(
-          initial_dist, name="initial_dist")
+      initial_dist_t = ops.convert_to_tensor(initial_dist, name="initial_dist")
       acceptance_dist, prob_of_original = (
           _calculate_acceptance_probs_with_mixing(initial_dist_t,
                                                   target_dist_t))
@@ -133,19 +78,51 @@ def rejection_resample_v2(class_func, target_dist, initial_dist=None,
           lambda accept_prob, _: accept_prob)
       prob_of_original_ds = acceptance_and_original_prob_ds.map(
           lambda _, prob_original: prob_original)
+      prob_of_original = None
     filtered_ds = _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds,
                              class_values_ds, seed)
     # Prefetch filtered dataset for speed.
     filtered_ds = filtered_ds.prefetch(3)
 
-    return interleave_ops.sample_from_datasets(
-        [dataset_ops.Dataset.zip((class_values_ds, dataset)), filtered_ds],
-        weights=prob_of_original_ds.map(lambda prob: [(prob, 1.0 - prob)]),
-        seed=seed)
+    prob_original_static = _get_prob_original_static(
+        initial_dist, target_dist_t) if initial_dist is not None else None
+    if prob_original_static == 1:
+      return dataset_ops.Dataset.zip((class_values_ds, dataset))
+    elif prob_original_static == 0:
+      return filtered_ds
+    else:
+      return interleave_ops.sample_from_datasets(
+          [dataset_ops.Dataset.zip((class_values_ds, dataset)), filtered_ds],
+          weights=prob_of_original_ds.map(lambda prob: [(prob, 1.0 - prob)]),
+          seed=seed)
 
   return _apply_fn
 
 
+def _get_prob_original_static(initial_dist_t, target_dist_t):
+  """Returns the static probability of sampling from the original.
+
+  For some reason, `tensor_util.constant_value(prob_of_original)` of a ratio
+  of two constant Tensors isn't a constant. We have some custom logic to avoid
+  this.
+
+  Args:
+    initial_dist_t: A tensor of the initial distribution.
+    target_dist_t: A tensor of the target distribution.
+
+  Returns:
+    The probability of sampling from the original distribution as a constant,
+    if it is a constant, or `None`.
+  """
+  init_static = tensor_util.constant_value(initial_dist_t)
+  target_static = tensor_util.constant_value(target_dist_t)
+
+  if init_static is None or target_static is None:
+    return None
+  else:
+    return np.min(target_static / init_static)
+
+
 def _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds, class_values_ds,
                seed):
   """Filters a dataset based on per-class acceptance probabilities.
@@ -216,16 +193,42 @@ def _get_target_to_initial_ratio(initial_probs, target_probs):
   return target_probs / denom
 
 
-def _calculate_acceptance_probs(initial_probs, target_probs):
-  """Calculate the per-class acceptance rates.
+def _estimate_data_distribution(c, num_examples_per_class_seen):
+  """Estimate data distribution as labels are seen.
 
   Args:
-    initial_probs: The class probabilities of the data.
-    target_probs: The desired class proportion in minibatches.
+    c: The class labels.  Type `int32`, shape `[batch_size]`.
+    num_examples_per_class_seen: Type `int64`, shape `[num_classes]`,
+      containing counts.
+
   Returns:
-    A list of the per-class acceptance probabilities.
+    num_examples_per_lass_seen: Updated counts.  Type `int64`, shape
+      `[num_classes]`.
+    dist: The updated distribution.  Type `float32`, shape `[num_classes]`.
+  """
+  num_classes = num_examples_per_class_seen.get_shape()[0].value
+  # Update the class-count based on what labels are seen in batch.
+  num_examples_per_class_seen = math_ops.add(
+      num_examples_per_class_seen, math_ops.reduce_sum(
+          array_ops.one_hot(c, num_classes, dtype=dtypes.int64), 0))
+  init_prob_estimate = math_ops.truediv(
+      num_examples_per_class_seen,
+      math_ops.reduce_sum(num_examples_per_class_seen))
+  dist = math_ops.cast(init_prob_estimate, dtypes.float32)
+  return num_examples_per_class_seen, dist
 
-  This method is based on solving the following analysis:
+
+def _calculate_acceptance_probs_with_mixing(initial_probs, target_probs):
+  """Calculates the acceptance probabilities and mixing ratio.
+
+  In this case, we assume that we can *either* sample from the original data
+  distribution with probability `m`, or sample from a reshaped distribution
+  that comes from rejection sampling on the original distribution. This
+  rejection sampling is done on a per-class basis, with `a_i` representing the
+  probability of accepting data from class `i`.
+
+  This method is based on solving the following analysis for the reshaped
+  distribution:
 
   Let F be the probability of a rejection (on any example).
   Let p_i be the proportion of examples in the data in class i (init_probs)
@@ -256,47 +259,6 @@ def _calculate_acceptance_probs(initial_probs, target_probs):
 
   A solution for a_i in terms of the other variables is the following:
     ```a_i = (t_i / p_i) / max_i[t_i / p_i]```
-  """
-  ratio_l = _get_target_to_initial_ratio(initial_probs, target_probs)
-
-  # Calculate list of acceptance probabilities.
-  max_ratio = math_ops.reduce_max(ratio_l)
-  return ratio_l / max_ratio
-
-
-def _estimate_data_distribution(c, num_examples_per_class_seen):
-  """Estimate data distribution as labels are seen.
-
-  Args:
-    c: The class labels.  Type `int32`, shape `[batch_size]`.
-    num_examples_per_class_seen: Type `int64`, shape `[num_classes]`,
-      containing counts.
-
-  Returns:
-    num_examples_per_lass_seen: Updated counts.  Type `int64`, shape
-      `[num_classes]`.
-    dist: The updated distribution.  Type `float32`, shape `[num_classes]`.
-  """
-  num_classes = num_examples_per_class_seen.get_shape()[0].value
-  # Update the class-count based on what labels are seen in batch.
-  num_examples_per_class_seen = math_ops.add(
-      num_examples_per_class_seen, math_ops.reduce_sum(
-          array_ops.one_hot(c, num_classes, dtype=dtypes.int64), 0))
-  init_prob_estimate = math_ops.truediv(
-      num_examples_per_class_seen,
-      math_ops.reduce_sum(num_examples_per_class_seen))
-  dist = math_ops.cast(init_prob_estimate, dtypes.float32)
-  return num_examples_per_class_seen, dist
-
-
-def _calculate_acceptance_probs_with_mixing(initial_probs, target_probs):
-  """Calculates the acceptance probabilities and mixing ratio.
-
-  In this case, we assume that we can *either* sample from the original data
-  distribution with probability `m`, or sample from a reshaped distribution
-  that comes from rejection sampling on the original distribution. This
-  rejection sampling is done on a per-class basis, with `a_i` representing the
-  probability of accepting data from class `i`.
 
   If we try to minimize the amount of data rejected, we get the following:
 
@@ -312,8 +274,6 @@ def _calculate_acceptance_probs_with_mixing(initial_probs, target_probs):
 
   m = M_min
 
-  See the docstring for `_calculate_acceptance_probs` for more details.
-
   Args:
     initial_probs: A Tensor of the initial probability distribution, given or
       estimated.
-- 
GitLab


From a13d0e527941f6affeeb8155a819a93f8b4ee0ba Mon Sep 17 00:00:00 2001
From: Sami Kama <skama@nvidia.com>
Date: Thu, 26 Apr 2018 16:34:59 -0700
Subject: [PATCH 021/395] Clang-format and version fix

---
 .../contrib/tensorrt/convert/convert_graph.cc |  3 ++-
 .../contrib/tensorrt/convert/convert_nodes.cc | 10 ++++-----
 .../contrib/tensorrt/convert/convert_nodes.h  |  1 -
 .../tensorrt/convert/trt_optimization_pass.cc |  4 ++--
 .../contrib/tensorrt/kernels/trt_engine_op.cc | 21 ++++++++++--------
 .../contrib/tensorrt/kernels/trt_engine_op.h  |  3 ++-
 .../contrib/tensorrt/test/test_tftrt.py       | 22 ++++++++++++-------
 .../tensorrt/test/tf_trt_integration_test.py  | 19 +++++-----------
 8 files changed, 43 insertions(+), 40 deletions(-)

diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.cc b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
index 785c33c4c4..b40a45ee78 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_graph.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
@@ -449,7 +449,8 @@ tensorflow::Status ConvertAfterShapes(
       Status s = GpuIdManager::TfToCudaGpuId(tf_gpu_id, &cuda_gpu_id);
       if (!s.ok()) {
         LOG(ERROR)
-            << "Cuda device identification failed, using device 0. Error= " << s;
+            << "Cuda device identification failed, using device 0. Error= "
+            << s;
       } else {
         cuda_device_id = cuda_gpu_id.value();
       }
diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
index b37c535736..8ed0ed7b7e 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
@@ -2249,8 +2249,8 @@ tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) {
   op_res->logger_ = new tensorflow::tensorrt::Logger();
   cudaSetDevice(s.cuda_device_id_);
   op_res->builder_ = nvinfer1::createInferBuilder(*(op_res->logger_));
-  op_res->allocator_=s.allocator_;
-#if NV_TENSORRT_MAJOR >4 
+  op_res->allocator_ = s.allocator_;
+#if NV_TENSORRT_MAJOR > 3
   op_res->builder_->setGpuAllocator(s.allocator_.get());
 #endif
   if (!op_res->builder_) {
@@ -2481,13 +2481,13 @@ tensorflow::Status ConvertSubGraphToTensorRTNodeDef(
   // Topological order is needed to build TRT network
 
   tensorflow::tensorrt::Logger trt_logger;
-cudaSetDevice(s.cuda_device_id_);
+  cudaSetDevice(s.cuda_device_id_);
   auto trt_builder = infer_object(nvinfer1::createInferBuilder(trt_logger));
   if (!trt_builder) {
     return tensorflow::errors::Internal(
         "Failed to create TensorRT builder object");
   }
-#if NV_TENSORRT_MAJOR >3 
+#if NV_TENSORRT_MAJOR > 3
   trt_builder->setGpuAllocator(s.allocator_.get());
 #endif
   auto trt_network = infer_object(trt_builder->createNetwork());
@@ -2718,7 +2718,7 @@ cudaSetDevice(s.cuda_device_id_);
                     .Finalize(s.trt_node);
 
   VLOG(0) << status.ToString() << " finished op building for " << engine_name
-          << " on device " << s.device_name_ ;
+          << " on device " << s.device_name_;
 
   return tensorflow::Status::OK();
 }
diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.h b/tensorflow/contrib/tensorrt/convert/convert_nodes.h
index ecccaf36e3..8e1d7c99b6 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_nodes.h
+++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.h
@@ -22,7 +22,6 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
-
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/grappler/costs/graph_properties.h"
diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
index 880ffe1b3a..5c08d5afdf 100644
--- a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
+++ b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
@@ -189,8 +189,8 @@ tensorflow::Status TRTOptimizationPass::Optimize(
     const auto& pname = dev->parsed_name();
     VLOG(1) << "Device name= " << dev->name()
             << " parsedname job= " << pname.job << " id= " << pname.id
-            << " has_id: " << pname.has_id << " has_job: " << pname.has_job<< 
-            "has_type: "<<pname.has_type<<" type ="<<pname.type;
+            << " has_id: " << pname.has_id << " has_job: " << pname.has_job
+            << "has_type: " << pname.has_type << " type =" << pname.type;
   }
   auto status = tensorflow::tensorrt::convert::ConvertAfterShapes(
       item.graph, item.fetch, maximum_batch_size_, maximum_workspace_size_,
diff --git a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
index 9c59fd973b..f05c2a1d22 100644
--- a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
+++ b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
@@ -15,12 +15,12 @@ limitations under the License.
 #include "tensorflow/contrib/tensorrt/kernels/trt_engine_op.h"
 
 #include "tensorflow/contrib/tensorrt/log/trt_logger.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h"
+#include "tensorflow/core/common_runtime/gpu/process_state.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/stream_executor.h"
 #include "tensorflow/core/platform/types.h"
-#include "tensorflow/core/common_runtime/gpu/process_state.h"
-#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
-#include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h"
 
 #if GOOGLE_CUDA
 #if GOOGLE_TENSORRT
@@ -67,17 +67,20 @@ TRTEngineOp::TRTEngineOp(OpKernelConstruction* context) : OpKernel(context) {
 }
 
 void TRTEngineOp::Compute(OpKernelContext* context) {
-  if(!trt_execution_context_ptr_){
-    tensorflow::TfGpuId tf_gpu_id(context->device()->tensorflow_gpu_device_info()->gpu_id);
+  if (!trt_execution_context_ptr_) {
+    IRuntime* infer = nvinfer1::createInferRuntime(logger);
+#if NV_TENSORRT_MAJOR > 3
+    tensorflow::TfGpuId tf_gpu_id(
+        context->device()->tensorflow_gpu_device_info()->gpu_id);
     tensorflow::GPUOptions gpuoptions;
     auto pm = tensorflow::ProcessState::singleton();
     auto dev_allocator = pm->GetGPUAllocator(gpuoptions, tf_gpu_id, 1);
-    IRuntime* infer = nvinfer1::createInferRuntime(logger);
-    if(!dev_allocator){
-      LOG(FATAL)<<"Can't find device allocator for gpu device"<<tf_gpu_id;
+    if (!dev_allocator) {
+      LOG(FATAL) << "Can't find device allocator for gpu device" << tf_gpu_id;
     }
     allocator_ = std::make_shared<TRTDeviceAllocator>(dev_allocator);
     infer->setGpuAllocator(allocator_.get());
+#endif
     trt_engine_ptr_.reset(infer->deserializeCudaEngine(
         serialized_engine_.c_str(), serialized_engine_.size(), nullptr));
     trt_execution_context_ptr_.reset(trt_engine_ptr_->createExecutionContext());
@@ -167,7 +170,7 @@ void TRTEngineOp::Compute(OpKernelContext* context) {
   VLOG(2) << "enqueue returns: " << ret;
   // sync should be done by TF.
 }
-TRTEngineOp::~TRTEngineOp(){
+TRTEngineOp::~TRTEngineOp() {
   // Order matters!
   trt_execution_context_ptr_.reset();
   trt_engine_ptr_.reset();
diff --git a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h
index 791bb6f583..38ceec4704 100644
--- a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h
+++ b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h
@@ -23,10 +23,10 @@ limitations under the License.
 #if GOOGLE_CUDA
 #if GOOGLE_TENSORRT
 #include "cuda/include/cuda_runtime_api.h"
+#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorrt/include/NvInfer.h"
-#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h"
 
 namespace tensorflow {
 namespace tensorrt {
@@ -38,6 +38,7 @@ class TRTEngineOp : public OpKernel {
 
   void Compute(OpKernelContext* context) override;
   ~TRTEngineOp();
+
  private:
   template <typename T>
   struct Destroyer {
diff --git a/tensorflow/contrib/tensorrt/test/test_tftrt.py b/tensorflow/contrib/tensorrt/test/test_tftrt.py
index aaaed0c30f..2295320117 100644
--- a/tensorflow/contrib/tensorrt/test/test_tftrt.py
+++ b/tensorflow/contrib/tensorrt/test/test_tftrt.py
@@ -78,8 +78,7 @@ def execute_graph(gdef, dumm_inp):
   # with csess.Session(
   #     config=cpb2.ConfigProto(gpu_options=gpu_options), graph=g) as sess:
   #   val = sess.run(out, {inp: dumm_inp})
-  with csess.Session(
-      config=sessconfig, graph=g) as sess:
+  with csess.Session(config=sessconfig, graph=g) as sess:
     val = sess.run(out, {inp: dumm_inp})
   return val
 
@@ -149,6 +148,7 @@ def user(run_graph=execute_graph, run_calibration=execute_calibration):
   assert np.allclose(o1, o5)
   print("Pass")
 
+
 def auto():
   """ Run the conversion as an optimization pass"""
   inp_dims = (100, 24, 24, 2)
@@ -165,8 +165,8 @@ def auto():
   print(custom_op)
   gpu_options = cpb2.GPUOptions(per_process_gpu_memory_fraction=0.50)
   graph_options = cpb2.GraphOptions(rewrite_options=opt_config)
-  sessconfig = cpb2.ConfigProto(gpu_options=gpu_options,
-                                graph_options=graph_options)
+  sessconfig = cpb2.ConfigProto(
+      gpu_options=gpu_options, graph_options=graph_options)
   print(sessconfig)
   g = ops.Graph()
   ops.reset_default_graph()
@@ -179,11 +179,17 @@ def auto():
       val = sess.run(out, {inp: dummy_input})
   print(val.shape)
 
+
 if "__main__" in __name__:
-  P = argparse.ArgumentParser(prog="tftrt_test",
-    description="Example utilization of TensorFlow-TensorRT integration")
-  P.add_argument("--automatic", "-a", action="store_true",
-                 help="Do TRT conversion automatically", default=False)
+  P = argparse.ArgumentParser(
+      prog="tftrt_test",
+      description="Example utilization of TensorFlow-TensorRT integration")
+  P.add_argument(
+      "--automatic",
+      "-a",
+      action="store_true",
+      help="Do TRT conversion automatically",
+      default=False)
   flags, unparsed = P.parse_known_args()
   if flags.automatic:
     auto()
diff --git a/tensorflow/contrib/tensorrt/test/tf_trt_integration_test.py b/tensorflow/contrib/tensorrt/test/tf_trt_integration_test.py
index 7a47328762..a5c00dd633 100644
--- a/tensorflow/contrib/tensorrt/test/tf_trt_integration_test.py
+++ b/tensorflow/contrib/tensorrt/test/tf_trt_integration_test.py
@@ -45,8 +45,7 @@ class IntegrationTest(test_util.TensorFlowTestCase):
     inp_dims = (100, 24, 24, 2)
     self._input = np.random.random_sample(inp_dims)
     self._original_graph = self.get_simple_graph_def()
-    self._gpu_options = cpb2.GPUOptions(
-        per_process_gpu_memory_fraction=0.50)
+    self._gpu_options = cpb2.GPUOptions(per_process_gpu_memory_fraction=0.50)
     self._config = cpb2.ConfigProto(gpu_options=self._gpu_options)
     self._reference = self.run_graph(self._original_graph, self._input)
 
@@ -61,11 +60,7 @@ class IntegrationTest(test_util.TensorFlowTestCase):
           name="weights",
           dtype=dtypes.float32)
       conv = nn.conv2d(
-          input=a,
-          filter=e,
-          strides=[1, 2, 2, 1],
-          padding="SAME",
-          name="conv")
+          input=a, filter=e, strides=[1, 2, 2, 1], padding="SAME", name="conv")
       b = cop.constant(
           [4., 1.5, 2., 3., 5., 7.], name="bias", dtype=dtypes.float32)
       t = nn.bias_add(conv, b, name="biasAdd")
@@ -86,8 +81,7 @@ class IntegrationTest(test_util.TensorFlowTestCase):
       inp = inp.outputs[0]
       out = out.outputs[0]
     with self.test_session(
-        graph=g, config=self._config, use_gpu=True,
-        force_gpu=True) as sess:
+        graph=g, config=self._config, use_gpu=True, force_gpu=True) as sess:
       val = sess.run(out, {inp: dumm_inp})
     return val
 
@@ -105,15 +99,14 @@ class IntegrationTest(test_util.TensorFlowTestCase):
       # run over real calibration data here, we are mimicking a calibration
       # set of 30 different batches. Use as much calibration data as you want
     with self.test_session(
-        graph=g, config=self._config, use_gpu=True,
-        force_gpu=True) as sess:
+        graph=g, config=self._config, use_gpu=True, force_gpu=True) as sess:
       for _ in range(30):
         val = sess.run(out, {inp: dumm_inp})
     return val
 
   def get_trt_graph(self, mode):
     """Return trt converted graph."""
-    if mode in  ["FP32", "FP16", "INT8"]:
+    if mode in ["FP32", "FP16", "INT8"]:
       return trt.create_inference_graph(
           input_graph_def=self._original_graph,
           outputs=["output"],
@@ -121,7 +114,7 @@ class IntegrationTest(test_util.TensorFlowTestCase):
           max_workspace_size_bytes=1 << 25,
           precision_mode=mode,  # TRT Engine precision "FP32","FP16" or "INT8"
           minimum_segment_size=2  # minimum number of nodes in an engine
-          )
+      )
     return None
 
   def testFP32(self):
-- 
GitLab


From 8838e2a84f98bd210147dc1a79e1037f2545dff9 Mon Sep 17 00:00:00 2001
From: Sami Kama <skama@nvidia.com>
Date: Thu, 26 Apr 2018 17:36:20 -0700
Subject: [PATCH 022/395] Remove some commented code and add a TODO

---
 .../contrib/tensorrt/convert/convert_graph.cc   |  6 +-----
 .../tensorrt/convert/trt_optimization_pass.cc   |  1 +
 .../tensorrt/convert/trt_optimization_pass.h    |  4 +---
 .../contrib/tensorrt/resources/trt_allocator.cc |  2 +-
 .../contrib/tensorrt/resources/trt_allocator.h  |  2 +-
 tensorflow/contrib/tensorrt/segment/segment.cc  | 17 ++++-------------
 6 files changed, 9 insertions(+), 23 deletions(-)

diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.cc b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
index 9d79c084ee..44b1a8f94c 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_graph.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
@@ -230,7 +230,7 @@ tensorflow::Status GetCalibNode(ConvertGraphParams* params) {
     auto src_output = in_edge->src_output();
     auto dst_node = in_edge->dst();
     auto dst_input = in_edge->dst_input();
-    VLOG(0) << " update edge " << trt_node->name() << ":" << src_output
+    VLOG(1) << " update edge " << trt_node->name() << ":" << src_output
             << " -> " << dst_node->name() << ":" << dst_input;
     TF_RETURN_IF_ERROR(
         params->graph.UpdateEdge(trt_node, src_output, dst_node, dst_input));
@@ -367,12 +367,8 @@ tensorflow::Status ConvertGraphDefToTensorRT(
   VLOG(2) << "gpus: " << num_gpus;
   tensorflow::RewriterConfig rw_cfg;
   tensorflow::grappler::MetaOptimizer meta_opt(nullptr, rw_cfg);
-  // TF_RETURN_IF_ERROR(optimizer.Optimize(cluster, item, &gdef));
   TF_RETURN_IF_ERROR(meta_opt.Optimize(cluster, item, &gdef));
-  // constant folding
   item.graph = gdef;
-  // tensorflow::grappler::ConstantFolding fold(nullptr);
-  // TF_RETURN_IF_ERROR(fold.Optimize(nullptr, item, &gdef));
 
   // AJ refactoring shape inference through grappler/GraphProperties.
   tensorflow::grappler::GraphProperties static_graph_properties(item);
diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
index 5c08d5afdf..999ad1274c 100644
--- a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
+++ b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
@@ -31,6 +31,7 @@ using tensorflow::strings::StrCat;
 namespace tensorflow {
 namespace tensorrt {
 namespace convert {
+// TODO(sami): Remove VLOG messages once the code matures
 tensorflow::Status TRTOptimizationPass::Init(
     const tensorflow::RewriterConfig_CustomGraphOptimizer* config) {
   VLOG(1) << "Called INIT for " << m_name_ << " with config = " << config;
diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
index 5b1462f573..81e3462a61 100644
--- a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
+++ b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
@@ -42,8 +42,6 @@ class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer {
         maximum_workspace_size_(-1) {
     VLOG(1) << "Constructing " << m_name_;
   };
-  // tensorflow::Status Run(const tensorflow::GraphOptimizationPassOptions
-  // &options) override;
   string name() const override { return m_name_; };
   tensorflow::Status Init(const tensorflow::RewriterConfig_CustomGraphOptimizer*
                               config = nullptr) override;
@@ -67,4 +65,4 @@ class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer {
 }  // namespace tensorflow
 #endif
 #endif
-#endif
\ No newline at end of file
+#endif
diff --git a/tensorflow/contrib/tensorrt/resources/trt_allocator.cc b/tensorflow/contrib/tensorrt/resources/trt_allocator.cc
index 4705f6d20f..9d40fea06b 100644
--- a/tensorflow/contrib/tensorrt/resources/trt_allocator.cc
+++ b/tensorflow/contrib/tensorrt/resources/trt_allocator.cc
@@ -54,4 +54,4 @@ void TRTDeviceAllocator::free(void* memory) {
 }  // namespace tensorflow
 #endif
 #endif
-#endif
\ No newline at end of file
+#endif
diff --git a/tensorflow/contrib/tensorrt/resources/trt_allocator.h b/tensorflow/contrib/tensorrt/resources/trt_allocator.h
index 8bdb0519ba..3001224b8d 100644
--- a/tensorflow/contrib/tensorrt/resources/trt_allocator.h
+++ b/tensorflow/contrib/tensorrt/resources/trt_allocator.h
@@ -62,4 +62,4 @@ class AllocatorFactory {};
 
 #endif
 #endif
-#endif
\ No newline at end of file
+#endif
diff --git a/tensorflow/contrib/tensorrt/segment/segment.cc b/tensorflow/contrib/tensorrt/segment/segment.cc
index 8f335f2bf1..ac0d782a2b 100644
--- a/tensorflow/contrib/tensorrt/segment/segment.cc
+++ b/tensorflow/contrib/tensorrt/segment/segment.cc
@@ -96,15 +96,6 @@ bool CanContractEdge(const Edge* edge, const Graph* graph) {
   }
 
   bool is_cycle = check_cycles(graph, src, dfs_start_nodes);
-  // if (!dfs_start_nodes.empty()) {
-  //   tensorflow::ReverseDFSFrom(graph, dfs_start_nodes, {},
-  //                              [&is_cycle, src](tensorflow::Node* node) {
-  //                                if (node == src) {
-  //                                  is_cycle = true;
-  //                                }
-  //                              });
-  // }
-
   return !is_cycle;
 }
 }  // namespace
@@ -140,7 +131,7 @@ Graph::Graph(const tensorflow::Graph* g) : g_(g) {
       auto dst = nodes_[tfdst->id()];
       auto edge =
           new Edge(i, src, e->src_output(), dst, e->dst_input(), is_control);
-      edges_[i]=edge;
+      edges_[i] = edge;
       src->out_edges_.push_back(edge);
       dst->in_edges_.push_back(edge);
     } else {
@@ -271,10 +262,10 @@ tensorflow::Status SegmentGraph(
     const std::function<bool(const tensorflow::Node*)>& candidate_fn,
     const SegmentOptions& options, SegmentNodesVector* segments) {
   // tensorflow::DumpGraph("Pre-Segment", &graph);
-  Graph* graph= new Graph(tf_graph);
+  Graph* graph = new Graph(tf_graph);
   // Use a union-find to collect the nodes that belong to the same
-  // segment. A node value of nullptr indicates that tusing
-  // ::tensorflow::strings::StrAppendhe node is not a candidate for TRT.
+  // segment. A node value of nullptr indicates that the node is not a candidate
+  // for TRT.
   std::vector<UnionFind<Node*>> node_segments;
   for (int i = 0; i < graph->num_node_ids(); ++i) {
     Node* node = graph->FindNodeId(i);
-- 
GitLab


From 09fc850e988e71983e9d0eb4e874f998b3a480e6 Mon Sep 17 00:00:00 2001
From: Yifei Feng <1192265+yifeif@users.noreply.github.com>
Date: Thu, 26 Apr 2018 18:47:07 -0700
Subject: [PATCH 023/395] Update build_pip_package.sh

---
 tensorflow/tools/pip_package/build_pip_package.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/tools/pip_package/build_pip_package.sh b/tensorflow/tools/pip_package/build_pip_package.sh
index 8f0cf8c3d1..3af79ee170 100755
--- a/tensorflow/tools/pip_package/build_pip_package.sh
+++ b/tensorflow/tools/pip_package/build_pip_package.sh
@@ -24,7 +24,7 @@ function real_path() {
 function cp_external() {
   local src_dir=$1
   local dest_dir=$2
-  for f in `find "$src_dir" -maxdepth 1 -mindepth 1 ! -name '*local_config_cuda*' ! -name '*org_tensorflow*'`; do
+  for f in `find "$src_dir" -maxdepth 1 -mindepth 1 ! -name '*local_config_cuda*' ! -name '*local_config_tensorrt*' ! -name '*org_tensorflow*'`; do
     cp -R "$f" "$dest_dir"
   done
   mkdir -p "${dest_dir}/local_config_cuda/cuda/cuda/"
-- 
GitLab


From f1e00684f14a9a2c50ca0e05710a1bd2bc2e734f Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Fri, 27 Apr 2018 13:01:10 +0300
Subject: [PATCH 024/395] [tf.data] Make documentation changes, and add correct
 import.

---
 .../contrib/data/python/kernel_tests/resample_test.py       | 2 +-
 tensorflow/contrib/data/python/ops/resampling.py            | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/contrib/data/python/kernel_tests/resample_test.py b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
index fc84301b17..b556525ce4 100644
--- a/tensorflow/contrib/data/python/kernel_tests/resample_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
@@ -59,7 +59,7 @@ def _time_resampling(
 class ResampleTest(test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(
-      ("InitialnDistributionKnown", True),
+      ("InitialDistributionKnown", True),
       ("InitialDistributionUnknown", False))
   def testDistribution(self, initial_known):
     classes = np.random.randint(5, size=(20000,))  # Uniformly sampled
diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index 66eaf9b69a..982ff66c13 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -25,6 +25,7 @@ from tensorflow.contrib.data.python.ops import scan_ops
 from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import logging_ops
@@ -102,9 +103,8 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
 def _get_prob_original_static(initial_dist_t, target_dist_t):
   """Returns the static probability of sampling from the original.
 
-  For some reason, `tensor_util.constant_value(prob_of_original)` of a ratio
-  of two constant Tensors isn't a constant. We have some custom logic to avoid
-  this.
+  `tensor_util.constant_value(prob_of_original)` returns `None` if it encounters
+  an Op that it isn't defined for. We have some custom logic to avoid this.
 
   Args:
     initial_dist_t: A tensor of the initial distribution.
-- 
GitLab


From 4dee7b57a47817ec8c972cbb117868463ef15cdf Mon Sep 17 00:00:00 2001
From: Yifei Feng <1192265+yifeif@users.noreply.github.com>
Date: Fri, 27 Apr 2018 07:46:49 -0700
Subject: [PATCH 025/395] Update tf_tests.cmake

---
 tensorflow/contrib/cmake/tf_tests.cmake | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/contrib/cmake/tf_tests.cmake b/tensorflow/contrib/cmake/tf_tests.cmake
index 92f2ab6dea..5942ff3363 100644
--- a/tensorflow/contrib/cmake/tf_tests.cmake
+++ b/tensorflow/contrib/cmake/tf_tests.cmake
@@ -267,6 +267,8 @@ if (tensorflow_BUILD_PYTHON_TESTS)
       "${tensorflow_source_dir}/tensorflow/python/kernel_tests/variable_scope_test.py"
       "${tensorflow_source_dir}/tensorflow/python/kernel_tests/functional_ops_test.py"
       "${tensorflow_source_dir}/tensorflow/python/kernel_tests/py_func_test.py"
+      # Flaky on Windows cpu with py36 (b/73556968)
+      "${tensorflow_source_dir}/tensorflow/python/kernel_tests/sparse_reshape_op_test.py"
       # Windows file management related issues.
       "${tensorflow_source_dir}/tensorflow/python/training/evaluation_test.py"
       # training tests
-- 
GitLab


From dd24a090971a68a42925b2d1276af165434c9913 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Fri, 27 Apr 2018 23:46:22 +0300
Subject: [PATCH 026/395] [tf.data] Pass a Tensor to
 `tensor_util.constant_value` instead of possible a python list.

---
 tensorflow/contrib/data/python/ops/resampling.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index 982ff66c13..f7ea44bec0 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -86,7 +86,7 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
     filtered_ds = filtered_ds.prefetch(3)
 
     prob_original_static = _get_prob_original_static(
-        initial_dist, target_dist_t) if initial_dist is not None else None
+          initial_dist_t, target_dist_t) if initial_dist is not None else None
     if prob_original_static == 1:
       return dataset_ops.Dataset.zip((class_values_ds, dataset))
     elif prob_original_static == 0:
-- 
GitLab


From ac2416120ddd13891486cce6135160cc2f412f92 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Sat, 28 Apr 2018 00:20:53 +0300
Subject: [PATCH 027/395] [tf.data] Fix indentation.

---
 tensorflow/contrib/data/python/ops/resampling.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index f7ea44bec0..1194b8447a 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -86,7 +86,7 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
     filtered_ds = filtered_ds.prefetch(3)
 
     prob_original_static = _get_prob_original_static(
-          initial_dist_t, target_dist_t) if initial_dist is not None else None
+        initial_dist_t, target_dist_t) if initial_dist is not None else None
     if prob_original_static == 1:
       return dataset_ops.Dataset.zip((class_values_ds, dataset))
     elif prob_original_static == 0:
-- 
GitLab


From 8753e2ebde6c58b56675cc19ab7ff83072824a62 Mon Sep 17 00:00:00 2001
From: Yifei Feng <1192265+yifeif@users.noreply.github.com>
Date: Fri, 27 Apr 2018 17:05:02 -0700
Subject: [PATCH 028/395] Fixing the mock import error for devel docker.
 (#18940)

* Fixing the mock import error for devel docker.

Same as #18843
---
 tensorflow/tools/docker/Dockerfile.devel     | 1 +
 tensorflow/tools/docker/Dockerfile.devel-gpu | 1 +
 2 files changed, 2 insertions(+)

diff --git a/tensorflow/tools/docker/Dockerfile.devel b/tensorflow/tools/docker/Dockerfile.devel
index 390d7442c3..5c49ac1d8d 100644
--- a/tensorflow/tools/docker/Dockerfile.devel
+++ b/tensorflow/tools/docker/Dockerfile.devel
@@ -31,6 +31,7 @@ RUN pip --no-cache-dir install \
         ipykernel \
         jupyter \
         matplotlib \
+        mock \
         numpy \
         scipy \
         sklearn \
diff --git a/tensorflow/tools/docker/Dockerfile.devel-gpu b/tensorflow/tools/docker/Dockerfile.devel-gpu
index 293028d229..196227861b 100644
--- a/tensorflow/tools/docker/Dockerfile.devel-gpu
+++ b/tensorflow/tools/docker/Dockerfile.devel-gpu
@@ -40,6 +40,7 @@ RUN pip --no-cache-dir install \
         ipykernel \
         jupyter \
         matplotlib \
+        mock \
         numpy \
         scipy \
         sklearn \
-- 
GitLab


From e276bf65e2f3ec452eb28d0a9d34849d65663788 Mon Sep 17 00:00:00 2001
From: Sami Kama <skama@nvidia.com>
Date: Fri, 27 Apr 2018 17:11:30 -0700
Subject: [PATCH 029/395] Fixes for review

---
 .../contrib/tensorrt/convert/convert_graph.cc | 10 +-
 .../contrib/tensorrt/convert/convert_nodes.cc |  4 +-
 .../contrib/tensorrt/convert/convert_nodes.h  |  6 +-
 .../tensorrt/convert/trt_optimization_pass.cc |  1 -
 .../tensorrt/convert/trt_optimization_pass.h  |  4 +-
 .../tensorrt/resources/trt_allocator.cc       |  4 +-
 .../tensorrt/resources/trt_allocator.h        |  7 +-
 .../contrib/tensorrt/segment/segment.cc       | 92 ++++++++++---------
 tensorflow/contrib/tensorrt/segment/segment.h | 82 +++++++++--------
 .../contrib/tensorrt/segment/segment_test.cc  | 10 +-
 .../contrib/tensorrt/test/test_tftrt.py       |  8 +-
 11 files changed, 116 insertions(+), 112 deletions(-)

diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.cc b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
index 44b1a8f94c..632908f078 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_graph.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
@@ -150,7 +150,7 @@ struct ConvertGraphParams {
       const tensorflow::grappler::GraphProperties& current_graph_properties,
       std::unordered_map<string, std::pair<int, string>>* output_edges,
       int engine_precision_mode, const string& device_name,
-      std::shared_ptr<nvinfer1::IGpuAllocator> allocator, int cuda_device_id)
+      std::shared_ptr<nvinfer1::IGpuAllocator> allocator, int cuda_gpu_id)
       : graph(inp_graph),
         output_names(output_node_names),
         subgraph_node_ids(subgraph_node_id_numbers),
@@ -161,7 +161,7 @@ struct ConvertGraphParams {
         precision_mode(engine_precision_mode),
         device_name_(device_name),
         allocator_(allocator),
-        cuda_device_id_(cuda_device_id) {}
+        cuda_gpu_id_(cuda_gpu_id) {}
   tensorflow::Graph& graph;
   const std::vector<string>& output_names;
   const std::set<int>& subgraph_node_ids;
@@ -172,7 +172,7 @@ struct ConvertGraphParams {
   int precision_mode;
   string device_name_;
   std::shared_ptr<nvinfer1::IGpuAllocator> allocator_;
-  int cuda_device_id_;
+  int cuda_gpu_id_;
   std::vector<std::pair<int, int>> subgraph_inputs;
   std::vector<std::pair<int, int>> subgraph_outputs;
   tensorflow::EdgeSet subgraph_incoming_edges;
@@ -216,7 +216,7 @@ tensorflow::Status GetCalibNode(ConvertGraphParams* params) {
                    params->max_batch_size, params->max_workspace_size_bytes,
                    params->graph_properties, params->output_edge_map,
                    &trt_node_def, params->precision_mode, params->device_name_,
-                   params->allocator_, params->cuda_device_id_);
+                   params->allocator_, params->cuda_gpu_id_);
   TF_RETURN_IF_ERROR(InjectCalibrationNode(s));
   tensorflow::Status status;
   tensorflow::Node* trt_node = params->graph.AddNode(trt_node_def, &status);
@@ -247,7 +247,7 @@ tensorflow::Status ConvertSubGraphToTensorRT(ConvertGraphParams* params) {
                    params->max_batch_size, params->max_workspace_size_bytes,
                    params->graph_properties, params->output_edge_map,
                    &trt_node_def, params->precision_mode, params->device_name_,
-                   params->allocator_, params->cuda_device_id_);
+                   params->allocator_, params->cuda_gpu_id_);
   TF_RETURN_IF_ERROR(ConvertSubGraphToTensorRTNodeDef(s));
   tensorflow::Status status;
   tensorflow::Node* trt_node = params->graph.AddNode(trt_node_def, &status);
diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
index 8ed0ed7b7e..ae0e861be5 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
@@ -2247,7 +2247,7 @@ tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) {
   auto op_res = new tensorflow::tensorrt::TRTCalibrationResource();
   TF_CHECK_OK(op_rmgr->Create(calib_op_name, calib_op_name, op_res));
   op_res->logger_ = new tensorflow::tensorrt::Logger();
-  cudaSetDevice(s.cuda_device_id_);
+  cudaSetDevice(s.cuda_gpu_id_);
   op_res->builder_ = nvinfer1::createInferBuilder(*(op_res->logger_));
   op_res->allocator_ = s.allocator_;
 #if NV_TENSORRT_MAJOR > 3
@@ -2481,7 +2481,7 @@ tensorflow::Status ConvertSubGraphToTensorRTNodeDef(
   // Topological order is needed to build TRT network
 
   tensorflow::tensorrt::Logger trt_logger;
-  cudaSetDevice(s.cuda_device_id_);
+  cudaSetDevice(s.cuda_gpu_id_);
   auto trt_builder = infer_object(nvinfer1::createInferBuilder(trt_logger));
   if (!trt_builder) {
     return tensorflow::errors::Internal(
diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.h b/tensorflow/contrib/tensorrt/convert/convert_nodes.h
index 8e1d7c99b6..50b0c37094 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_nodes.h
+++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.h
@@ -50,7 +50,7 @@ struct SubGraphParams {
       tensorflow::NodeDef* constructed_trt_node,
       int engine_precision_mode = FP32MODE, const string& device_name = "",
       std::shared_ptr<nvinfer1::IGpuAllocator> allocator = 0,
-      int cuda_device_id = 0)
+      int cuda_gpu_id = 0)
       : graph(inp_graph),
         subgraph_node_ids(subgraph_node_id_numbers),
         input_inds(input_indices),
@@ -63,7 +63,7 @@ struct SubGraphParams {
         precision_mode(engine_precision_mode),
         device_name_(device_name),
         allocator_(allocator),
-        cuda_device_id_(cuda_device_id) {}
+        cuda_gpu_id_(cuda_gpu_id) {}
 
   tensorflow::Graph& graph;
   const std::set<int>& subgraph_node_ids;
@@ -77,7 +77,7 @@ struct SubGraphParams {
   const int precision_mode;
   const string device_name_;
   std::shared_ptr<nvinfer1::IGpuAllocator> allocator_;
-  const int cuda_device_id_;
+  const int cuda_gpu_id_;
 };
 
 // TODO(sami): Replace references with const reference or pointers
diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
index 999ad1274c..743750998c 100644
--- a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
+++ b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
@@ -1,5 +1,4 @@
 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
-1;4804;0c
 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
diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
index 81e3462a61..aa9f289550 100644
--- a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
+++ b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
@@ -34,8 +34,8 @@ namespace tensorrt {
 namespace convert {
 class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer {
  public:
-  TRTOptimizationPass(string optName = "TRTOptimizationPass")
-      : m_name_(optName),
+  TRTOptimizationPass(const string& name = "TRTOptimizationPass")
+      : m_name_(name),
         minimum_segment_size_(3),
         precision_mode_(0),
         maximum_batch_size_(-1),
diff --git a/tensorflow/contrib/tensorrt/resources/trt_allocator.cc b/tensorflow/contrib/tensorrt/resources/trt_allocator.cc
index 9d40fea06b..b94f8a2da7 100644
--- a/tensorflow/contrib/tensorrt/resources/trt_allocator.cc
+++ b/tensorflow/contrib/tensorrt/resources/trt_allocator.cc
@@ -30,6 +30,7 @@ void* TRTCudaAllocator::allocate(uint64_t size, uint64_t alignment,
   cudaMalloc(&memory, size);
   return memory;
 }
+
 void TRTCudaAllocator::free(void* memory) { cudaFree(memory); }
 
 void* TRTDeviceAllocator::allocate(uint64_t size, uint64_t alignment,
@@ -44,7 +45,8 @@ void* TRTDeviceAllocator::allocate(uint64_t size, uint64_t alignment,
 TRTDeviceAllocator::TRTDeviceAllocator(tensorflow::Allocator* allocator)
     : allocator_(allocator) {
   VLOG(1) << "Using " << allocator->Name() << " allocator from TensorFlow";
-};
+}
+
 void TRTDeviceAllocator::free(void* memory) {
   VLOG(2) << "Deallocating " << memory;
   allocator_->DeallocateRaw(memory);
diff --git a/tensorflow/contrib/tensorrt/resources/trt_allocator.h b/tensorflow/contrib/tensorrt/resources/trt_allocator.h
index 3001224b8d..05dcb7cde6 100644
--- a/tensorflow/contrib/tensorrt/resources/trt_allocator.h
+++ b/tensorflow/contrib/tensorrt/resources/trt_allocator.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include <string>
 #include <thread>
 #include <vector>
+
 #include "tensorflow/contrib/tensorrt/log/trt_logger.h"
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/framework/resource_mgr.h"
@@ -41,21 +42,21 @@ namespace tensorrt {
 class TRTCudaAllocator : public nvinfer1::IGpuAllocator {
  public:
   TRTCudaAllocator() {}
-  virtual ~TRTCudaAllocator(){};
+  virtual ~TRTCudaAllocator() {};
   void* allocate(uint64_t size, uint64_t alignment, uint32_t flags) override;
   void free(void* memory) override;
 };
+
 class TRTDeviceAllocator : public nvinfer1::IGpuAllocator {
  public:
   TRTDeviceAllocator(tensorflow::Allocator* allocator);
-  virtual ~TRTDeviceAllocator(){};
+  virtual ~TRTDeviceAllocator() {};
   void* allocate(uint64_t size, uint64_t alignment, uint32_t flags) override;
   void free(void* memory) override;
 
  private:
   tensorflow::Allocator* allocator_;
 };
-class AllocatorFactory {};
 
 }  // namespace tensorrt
 }  // namespace tensorflow
diff --git a/tensorflow/contrib/tensorrt/segment/segment.cc b/tensorflow/contrib/tensorrt/segment/segment.cc
index ac0d782a2b..a76d170236 100644
--- a/tensorflow/contrib/tensorrt/segment/segment.cc
+++ b/tensorflow/contrib/tensorrt/segment/segment.cc
@@ -34,10 +34,11 @@ namespace segment {
 using ::tensorflow::strings::StrAppend;
 namespace {
 
-bool check_cycles(const Graph* g, const Node* src,
-                  const std::vector<Node*>& start) {
+bool CheckCycles(const SimpleGraph* g, const SimpleNode* src,
+                 const std::vector<SimpleNode*>& start) {
+  //  copied from TF ReverseDFS
   struct Work {
-    Node* node;
+    SimpleNode* node;
     bool leave;  // Are we entering or leaving n?
   };
 
@@ -74,7 +75,7 @@ bool check_cycles(const Graph* g, const Node* src,
   return false;
 }
 
-bool CanContractEdge(const Edge* edge, const Graph* graph) {
+bool CanContractEdge(const SimpleEdge* edge, const SimpleGraph* graph) {
   const auto src = edge->src();
   const auto dst = edge->dst();
 
@@ -88,35 +89,36 @@ bool CanContractEdge(const Edge* edge, const Graph* graph) {
   //   1. Get all nodes incoming to 'dst', excluding 'src'
   //   2. Reverse DFS from those nodes
   //   3. If reverse DFS reaches 'src' then we have a cycle
-  std::vector<Node*> dfs_start_nodes;
-  for (Node* node : dst->in_nodes()) {
+  std::vector<SimpleNode*> dfs_start_nodes;
+  for (SimpleNode* node : dst->in_nodes()) {
     if (node != src) {
       dfs_start_nodes.push_back(node);
     }
   }
 
-  bool is_cycle = check_cycles(graph, src, dfs_start_nodes);
+  bool is_cycle = CheckCycles(graph, src, dfs_start_nodes);
   return !is_cycle;
 }
 }  // namespace
-Node::Node(const tensorflow::Node* node, const int id) : node_(node), id_(id) {
+SimpleNode::SimpleNode(const tensorflow::Node* node, const int id)
+    : node_(node), id_(id) {
   if (node_) {
     in_edges_.reserve(node_->in_edges().size());
     out_edges_.reserve(node_->out_edges().size());
   }
 }
 
-Graph::Graph(const tensorflow::Graph* g) : g_(g) {
+SimpleGraph::SimpleGraph(const tensorflow::Graph* g) : g_(g) {
   int n_nodes = g_->num_node_ids();
   nodes_.resize(n_nodes, nullptr);
-  nodes_[g->kSourceId] = new Node(g->source_node(), g->kSourceId);
-  nodes_[g->kSinkId] = new Node(g->sink_node(), g->kSinkId);
+  nodes_[g->kSourceId] = new SimpleNode(g->source_node(), g->kSourceId);
+  nodes_[g->kSinkId] = new SimpleNode(g->sink_node(), g->kSinkId);
   int n_edges = g->num_edge_ids();
   edges_.resize(n_edges, nullptr);
   for (int i = 2; i < n_nodes; i++) {
     const auto n = g->FindNodeId(i);
     if (n) {
-      nodes_[i] = new Node(n, i);
+      nodes_[i] = new SimpleNode(n, i);
     } else {
       node_ids_.insert(i);
     }
@@ -129,8 +131,8 @@ Graph::Graph(const tensorflow::Graph* g) : g_(g) {
       bool is_control = e->IsControlEdge();
       auto src = nodes_[tfsrc->id()];
       auto dst = nodes_[tfdst->id()];
-      auto edge =
-          new Edge(i, src, e->src_output(), dst, e->dst_input(), is_control);
+      auto edge = new SimpleEdge(i, src, e->src_output(), dst, e->dst_input(),
+                                 is_control);
       edges_[i] = edge;
       src->out_edges_.push_back(edge);
       dst->in_edges_.push_back(edge);
@@ -140,7 +142,8 @@ Graph::Graph(const tensorflow::Graph* g) : g_(g) {
   }
 }
 
-void Graph::AddEdge(Node* src, int out_port, Node* dst, int in_port) {
+void SimpleGraph::AddEdge(SimpleNode* src, int out_port, SimpleNode* dst,
+                          int in_port) {
   int i = edges_.size();
   if (edge_ids_.size()) {
     auto it = edge_ids_.begin();
@@ -151,18 +154,18 @@ void Graph::AddEdge(Node* src, int out_port, Node* dst, int in_port) {
   }
   bool is_control = (out_port == tensorflow::Graph::kControlSlot);
   is_control |= (in_port == tensorflow::Graph::kControlSlot);
-  auto edge = new Edge(i, src, out_port, dst, in_port, is_control);
+  auto edge = new SimpleEdge(i, src, out_port, dst, in_port, is_control);
   edges_[i] = edge;
   src->out_edges_.push_back(edge);
   dst->in_edges_.push_back(edge);
 }
 
-void Graph::AddControlEdge(Node* src, Node* dst) {
+void SimpleGraph::AddControlEdge(SimpleNode* src, SimpleNode* dst) {
   AddEdge(src, tensorflow::Graph::kControlSlot, dst,
           tensorflow::Graph::kControlSlot);
 }
 
-void Graph::RemoveEdge(const Edge* edge) {
+void SimpleGraph::RemoveEdge(const SimpleEdge* edge) {
   auto src = edge->src();
   auto dst = edge->dst();
   for (auto it = src->out_edges_.begin(); it != src->out_edges_.end(); ++it) {
@@ -179,13 +182,13 @@ void Graph::RemoveEdge(const Edge* edge) {
   }
 }
 
-Graph::~Graph() {
+SimpleGraph::~SimpleGraph() {
   for (auto x : nodes_) delete x;
   for (auto x : edges_) delete x;
 }
 
-void ContractEdge(Edge* edge, Graph* graph,
-                  std::vector<const Edge*>* remove_edges) {
+void ContractEdge(SimpleEdge* edge, SimpleGraph* graph,
+                  std::vector<const SimpleEdge*>* remove_edges) {
   // Transfer all inputs and outputs of 'dst' to 'src' except edges
   // connecting the two.
   auto src = edge->src();
@@ -193,17 +196,17 @@ void ContractEdge(Edge* edge, Graph* graph,
 
   // We can use '0' for input/output index because we don't need them
   // to be accurate for the way we are using the graph.
-  std::vector<const Edge*> in_edges(dst->in_edges().begin(),
-                                    dst->in_edges().end());
-  for (const Edge* in_edge : in_edges) {
+  std::vector<const SimpleEdge*> in_edges(dst->in_edges().begin(),
+                                          dst->in_edges().end());
+  for (const SimpleEdge* in_edge : in_edges) {
     if (in_edge->IsControlEdge()) {
       if (in_edge->src() != src) {
-        Edge* e = const_cast<Edge*>(in_edge);
+        SimpleEdge* e = const_cast<SimpleEdge*>(in_edge);
         graph->AddControlEdge(e->src(), src);
       }
     } else {
       if (in_edge->src() != src) {
-        Edge* e = const_cast<Edge*>(in_edge);
+        SimpleEdge* e = const_cast<SimpleEdge*>(in_edge);
         if (e->src() == graph->source_node()) {
           graph->AddEdge(e->src(), e->src_output(), src,
                          tensorflow::Graph::kControlSlot);
@@ -214,14 +217,14 @@ void ContractEdge(Edge* edge, Graph* graph,
     }
   }
 
-  std::vector<const Edge*> out_edges(dst->out_edges().begin(),
-                                     dst->out_edges().end());
-  for (const Edge* out_edge : out_edges) {
+  std::vector<const SimpleEdge*> out_edges(dst->out_edges().begin(),
+                                           dst->out_edges().end());
+  for (const SimpleEdge* out_edge : out_edges) {
     if (out_edge->IsControlEdge()) {
-      Edge* e = const_cast<Edge*>(out_edge);
+      SimpleEdge* e = const_cast<SimpleEdge*>(out_edge);
       graph->AddControlEdge(src, e->dst());
     } else {
-      Edge* e = const_cast<Edge*>(out_edge);
+      SimpleEdge* e = const_cast<SimpleEdge*>(out_edge);
       if (e->dst() == graph->sink_node()) {
         VLOG(1) << " edge to sink node " << src->name() << " -> "
                 << e->dst()->name();
@@ -262,13 +265,13 @@ tensorflow::Status SegmentGraph(
     const std::function<bool(const tensorflow::Node*)>& candidate_fn,
     const SegmentOptions& options, SegmentNodesVector* segments) {
   // tensorflow::DumpGraph("Pre-Segment", &graph);
-  Graph* graph = new Graph(tf_graph);
+  SimpleGraph* graph = new SimpleGraph(tf_graph);
   // Use a union-find to collect the nodes that belong to the same
   // segment. A node value of nullptr indicates that the node is not a candidate
   // for TRT.
-  std::vector<UnionFind<Node*>> node_segments;
+  std::vector<UnionFind<SimpleNode*>> node_segments;
   for (int i = 0; i < graph->num_node_ids(); ++i) {
-    Node* node = graph->FindNodeId(i);
+    SimpleNode* node = graph->FindNodeId(i);
     if (options.exclude_node_list.count(node->name()) != 0 ||
         !candidate_fn(node->tf_node())) {
       node = nullptr;
@@ -288,12 +291,12 @@ tensorflow::Status SegmentGraph(
   tensorflow::GetPostOrder(*tf_graph, &tforder);
   // use postorder implementation from tensorflow and construct mirror in
   // internal format
-  std::vector<Node*> order;
+  std::vector<SimpleNode*> order;
   order.reserve(tforder.size());
   for (const auto tfnode : tforder) {
     order.push_back(graph->FindNodeId(tfnode->id()));
   }
-  for (const Node* node : order) {
+  for (const SimpleNode* node : order) {
     // All output nodes of 'node' have been visited...
     VLOG(2) << "Trying node " << node->name() << " id=" << node->id();
 
@@ -307,8 +310,8 @@ tensorflow::Status SegmentGraph(
     // nodes. Iterate since combining two nodes may unblock other
     // combining.
     while (true) {
-      std::set<const Edge*> contract_edges;
-      for (const Edge* out_edge : node->out_edges()) {
+      std::set<const SimpleEdge*> contract_edges;
+      for (const SimpleEdge* out_edge : node->out_edges()) {
         VLOG(2) << "... out node " << out_edge->dst()->name() << " ( "
                 << out_edge->dst()->id() << " <- " << node->id() << " )";
         if (out_edge->IsControlEdge()) {
@@ -336,9 +339,9 @@ tensorflow::Status SegmentGraph(
       // Contract edges and collect the adjacent nodes into the same
       // segment/subgraph.
       while (!contract_edges.empty()) {
-        const Edge* contract_edge = *contract_edges.begin();
-        const Node* src = contract_edge->src();
-        const Node* dst = contract_edge->dst();
+        const SimpleEdge* contract_edge = *contract_edges.begin();
+        const SimpleNode* src = contract_edge->src();
+        const SimpleNode* dst = contract_edge->dst();
 
         VLOG(2) << "Merge " << src->name() << " <- " << dst->name() << " ("
                 << src->id() << " <- " << dst->id();
@@ -347,11 +350,11 @@ tensorflow::Status SegmentGraph(
         // Contracting the edge leaves disconnected graph edges.
         // Remove these from the graph and from 'contract_edges' so we
         // don't visit them again.
-        Edge* e = const_cast<Edge*>(contract_edge);
-        std::vector<const Edge*> remove_edges;
+        SimpleEdge* e = const_cast<SimpleEdge*>(contract_edge);
+        std::vector<const SimpleEdge*> remove_edges;
         ContractEdge(e, graph, &remove_edges);
 
-        for (const Edge* r : remove_edges) {
+        for (const SimpleEdge* r : remove_edges) {
           contract_edges.erase(r);
           graph->RemoveEdge(r);
         }
@@ -399,6 +402,7 @@ tensorflow::Status SegmentGraph(
               << segment_node_names.size() << " nodes, dropping";
       continue;
     }
+    // TODO(sami): Make segmenter placement aware once trtscopes are in place
     const auto& dev_itr = device_maps.find(itr.first);
     if (dev_itr == device_maps.end() || dev_itr->second.size() == 0) {
       VLOG(1) << "No device assigned to segment " << segments->size();
diff --git a/tensorflow/contrib/tensorrt/segment/segment.h b/tensorflow/contrib/tensorrt/segment/segment.h
index 659fea1859..44a84cbd38 100644
--- a/tensorflow/contrib/tensorrt/segment/segment.h
+++ b/tensorflow/contrib/tensorrt/segment/segment.h
@@ -30,43 +30,41 @@ namespace tensorrt {
 namespace segment {
 
 using SegmentNodesVector = std::vector<std::pair<std::set<string>, string>>;
-class Node;
-class Graph;
-class Edge {
+class SimpleNode;
+class SimpleGraph;
+class SimpleEdge {
  public:
-  Edge(int id, Node* src, int src_port, Node* dst, int dst_port,
-       bool is_control = false)
+  SimpleEdge(int id, SimpleNode* src, int src_port, SimpleNode* dst,
+             int dst_port, bool is_control = false)
       : id_(id),
         src_(src),
         src_port_(src_port),
         dst_(dst),
         dst_port_(dst_port),
         control_(is_control){};
-  Node* src() const { return src_; }
-  Node* dst() const { return dst_; }
+  SimpleNode* src() const { return src_; }
+  SimpleNode* dst() const { return dst_; }
   int src_output() const { return src_port_; }
   int dst_input() const { return dst_port_; }
   int id() const { return id_; }
   bool IsControlEdge() const { return control_; }
-  ~Edge() {}
+  ~SimpleEdge() {}
 
  private:
   int id_;
-  Node* src_;
+  SimpleNode* src_;
   int src_port_;
-  Node* dst_;
+  SimpleNode* dst_;
   int dst_port_;
   bool control_;
 };
-class Node {
-  friend class Graph;
-
+class SimpleNode {
  public:
-  Node(const tensorflow::Node* node, const int id);
-  const std::vector<Edge*>& in_edges() const { return in_edges_; };
-  const std::vector<Edge*>& out_edges() const { return out_edges_; };
-  std::vector<Node*> in_nodes() const {
-    std::vector<Node*> res;
+  SimpleNode(const tensorflow::Node* node, const int id);
+  const std::vector<SimpleEdge*>& in_edges() const { return in_edges_; };
+  const std::vector<SimpleEdge*>& out_edges() const { return out_edges_; };
+  std::vector<SimpleNode*> in_nodes() const {
+    std::vector<SimpleNode*> res;
     res.reserve(in_edges_.size());
     for (const auto e : in_edges_) {
       if (e) res.push_back(e->src());
@@ -79,32 +77,36 @@ class Node {
 
  private:
   const tensorflow::Node* node_;
-  std::vector<Edge*> in_edges_;
-  std::vector<Edge*> out_edges_;
+  std::vector<SimpleEdge*> in_edges_;
+  std::vector<SimpleEdge*> out_edges_;
   int id_;
+
+  friend class SimpleGraph;
 };
 
-class Graph {
+class SimpleGraph {
  public:
-  Graph(const tensorflow::Graph* g);
-  void AddControlEdge(Node* src, Node* dst);
-  void AddEdge(Node* src, int out_port, Node* dst, int in_port);
-  void RemoveEdge(const Edge*);
-  Node* FindNodeId(int node_id) {
+  SimpleGraph(const tensorflow::Graph* g);
+  void AddControlEdge(SimpleNode* src, SimpleNode* dst);
+  void AddEdge(SimpleNode* src, int out_port, SimpleNode* dst, int in_port);
+  void RemoveEdge(const SimpleEdge*);
+  SimpleNode* FindNodeId(int node_id) {
     if (node_id < 0 || node_id > (int)nodes_.size()) return nullptr;
     return nodes_[node_id];
   }
-  ~Graph();
+  ~SimpleGraph();
   int num_node_ids() const { return nodes_.size(); }
-  const Node* source_node() const {
+  const SimpleNode* source_node() const {
     return nodes_[tensorflow::Graph::kSourceId];
   }
-  const Node* sink_node() const { return nodes_[tensorflow::Graph::kSinkId]; }
+  const SimpleNode* sink_node() const {
+    return nodes_[tensorflow::Graph::kSinkId];
+  }
 
  private:
   const tensorflow::Graph* g_;
-  std::vector<Node*> nodes_;
-  std::vector<Edge*> edges_;
+  std::vector<SimpleNode*> nodes_;
+  std::vector<SimpleEdge*> edges_;
   std::set<int> edge_ids_;
   std::set<int> node_ids_;
 };
@@ -114,15 +116,15 @@ struct SegmentOptions {
   std::set<string> exclude_node_list;
 };
 
-// // Get the subgraphs of a graph that can be handled by TensorRT.
-// //
-// // @param gdef The GraphDef describing the network
-// // @param candidate_fn A function that returns true for a NodeDef if
-// // that node can be handled by TensorRT.
-// // @param segments Returns the TensorRT segments/subgraphs. Each entry
-// // in the vector describes a subgraph by giving a set of the names of
-// // all the NodeDefs in that subgraph.
-// // @return the status.
+// Get the subgraphs of a graph that can be handled by TensorRT.
+//
+// @param gdef The GraphDef describing the network
+// @param candidate_fn A function that returns true for a NodeDef if
+// that node can be handled by TensorRT.
+// @param segments Returns the TensorRT segments/subgraphs. Each entry
+// in the vector describes a subgraph by giving a set of the names of
+// all the NodeDefs in that subgraph.
+// @return the status.
 tensorflow::Status SegmentGraph(
     const tensorflow::GraphDef& gdef,
     const std::function<bool(const tensorflow::Node*)>& candidate_fn,
diff --git a/tensorflow/contrib/tensorrt/segment/segment_test.cc b/tensorflow/contrib/tensorrt/segment/segment_test.cc
index 7fe824b12f..8038085a06 100644
--- a/tensorflow/contrib/tensorrt/segment/segment_test.cc
+++ b/tensorflow/contrib/tensorrt/segment/segment_test.cc
@@ -165,7 +165,7 @@ TEST_F(SegmentTest, Simple) {
   ASSERT_EQ(segments.size(), 1);
   std::vector<string> expected{"add0", "add1", "add2", "add3", "add4"};
   for (const auto& ex : expected) {
-    EXPECT_TRUE(segments[0].find(ex) != segments[0].end())
+    EXPECT_TRUE(segments[0].first.find(ex) != segments[0].first.end())
         << "Missing expected node " << ex;
   }
   TF_DeleteGraph(graph);
@@ -278,13 +278,13 @@ TEST_F(SegmentTest, Multiple) {
 
   std::vector<string> expected0{"add0", "add1", "add2", "add3"};
   for (const auto& ex : expected0) {
-    EXPECT_TRUE(segments[0].find(ex) != segments[0].end())
+    EXPECT_TRUE(segments[0].first.find(ex) != segments[0].first.end())
         << "Missing expected node " << ex;
   }
 
   std::vector<string> expected1{"add6", "add8"};
   for (const auto& ex : expected1) {
-    EXPECT_TRUE(segments[1].find(ex) != segments[1].end())
+    EXPECT_TRUE(segments[1].first.find(ex) != segments[1].first.end())
         << "Missing expected node " << ex;
   }
   TF_DeleteGraph(graph);
@@ -348,13 +348,13 @@ TEST_F(SegmentTest, BigIfElse) {
 
   std::vector<string> expected0{"add3", "add4", "add5", "add6", "add7"};
   for (const auto& ex : expected0) {
-    EXPECT_TRUE(segments[0].find(ex) != segments[0].end())
+    EXPECT_TRUE(segments[0].first.find(ex) != segments[0].first.end())
         << "Missing expected node " << ex;
   }
 
   std::vector<string> expected1{"add0", "add1"};
   for (const auto& ex : expected1) {
-    EXPECT_TRUE(segments[1].find(ex) != segments[1].end())
+    EXPECT_TRUE(segments[1].first.find(ex) != segments[1].first.end())
         << "Missing expected node " << ex;
   }
   TF_DeleteGraph(graph);
diff --git a/tensorflow/contrib/tensorrt/test/test_tftrt.py b/tensorflow/contrib/tensorrt/test/test_tftrt.py
index 2295320117..175ccd8006 100644
--- a/tensorflow/contrib/tensorrt/test/test_tftrt.py
+++ b/tensorflow/contrib/tensorrt/test/test_tftrt.py
@@ -66,7 +66,6 @@ def execute_graph(gdef, dumm_inp):
   """Run given graphdef once."""
   print("executing")
   gpu_options = cpb2.GPUOptions(per_process_gpu_memory_fraction=0.50)
-  #graph_options = cpb2.GraphOptions(rewrite_options=opt_config)
   sessconfig = cpb2.ConfigProto(gpu_options=gpu_options)
   ops.reset_default_graph()
   g = ops.Graph()
@@ -75,9 +74,6 @@ def execute_graph(gdef, dumm_inp):
         graph_def=gdef, return_elements=["input", "output"])
     inp = inp.outputs[0]
     out = out.outputs[0]
-  # with csess.Session(
-  #     config=cpb2.ConfigProto(gpu_options=gpu_options), graph=g) as sess:
-  #   val = sess.run(out, {inp: dumm_inp})
   with csess.Session(config=sessconfig, graph=g) as sess:
     val = sess.run(out, {inp: dumm_inp})
   return val
@@ -105,7 +101,7 @@ def execute_calibration(gdef, dumm_inp):
 
 
 def user(run_graph=execute_graph, run_calibration=execute_calibration):
-  """ Example function that converts a graph to TFTRT graph """
+  """Example function that converts a graph to TFTRT graph."""
 
   inp_dims = (100, 24, 24, 2)
   dummy_input = np.random.random_sample(inp_dims)
@@ -150,7 +146,7 @@ def user(run_graph=execute_graph, run_calibration=execute_calibration):
 
 
 def auto():
-  """ Run the conversion as an optimization pass"""
+  """Run the conversion as an optimization pass."""
   inp_dims = (100, 24, 24, 2)
   dummy_input = np.random.random_sample(inp_dims)
   orig_graph = get_simple_graph_def()
-- 
GitLab


From c45b05197623b375a056dd9577a778c5d5cc7d03 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Sat, 28 Apr 2018 23:30:22 +0300
Subject: [PATCH 030/395] [tf.data] A change to use Jenkins to test the Winsows
 build.

don't submit with this change!
---
 tensorflow/contrib/data/python/kernel_tests/resample_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/data/python/kernel_tests/resample_test.py b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
index b556525ce4..c08283a041 100644
--- a/tensorflow/contrib/data/python/kernel_tests/resample_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
@@ -60,7 +60,7 @@ class ResampleTest(test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(
       ("InitialDistributionKnown", True),
-      ("InitialDistributionUnknown", False))
+      ("InitialDistributionUnknown", True))  # THIS IS TO TEST THE WINDOWS BUILD DONT SUBMIT
   def testDistribution(self, initial_known):
     classes = np.random.randint(5, size=(20000,))  # Uniformly sampled
     target_dist = [0.9, 0.05, 0.05, 0.0, 0.0]
-- 
GitLab


From b384c339ee7d8440b6d4e39c09202c19f900aebe Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Sun, 29 Apr 2018 01:16:16 +0300
Subject: [PATCH 031/395] [tf.data] Possible bug fix to fix Winsows build.

---
 tensorflow/contrib/data/python/kernel_tests/resample_test.py | 4 ++--
 tensorflow/contrib/data/python/ops/resampling.py             | 1 -
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensorflow/contrib/data/python/kernel_tests/resample_test.py b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
index c08283a041..bbb8ca22f6 100644
--- a/tensorflow/contrib/data/python/kernel_tests/resample_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
@@ -60,9 +60,9 @@ class ResampleTest(test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(
       ("InitialDistributionKnown", True),
-      ("InitialDistributionUnknown", True))  # THIS IS TO TEST THE WINDOWS BUILD DONT SUBMIT
+      ("InitialDistributionUnknown", False))
   def testDistribution(self, initial_known):
-    classes = np.random.randint(5, size=(20000,))  # Uniformly sampled
+    classes = np.random.randint(5, size=(20000,), dtype=np.int64)
     target_dist = [0.9, 0.05, 0.05, 0.0, 0.0]
     initial_dist = [0.2] * 5 if initial_known else None
     dataset = dataset_ops.Dataset.from_tensor_slices(classes).shuffle(
diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index 1194b8447a..bad6edd514 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -79,7 +79,6 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
           lambda accept_prob, _: accept_prob)
       prob_of_original_ds = acceptance_and_original_prob_ds.map(
           lambda _, prob_original: prob_original)
-      prob_of_original = None
     filtered_ds = _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds,
                              class_values_ds, seed)
     # Prefetch filtered dataset for speed.
-- 
GitLab


From 9033bb2a175e344448772f5641020023badeacd8 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Sun, 29 Apr 2018 03:01:12 +0300
Subject: [PATCH 032/395] [tf.data] Undo previously unsuccessful bugfix, and
 try another one to fix the Windows build.

don't submit with this change, because it includes some debugging!
---
 tensorflow/contrib/data/python/kernel_tests/resample_test.py | 2 +-
 tensorflow/contrib/data/python/ops/resampling.py             | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/contrib/data/python/kernel_tests/resample_test.py b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
index bbb8ca22f6..b556525ce4 100644
--- a/tensorflow/contrib/data/python/kernel_tests/resample_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
@@ -62,7 +62,7 @@ class ResampleTest(test.TestCase, parameterized.TestCase):
       ("InitialDistributionKnown", True),
       ("InitialDistributionUnknown", False))
   def testDistribution(self, initial_known):
-    classes = np.random.randint(5, size=(20000,), dtype=np.int64)
+    classes = np.random.randint(5, size=(20000,))  # Uniformly sampled
     target_dist = [0.9, 0.05, 0.05, 0.0, 0.0]
     initial_dist = [0.2] * 5 if initial_known else None
     dataset = dataset_ops.Dataset.from_tensor_slices(classes).shuffle(
diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index bad6edd514..e65207f675 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -59,7 +59,7 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
 
     # Get initial distribution.
     if initial_dist is not None:
-      initial_dist_t = ops.convert_to_tensor(initial_dist, name="initial_dist")
+      initial_dist_t = math_ops.to_float(ops.convert_to_tensor(initial_dist, name="initial_dist"))
       acceptance_dist, prob_of_original = (
           _calculate_acceptance_probs_with_mixing(initial_dist_t,
                                                   target_dist_t))
@@ -291,4 +291,4 @@ def _calculate_acceptance_probs_with_mixing(initial_probs, target_probs):
 
   # TODO(joelshor): Simplify fraction, if possible.
   a_i = (ratio_l - m) / (max_ratio - m)
-  return a_i, m
\ No newline at end of file
+  return math_ops.to_float(a_i), math_ops.to_float(m)
\ No newline at end of file
-- 
GitLab


From 87f7d4b894c08031ba5942c1c391199de793eb88 Mon Sep 17 00:00:00 2001
From: ManHyuk <manhyuk@kw.ac.kr>
Date: Sun, 29 Apr 2018 16:07:33 +0900
Subject: [PATCH 033/395] fix typo

---
 .../tools/ci_build/windows/cpu/bazel/run_cc_test_windows.sh     | 2 +-
 .../tools/ci_build/windows/gpu/bazel/run_cc_test_windows.sh     | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/tools/ci_build/windows/cpu/bazel/run_cc_test_windows.sh b/tensorflow/tools/ci_build/windows/cpu/bazel/run_cc_test_windows.sh
index 748a961e44..dc9af221ec 100644
--- a/tensorflow/tools/ci_build/windows/cpu/bazel/run_cc_test_windows.sh
+++ b/tensorflow/tools/ci_build/windows/cpu/bazel/run_cc_test_windows.sh
@@ -44,7 +44,7 @@ source "tensorflow/tools/ci_build/windows/bazel/bazel_test_lib.sh" \
 
 run_configure_for_cpu_build
 
-# Compliling the following test is extremely slow with -c opt
+# Compiling the following test is extremely slow with -c opt
 slow_compiling_test="//tensorflow/core/kernels:eigen_backward_spatial_convolutions_test"
 
 # Find all the passing cc_tests on Windows and store them in a variable
diff --git a/tensorflow/tools/ci_build/windows/gpu/bazel/run_cc_test_windows.sh b/tensorflow/tools/ci_build/windows/gpu/bazel/run_cc_test_windows.sh
index f26f8727e5..f1114f4ffa 100644
--- a/tensorflow/tools/ci_build/windows/gpu/bazel/run_cc_test_windows.sh
+++ b/tensorflow/tools/ci_build/windows/gpu/bazel/run_cc_test_windows.sh
@@ -46,7 +46,7 @@ clean_output_base
 
 run_configure_for_gpu_build
 
-# Compliling the following test is extremely slow with -c opt
+# Compiling the following test is extremely slow with -c opt
 slow_compiling_test="//tensorflow/core/kernels:eigen_backward_spatial_convolutions_test"
 
 # Find all the passing cc_tests on Windows and store them in a variable
-- 
GitLab


From 9310de4af4816e5820d1907a9550ed427321eb33 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Mon, 30 Apr 2018 00:52:37 +0300
Subject: [PATCH 034/395] [tf.data] Add a bunch of debugging for Jenkins to run
 on the Windows build.

---
 tensorflow/contrib/data/python/ops/resampling.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index e65207f675..4caa25197e 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -91,9 +91,18 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
     elif prob_original_static == 0:
       return filtered_ds
     else:
+      print('class_values_ds.output_shapes: %s', class_values_ds.output_shapes)
+      print('class_values_ds.output_types: %s', class_values_ds.output_types)
+      print('dataset.output_shapes: %s', dataset.output_shapes)
+      print('dataset.output_types: %s', dataset.output_types)
+      print('filtered_ds.output_shapes: %s', filtered_ds.output_shapes)
+      print('filtered_ds.output_types: %s', filtered_ds.output_types)
+      weights = prob_of_original_ds.map(lambda prob: [(prob, 1.0 - prob)])
+      print('weights.output_shapes: %s', weights.output_shapes)
+      print('weights.output_types: %s', weights.output_types)
       return interleave_ops.sample_from_datasets(
           [dataset_ops.Dataset.zip((class_values_ds, dataset)), filtered_ds],
-          weights=prob_of_original_ds.map(lambda prob: [(prob, 1.0 - prob)]),
+          weights=weights,
           seed=seed)
 
   return _apply_fn
-- 
GitLab


From fc23d94b4c9c48c5abef87641cb6586fb9124d21 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Mon, 30 Apr 2018 11:45:27 +0300
Subject: [PATCH 035/395] [tf.data] Add a bunch of debugging for Jenkins to run
 on the Windows build.

---
 tensorflow/contrib/data/python/ops/BUILD      |  1 +
 .../contrib/data/python/ops/resampling.py     | 19 ++++++++++---------
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/tensorflow/contrib/data/python/ops/BUILD b/tensorflow/contrib/data/python/ops/BUILD
index 7a3e42cc72..299062212d 100644
--- a/tensorflow/contrib/data/python/ops/BUILD
+++ b/tensorflow/contrib/data/python/ops/BUILD
@@ -204,6 +204,7 @@ py_library(
         "//tensorflow/python:random_ops",
         "//tensorflow/python/data/ops:dataset_ops",
         "//third_party/py/numpy",
+        "//third_party/tensorflow/python:platform",
     ],
 )
 
diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index 4caa25197e..6b9ae772dc 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -31,6 +31,7 @@ from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import logging_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import random_ops
+from google3.third_party.tensorflow.python.platform import tf_logging as logging
 
 
 def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
@@ -91,15 +92,15 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
     elif prob_original_static == 0:
       return filtered_ds
     else:
-      print('class_values_ds.output_shapes: %s', class_values_ds.output_shapes)
-      print('class_values_ds.output_types: %s', class_values_ds.output_types)
-      print('dataset.output_shapes: %s', dataset.output_shapes)
-      print('dataset.output_types: %s', dataset.output_types)
-      print('filtered_ds.output_shapes: %s', filtered_ds.output_shapes)
-      print('filtered_ds.output_types: %s', filtered_ds.output_types)
+      logging.warn('class_values_ds.output_shapes: %s'% class_values_ds.output_shapes)
+      logging.warn('class_values_ds.output_types: %s'% class_values_ds.output_types)
+      logging.warn('dataset.output_shapes: %s'% dataset.output_shapes)
+      logging.warn('dataset.output_types: %s'% dataset.output_types)
+      logging.warn('filtered_ds.output_shapes: %s'% filtered_ds.output_shapes)
+      logging.warn('filtered_ds.output_types: %s'% filtered_ds.output_types)
       weights = prob_of_original_ds.map(lambda prob: [(prob, 1.0 - prob)])
-      print('weights.output_shapes: %s', weights.output_shapes)
-      print('weights.output_types: %s', weights.output_types)
+      logging.warn('weights.output_shapes: %s'% weights.output_shapes)
+      logging.warn('weights.output_types: %s'% weights.output_types)
       return interleave_ops.sample_from_datasets(
           [dataset_ops.Dataset.zip((class_values_ds, dataset)), filtered_ds],
           weights=weights,
@@ -151,7 +152,7 @@ def _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds, class_values_ds,
     return control_flow_ops.cond(
         math_ops.less(proportion_rejected, .5),
         lambda: accept_dist,
-        lambda: logging_ops.Print(  # pylint: disable=g-long-lambda
+        lambda: logging_ops.logging.warn(  # pylint: disable=g-long-lambda
             accept_dist, [proportion_rejected, initial_dist, accept_dist],
             message="Proportion of examples rejected by sampler is high: ",
             summarize=100,
-- 
GitLab


From 44ecd94792574be012d0a803c0b57ffec637c3e2 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Mon, 30 Apr 2018 11:54:56 +0300
Subject: [PATCH 036/395] [tf.data] Add a bunch of debugging for Jenkins to run
 on the Windows build.

---
 tensorflow/contrib/data/python/ops/BUILD         | 2 +-
 tensorflow/contrib/data/python/ops/resampling.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/contrib/data/python/ops/BUILD b/tensorflow/contrib/data/python/ops/BUILD
index 299062212d..6d94a2bd82 100644
--- a/tensorflow/contrib/data/python/ops/BUILD
+++ b/tensorflow/contrib/data/python/ops/BUILD
@@ -204,7 +204,7 @@ py_library(
         "//tensorflow/python:random_ops",
         "//tensorflow/python/data/ops:dataset_ops",
         "//third_party/py/numpy",
-        "//third_party/tensorflow/python:platform",
+        "//tensorflow/python:platform",
     ],
 )
 
diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index 6b9ae772dc..47bf6ecb58 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -31,7 +31,7 @@ from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import logging_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import random_ops
-from google3.third_party.tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.platform import tf_logging as logging
 
 
 def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
-- 
GitLab


From d4aa90c5eeb00bd46a2c7a5ee99d8eff04407e38 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Mon, 30 Apr 2018 12:20:20 +0300
Subject: [PATCH 037/395] [tf.data] Fix logging ops debug statement.

---
 tensorflow/contrib/data/python/ops/resampling.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index 47bf6ecb58..6be41985bb 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -152,7 +152,7 @@ def _filter_ds(dataset, acceptance_dist_ds, initial_dist_ds, class_values_ds,
     return control_flow_ops.cond(
         math_ops.less(proportion_rejected, .5),
         lambda: accept_dist,
-        lambda: logging_ops.logging.warn(  # pylint: disable=g-long-lambda
+        lambda: logging_ops.Print(  # pylint: disable=g-long-lambda
             accept_dist, [proportion_rejected, initial_dist, accept_dist],
             message="Proportion of examples rejected by sampler is high: ",
             summarize=100,
-- 
GitLab


From 19e7b123408fe6085294fe62479ddf0b31060ab2 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Mon, 30 Apr 2018 12:36:32 +0300
Subject: [PATCH 038/395] [tf.data] Properly format debug statements.

---
 tensorflow/contrib/data/python/ops/resampling.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index 6be41985bb..f041b7bcbf 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -92,15 +92,15 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
     elif prob_original_static == 0:
       return filtered_ds
     else:
-      logging.warn('class_values_ds.output_shapes: %s'% class_values_ds.output_shapes)
-      logging.warn('class_values_ds.output_types: %s'% class_values_ds.output_types)
-      logging.warn('dataset.output_shapes: %s'% dataset.output_shapes)
-      logging.warn('dataset.output_types: %s'% dataset.output_types)
-      logging.warn('filtered_ds.output_shapes: %s'% filtered_ds.output_shapes)
-      logging.warn('filtered_ds.output_types: %s'% filtered_ds.output_types)
+      logging.warn('class_values_ds.output_shapes: %s'% str(class_values_ds.output_shapes))
+      logging.warn('class_values_ds.output_types: %s'% str(class_values_ds.output_types))
+      logging.warn('dataset.output_shapes: %s'% str(dataset.output_shapes))
+      logging.warn('dataset.output_types: %s'% str(dataset.output_types))
+      logging.warn('filtered_ds.output_shapes: %s'% str(filtered_ds.output_shapes))
+      logging.warn('filtered_ds.output_types: %s'% str(filtered_ds.output_types))
       weights = prob_of_original_ds.map(lambda prob: [(prob, 1.0 - prob)])
-      logging.warn('weights.output_shapes: %s'% weights.output_shapes)
-      logging.warn('weights.output_types: %s'% weights.output_types)
+      logging.warn('weights.output_shapes: %s'% str(weights.output_shapes))
+      logging.warn('weights.output_types: %s'% str(weights.output_types))
       return interleave_ops.sample_from_datasets(
           [dataset_ops.Dataset.zip((class_values_ds, dataset)), filtered_ds],
           weights=weights,
-- 
GitLab


From a5a51ad3a1200e2e5ef46c140bab717422e41ca2 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 06:59:23 -0700
Subject: [PATCH 039/395] Adding a depthwise convolution kernel op (with label
 'cudnn_grouped_convolution') which forwards to cuDNN grouped convolutions.

PiperOrigin-RevId: 194780352
---
 tensorflow/core/kernels/BUILD                 |  10 +-
 .../core/kernels/conv_grad_filter_ops.cc      |  71 +++--
 .../core/kernels/conv_grad_input_ops.cc       |  74 +++--
 tensorflow/core/kernels/conv_grad_ops.cc      |   7 +-
 tensorflow/core/kernels/conv_ops.cc           |  85 +++---
 .../core/kernels/depthwise_conv_grad_op.cc    | 263 +++++++++++++++---
 tensorflow/core/kernels/depthwise_conv_op.cc  | 118 +++++---
 .../kernel_tests/depthwise_conv_op_test.py    | 222 +++++++++------
 tensorflow/stream_executor/cuda/cuda_dnn.cc   |  18 +-
 tensorflow/stream_executor/dnn.cc             |   1 +
 tensorflow/stream_executor/dnn.h              |   6 +
 11 files changed, 637 insertions(+), 238 deletions(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 6355f13654..3fb03cd5bd 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -3299,7 +3299,10 @@ tf_kernel_library(
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:nn_ops_op_lib",
-    ] + if_cuda(["@cub_archive//:cub"]),
+    ] + if_cuda([
+        "@cub_archive//:cub",
+        "@local_config_cuda//cuda:cudnn",
+    ]),
 )
 
 tf_kernel_library(
@@ -3310,12 +3313,15 @@ tf_kernel_library(
     prefix = "depthwise_conv_grad_op",
     deps = [
         ":bounds_check",
+        ":conv_ops",
         ":ops_util",
         "//tensorflow/core:core_cpu",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:nn_ops_op_lib",
-    ],
+    ] + if_cuda([
+        "@local_config_cuda//cuda:cudnn",
+    ]),
 )
 
 cc_library(
diff --git a/tensorflow/core/kernels/conv_grad_filter_ops.cc b/tensorflow/core/kernels/conv_grad_filter_ops.cc
index ef1e73e5ab..aca75176a5 100644
--- a/tensorflow/core/kernels/conv_grad_filter_ops.cc
+++ b/tensorflow/core/kernels/conv_grad_filter_ops.cc
@@ -96,7 +96,8 @@ template <typename T>
 struct LaunchConv2DBackpropFilterOp<CPUDevice, T> {
   void operator()(OpKernelContext* ctx, bool use_cudnn, bool cudnn_use_autotune,
                   const Tensor& out_backprop, const Tensor& input,
-                  int row_stride, int col_stride, const Padding& padding,
+                  int row_dilation, int col_dilation, int row_stride,
+                  int col_stride, const Padding& padding,
                   Tensor* filter_backprop, TensorFormat data_format) {
     const CPUDevice& d = ctx->eigen_device<CPUDevice>();
     functor::SpatialConvolutionBackwardFilter<CPUDevice, T>()(
@@ -275,7 +276,8 @@ class Conv2DFastBackpropFilterOp : public OpKernel {
 #endif
 
     LaunchConv2DBackpropFilterOp<Device, T>()(
-        context, false, false, out_backprop, input, dims.spatial_dims[0].stride,
+        context, false, false, out_backprop, input,
+        /*row_dilation=*/1, /*col_dilation=*/1, dims.spatial_dims[0].stride,
         dims.spatial_dims[1].stride, padding_, filter_backprop, data_format_);
   }
 
@@ -523,6 +525,11 @@ TF_CALL_float(REGISTER_CPU_KERNELS);
 TF_CALL_double(REGISTER_CPU_KERNELS);
 #undef REGISTER_CPU_KERNELS
 
+// To be used inside depthwise_conv_grad_op.cc.
+template struct LaunchConv2DBackpropFilterOp<CPUDevice, Eigen::half>;
+template struct LaunchConv2DBackpropFilterOp<CPUDevice, float>;
+template struct LaunchConv2DBackpropFilterOp<CPUDevice, double>;
+
 // GPU definitions.
 #if GOOGLE_CUDA
 // The slow version (but compiles for GPU)
@@ -690,10 +697,15 @@ void LaunchConv2DBackpropFilterOp<Eigen::GpuDevice, T>::operator()(
     return;
   }
 
+  // If the filter in-depth (filter_shape.dim_size(2)) is 1 and smaller than the
+  // input depth, it's a depthwise convolution. More generally, if the filter
+  // in-depth divides but is smaller than the input depth, it is a grouped
+  // convolution.
+  bool is_grouped_convolution = filter_shape.dim_size(2) != dims.in_depth;
   bool cudnn_disable_conv_1x1_optimization_ = CudnnDisableConv1x1Optimization();
   if (!cudnn_disable_conv_1x1_optimization_ &&
       dims.spatial_dims[0].filter_size == 1 &&
-      dims.spatial_dims[1].filter_size == 1 &&
+      dims.spatial_dims[1].filter_size == 1 && !is_grouped_convolution &&
       dims.spatial_dims[0].stride == 1 && dims.spatial_dims[1].stride == 1 &&
       data_format == FORMAT_NHWC) {
     const uint64 m = dims.in_depth;
@@ -734,9 +746,10 @@ void LaunchConv2DBackpropFilterOp<Eigen::GpuDevice, T>::operator()(
                  dims.spatial_dims[0].input_size &&
              dims.spatial_dims[1].filter_size ==
                  dims.spatial_dims[1].input_size &&
-             padding == VALID && data_format == FORMAT_NHWC) {
-    // The input data and filter have the same height/width, so call cublas
-    // directly.
+             !is_grouped_convolution && padding == VALID &&
+             data_format == FORMAT_NHWC) {
+    // The input data and filter have the same height/width, and we are not
+    // using grouped convolution, so call cublas directly.
     const uint64 m = dims.spatial_dims[0].input_size *
                      dims.spatial_dims[1].input_size * dims.in_depth;
     const uint64 k = dims.batch_size;
@@ -802,15 +815,16 @@ void LaunchConv2DBackpropFilterOp<Eigen::GpuDevice, T>::operator()(
   se::dnn::FilterDescriptor filter_desc;
   filter_desc.set_input_filter_height(dims.spatial_dims[0].filter_size)
       .set_input_filter_width(dims.spatial_dims[1].filter_size)
-      .set_input_feature_map_count(dims.in_depth)
-      .set_output_feature_map_count(dims.out_depth);
+      .set_input_feature_map_count(filter_shape.dim_size(2))
+      .set_output_feature_map_count(filter_shape.dim_size(3));
   se::dnn::ConvolutionDescriptor conv_desc;
   conv_desc.set_vertical_dilation_rate(dims.spatial_dims[0].dilation)
       .set_horizontal_dilation_rate(dims.spatial_dims[1].dilation)
       .set_vertical_filter_stride(dims.spatial_dims[0].stride)
       .set_horizontal_filter_stride(dims.spatial_dims[1].stride)
       .set_zero_padding_height(padding_rows / 2)
-      .set_zero_padding_width(padding_cols / 2);
+      .set_zero_padding_width(padding_cols / 2)
+      .set_group_count(dims.in_depth / filter_shape.dim_size(2));
 
   // NOTE(zhengxq):
   // cuDNN only supports the following layouts :
@@ -891,21 +905,22 @@ void LaunchConv2DBackpropFilterOp<Eigen::GpuDevice, T>::operator()(
   int device_id = stream->parent()->device_ordinal();
   DataType dtype = input.dtype();
   ConvParameters conv_parameters = {
-      dims.batch_size,                       // batch
-      dims.in_depth,                         // in_depths
-      {{input_desc.height(),                 // in_rows
-        input_desc.width()}},                // in_cols
-      dims.out_depth,                        // out_depths
-      {{dims.spatial_dims[0].filter_size,    // filter_rows
-        dims.spatial_dims[1].filter_size}},  // filter_cols
-      {{dims.spatial_dims[0].dilation,       // dilation_rows
-        dims.spatial_dims[1].dilation}},     // dilation_cols
-      {{dims.spatial_dims[0].stride,         // stride_rows
-        dims.spatial_dims[1].stride}},       // stride_cols
-      {{padding_rows,                        // padding_rows
-        padding_cols}},                      // padding_cols
-      dtype,                                 // tensor datatype
-      device_id,                             // device_id
+      dims.batch_size,                     // batch
+      dims.in_depth,                       // in_depths
+      {{input_desc.height(),               // in_rows
+        input_desc.width()}},              // in_cols
+      dims.out_depth,                      // out_depths
+      {{dims.spatial_dims[0].filter_size,  // filter_rows
+        dims.spatial_dims[1].filter_size,  // filter_cols
+        filter_shape.dim_size(2)}},        // filter_depth
+      {{dims.spatial_dims[0].dilation,     // dilation_rows
+        dims.spatial_dims[1].dilation}},   // dilation_cols
+      {{dims.spatial_dims[0].stride,       // stride_rows
+        dims.spatial_dims[1].stride}},     // stride_cols
+      {{padding_rows,                      // padding_rows
+        padding_cols}},                    // padding_cols
+      dtype,                               // tensor datatype
+      device_id,                           // device_id
   };
   AlgorithmConfig algorithm_config;
   if (cudnn_use_autotune && !AutoTuneConvBwdFilter::GetInstance()->Find(
@@ -1019,9 +1034,9 @@ namespace functor {
       typename TTypes<T, 4, int>::Tensor out, TensorFormat data_format); \
   extern template struct PadInput<GPUDevice, T, int, 4>;
 
-DECLARE_GPU_SPEC(double);
 DECLARE_GPU_SPEC(float);
 DECLARE_GPU_SPEC(Eigen::half);
+DECLARE_GPU_SPEC(double);
 #undef DECLARE_GPU_SPEC
 }  // namespace functor
 
@@ -1040,6 +1055,12 @@ REGISTER_KERNEL_BUILDER(Name("Conv2DBackpropFilter")
                             .TypeConstraint<Eigen::half>("T")
                             .HostMemory("filter_sizes"),
                         Conv2DSlowBackpropFilterOp<GPUDevice, Eigen::half>);
+
+// To be used inside depthwise_conv_grad_op.cc.
+template struct LaunchConv2DBackpropFilterOp<GPUDevice, float>;
+template struct LaunchConv2DBackpropFilterOp<GPUDevice, Eigen::half>;
+template struct LaunchConv2DBackpropFilterOp<GPUDevice, double>;
+
 #endif  // GOOGLE_CUDA
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/kernels/conv_grad_input_ops.cc b/tensorflow/core/kernels/conv_grad_input_ops.cc
index 35f2676023..63a775afa8 100644
--- a/tensorflow/core/kernels/conv_grad_input_ops.cc
+++ b/tensorflow/core/kernels/conv_grad_input_ops.cc
@@ -101,8 +101,9 @@ template <typename T>
 struct LaunchConv2DBackpropInputOp<CPUDevice, T> {
   void operator()(OpKernelContext* ctx, bool use_cudnn, bool cudnn_use_autotune,
                   const Tensor& out_backprop, const Tensor& filter,
-                  int row_stride, int col_stride, const Padding& padding,
-                  Tensor* in_backprop, TensorFormat data_format) {
+                  int row_dilation, int col_dilation, int row_stride,
+                  int col_stride, const Padding& padding, Tensor* in_backprop,
+                  TensorFormat data_format) {
     const CPUDevice& d = ctx->eigen_device<CPUDevice>();
     functor::SpatialConvolutionBackwardInput<CPUDevice, T>()(
         d, in_backprop->tensor<T, 4>(), filter.tensor<T, 4>(),
@@ -280,8 +281,8 @@ class Conv2DFastBackpropInputOp : public OpKernel {
 
     LaunchConv2DBackpropInputOp<Device, T>()(
         context, false, false, out_backprop, filter,
-        dims.spatial_dims[0].stride, dims.spatial_dims[1].stride, padding_,
-        in_backprop, data_format_);
+        /*row_dilation=*/1, /*col_dilation=*/1, dims.spatial_dims[0].stride,
+        dims.spatial_dims[1].stride, padding_, in_backprop, data_format_);
   }
 
  private:
@@ -595,6 +596,11 @@ TF_CALL_float(REGISTER_CPU_KERNELS);
 TF_CALL_double(REGISTER_CPU_KERNELS);
 #undef REGISTER_CPU_KERNELS
 
+// To be used inside depthwise_conv_grad_op.cc.
+template struct LaunchConv2DBackpropInputOp<CPUDevice, Eigen::half>;
+template struct LaunchConv2DBackpropInputOp<CPUDevice, float>;
+template struct LaunchConv2DBackpropInputOp<CPUDevice, double>;
+
 // GPU definitions.
 #if GOOGLE_CUDA
 // The slow version (but compiles for GPU)
@@ -761,8 +767,13 @@ void LaunchConv2DBackpropInputOp<GPUDevice, T>::operator()(
     return;
   }
 
+  // If the filter in-depth (filter_shape.dim_size(2)) is 1 and smaller than the
+  // input depth, it's a depthwise convolution. More generally, if the filter
+  // in-depth divides but is smaller than the input depth, it is a grouped
+  // convolution.
+  bool is_grouped_convolution = filter_shape.dim_size(2) != dims.in_depth;
   if (dims.spatial_dims[0].filter_size == 1 &&
-      dims.spatial_dims[1].filter_size == 1 &&
+      dims.spatial_dims[1].filter_size == 1 && !is_grouped_convolution &&
       dims.spatial_dims[0].stride == 1 && dims.spatial_dims[1].stride == 1 &&
       data_format == FORMAT_NHWC) {
     // 1x1 filter, so call cublas directly.
@@ -795,9 +806,10 @@ void LaunchConv2DBackpropInputOp<GPUDevice, T>::operator()(
                  dims.spatial_dims[0].input_size &&
              dims.spatial_dims[1].filter_size ==
                  dims.spatial_dims[1].input_size &&
-             padding == VALID && data_format == FORMAT_NHWC) {
-    // The input data and filter have the same height/width, so call cublas
-    // directly.
+             !is_grouped_convolution && padding == VALID &&
+             data_format == FORMAT_NHWC) {
+    // The input data and filter have the same height/width, and we are not
+    // using grouped convolution, so call cublas directly.
     const uint64 m = dims.batch_size;
     const uint64 k = dims.out_depth;
     const uint64 n = dims.spatial_dims[0].input_size *
@@ -856,15 +868,16 @@ void LaunchConv2DBackpropInputOp<GPUDevice, T>::operator()(
   se::dnn::FilterDescriptor filter_desc;
   filter_desc.set_input_filter_height(dims.spatial_dims[0].filter_size)
       .set_input_filter_width(dims.spatial_dims[1].filter_size)
-      .set_input_feature_map_count(dims.in_depth)
-      .set_output_feature_map_count(dims.out_depth);
+      .set_input_feature_map_count(filter_shape.dim_size(2))
+      .set_output_feature_map_count(filter_shape.dim_size(3));
   se::dnn::ConvolutionDescriptor conv_desc;
   conv_desc.set_vertical_dilation_rate(dims.spatial_dims[0].dilation)
       .set_horizontal_dilation_rate(dims.spatial_dims[1].dilation)
       .set_vertical_filter_stride(dims.spatial_dims[0].stride)
       .set_horizontal_filter_stride(dims.spatial_dims[1].stride)
       .set_zero_padding_height(padding_rows / 2)
-      .set_zero_padding_width(padding_cols / 2);
+      .set_zero_padding_width(padding_cols / 2)
+      .set_group_count(dims.in_depth / filter_shape.dim_size(2));
 
   // NOTE(keveman):
   // cuDNN only supports the following layouts :
@@ -940,21 +953,22 @@ void LaunchConv2DBackpropInputOp<GPUDevice, T>::operator()(
   int device_id = stream->parent()->device_ordinal();
   DataType dtype = out_backprop.dtype();
   ConvParameters conv_parameters = {
-      dims.batch_size,                       // batch
-      dims.in_depth,                         // in_depths
-      {{input_desc.height(),                 // in_rows
-        input_desc.width()}},                // in_cols
-      dims.out_depth,                        // out_depths
-      {{dims.spatial_dims[0].filter_size,    // filter_rows
-        dims.spatial_dims[1].filter_size}},  // filter_cols
-      {{dims.spatial_dims[0].dilation,       // dilation_rows
-        dims.spatial_dims[1].dilation}},     // dilation_cols
-      {{dims.spatial_dims[0].stride,         // stride_rows
-        dims.spatial_dims[1].stride}},       // stride_cols
-      {{padding_rows,                        // padding_rows
-        padding_cols}},                      // padding_cols
-      dtype,                                 // tensor data type
-      device_id,                             // device_id
+      dims.batch_size,                     // batch
+      dims.in_depth,                       // in_depths
+      {{input_desc.height(),               // in_rows
+        input_desc.width()}},              // in_cols
+      dims.out_depth,                      // out_depths
+      {{dims.spatial_dims[0].filter_size,  // filter_rows
+        dims.spatial_dims[1].filter_size,  // filter_cols
+        filter_shape.dim_size(2)}},        // filter_depths
+      {{dims.spatial_dims[0].dilation,     // dilation_rows
+        dims.spatial_dims[1].dilation}},   // dilation_cols
+      {{dims.spatial_dims[0].stride,       // stride_rows
+        dims.spatial_dims[1].stride}},     // stride_cols
+      {{padding_rows,                      // padding_rows
+        padding_cols}},                    // padding_cols
+      dtype,                               // tensor data type
+      device_id,                           // device_id
   };
   AlgorithmConfig algorithm_config;
   if (cudnn_use_autotune && !AutoTuneConvBwdData::GetInstance()->Find(
@@ -1092,9 +1106,9 @@ namespace functor {
       typename TTypes<T, 4, int>::Tensor out, TensorFormat data_format); \
   extern template struct PadInput<GPUDevice, T, int, 4>;
 
-DECLARE_GPU_SPEC(double);
 DECLARE_GPU_SPEC(float);
 DECLARE_GPU_SPEC(Eigen::half);
+DECLARE_GPU_SPEC(double);
 #undef DECLARE_GPU_SPEC
 }  // namespace functor
 
@@ -1113,6 +1127,12 @@ REGISTER_KERNEL_BUILDER(Name("Conv2DBackpropInput")
                             .TypeConstraint<Eigen::half>("T")
                             .HostMemory("input_sizes"),
                         Conv2DSlowBackpropInputOp<GPUDevice, Eigen::half>);
+
+// To be used inside depthwise_conv_grad_op.cc.
+template struct LaunchConv2DBackpropInputOp<GPUDevice, float>;
+template struct LaunchConv2DBackpropInputOp<GPUDevice, Eigen::half>;
+template struct LaunchConv2DBackpropInputOp<GPUDevice, double>;
+
 #endif  // GOOGLE_CUDA
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/kernels/conv_grad_ops.cc b/tensorflow/core/kernels/conv_grad_ops.cc
index 170ce31d17..5bf709af08 100644
--- a/tensorflow/core/kernels/conv_grad_ops.cc
+++ b/tensorflow/core/kernels/conv_grad_ops.cc
@@ -127,16 +127,17 @@ Status ConvBackpropComputeDimensionsV2(
   dims->in_depth = input_shape.dim_size(feature_dim);
   // The input and output feature dimensions are the second last and last
   // dimensions of the filter Tensor.
-  if (dims->in_depth != filter_shape.dim_size(num_dims - 2)) {
+  VLOG(2) << "input vs filter_in depth " << dims->in_depth << " "
+          << filter_shape.dim_size(num_dims - 2);
+  if (dims->in_depth % filter_shape.dim_size(num_dims - 2)) {
     return errors::InvalidArgument(
-        label, ": input and filter must have the same depth");
+        label, ": input depth must be evenly divisible by filter depth");
   }
   dims->out_depth = filter_shape.dim_size(num_dims - 1);
   if (dims->out_depth != out_backprop_shape.dim_size(feature_dim)) {
     return errors::InvalidArgument(
         label, ": filter and out_backprop must have the same out_depth");
   }
-
   dims->spatial_dims.resize(num_spatial_dims);
   for (int i = 0; i < num_spatial_dims; ++i) {
     int image_dim = GetTensorSpatialDimIndex(num_dims, data_format, i);
diff --git a/tensorflow/core/kernels/conv_ops.cc b/tensorflow/core/kernels/conv_ops.cc
index c6d36b40fe..3b9886eece 100644
--- a/tensorflow/core/kernels/conv_ops.cc
+++ b/tensorflow/core/kernels/conv_ops.cc
@@ -18,10 +18,16 @@ limitations under the License.
 #define USE_EIGEN_TENSOR
 #define EIGEN_USE_THREADS
 
+#if GOOGLE_CUDA
+#define EIGEN_USE_GPU
+#endif  // GOOGLE_CUDA
+
 #include "tensorflow/core/kernels/conv_ops.h"
+
 #include <string.h>
 #include <map>
 #include <vector>
+
 #include "tensorflow/core/framework/numeric_op.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
@@ -32,9 +38,6 @@ limitations under the License.
 #include "tensorflow/core/kernels/conv_2d.h"
 #include "tensorflow/core/kernels/deep_conv2d.h"
 #include "tensorflow/core/kernels/ops_util.h"
-#ifdef TENSORFLOW_USE_LIBXSMM_CONVOLUTIONS
-#include "tensorflow/core/kernels/xsmm_conv2d.h"
-#endif
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/gtl/array_slice.h"
 #include "tensorflow/core/lib/strings/numbers.h"
@@ -45,6 +48,10 @@ limitations under the License.
 #include "tensorflow/core/util/tensor_format.h"
 #include "tensorflow/core/util/use_cudnn.h"
 
+#ifdef TENSORFLOW_USE_LIBXSMM_CONVOLUTIONS
+#include "tensorflow/core/kernels/xsmm_conv2d.h"
+#endif
+
 #if GOOGLE_CUDA
 #include "tensorflow/core/kernels/conv_ops_gpu.h"
 #include "tensorflow/core/platform/stream_executor.h"
@@ -123,6 +130,10 @@ struct LaunchConv2DOp<CPUDevice, T> {
                                 "NHWC tensor format for now."));
       return;
     }
+    const int64 in_depth = GetTensorDim(input, data_format, 'C');
+    OP_REQUIRES(ctx, in_depth == filter.dim_size(2),
+                errors::Unimplemented("Generic conv implementation does not "
+                                      "support grouped convolutions for now."));
     LaunchGeneric<CPUDevice, T>()(ctx, input, filter, row_stride, col_stride,
                                   row_dilation, col_dilation, padding, output,
                                   data_format);
@@ -324,12 +335,13 @@ class Conv2DOp : public BinaryOp<T> {
     }
 
     // The last dimension for input is in_depth. It must be the same as the
-    // filter's in_depth.
+    // filter's in_depth or be evenly divisible by filter's in_depth.
     const int64 in_depth = GetTensorDim(input, data_format_, 'C');
-    OP_REQUIRES(context, in_depth == filter.dim_size(2),
+    const int64 patch_depth = filter.dim_size(2);
+    OP_REQUIRES(context, in_depth % patch_depth == 0,
                 errors::InvalidArgument(
-                    "input and filter must have the same depth: ", in_depth,
-                    " vs ", filter.dim_size(2)));
+                    "input depth must be evenly divisible by filter depth: ",
+                    in_depth, " vs ", patch_depth));
 
     // The last dimension for filter is out_depth.
     const int out_depth = static_cast<int>(filter.dim_size(3));
@@ -386,6 +398,7 @@ class Conv2DOp : public BinaryOp<T> {
     OP_REQUIRES_OK(context, context->allocate_output(0, out_shape, &output));
 
     VLOG(2) << "Conv2D: in_depth = " << in_depth
+            << ", patch_depth = " << patch_depth
             << ", input_cols = " << input_cols
             << ", filter_cols = " << filter_cols
             << ", input_rows = " << input_rows
@@ -450,7 +463,9 @@ TF_CALL_double(REGISTER_CPU);
 #endif  // USE_GEMM_FOR_CONV
 
 // To be used inside depthwise_conv_op.cc.
+template struct LaunchConv2DOp<CPUDevice, Eigen::half>;
 template struct LaunchConv2DOp<CPUDevice, float>;
+template struct LaunchConv2DOp<CPUDevice, double>;
 
 #if GOOGLE_CUDA
 int64 GetCudnnWorkspaceLimit(const string& envvar_in_mb,
@@ -498,13 +513,24 @@ void LaunchConv2DOp<GPUDevice, T>::operator()(
   }
 
   Tensor input = input_param;
-
-  if (filter.dim_size(0) == 1 && filter.dim_size(1) == 1 && row_dilation == 1 &&
-      col_dilation == 1 && row_stride == 1 && col_stride == 1 &&
-      data_format == FORMAT_NHWC) {
+  const int64 in_batch = GetTensorDim(input, data_format, 'N');
+  int64 in_rows = GetTensorDim(input, data_format, 'H');
+  int64 in_cols = GetTensorDim(input, data_format, 'W');
+  const int64 in_depths = GetTensorDim(input, data_format, 'C');
+  const int64 patch_rows = filter.dim_size(0);
+  const int64 patch_cols = filter.dim_size(1);
+  const int64 patch_depths = filter.dim_size(2);
+
+  // If the filter in-depth (patch_depths) is 1 and smaller than the input
+  // depth, it's a depthwise convolution. More generally, if the filter in-depth
+  // divides but is smaller than the input depth, it is a grouped convolution.
+  bool is_grouped_convolution = patch_depths != in_depths;
+  if (patch_rows == 1 && patch_cols == 1 && !is_grouped_convolution &&
+      row_dilation == 1 && col_dilation == 1 && row_stride == 1 &&
+      col_stride == 1 && data_format == FORMAT_NHWC) {
     // 1x1 filter, so call cublas directly.
-    const uint64 m = input.dim_size(0) * input.dim_size(1) * input.dim_size(2);
-    const uint64 k = filter.dim_size(2);
+    const uint64 m = in_batch * in_rows * in_cols;
+    const uint64 k = patch_depths;
     const uint64 n = filter.dim_size(3);
 
     auto a_ptr = AsDeviceMemory(input.template flat<T>().data(),
@@ -525,15 +551,14 @@ void LaunchConv2DOp<GPUDevice, T>::operator()(
                                       ", n=", n, ", k=", k));
     }
     return;
-  } else if (filter.dim_size(0) == input.dim_size(1) &&
-             filter.dim_size(1) == input.dim_size(2) && row_dilation == 1 &&
+  } else if (patch_rows == in_rows && patch_cols == in_cols &&
+             !is_grouped_convolution && row_dilation == 1 &&
              col_dilation == 1 && padding == VALID &&
              data_format == FORMAT_NHWC) {
     // The input data and filter have the same height/width, so call cublas
     // directly.
-    const uint64 m = input.dim_size(0);
-    const uint64 k =
-        filter.dim_size(0) * filter.dim_size(1) * filter.dim_size(2);
+    const uint64 m = in_batch;
+    const uint64 k = patch_rows * patch_cols * patch_depths;
     const uint64 n = filter.dim_size(3);
 
     auto a_ptr = AsDeviceMemory(input.template flat<T>().data(),
@@ -558,16 +583,10 @@ void LaunchConv2DOp<GPUDevice, T>::operator()(
 
   int padding_rows = 0;
   int padding_cols = 0;
-  const int64 in_batch = GetTensorDim(input, data_format, 'N');
-  int64 in_rows = GetTensorDim(input, data_format, 'H');
-  int64 in_cols = GetTensorDim(input, data_format, 'W');
-  const int64 in_depths = GetTensorDim(input, data_format, 'C');
   const int64 out_batch = GetTensorDim(*output, data_format, 'N');
   const int64 out_rows = GetTensorDim(*output, data_format, 'H');
   const int64 out_cols = GetTensorDim(*output, data_format, 'W');
   const int64 out_depths = GetTensorDim(*output, data_format, 'C');
-  const int64 patch_rows = filter.dim_size(0);
-  const int64 patch_cols = filter.dim_size(1);
   if (padding == SAME) {
     // Total padding on rows and cols is
     // Pr = (R' - 1) * S + (Kr - 1) * Dr + 1 - R
@@ -642,9 +661,9 @@ void LaunchConv2DOp<GPUDevice, T>::operator()(
       .set_feature_map_count(out_depths)
       .set_layout(se::dnn::DataLayout::kBatchDepthYX);
   se::dnn::FilterDescriptor filter_desc;
-  filter_desc.set_input_filter_height(filter.dim_size(0))
-      .set_input_filter_width(filter.dim_size(1))
-      .set_input_feature_map_count(filter.dim_size(2))
+  filter_desc.set_input_filter_height(patch_rows)
+      .set_input_filter_width(patch_cols)
+      .set_input_feature_map_count(patch_depths)
       .set_output_feature_map_count(filter.dim_size(3));
   se::dnn::ConvolutionDescriptor conv_desc;
   conv_desc.set_vertical_dilation_rate(row_dilation)
@@ -652,7 +671,8 @@ void LaunchConv2DOp<GPUDevice, T>::operator()(
       .set_vertical_filter_stride(row_stride)
       .set_horizontal_filter_stride(col_stride)
       .set_zero_padding_height(padding_rows / 2)
-      .set_zero_padding_width(padding_cols / 2);
+      .set_zero_padding_width(padding_cols / 2)
+      .set_group_count(in_depths / patch_depths);
 
   Tensor transformed_filter;
   OP_REQUIRES_OK(ctx, ctx->allocate_temp(
@@ -695,7 +715,8 @@ void LaunchConv2DOp<GPUDevice, T>::operator()(
         in_cols}},       // in_cols
       out_depths,        // out_depths
       {{patch_rows,      // filter_rows
-        patch_cols}},    // filter_cols
+        patch_cols,      // filter_cols
+        patch_depths}},  // filter_depths
       {{row_dilation,    // dilation_rows
         col_dilation}},  // dilation_cols
       {{row_stride,      // stride_rows
@@ -812,9 +833,9 @@ namespace functor {
       typename TTypes<T, 4, int>::Tensor out, TensorFormat data_format);     \
   extern template struct PadInput<GPUDevice, T, int, 4>
 
-DECLARE_GPU_SPEC(double);
 DECLARE_GPU_SPEC(float);
 DECLARE_GPU_SPEC(Eigen::half);
+DECLARE_GPU_SPEC(double);
 #undef DECLARE_GPU_SPEC
 }  // namespace functor
 
@@ -830,7 +851,9 @@ REGISTER_KERNEL_BUILDER(
     Conv2DOp<GPUDevice, double>);
 
 // To be used inside depthwise_conv_op.cc.
-template class LaunchConv2DOp<GPUDevice, float>;
+template struct LaunchConv2DOp<GPUDevice, float>;
+template struct LaunchConv2DOp<GPUDevice, Eigen::half>;
+template struct LaunchConv2DOp<GPUDevice, double>;
 
 #endif  // GOOGLE_CUDA
 
diff --git a/tensorflow/core/kernels/depthwise_conv_grad_op.cc b/tensorflow/core/kernels/depthwise_conv_grad_op.cc
index 91a9587174..7afa21acb9 100644
--- a/tensorflow/core/kernels/depthwise_conv_grad_op.cc
+++ b/tensorflow/core/kernels/depthwise_conv_grad_op.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_types.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/kernels/bounds_check.h"
+#include "tensorflow/core/kernels/conv_grad_ops.h"
 #include "tensorflow/core/kernels/depthwise_conv_op.h"
 #include "tensorflow/core/kernels/ops_util.h"
 #include "tensorflow/core/lib/core/status.h"
@@ -33,9 +34,11 @@ limitations under the License.
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/util/padding.h"
 #include "tensorflow/core/util/tensor_format.h"
+#include "tensorflow/core/util/use_cudnn.h"
 #include "tensorflow/core/util/work_sharder.h"
 
 #if GOOGLE_CUDA
+#include "cuda/include/cudnn.h"
 #include "tensorflow/core/platform/stream_executor.h"
 #endif  // GOOGLE_CUDA
 
@@ -509,8 +512,19 @@ static void DepthwiseConvBackpropInputReference(const DepthwiseArgs& args,
   }
 }
 
+// Extern template instantiated in conv_grad_input_ops.cc.
+extern template struct LaunchConv2DBackpropInputOp<CPUDevice, Eigen::half>;
+extern template struct LaunchConv2DBackpropInputOp<CPUDevice, float>;
+extern template struct LaunchConv2DBackpropInputOp<CPUDevice, double>;
+
 #if GOOGLE_CUDA
 
+// Extern template instantiated in conv_grad_input_ops.cc.
+extern template struct LaunchConv2DBackpropInputOp<GPUDevice, Eigen::half>;
+extern template struct LaunchConv2DBackpropInputOp<GPUDevice, float>;
+extern template struct LaunchConv2DBackpropInputOp<GPUDevice, double>;
+
+// Extern template instantiated in depthwise_conv_op_gpu.cu.cc.
 extern template struct LaunchDepthwiseConvBackpropInputOp<GPUDevice,
                                                           Eigen::half>;
 extern template struct LaunchDepthwiseConvBackpropInputOp<GPUDevice, float>;
@@ -548,6 +562,12 @@ class DepthwiseConv2dNativeBackpropInputOp : public OpKernel {
         errors::InvalidArgument("Current implementation does not yet support "
                                 "strides in the batch and depth dimensions."));
     OP_REQUIRES_OK(context, context->GetAttr("padding", &padding_));
+
+    // For in_depth == 1 and grouped convolutions.
+    use_cudnn_ = CanUseCudnn();
+    cudnn_use_autotune_ = CudnnUseAutotune();
+    use_cudnn_grouped_conv_ = false;
+    dtype_ = DataTypeToEnum<T>::value;
   }
 
   void Compute(OpKernelContext* context) override {
@@ -560,6 +580,7 @@ class DepthwiseConv2dNativeBackpropInputOp : public OpKernel {
             input_sizes.dims()));
     TensorShape input_shape;
     const int32* in_sizes_data = input_sizes.template flat<int32>().data();
+
     for (int i = 0; i < input_sizes.NumElements(); ++i) {
       OP_REQUIRES(context, in_sizes_data[i] >= 0,
                   errors::InvalidArgument("Dimension ", i,
@@ -568,27 +589,77 @@ class DepthwiseConv2dNativeBackpropInputOp : public OpKernel {
     }
     const TensorShape& filter_shape = filter.shape();
     EXTRACT_AND_VERIFY_DIMENSIONS("DepthwiseConv2DBackpropInput");
+
     Tensor* in_backprop = nullptr;
     OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
                                 {0}, 0, input_shape, &in_backprop));
-    auto out_backprop_ptr = out_backprop.template flat<T>().data();
-    auto filter_ptr = filter.template flat<T>().data();
-    auto in_backprop_ptr = in_backprop->template flat<T>().data();
+
     // If there is nothing to compute, return.
     if (input_shape.num_elements() == 0) {
       return;
     }
+
+    // If in_depth==1, this operation is just a standard convolution.
+    // Depthwise convolution is a special case of cuDNN's grouped convolution.
+    bool use_cudnn = use_cudnn_ && (in_depth == 1 || use_cudnn_grouped_conv_);
+
+    VLOG(2) << "DepthwiseConv2dNativeBackpropInput: "
+            << " Input: [" << batch << ", " << input_rows << ", " << input_cols
+            << ", " << in_depth << "]; Filter: [" << filter_rows << ", "
+            << filter_cols << ", " << in_depth << ", " << depth_multiplier
+            << "]; Output: [" << batch << ", " << out_rows << ", " << out_cols
+            << ", " << out_depth << "], stride = " << stride_
+            << ", pad_rows = " << pad_rows << ", pad_cols = " << pad_cols
+            << ", Use cuDNN: " << use_cudnn;
+
+    if (use_cudnn) {
+      // Reshape from TF depthwise filter to cuDNN grouped convolution filter:
+      //
+      //                  | TensorFlow       | cuDNN
+      // --------------------------------------------------------------------
+      // filter_out_depth | depth_multiplier | depth_multiplier * group_count
+      // filter_in_depth  | in_depth         | in_depth / group_count
+      //
+      // For depthwise convolution, we have group_count == in_depth.
+      int32 filter_in_depth = 1;
+      TensorShape shape =
+          TensorShape{filter_rows, filter_cols, filter_in_depth, out_depth};
+      Tensor reshaped_filter(/*type=*/dtype_);
+      OP_REQUIRES(
+          context, reshaped_filter.CopyFrom(filter, shape),
+          errors::Internal(
+              "Failed to reshape filter tensor for grouped convolution."));
+      // TODO(yangzihao): Send in arbitrary dilation rates after the dilated
+      // conv is supported.
+      launcher_(context, use_cudnn_, cudnn_use_autotune_, out_backprop,
+                reshaped_filter, /*row_dilation=*/1, /*col_dilation=*/1,
+                stride_, stride_, padding_, in_backprop, data_format_);
+      return;
+    }
+
+    auto out_backprop_ptr = out_backprop.template flat<T>().data();
+    auto filter_ptr = filter.template flat<T>().data();
+    auto in_backprop_ptr = in_backprop->template flat<T>().data();
     LaunchDepthwiseConvBackpropInputOp<Device, T>()(
         context, args, out_backprop_ptr, filter_ptr, in_backprop_ptr,
         data_format_);
   }
 
+ protected:
+  bool use_cudnn_grouped_conv_;
+
  private:
   std::vector<int32> strides_;
   Padding padding_;
   TensorFormat data_format_;
   int64 stride_;
 
+  // For in_depth == 1 and grouped convolutions.
+  LaunchConv2DBackpropInputOp<Device, T> launcher_;
+  bool use_cudnn_;
+  bool cudnn_use_autotune_;
+  DataType dtype_;
+
   TF_DISALLOW_COPY_AND_ASSIGN(DepthwiseConv2dNativeBackpropInputOp);
 };
 
@@ -597,23 +668,52 @@ class DepthwiseConv2dNativeBackpropInputOp : public OpKernel {
                               .Device(DEVICE_CPU)                    \
                               .TypeConstraint<T>("T"),               \
                           DepthwiseConv2dNativeBackpropInputOp<CPUDevice, T>);
+
+TF_CALL_half(REGISTER_CPU_KERNEL);
 TF_CALL_float(REGISTER_CPU_KERNEL);
+#if !defined(PLATFORM_WINDOWS) || !defined(_DEBUG)
 TF_CALL_double(REGISTER_CPU_KERNEL);
+#endif
 #undef REGISTER_CPU_KERNEL
 
 #if GOOGLE_CUDA
-REGISTER_KERNEL_BUILDER(Name("DepthwiseConv2dNativeBackpropInput")
-                            .Device(DEVICE_GPU)
-                            .TypeConstraint<float>("T")
-                            .HostMemory("input_sizes"),
-                        DepthwiseConv2dNativeBackpropInputOp<GPUDevice, float>);
-
-REGISTER_KERNEL_BUILDER(
-    Name("DepthwiseConv2dNativeBackpropInput")
-        .Device(DEVICE_GPU)
-        .TypeConstraint<double>("T")
-        .HostMemory("input_sizes"),
-    DepthwiseConv2dNativeBackpropInputOp<GPUDevice, double>);
+
+#define REGISTER_GPU_KERNEL(T)                                       \
+  REGISTER_KERNEL_BUILDER(Name("DepthwiseConv2dNativeBackpropInput") \
+                              .Device(DEVICE_GPU)                    \
+                              .TypeConstraint<T>("T")                \
+                              .HostMemory("input_sizes"),            \
+                          DepthwiseConv2dNativeBackpropInputOp<GPUDevice, T>)
+
+TF_CALL_half(REGISTER_GPU_KERNEL);
+TF_CALL_float(REGISTER_GPU_KERNEL);
+TF_CALL_double(REGISTER_GPU_KERNEL);
+#undef REGISTER_GPU_KERNEL
+
+#if CUDNN_VERSION >= 7000
+template <typename T>
+class DepthwiseConv2dGroupedConvBackpropInputOp
+    : public DepthwiseConv2dNativeBackpropInputOp<GPUDevice, T> {
+ public:
+  DepthwiseConv2dGroupedConvBackpropInputOp(OpKernelConstruction* context)
+      : DepthwiseConv2dNativeBackpropInputOp<GPUDevice, T>(context) {
+    this->use_cudnn_grouped_conv_ = true;
+  }
+};
+
+#define REGISTER_GROUPED_CONV_KERNEL(T)                              \
+  REGISTER_KERNEL_BUILDER(Name("DepthwiseConv2dNativeBackpropInput") \
+                              .Device(DEVICE_GPU)                    \
+                              .TypeConstraint<T>("T")                \
+                              .HostMemory("input_sizes")             \
+                              .Label("cudnn_grouped_convolution"),   \
+                          DepthwiseConv2dGroupedConvBackpropInputOp<T>)
+
+TF_CALL_half(REGISTER_GROUPED_CONV_KERNEL);
+TF_CALL_float(REGISTER_GROUPED_CONV_KERNEL);
+TF_CALL_double(REGISTER_GROUPED_CONV_KERNEL);
+#undef REGISTER_GROUPED_CONV_KERNEL
+#endif  // CUDNN_VERSION
 #endif  // GOOGLE_CUDA
 
 // Kernels to compute the gradients of the filters for depthwise convolution.
@@ -885,8 +985,19 @@ static void DepthwiseConvBackpropFilterReference(const DepthwiseArgs& args,
   }
 }
 
+// Extern template instantiated in conv_grad_filter_ops.cc.
+extern template struct LaunchConv2DBackpropFilterOp<CPUDevice, Eigen::half>;
+extern template struct LaunchConv2DBackpropFilterOp<CPUDevice, float>;
+extern template struct LaunchConv2DBackpropFilterOp<CPUDevice, double>;
+
 #if GOOGLE_CUDA
 
+// Extern template instantiated in conv_grad_filter_ops.cc.
+extern template struct LaunchConv2DBackpropFilterOp<GPUDevice, Eigen::half>;
+extern template struct LaunchConv2DBackpropFilterOp<GPUDevice, float>;
+extern template struct LaunchConv2DBackpropFilterOp<GPUDevice, double>;
+
+// Extern template instantiated in depthwise_conv_op_gpu.cu.cc.
 extern template struct LaunchDepthwiseConvBackpropFilterOp<GPUDevice,
                                                            Eigen::half>;
 extern template struct LaunchDepthwiseConvBackpropFilterOp<GPUDevice, float>;
@@ -924,6 +1035,21 @@ class DepthwiseConv2dNativeBackpropFilterOp : public OpKernel {
         errors::InvalidArgument("Current implementation does not yet support "
                                 "strides in the batch and depth dimensions."));
     OP_REQUIRES_OK(context, context->GetAttr("padding", &padding_));
+
+    // For in_depth == 1 and grouped convolutions.
+    use_cudnn_ = CanUseCudnn();
+    cudnn_use_autotune_ = CudnnUseAutotune();
+    use_cudnn_grouped_conv_ = false;
+
+    if (std::is_same<T, Eigen::half>::value) {
+      dtype_ = DT_HALF;
+    } else if (std::is_same<T, float>::value) {
+      dtype_ = DT_FLOAT;
+    } else if (std::is_same<T, double>::value) {
+      dtype_ = DT_DOUBLE;
+    } else {
+      LOG(ERROR) << "Only half, float, and double are supported.";
+    }
   }
 
   void Compute(OpKernelContext* context) override {
@@ -949,24 +1075,73 @@ class DepthwiseConv2dNativeBackpropFilterOp : public OpKernel {
     OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
                                 {1}, 0, filter_shape, &filter_backprop));
 
-    auto out_backprop_ptr = out_backprop.template flat<T>().data();
-    auto input_ptr = input.template flat<T>().data();
-    auto filter_backprop_ptr = filter_backprop->template flat<T>().data();
     // If there is nothing to compute, return.
     if (filter_shape.num_elements() == 0) {
       return;
     }
+
+    // If in_depth==1, this operation is just a standard convolution.
+    // Depthwise convolution is a special case of cuDNN's grouped convolution.
+    bool use_cudnn = use_cudnn_ && (in_depth == 1 || use_cudnn_grouped_conv_);
+
+    VLOG(2) << "DepthwiseConv2dNativeBackpropFilter: "
+            << " Input: [" << batch << ", " << input_rows << ", " << input_cols
+            << ", " << in_depth << "]; Filter: [" << filter_rows << ", "
+            << filter_cols << ", " << in_depth << ", " << depth_multiplier
+            << "]; Output: [" << batch << ", " << out_rows << ", " << out_cols
+            << ", " << out_depth << "], stride = " << stride_
+            << ", pad_rows = " << pad_rows << ", pad_cols = " << pad_cols
+            << ", Use cuDNN: " << use_cudnn;
+
+    if (use_cudnn) {
+      // Reshape from TF depthwise filter to cuDNN grouped convolution filter:
+      //
+      //                  | TensorFlow       | cuDNN
+      // --------------------------------------------------------------------
+      // filter_out_depth | depth_multiplier | depth_multiplier * group_count
+      // filter_in_depth  | in_depth         | in_depth / group_count
+      //
+      // For depthwise convolution, we have group_count == in_depth.
+      int32 filter_in_depth = 1;
+      TensorShape shape =
+          TensorShape{filter_rows, filter_cols, filter_in_depth, out_depth};
+      Tensor reshaped_filter(/*type=*/dtype_);
+      OP_REQUIRES(
+          context, reshaped_filter.CopyFrom(*filter_backprop, shape),
+          errors::Internal(
+              "Failed to reshape filter tensor for grouped convolution."));
+
+      // TODO(yangzihao): Send in arbitrary dilation rates after the dilated
+      // conv is supported.
+      launcher_(context, use_cudnn_, cudnn_use_autotune_, out_backprop, input,
+                /*row_dilation=*/1, /*col_dilation=*/1, stride_, stride_,
+                padding_, &reshaped_filter, data_format_);
+      return;
+    }
+
+    auto out_backprop_ptr = out_backprop.template flat<T>().data();
+    auto input_ptr = input.template flat<T>().data();
+    auto filter_backprop_ptr = filter_backprop->template flat<T>().data();
     LaunchDepthwiseConvBackpropFilterOp<Device, T>()(
         context, args, out_backprop_ptr, input_ptr, filter_backprop_ptr,
         data_format_);
   }
 
+ protected:
+  bool use_cudnn_grouped_conv_;
+
  private:
   std::vector<int32> strides_;
   Padding padding_;
   TensorFormat data_format_;
   int64 stride_;
 
+  // For in_depth == 1 and grouped convolutions.
+  LaunchConv2DBackpropFilterOp<Device, T> launcher_;
+  bool use_cudnn_;
+  bool cudnn_use_autotune_;
+  DataType dtype_;
+
   TF_DISALLOW_COPY_AND_ASSIGN(DepthwiseConv2dNativeBackpropFilterOp);
 };
 
@@ -976,24 +1151,50 @@ class DepthwiseConv2dNativeBackpropFilterOp : public OpKernel {
           .Device(DEVICE_CPU)                     \
           .TypeConstraint<T>("T"),                \
       DepthwiseConv2dNativeBackpropFilterOp<CPUDevice, T>);
+TF_CALL_half(REGISTER_CPU_KERNEL);
 TF_CALL_float(REGISTER_CPU_KERNEL);
+#if !defined(PLATFORM_WINDOWS) || !defined(_DEBUG)
 TF_CALL_double(REGISTER_CPU_KERNEL);
+#endif
 #undef REGISTER_CPU_KERNEL
 
 #if GOOGLE_CUDA
-REGISTER_KERNEL_BUILDER(
-    Name("DepthwiseConv2dNativeBackpropFilter")
-        .Device(DEVICE_GPU)
-        .TypeConstraint<float>("T")
-        .HostMemory("filter_sizes"),
-    DepthwiseConv2dNativeBackpropFilterOp<GPUDevice, float>);
-
-REGISTER_KERNEL_BUILDER(
-    Name("DepthwiseConv2dNativeBackpropFilter")
-        .Device(DEVICE_GPU)
-        .TypeConstraint<double>("T")
-        .HostMemory("filter_sizes"),
-    DepthwiseConv2dNativeBackpropFilterOp<GPUDevice, double>);
+#define REGISTER_GPU_KERNEL(T)                                        \
+  REGISTER_KERNEL_BUILDER(Name("DepthwiseConv2dNativeBackpropFilter") \
+                              .Device(DEVICE_GPU)                     \
+                              .TypeConstraint<T>("T")                 \
+                              .HostMemory("filter_sizes"),            \
+                          DepthwiseConv2dNativeBackpropFilterOp<GPUDevice, T>)
+
+TF_CALL_half(REGISTER_GPU_KERNEL);
+TF_CALL_float(REGISTER_GPU_KERNEL);
+TF_CALL_double(REGISTER_GPU_KERNEL);
+#undef REGISTER_GPU_KERNEL
+
+#if CUDNN_VERSION >= 7000
+template <typename T>
+class DepthwiseConv2dGroupedConvBackpropFilterOp
+    : public DepthwiseConv2dNativeBackpropFilterOp<GPUDevice, T> {
+ public:
+  DepthwiseConv2dGroupedConvBackpropFilterOp(OpKernelConstruction* context)
+      : DepthwiseConv2dNativeBackpropFilterOp<GPUDevice, T>(context) {
+    this->use_cudnn_grouped_conv_ = true;
+  }
+};
+
+#define REGISTER_GROUPED_CONV_KERNEL(T)                               \
+  REGISTER_KERNEL_BUILDER(Name("DepthwiseConv2dNativeBackpropFilter") \
+                              .Device(DEVICE_GPU)                     \
+                              .TypeConstraint<T>("T")                 \
+                              .HostMemory("filter_sizes")             \
+                              .Label("cudnn_grouped_convolution"),    \
+                          DepthwiseConv2dGroupedConvBackpropFilterOp<T>)
+
+TF_CALL_half(REGISTER_GROUPED_CONV_KERNEL);
+TF_CALL_float(REGISTER_GROUPED_CONV_KERNEL);
+TF_CALL_double(REGISTER_GROUPED_CONV_KERNEL);
+#undef REGISTER_GROUPED_CONV_KERNEL
+#endif  // CUDNN_VERSION
 #endif  // GOOGLE_CUDA
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/kernels/depthwise_conv_op.cc b/tensorflow/core/kernels/depthwise_conv_op.cc
index 6dedb1a61e..d5f4a68120 100644
--- a/tensorflow/core/kernels/depthwise_conv_op.cc
+++ b/tensorflow/core/kernels/depthwise_conv_op.cc
@@ -39,6 +39,7 @@ limitations under the License.
 #include "tensorflow/core/util/work_sharder.h"
 
 #if GOOGLE_CUDA
+#include "cuda/include/cudnn.h"
 #include "tensorflow/core/platform/stream_executor.h"
 #endif  // GOOGLE_CUDA
 
@@ -241,18 +242,22 @@ struct LaunchDepthwiseConvOp<CPUDevice, T> {
 };
 
 // Extern template instantiated in conv_ops.cc.
+extern template struct LaunchConv2DOp<CPUDevice, Eigen::half>;
 extern template struct LaunchConv2DOp<CPUDevice, float>;
+extern template struct LaunchConv2DOp<CPUDevice, double>;
 
 #if GOOGLE_CUDA
 
+// Extern template instantiated in conv_ops.cc.
+extern template struct LaunchConv2DOp<GPUDevice, Eigen::half>;
+extern template struct LaunchConv2DOp<GPUDevice, float>;
+extern template struct LaunchConv2DOp<GPUDevice, double>;
+
 // Extern template instantiated in depthwise_conv_op_gpu.cc.
 extern template struct LaunchDepthwiseConvOp<GPUDevice, Eigen::half>;
 extern template struct LaunchDepthwiseConvOp<GPUDevice, float>;
 extern template struct LaunchDepthwiseConvOp<GPUDevice, double>;
 
-// Extern template instantiated in conv_ops.cc.
-extern template struct LaunchConv2DOp<GPUDevice, float>;
-
 #endif
 
 template <typename Device, typename T>
@@ -284,9 +289,11 @@ class DepthwiseConv2dNativeOp : public BinaryOp<T> {
                                 "strides in the batch and depth dimensions."));
     OP_REQUIRES_OK(context, context->GetAttr("padding", &padding_));
 
-    // For special case when in_depth == 1.
+    // For in_depth == 1 and grouped convolutions.
     use_cudnn_ = CanUseCudnn();
     cudnn_use_autotune_ = CudnnUseAutotune();
+    use_cudnn_grouped_conv_ = false;
+    dtype_ = DataTypeToEnum<T>::value;
   }
 
   void Compute(OpKernelContext* context) override {
@@ -357,27 +364,47 @@ class DepthwiseConv2dNativeOp : public BinaryOp<T> {
     Tensor* output = nullptr;
     OP_REQUIRES_OK(context, context->allocate_output(0, out_shape, &output));
 
-    VLOG(2) << "DepthwiseConv2dNative: "
-            << " Input: [" << batch << ", " << input_rows << ", " << input_cols
-            << ", " << in_depth << "]; Filter: [" << filter_rows << ", "
-            << filter_cols << ", " << in_depth << ", " << depth_multiplier
-            << "]; stride = " << stride_ << ", pad_rows = " << pad_rows
-            << ", pad_cols = " << pad_cols << ", output: [" << batch << ", "
-            << out_rows << ", " << out_cols << ", " << out_depth << "]";
-
     // If there is nothing to compute, return.
     if (out_shape.num_elements() == 0) {
       return;
     }
 
-    // If in_depth==1, this operation is just a standard convolution, so
-    // invoke that op.
-    if (std::is_same<T, float>::value && in_depth == 1) {
+    // TODO(csigg): Have autotune decide if native is faster than cuDNN.
+    // If in_depth==1, this operation is just a standard convolution.
+    // Depthwise convolution is a special case of cuDNN's grouped convolution.
+    bool use_cudnn = use_cudnn_ && (in_depth == 1 || use_cudnn_grouped_conv_);
+
+    VLOG(2) << "DepthwiseConv2dNative: "
+            << " Input: [" << batch << ", " << input_rows << ", " << input_cols
+            << ", " << in_depth << "]; Filter: [" << filter_rows << ", "
+            << filter_cols << ", " << in_depth << ", " << depth_multiplier
+            << "]; Output: [" << batch << ", " << out_rows << ", " << out_cols
+            << ", " << out_depth << "], stride = " << stride_
+            << ", pad_rows = " << pad_rows << ", pad_cols = " << pad_cols
+            << ", Use cuDNN: " << use_cudnn;
+
+    if (use_cudnn) {
+      // Reshape from TF depthwise filter to cuDNN grouped convolution filter:
+      //
+      //                  | TensorFlow       | cuDNN
+      // --------------------------------------------------------------------
+      // filter_out_depth | depth_multiplier | depth_multiplier * group_count
+      // filter_in_depth  | in_depth         | in_depth / group_count
+      //
+      // For depthwise convolution, we have group_count == in_depth.
+      int32 filter_in_depth = 1;
+      TensorShape shape =
+          TensorShape{filter_rows, filter_cols, filter_in_depth, out_depth};
+      Tensor reshaped_filter(/*type=*/dtype_);
+      OP_REQUIRES(
+          context, reshaped_filter.CopyFrom(filter, shape),
+          errors::Internal(
+              "Failed to reshape filter tensor for grouped convolution."));
       // TODO(yangzihao): Send in arbitrary dilation rates after the dilated
       // conv is supported.
-      launcher_(context, use_cudnn_, cudnn_use_autotune_, input, filter,
-                /*row_dilation=*/1, /*col_dilation=*/1, stride_, stride_,
-                padding_, output, data_format_);
+      launcher_(context, use_cudnn_, cudnn_use_autotune_, input,
+                reshaped_filter, /*row_dilation=*/1, /*col_dilation=*/1,
+                stride_, stride_, padding_, output, data_format_);
       return;
     }
 
@@ -403,6 +430,9 @@ class DepthwiseConv2dNativeOp : public BinaryOp<T> {
                                        output_ptr, data_format_);
   }
 
+ protected:
+  bool use_cudnn_grouped_conv_;
+
  private:
   std::vector<int32> strides_;
   Padding padding_;
@@ -410,10 +440,11 @@ class DepthwiseConv2dNativeOp : public BinaryOp<T> {
 
   int64 stride_;  // in height/width dimension.
 
-  // For the case in_depth == 1.
+  // For in_depth == 1 and grouped convolutions.
   LaunchConv2DOp<Device, T> launcher_;
   bool use_cudnn_;
   bool cudnn_use_autotune_;
+  DataType dtype_;
 
   TF_DISALLOW_COPY_AND_ASSIGN(DepthwiseConv2dNativeOp);
 };
@@ -421,7 +452,7 @@ class DepthwiseConv2dNativeOp : public BinaryOp<T> {
 #define REGISTER_CPU_KERNEL(T)                                                 \
   REGISTER_KERNEL_BUILDER(                                                     \
       Name("DepthwiseConv2dNative").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
-      DepthwiseConv2dNativeOp<CPUDevice, T>);
+      DepthwiseConv2dNativeOp<CPUDevice, T>)
 
 TF_CALL_half(REGISTER_CPU_KERNEL);
 TF_CALL_float(REGISTER_CPU_KERNEL);
@@ -430,19 +461,38 @@ TF_CALL_double(REGISTER_CPU_KERNEL);
 #endif
 
 #if GOOGLE_CUDA
-REGISTER_KERNEL_BUILDER(Name("DepthwiseConv2dNative")
-                            .Device(DEVICE_GPU)
-                            .TypeConstraint<Eigen::half>("T"),
-                        DepthwiseConv2dNativeOp<GPUDevice, Eigen::half>);
-
-REGISTER_KERNEL_BUILDER(
-    Name("DepthwiseConv2dNative").Device(DEVICE_GPU).TypeConstraint<float>("T"),
-    DepthwiseConv2dNativeOp<GPUDevice, float>);
-
-REGISTER_KERNEL_BUILDER(Name("DepthwiseConv2dNative")
-                            .Device(DEVICE_GPU)
-                            .TypeConstraint<double>("T"),
-                        DepthwiseConv2dNativeOp<GPUDevice, double>);
-#endif
+
+#define REGISTER_GPU_KERNEL(T)                                                 \
+  REGISTER_KERNEL_BUILDER(                                                     \
+      Name("DepthwiseConv2dNative").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
+      DepthwiseConv2dNativeOp<GPUDevice, T>)
+
+TF_CALL_half(REGISTER_GPU_KERNEL);
+TF_CALL_float(REGISTER_GPU_KERNEL);
+TF_CALL_double(REGISTER_GPU_KERNEL);
+
+#if CUDNN_VERSION >= 7000
+template <typename T>
+class DepthwiseConv2dGroupedConvOp
+    : public DepthwiseConv2dNativeOp<GPUDevice, T> {
+ public:
+  DepthwiseConv2dGroupedConvOp(OpKernelConstruction* context)
+      : DepthwiseConv2dNativeOp<GPUDevice, T>(context) {
+    this->use_cudnn_grouped_conv_ = true;
+  }
+};
+
+#define REGISTER_GROUPED_CONV_KERNEL(T)                            \
+  REGISTER_KERNEL_BUILDER(Name("DepthwiseConv2dNative")            \
+                              .Device(DEVICE_GPU)                  \
+                              .TypeConstraint<T>("T")              \
+                              .Label("cudnn_grouped_convolution"), \
+                          DepthwiseConv2dGroupedConvOp<T>)
+
+TF_CALL_half(REGISTER_GROUPED_CONV_KERNEL);
+TF_CALL_float(REGISTER_GROUPED_CONV_KERNEL);
+TF_CALL_double(REGISTER_GROUPED_CONV_KERNEL);
+#endif  // CUDNN_VERSION
+#endif  // GOOGLE_CUDA
 
 }  // namespace tensorflow
diff --git a/tensorflow/python/kernel_tests/depthwise_conv_op_test.py b/tensorflow/python/kernel_tests/depthwise_conv_op_test.py
index f7ae1a0f37..659dc0419a 100644
--- a/tensorflow/python/kernel_tests/depthwise_conv_op_test.py
+++ b/tensorflow/python/kernel_tests/depthwise_conv_op_test.py
@@ -22,12 +22,15 @@ import numpy as np
 
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import gradient_checker
 from tensorflow.python.ops import nn_impl
 from tensorflow.python.ops import nn_ops
 import tensorflow.python.ops.nn_grad  # pylint: disable=unused-import
 from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging
 
 
 def ConfigsToTest():
@@ -98,6 +101,7 @@ class DepthwiseConv2DTest(test.TestCase):
                     padding,
                     data_type,
                     use_gpu,
+                    grouped_conv=False,
                     data_format="NHWC"):
     """Verifies the output values of the convolution function.
 
@@ -110,25 +114,26 @@ class DepthwiseConv2DTest(test.TestCase):
       padding: Padding type.
       data_type: The data type to use.
       use_gpu: Whether to use GPU.
+      grouped_conv: Whether to use cuDNN 7's grouped convolution.
       data_format: The data_format of the input. "NHWC" or "NCHW".
     """
-    total_size_1 = 1
-    total_size_2 = 1
+    input_size = 1
+    filter_size = 1
     for s in tensor_in_sizes:
-      total_size_1 *= s
+      input_size *= s
     for s in filter_in_sizes:
-      total_size_2 *= s
+      filter_size *= s
     # Initializes the input and filter tensor with numbers incrementing from 1.
-    x1 = [f * 1.0 for f in range(1, total_size_1 + 1)]
-    x2 = [f * 1.0 for f in range(1, total_size_2 + 1)]
-    with self.test_session(use_gpu=use_gpu) as sess:
-      if data_type == dtypes.float16:
-        tolerance = 1e-5
-      elif data_type == dtypes.float32:
-        tolerance = 1e-5
-      else:
-        self.assertEqual(data_type, dtypes.float64)
-        tolerance = 1e-8
+    x1 = [f * 1.0 / input_size for f in range(1, input_size + 1)]
+    x2 = [f * 1.0 / filter_size for f in range(1, filter_size + 1)]
+    ops.reset_default_graph()
+    graph = ops.get_default_graph()
+    with self.test_session(graph=graph, use_gpu=use_gpu) as sess:
+      tolerance = {
+          dtypes.float16: 4e-2,
+          dtypes.float32: 1e-8,
+          dtypes.float64: 1e-13,
+      }[data_type]
 
       t1 = constant_op.constant(x1, shape=tensor_in_sizes, dtype=data_type)
       t1.set_shape(tensor_in_sizes)
@@ -142,25 +147,39 @@ class DepthwiseConv2DTest(test.TestCase):
         native_t1 = array_ops.transpose(t1, [0, 3, 1, 2])
         strides = [1, 1, stride, stride]
 
-      conv_native = nn_ops.depthwise_conv2d_native(
-          native_t1,
-          t2,
-          strides=strides,
-          data_format=data_format,
-          padding=padding)
+      with sess.graph._kernel_label_map({
+          "DepthwiseConv2dNative": "cudnn_grouped_convolution"
+      } if grouped_conv else {}):
+        conv_native = nn_ops.depthwise_conv2d_native(
+            native_t1,
+            t2,
+            strides=strides,
+            data_format=data_format,
+            padding=padding)
 
       if data_format == "NCHW":
         # Transpose back from NCHW to NHWC
         conv_native = array_ops.transpose(conv_native, [0, 2, 3, 1])
 
+      try:
+        native_result = sess.run(conv_native)
+      except errors.InvalidArgumentError as e:
+        # Grouped convolution kernel is only registered for cuDNN 7. Silently
+        # return when we are running on an earlier version or without GPU.
+        if e.message.startswith(
+            "No OpKernel was registered to support Op 'DepthwiseConv2dNative'"):
+          tf_logging.warn("Skipping grouped convolution test")
+          return
+        raise e
+
       conv_interface = nn_impl.depthwise_conv2d(
           t1, t2, strides=[1, stride, stride, 1], padding=padding)
-
-      native_result = sess.run(conv_native)
       interface_result = sess.run(conv_interface)
 
-    print("data_type:", data_type, "use_gpu:", use_gpu, "max diff = ",
-          np.amax(np.absolute(native_result - interface_result)))
+    tf_logging.info(
+        "data_type: %r, use_gpu: %r, grouped_conv: %r, max diff = %f",
+        data_type, use_gpu, grouped_conv,
+        np.amax(np.absolute(native_result - interface_result)))
     self.assertArrayNear(
         np.ravel(native_result), np.ravel(interface_result), tolerance)
     self.assertShapeEqual(native_result, conv_native)
@@ -169,11 +188,22 @@ class DepthwiseConv2DTest(test.TestCase):
   def testDepthwiseConv2D(self):
     for index, (input_size, filter_size, _, stride,
                 padding) in enumerate(ConfigsToTest()):
-      print("Testing DepthwiseConv2D,", index, "th config:", input_size, "*",
-            filter_size, "stride:", stride, "padding:", padding)
+      tf_logging.info(
+          "Testing DepthwiseConv2D, %dth config: %r * %r, stride: %d, padding: "
+          "%s", index, input_size, filter_size, stride, padding)
       for data_type in [dtypes.float16, dtypes.float32, dtypes.float64]:
+        tf_logging.info("Testing without grouped_conv")
         self._VerifyValues(
             input_size, filter_size, stride, padding, data_type, use_gpu=True)
+        tf_logging.info("Testing with grouped_conv")
+        self._VerifyValues(
+            input_size,
+            filter_size,
+            stride,
+            padding,
+            data_type,
+            use_gpu=True,
+            grouped_conv=True)
 
   def testDepthwiseConv2DFormat(self):
     if not test.is_gpu_available():
@@ -181,8 +211,9 @@ class DepthwiseConv2DTest(test.TestCase):
 
     for index, (input_size, filter_size, _, stride,
                 padding) in enumerate(ConfigsToTest()):
-      print("Testing DepthwiseConv2DFormat,", index, "th config:", input_size,
-            "*", filter_size, "stride:", stride, "padding:", padding)
+      tf_logging.info(
+          "Testing DepthwiseConv2DFormat, %dth config: %r * %r, stride: %d, "
+          "padding: %s", index, input_size, filter_size, stride, padding)
       for data_type in [dtypes.float16, dtypes.float32, dtypes.float64]:
         self._VerifyValues(
             input_size,
@@ -226,7 +257,7 @@ class DepthwiseConv2DTest(test.TestCase):
       conv = nn_ops.depthwise_conv2d_native(
           t1, t2, strides=[1, stride, stride, 1], padding=padding)
       value = sess.run(conv)
-    print("value = ", value)
+    tf_logging.info("value = %r", value)
     self.assertArrayNear(expected, np.ravel(value), 1e-5)
     self.assertShapeEqual(value, conv)
 
@@ -296,7 +327,7 @@ class DepthwiseConv2DTest(test.TestCase):
         expected=expected_output,
         use_gpu=True)
 
-  # Gradient checkers.This tests depthwise gradient computations for both
+  # Gradient checkers. This tests depthwise gradient computations for both
   # BackpropFilter and BackpropInput by comparing gradients computed by the
   # depthwise gradient ops with the gradients computed numerically (details can
   # be found in the compute_gradient_error().
@@ -310,6 +341,7 @@ class DepthwiseConv2DTest(test.TestCase):
                                 data_type,
                                 test_input,
                                 use_gpu,
+                                grouped_conv=False,
                                 data_format="NHWC"):
     input_size = 1
     for x in input_shape:
@@ -319,14 +351,14 @@ class DepthwiseConv2DTest(test.TestCase):
       filter_size *= x
     input_data = [x * 1.0 / input_size for x in range(0, input_size)]
     filter_data = [x * 1.0 / filter_size for x in range(0, filter_size)]
-    with self.test_session(use_gpu=use_gpu):
-      if data_type == dtypes.float16:
-        tolerance = 0.002
-      elif data_type == dtypes.float32:
-        tolerance = 0.002
-      else:
-        self.assertEqual(data_type, dtypes.float64)
-        tolerance = 1e-8
+    ops.reset_default_graph()
+    graph = ops.get_default_graph()
+    with self.test_session(graph=graph, use_gpu=use_gpu) as sess:
+      tolerance = {
+          dtypes.float16: 2e-0,
+          dtypes.float32: 5e-4,
+          dtypes.float64: 1e-12,
+      }[data_type]
 
       input_tensor = constant_op.constant(
           input_data, shape=input_shape, dtype=data_type, name="input")
@@ -347,35 +379,49 @@ class DepthwiseConv2DTest(test.TestCase):
         ]
         strides = [1, 1, stride, stride]
 
-      depthwise_conv2d = nn_ops.depthwise_conv2d_native(
-          native_input,
-          filter_tensor,
-          strides,
-          padding,
-          data_format=data_format,
-          name="depthwise_conv2d")
+      with sess.graph._kernel_label_map({
+          "DepthwiseConv2dNative": "cudnn_grouped_convolution",
+          "DepthwiseConv2dNativeBackpropInput": "cudnn_grouped_convolution",
+          "DepthwiseConv2dNativeBackpropFilter": "cudnn_grouped_convolution",
+      } if grouped_conv else {}):
+        depthwise_conv2d = nn_ops.depthwise_conv2d_native(
+            native_input,
+            filter_tensor,
+            strides,
+            padding,
+            data_format=data_format,
+            name="depthwise_conv2d")
 
       self.assertEqual(output_shape, depthwise_conv2d.get_shape())
-      if test_input:
-        err = gradient_checker.compute_gradient_error(
-            native_input, input_shape, depthwise_conv2d, output_shape)
-      else:
-        err = gradient_checker.compute_gradient_error(filter_tensor,
-                                                      filter_shape,
-                                                      depthwise_conv2d,
-                                                      output_shape)
-      print("data_type:", data_type, "use_gpu:", use_gpu, ", error = ", err)
+
+      try:
+        if test_input:
+          err = gradient_checker.compute_gradient_error(
+              native_input, input_shape, depthwise_conv2d, output_shape)
+        else:
+          err = gradient_checker.compute_gradient_error(
+              filter_tensor, filter_shape, depthwise_conv2d, output_shape)
+      except errors.InvalidArgumentError as e:
+        # Grouped convolution kernel is only registered for cuDNN 7. Silently
+        # return when we are running on an earlier version or without GPU.
+        if grouped_conv and e.message.startswith(
+            "No OpKernel was registered to support Op 'DepthwiseConv2dNative'"):
+          tf_logging.warn("Skipping grouped convolution test")
+          return
+        raise e
+
+      tf_logging.info(
+          "data_type: %r, use_gpu: %r, grouped_conv: %r, error = %f", data_type,
+          use_gpu, grouped_conv, err)
       self.assertLess(err, tolerance)
 
   def testDepthwiseConv2DInputGrad(self):
     for index, (input_size, filter_size, output_size, stride,
                 padding) in enumerate(CheckGradConfigsToTest()):
-      print("Testing DepthwiseConv2DInputGrad,", index, "th config:",
-            input_size, "*", filter_size, "stride:", stride, "padding:",
-            padding)
-      # Note: float16 test for DepthwiseConv2DInputGrad is not enabled,
-      # calculations are not very precise.
-      for data_type in [dtypes.float32, dtypes.float64]:
+      tf_logging.info(
+          "Testing DepthwiseConv2DInputGrad, %dth config: %r * %r, stride: %d, "
+          "padding: %s", index, input_size, filter_size, stride, padding)
+      for data_type in [dtypes.float16, dtypes.float32, dtypes.float64]:
         self._ConstructAndTestGradient(
             input_size,
             filter_size,
@@ -385,6 +431,16 @@ class DepthwiseConv2DTest(test.TestCase):
             data_type,
             test_input=True,
             use_gpu=True)
+        self._ConstructAndTestGradient(
+            input_size,
+            filter_size,
+            output_size,
+            stride,
+            padding,
+            data_type,
+            test_input=True,
+            use_gpu=True,
+            grouped_conv=True)
 
   def testDepthwiseConv2DInputGradFormat(self):
     if not test.is_gpu_available():
@@ -392,12 +448,11 @@ class DepthwiseConv2DTest(test.TestCase):
 
     for index, (input_size, filter_size, output_size, stride,
                 padding) in enumerate(CheckGradConfigsToTest()):
-      print("Testing DepthwiseConv2DInputGradFormat,", index, "th config:",
-            input_size, "*", filter_size, "stride:", stride, "padding:",
-            padding)
-      # Note: float16 test for DepthwiseConv2DInputGradFormat is not enabled,
-      # calculations are not very precise.
-      for data_type in [dtypes.float32, dtypes.float64]:
+      tf_logging.info(
+          "Testing DepthwiseConv2DInputGradFormat, %dth config: %r * %r, "
+          "stride: %d, padding: %s", index, input_size, filter_size, stride,
+          padding)
+      for data_type in [dtypes.float16, dtypes.float32, dtypes.float64]:
         self._ConstructAndTestGradient(
             input_size,
             filter_size,
@@ -412,12 +467,10 @@ class DepthwiseConv2DTest(test.TestCase):
   def testDepthwiseConv2DFilterGrad(self):
     for index, (input_size, filter_size, output_size, stride,
                 padding) in enumerate(CheckGradConfigsToTest()):
-      print("Testing DepthwiseConv2DFilterGrad,", index, "th config:",
-            input_size, "*", filter_size, "stride:", stride, "padding:",
-            padding)
-      # Note: float16 test for DepthwiseConv2DFilterGrad is not enabled,
-      # calculations are not very precise.
-      for data_type in [dtypes.float32, dtypes.float64]:
+      tf_logging.info(
+          "Testing DepthwiseConv2DFilterGrad, %dth config: %r * %r, stride: "
+          "%d, padding: %s", index, input_size, filter_size, stride, padding)
+      for data_type in [dtypes.float16, dtypes.float32, dtypes.float64]:
         self._ConstructAndTestGradient(
             input_size,
             filter_size,
@@ -434,12 +487,11 @@ class DepthwiseConv2DTest(test.TestCase):
 
     for index, (input_size, filter_size, output_size, stride,
                 padding) in enumerate(CheckGradConfigsToTest()):
-      print("Testing DepthwiseConv2DFilterGradFormat,", index, "th config:",
-            input_size, "*", filter_size, "stride:", stride, "padding:",
-            padding)
-      # Note: float16 test for DepthwiseConv2DFilterGradFormat is not enabled,
-      # calculations are not very precise.
-      for data_type in [dtypes.float32, dtypes.float64]:
+      tf_logging.info(
+          "Testing DepthwiseConv2DFilterGradFormat, %dth config: %r * %r, "
+          "stride: %d, padding: %s", index, input_size, filter_size, stride,
+          padding)
+      for data_type in [dtypes.float16, dtypes.float32, dtypes.float64]:
         self._ConstructAndTestGradient(
             input_size,
             filter_size,
@@ -494,9 +546,10 @@ class DepthwiseConv2DTest(test.TestCase):
   def testDepthwiseConv2DInputGradCompare(self):
     for index, (input_size, filter_size, output_size, stride,
                 padding) in enumerate(ConfigsToTest()):
-      print("Testing DepthwiseConv2DInputGradCompare,", index, "th config:",
-            input_size, "*", filter_size, "stride:", stride, "padding:",
-            padding)
+      tf_logging.info(
+          "Testing DepthwiseConv2DInputGradCompare, %dth config: %r * %r, "
+          "stride: %d, padding: %s", index, input_size, filter_size, stride,
+          padding)
       self._CompareBackpropInputFloat(input_size, filter_size, output_size,
                                       stride, padding)
       self._CompareBackpropInputDouble(input_size, filter_size, output_size,
@@ -545,9 +598,10 @@ class DepthwiseConv2DTest(test.TestCase):
   def testDepthwiseConv2DFilterGradCompare(self):
     for index, (input_size, filter_size, output_size, stride,
                 padding) in enumerate(ConfigsToTest()):
-      print("Testing DepthwiseConv2DFilterGradCompare,", index, "th config:",
-            input_size, "*", filter_size, "stride:", stride, "padding:",
-            padding)
+      tf_logging.info(
+          "Testing DepthwiseConv2DFilterGradCompare, %dth config: %r * %r, "
+          "stride: %d, padding: %s", index, input_size, filter_size, stride,
+          padding)
       self._CompareBackpropFilterFloat(input_size, filter_size, output_size,
                                        stride, padding)
       self._CompareBackpropFilterDouble(input_size, filter_size, output_size,
diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc
index 42a77aa3f8..773cac2c40 100644
--- a/tensorflow/stream_executor/cuda/cuda_dnn.cc
+++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc
@@ -337,7 +337,9 @@ CUDNN_DNN_ROUTINE_EACH_R6_WITH_STREAM(
 #if CUDNN_VERSION >= 7000
 #define CUDNN_DNN_ROUTINE_EACH_R7(__macro)                    \
   __macro(cudnnSetConvolutionMathType)                        \
-  __macro(cudnnSetRNNMatrixMathType)
+  __macro(cudnnSetRNNMatrixMathType)                          \
+  __macro(cudnnSetConvolutionGroupCount)                      \
+  __macro(cudnnGetConvolutionGroupCount)
 
 // clang-format on
 CUDNN_DNN_ROUTINE_EACH_R7(STREAM_EXECUTOR_CUDNN_WRAP)
@@ -779,6 +781,20 @@ class ScopedConvolutionDescriptor {
     // NOTE(benbarsdell): This only applies if tensor op math is enabled
     //                      and algo selection is set to Default.
     this->set_use_tensor_op_math(true);
+
+#if CUDNN_MAJOR >= 7
+    VLOG(2) << "Requesting grouped convolution: "
+            << convolution_descriptor.group_count();
+    status = wrap::cudnnSetConvolutionGroupCount(
+        parent_, handle_, convolution_descriptor.group_count());
+    if (status != CUDNN_STATUS_SUCCESS) {
+      LOG(FATAL) << "could not set cudnn convolution group count: "
+                 << ToString(status);
+    }
+#else
+    CHECK_EQ(convolution_descriptor.group_count(), 1)
+        << "Requested grouped convolution for cuDNN version < 7";
+#endif
   }
 
   void set_use_tensor_op_math(bool use_tensor_op_math) {
diff --git a/tensorflow/stream_executor/dnn.cc b/tensorflow/stream_executor/dnn.cc
index 031c82d3f4..eed93efc8d 100644
--- a/tensorflow/stream_executor/dnn.cc
+++ b/tensorflow/stream_executor/dnn.cc
@@ -434,6 +434,7 @@ ConvolutionDescriptor::ConvolutionDescriptor(int ndims)
       filter_strides_(ndims, 1),
       dilation_rates_(ndims, 1),
       pad_alignment_(PadAlignment::kDefault),
+      group_count_(1),
       ndims_(ndims) {}
 
 ConvolutionDescriptor::ConvolutionDescriptor()
diff --git a/tensorflow/stream_executor/dnn.h b/tensorflow/stream_executor/dnn.h
index 0c2e083b39..18606eb717 100644
--- a/tensorflow/stream_executor/dnn.h
+++ b/tensorflow/stream_executor/dnn.h
@@ -543,6 +543,10 @@ class ConvolutionDescriptor {
     pad_alignment_ = pad_alignment;
     return *this;
   }
+  ConvolutionDescriptor& set_group_count(int group_count) {
+    group_count_ = group_count;
+    return *this;
+  }
   int64 zero_padding_height() const {
     return GetDim(zero_padding_, DimIndex::Y);
   }
@@ -566,6 +570,7 @@ class ConvolutionDescriptor {
   int filter_stride(DimIndex dim) const { return GetDim(filter_strides_, dim); }
   int dilation_rate(DimIndex dim) const { return GetDim(dilation_rates_, dim); }
   PadAlignment pad_alignment() const { return pad_alignment_; }
+  int group_count() const { return group_count_; }
   int ndims() const { return ndims_; }
 
   std::vector<int64> strides() const { return filter_strides_; }
@@ -578,6 +583,7 @@ class ConvolutionDescriptor {
   std::vector<int64> filter_strides_;
   std::vector<int64> dilation_rates_;
   PadAlignment pad_alignment_;
+  int group_count_;
   int ndims_;
   // TODO(leary) cudnn provides these fields, but need to characterize what
   // their effect is -- they may be boolean rather than integral.
-- 
GitLab


From 541bd480c43ca48fcb1f4353d92687019b4cb765 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Mon, 30 Apr 2018 18:06:21 +0300
Subject: [PATCH 040/395] [tf.data] Explicitly make test's dataset int64.

---
 tensorflow/contrib/data/python/kernel_tests/resample_test.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/contrib/data/python/kernel_tests/resample_test.py b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
index b556525ce4..bdc003a8a5 100644
--- a/tensorflow/contrib/data/python/kernel_tests/resample_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/resample_test.py
@@ -65,6 +65,7 @@ class ResampleTest(test.TestCase, parameterized.TestCase):
     classes = np.random.randint(5, size=(20000,))  # Uniformly sampled
     target_dist = [0.9, 0.05, 0.05, 0.0, 0.0]
     initial_dist = [0.2] * 5 if initial_known else None
+    classes = math_ops.to_int64(classes)  # needed for Windows build.
     dataset = dataset_ops.Dataset.from_tensor_slices(classes).shuffle(
         200, seed=21).map(lambda c: (c, string_ops.as_string(c))).repeat()
 
-- 
GitLab


From 5da0d0022e08e60a30b88e4ef28c7f864e50fd1e Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Mon, 30 Apr 2018 18:26:08 +0300
Subject: [PATCH 041/395] [tf.data] Removed debug code.

---
 tensorflow/contrib/data/python/ops/BUILD         |  1 -
 tensorflow/contrib/data/python/ops/resampling.py | 16 +++-------------
 2 files changed, 3 insertions(+), 14 deletions(-)

diff --git a/tensorflow/contrib/data/python/ops/BUILD b/tensorflow/contrib/data/python/ops/BUILD
index 6d94a2bd82..7a3e42cc72 100644
--- a/tensorflow/contrib/data/python/ops/BUILD
+++ b/tensorflow/contrib/data/python/ops/BUILD
@@ -204,7 +204,6 @@ py_library(
         "//tensorflow/python:random_ops",
         "//tensorflow/python/data/ops:dataset_ops",
         "//third_party/py/numpy",
-        "//tensorflow/python:platform",
     ],
 )
 
diff --git a/tensorflow/contrib/data/python/ops/resampling.py b/tensorflow/contrib/data/python/ops/resampling.py
index f041b7bcbf..bad6edd514 100644
--- a/tensorflow/contrib/data/python/ops/resampling.py
+++ b/tensorflow/contrib/data/python/ops/resampling.py
@@ -31,7 +31,6 @@ from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import logging_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import random_ops
-from tensorflow.python.platform import tf_logging as logging
 
 
 def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
@@ -60,7 +59,7 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
 
     # Get initial distribution.
     if initial_dist is not None:
-      initial_dist_t = math_ops.to_float(ops.convert_to_tensor(initial_dist, name="initial_dist"))
+      initial_dist_t = ops.convert_to_tensor(initial_dist, name="initial_dist")
       acceptance_dist, prob_of_original = (
           _calculate_acceptance_probs_with_mixing(initial_dist_t,
                                                   target_dist_t))
@@ -92,18 +91,9 @@ def rejection_resample(class_func, target_dist, initial_dist=None, seed=None):
     elif prob_original_static == 0:
       return filtered_ds
     else:
-      logging.warn('class_values_ds.output_shapes: %s'% str(class_values_ds.output_shapes))
-      logging.warn('class_values_ds.output_types: %s'% str(class_values_ds.output_types))
-      logging.warn('dataset.output_shapes: %s'% str(dataset.output_shapes))
-      logging.warn('dataset.output_types: %s'% str(dataset.output_types))
-      logging.warn('filtered_ds.output_shapes: %s'% str(filtered_ds.output_shapes))
-      logging.warn('filtered_ds.output_types: %s'% str(filtered_ds.output_types))
-      weights = prob_of_original_ds.map(lambda prob: [(prob, 1.0 - prob)])
-      logging.warn('weights.output_shapes: %s'% str(weights.output_shapes))
-      logging.warn('weights.output_types: %s'% str(weights.output_types))
       return interleave_ops.sample_from_datasets(
           [dataset_ops.Dataset.zip((class_values_ds, dataset)), filtered_ds],
-          weights=weights,
+          weights=prob_of_original_ds.map(lambda prob: [(prob, 1.0 - prob)]),
           seed=seed)
 
   return _apply_fn
@@ -301,4 +291,4 @@ def _calculate_acceptance_probs_with_mixing(initial_probs, target_probs):
 
   # TODO(joelshor): Simplify fraction, if possible.
   a_i = (ratio_l - m) / (max_ratio - m)
-  return math_ops.to_float(a_i), math_ops.to_float(m)
\ No newline at end of file
+  return a_i, m
\ No newline at end of file
-- 
GitLab


From aa2405ee79dbcfabb8862ef3e1f8ca60e52159a0 Mon Sep 17 00:00:00 2001
From: Alexandre Passos <apassos@google.com>
Date: Mon, 30 Apr 2018 09:29:31 -0700
Subject: [PATCH 042/395] Fixes to tape gradient for providing outputs and
 having multiple targets.

PiperOrigin-RevId: 194796304
---
 tensorflow/c/eager/tape.h                | 65 ++++++++++--------------
 tensorflow/python/eager/backprop.py      |  8 ++-
 tensorflow/python/eager/backprop_test.py | 20 ++++++++
 3 files changed, 53 insertions(+), 40 deletions(-)

diff --git a/tensorflow/c/eager/tape.h b/tensorflow/c/eager/tape.h
index 97c323b872..8026076b9e 100644
--- a/tensorflow/c/eager/tape.h
+++ b/tensorflow/c/eager/tape.h
@@ -380,49 +380,39 @@ Status InitialGradients(const VSpace<Gradient, BackwardFunction>& vspace,
                         gtl::ArraySlice<Gradient*> output_gradients,
                         const TensorTape& tensor_tape,
                         const OpTape<BackwardFunction>& op_tape,
-                        const gtl::FlatMap<int64, int64>& tensor_usage_counts,
                         gtl::FlatMap<int64, std::vector<Gradient*>>* result) {
   for (int i = 0; i < target_tensor_ids.size(); ++i) {
     const int64 id = target_tensor_ids[i];
-    if (tensor_usage_counts.find(id) != tensor_usage_counts.end()) {
-      if (!output_gradients.empty() && output_gradients[i] != nullptr) {
-        // TODO(apassos) figure out how to print debugging information here.
-        return errors::InvalidArgument(
-            "A gradient was provided for a tensor which is used as part of the "
-            "computation.");
-      }
-    } else {
-      if (output_gradients.empty() || output_gradients[i] == nullptr) {
-        auto tensor_it = tensor_tape.find(id);
-        if (tensor_it != tensor_tape.end() && tensor_it->second != -1) {
-          auto op_it = op_tape.find(tensor_it->second);
-          if (op_it == op_tape.end()) {
-            return errors::Internal(
-                "Internal state of the gradient tape is invalid: "
-                "failed to find operation producing a tensor");
-          }
-          bool found = false;
-          for (int j = 0; j < op_it->second.output_tensor_info.size(); ++j) {
-            if (op_it->second.output_tensor_info[j].id == id) {
-              found = true;
-              (*result)[id].push_back(
-                  vspace.Ones(op_it->second.output_tensor_info[j].shape,
-                              op_it->second.output_tensor_info[j].dtype));
-              break;
-            }
-          }
-          if (!found) {
-            return errors::Internal(
-                "Internal state of the gradient tape is invalid: "
-                "none of operations outputs match expected tensor");
+    if (output_gradients.empty() || output_gradients[i] == nullptr) {
+      auto tensor_it = tensor_tape.find(id);
+      if (tensor_it != tensor_tape.end() && tensor_it->second != -1) {
+        auto op_it = op_tape.find(tensor_it->second);
+        if (op_it == op_tape.end()) {
+          return errors::Internal(
+              "Internal state of the gradient tape is invalid: "
+              "failed to find operation producing a tensor");
+        }
+        bool found = false;
+        for (int j = 0; j < op_it->second.output_tensor_info.size(); ++j) {
+          if (op_it->second.output_tensor_info[j].id == id) {
+            found = true;
+            (*result)[id].push_back(
+                vspace.Ones(op_it->second.output_tensor_info[j].shape,
+                            op_it->second.output_tensor_info[j].dtype));
+            break;
           }
-        } else {
-          // No record of the target tensor found on the tape, so no gradient
-          // needs to be computed from it. Do nothing.
+        }
+        if (!found) {
+          return errors::Internal(
+              "Internal state of the gradient tape is invalid: "
+              "none of operations outputs match expected tensor");
         }
       } else {
-        (*result)[id].push_back(output_gradients[i]);
+        // No record of the target tensor found on the tape, so no gradient
+        // needs to be computed from it. Do nothing.
       }
+    } else {
+      (*result)[id].push_back(output_gradients[i]);
     }
   }
   return Status::OK();
@@ -451,8 +441,7 @@ Status GradientTape<Gradient, BackwardFunction>::ComputeGradient(
       InitialStack(state.op_tape, state.op_missing_tensor);
   gtl::FlatMap<int64, std::vector<Gradient*>> gradients;
   Status s = InitialGradients(vspace, target_tensor_ids, output_gradients,
-                              tensor_tape_, state.op_tape,
-                              state.tensor_usage_counts, &gradients);
+                              tensor_tape_, state.op_tape, &gradients);
   auto cleanup = [this, &state]() {
     if (!persistent_) {
       // Release all backprop functions
diff --git a/tensorflow/python/eager/backprop.py b/tensorflow/python/eager/backprop.py
index 92774d4d50..07aec59cc8 100644
--- a/tensorflow/python/eager/backprop.py
+++ b/tensorflow/python/eager/backprop.py
@@ -740,7 +740,7 @@ class GradientTape(object):
     """Computes the gradient using operations recorded in context of this tape.
 
     Args:
-      target: Tensor to be differentiated.
+      target: Tensor (or list of tensors) to be differentiated.
       sources: a list or nested structure of Tensors or Variables. `target`
         will be differentiated against elements in `sources`.
       output_gradients: a list of gradients, one for each element of
@@ -762,8 +762,12 @@ class GradientTape(object):
     flat_sources = nest.flatten(sources)
     flat_sources = [_handle_or_self(x) for x in flat_sources]
 
+    if output_gradients is not None:
+      output_gradients = [None if x is None else ops.convert_to_tensor(x)
+                          for x in nest.flatten(output_gradients)]
+
     flat_grad = imperative_grad.imperative_grad(
-        _default_vspace, self._tape, [target], flat_sources,
+        _default_vspace, self._tape, nest.flatten(target), flat_sources,
         output_gradients=output_gradients)
 
     if not self._persistent:
diff --git a/tensorflow/python/eager/backprop_test.py b/tensorflow/python/eager/backprop_test.py
index 991b4dbe7a..8d9959fe20 100644
--- a/tensorflow/python/eager/backprop_test.py
+++ b/tensorflow/python/eager/backprop_test.py
@@ -96,6 +96,26 @@ class BackpropTest(test.TestCase):
     self.assertAllEqual(grads_and_vars[0][0], 1.0)
     self.assertAllEqual(id(grads_and_vars[0][1]), id(x))
 
+  def testTwoTargets(self):
+    with backprop.GradientTape() as t:
+      x = constant_op.constant(3.0)
+      y = constant_op.constant(2.0)
+      t.watch([x, y])
+      xx = 2 * x
+      yy = 3 * y
+    dx, dy = t.gradient([xx, yy], [x, y])
+    self.assertAllEqual(dx, 2.0)
+    self.assertAllEqual(dy, 3.0)
+
+  def testOutputGradUsedInComputation(self):
+    with backprop.GradientTape() as t:
+      x = constant_op.constant(3.0)
+      y = constant_op.constant(2.0)
+      t.watch([x, y])
+      loss = x * y
+    dx, = t.gradient([loss, x], [x], output_gradients=[1.0, 2.0])
+    self.assertAllEqual(dx, 4.0)
+
   def testDy(self):
 
     def f(x):
-- 
GitLab


From 1872f29b52d4bc4e32502715f461c4150e8c66a9 Mon Sep 17 00:00:00 2001
From: Tom Hennigan <tomhennigan@google.com>
Date: Mon, 30 Apr 2018 09:31:54 -0700
Subject: [PATCH 043/395] Clarify return type for defun as zero or more
 `tf.Tensor`s.

PiperOrigin-RevId: 194796621
---
 tensorflow/python/eager/function.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py
index 426ee4c215..741bd2ac9c 100644
--- a/tensorflow/python/eager/function.py
+++ b/tensorflow/python/eager/function.py
@@ -716,8 +716,7 @@ def defun(func):
   objects. Non-Tensor python objects are treated as constants, and new function
   definitions are created internally based on their values.
 
-  func must return a tf.Tensor (NOT a Tensor) or a list of tf.Tensor (NOT a
-  Tensor).
+  func must return zero or more `tf.Tensor`.
 
   Control flow constructs (e.g., `if`, `while`) are not yet compatible with
   `defun`.
@@ -748,7 +747,7 @@ def defun(func):
 
   Returns:
      A callable that will execute the compiled function (and return zero
-     or more Tensor objects).
+     or more `tf.Tensor` objects).
   """
   # TODO(apassos): deal with captured global state. Deal with control flow.
   try:
-- 
GitLab


From a3ae05d256a9fe82d9a2e50d3f72c3361c1162e4 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 09:46:59 -0700
Subject: [PATCH 044/395] Remove manifest_merger that is being removed from
 Bazel 0.13.0.

PiperOrigin-RevId: 194798790
---
 tensorflow/contrib/lite/examples/android/BUILD | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/contrib/lite/examples/android/BUILD b/tensorflow/contrib/lite/examples/android/BUILD
index 4928012997..5700007256 100644
--- a/tensorflow/contrib/lite/examples/android/BUILD
+++ b/tensorflow/contrib/lite/examples/android/BUILD
@@ -42,7 +42,6 @@ android_binary(
     custom_package = "org.tensorflow.lite.demo",
     inline_constants = 1,
     manifest = "AndroidManifest.xml",
-    manifest_merger = "android",
     nocompress_extensions = [
         ".tflite",
     ],
-- 
GitLab


From 09e529ff5adb916e40481563698dee72e8a15162 Mon Sep 17 00:00:00 2001
From: Ayush Dubey <ayushd@google.com>
Date: Mon, 30 Apr 2018 10:36:00 -0700
Subject: [PATCH 045/395] Prepare nodes that will be allocated using
 ScopedAllocator.

This includes changes to Executor that (1) set scope_id on nodes that are
decorated with _scoped_allocator attribute, (2) mark such nodes to never
forward input.

PiperOrigin-RevId: 194807086
---
 tensorflow/core/common_runtime/executor.cc | 114 +++++++++++++++++++--
 tensorflow/core/graph/graph.cc             |   1 +
 tensorflow/core/graph/graph.h              |  10 +-
 3 files changed, 115 insertions(+), 10 deletions(-)

diff --git a/tensorflow/core/common_runtime/executor.cc b/tensorflow/core/common_runtime/executor.cc
index 0c461a9ee9..e389eb9b2a 100644
--- a/tensorflow/core/common_runtime/executor.cc
+++ b/tensorflow/core/common_runtime/executor.cc
@@ -322,6 +322,7 @@ class GraphView {
 
   void Initialize(const Graph* g);
   Status SetAllocAttrs(const Graph* g, const Device* device);
+  void SetScopedAllocatorAttrs(const std::vector<const Node*>& sa_nodes);
 
   NodeItem* node(size_t id) const {
     DCHECK_GE(id, 0);
@@ -566,11 +567,46 @@ char* GraphView::InitializeNode(char* ptr, const Node* n) {
     DCHECK_EQ(item->input_type(i), n->input_type(i));
   }
 
-  uint8* output_types = item->output_type_base();
-  for (int i = 0; i < num_outputs; i++) {
-    output_types[i] = static_cast<uint8>(n->output_type(i));
-    DCHECK_EQ(item->output_type(i), n->output_type(i));
+  // Check ScopedAllocatorAttrs and forward_from.  Also assign output_types.
+  {
+    std::vector<int> forward_input;
+    Status fwd_status =
+        GetNodeAttr(n->attrs(), "_forward_input", &forward_input);
+    std::vector<int> scoped_allocator_attrs;
+    Status sa_status =
+        GetNodeAttr(n->attrs(), "_scoped_allocator", &scoped_allocator_attrs);
+
+    int* forward_from = item->forward_from_base();
+    uint8* output_types = item->output_type_base();
+    for (int i = 0; i < num_outputs; ++i) {
+      output_types[i] = static_cast<uint8>(n->output_type(i));
+      DCHECK_EQ(item->output_type(i), n->output_type(i));
+
+      forward_from[i] = OpKernelContext::Params::kNoReservation;
+      if (sa_status.ok()) {
+        for (int j = 0; j < scoped_allocator_attrs.size(); j += 2) {
+          if (scoped_allocator_attrs[j] == i) {
+            // This output slot must be explicitly allocated from a
+            // ScopedAllocator.
+            forward_from[i] = OpKernelContext::Params::kNeverForward;
+            DCHECK_EQ(output_attrs[i].scope_id, 0);
+            output_attrs[i].scope_id = scoped_allocator_attrs[j + 1];
+          }
+        }
+      }
+      if (fwd_status.ok() && forward_from[i] == -1) {
+        DCHECK_EQ(forward_input.size() % 2, 0);
+        for (int j = 0; j < forward_input.size(); j += 2) {
+          if (forward_input[j + 1] == i) {
+            DCHECK_EQ(forward_from[i], OpKernelContext::Params::kNoReservation);
+            forward_from[i] = forward_input[j];
+            break;
+          }
+        }
+      }
+    }
   }
+
   return ptr;
 }
 
@@ -696,22 +732,85 @@ Status ExecutorImpl::Initialize() {
   return gview_.SetAllocAttrs(graph_.get(), params_.device);
 }
 
+// If a Node has been marked to use a ScopedAllocator x for output i, then
+// sc_attr will contain the subsequence (i, x) at an even offset.  This function
+// extracts and transfers that ScopedAllocator id to alloc_attr.  For now, we
+// only allow one ScopedAllocator use per Node.
+bool ExtractScopedAllocatorAttr(const std::vector<int>& sc_attr,
+                                int output_index,
+                                AllocatorAttributes* alloc_attr) {
+  DCHECK_LE(2, sc_attr.size());
+  for (int i = 0; i < sc_attr.size(); i += 2) {
+    if (sc_attr[i] == output_index) {
+      CHECK_EQ(alloc_attr->scope_id, 0);
+      alloc_attr->scope_id = sc_attr[i + 1];
+      return true;
+    }
+  }
+  return false;
+}
+
+void GraphView::SetScopedAllocatorAttrs(
+    const std::vector<const Node*>& sa_nodes) {
+  for (const Node* sa : sa_nodes) {
+    NodeItem* sa_item = node(sa->id());
+    AllocatorAttributes* sa_attrs = sa_item->output_attr_base();
+    // Control edges out of the ScopedAllocator should be use instances, but may
+    // include a few other nodes.
+    for (const auto& e : sa->out_edges()) {
+      if (!e->IsControlEdge()) {
+        continue;
+      }
+      Node* use_node = e->dst();
+      NodeItem* item = node(use_node->id());
+      AllocatorAttributes* use_attrs = item->output_attr_base();
+      std::vector<int> scoped_allocator_attrs;
+      Status s = GetNodeAttr(use_node->attrs(), "_scoped_allocator",
+                             &scoped_allocator_attrs);
+      if (!s.ok()) {
+        VLOG(2) << "Failed to find expected ScopedAllocator attr on "
+                << use_node->name();
+        continue;
+      }
+      // There should be exactly one output using ScopedAllocation.
+      for (const auto& e : use_node->out_edges()) {
+        if (!e->IsControlEdge()) {
+          AllocatorAttributes attr;
+          if (ExtractScopedAllocatorAttr(scoped_allocator_attrs,
+                                         e->src_output(), &attr)) {
+            // Set the scope_id on this use instance node.
+            (use_attrs + e->src_output())->Merge(attr);
+            // Propagate the other attributes of this node back to the SA node.
+            attr = *(use_attrs + e->src_output());
+            attr.scope_id = 0;
+            sa_attrs->Merge(attr);
+          }
+        }
+      }
+    }
+  }
+}
+
 Status GraphView::SetAllocAttrs(const Graph* g, const Device* device) {
   Status s;
   DeviceNameUtils::ParsedName local_dev_name = device->parsed_name();
 
+  std::vector<const Node*> scoped_allocator_instances;
   for (const Node* n : g->nodes()) {
     NodeItem* item = node(n->id());
     AllocatorAttributes* attrs = item->output_attr_base();
+    if (IsScopedAllocator(n)) {
+      scoped_allocator_instances.push_back(n);
+    }
 
     // Examine the out edges of each node looking for special use
     // cases that may affect memory allocation attributes.
-    for (auto e : n->out_edges()) {
+    for (const auto& e : n->out_edges()) {
       if (!e->IsControlEdge()) {
         AllocatorAttributes attr;
         s = InferAllocAttr(n, e->dst(), local_dev_name, &attr);
         if (!s.ok()) return s;
-        if (attr.value != 0) {
+        if (attr.value != 0 || attr.scope_id != 0) {
           attrs[e->src_output()].Merge(attr);
         }
       }
@@ -728,6 +827,7 @@ Status GraphView::SetAllocAttrs(const Graph* g, const Device* device) {
       }
     }
   }
+  SetScopedAllocatorAttrs(scoped_allocator_instances);
   return s;
 }
 
@@ -1614,7 +1714,7 @@ void ExecutorState::Process(TaggedNode tagged_node, int64 scheduled_usec) {
       params.frame_iter = FrameAndIter(input_frame->frame_id, input_iter);
       params.is_input_dead = is_input_dead;
       params.output_attr_array = item.output_attrs();
-      params.forward_from_array = nullptr;  // later: item.forward_from();
+      params.forward_from_array = item.forward_from();
 
       if (item.kernel_is_async) {
         // Asynchronous computes.
diff --git a/tensorflow/core/graph/graph.cc b/tensorflow/core/graph/graph.cc
index fb8a6c39e6..eeb6c60f71 100644
--- a/tensorflow/core/graph/graph.cc
+++ b/tensorflow/core/graph/graph.cc
@@ -79,6 +79,7 @@ const std::unordered_map<string, Node::NodeClass>& Node::kNodeClassTable =
         {"Size", NC_METADATA},
         {"Shape", NC_METADATA},
         {"Rank", NC_METADATA},
+        {"_ScopedAllocator", NC_SCOPED_ALLOCATOR},
     });
 
 #undef REF_CLASS
diff --git a/tensorflow/core/graph/graph.h b/tensorflow/core/graph/graph.h
index f7ca7d0620..83a69e6b2d 100644
--- a/tensorflow/core/graph/graph.h
+++ b/tensorflow/core/graph/graph.h
@@ -34,8 +34,8 @@ limitations under the License.
 // between output O of layer A and input I of layer B using
 // "input index" and "output index" labels per edge.
 
-#ifndef TENSORFLOW_GRAPH_GRAPH_H_
-#define TENSORFLOW_GRAPH_GRAPH_H_
+#ifndef TENSORFLOW_CORE_GRAPH_GRAPH_H_
+#define TENSORFLOW_CORE_GRAPH_GRAPH_H_
 
 #include <functional>
 #include <string>
@@ -162,6 +162,7 @@ class Node {
   }
   bool IsHostSend() const { return class_ == NC_HOST_SEND; }
   bool IsHostRecv() const { return class_ == NC_HOST_RECV; }
+  bool IsScopedAllocator() const { return class_ == NC_SCOPED_ALLOCATOR; }
 
   bool IsMetadata() const { return class_ == NC_METADATA; }
 
@@ -233,6 +234,7 @@ class Node {
     NC_GET_SESSION_TENSOR,
     NC_DELETE_SESSION_TENSOR,
     NC_METADATA,
+    NC_SCOPED_ALLOCATOR,
     NC_OTHER  // Not a special kind of node
   };
 
@@ -696,6 +698,8 @@ inline bool IsControlFlow(const Node* n) { return n->IsControlFlow(); }
 // (shape).  Specifically, returns true for "Size", "Shape" and "Rank" ops.
 inline bool IsMetadata(const Node* n) { return n->IsMetadata(); }
 
+inline bool IsScopedAllocator(const Node* n) { return n->IsScopedAllocator(); }
+
 inline bool IsHostMemoryPreserving(const Node* node) {
   return IsIdentity(node) || IsControlFlow(node);
 }
@@ -827,4 +831,4 @@ inline const string& Node::assigned_device_name() const {
 
 }  // namespace tensorflow
 
-#endif  // TENSORFLOW_GRAPH_GRAPH_H_
+#endif  // TENSORFLOW_CORE_GRAPH_GRAPH_H_
-- 
GitLab


From 83e3c466b41f0235a19d5a511822b376a19cd982 Mon Sep 17 00:00:00 2001
From: ctiijima <ctiijima@us.ibm.com>
Date: Mon, 30 Apr 2018 10:55:26 -0700
Subject: [PATCH 046/395] Fixed some grammar errors.

---
 tensorflow/docs_src/community/benchmarks.md | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/tensorflow/docs_src/community/benchmarks.md b/tensorflow/docs_src/community/benchmarks.md
index 67856ce869..153ef4a015 100644
--- a/tensorflow/docs_src/community/benchmarks.md
+++ b/tensorflow/docs_src/community/benchmarks.md
@@ -1,14 +1,14 @@
 # Defining and Running Benchmarks
 
-This guide contains instructions for defining and running a TensorFlow benchmark. These benchmarks store output in [TestResults](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/util/test_log.proto) format. If these benchmarks are added to TensorFlow github repo, then we will run them daily with our continuous build and display a graph on our dashboard: https://benchmarks-dot-tensorflow-testing.appspot.com/.
+This guide contains instructions for defining and running a TensorFlow benchmark. These benchmarks store output in [TestResults](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/util/test_log.proto) format. If these benchmarks are added to the TensorFlow github repo, we will run them daily with our continuous build and display a graph on our dashboard: https://benchmarks-dot-tensorflow-testing.appspot.com/.
 
 [TOC]
 
 
 ## Defining a Benchmark
 
-Defining a TensorFlow benchmark requires extending from `tf.test.Benchmark`
-class and calling `self.report_benchmark` method. For example, take a look at the sample benchmark code below:
+Defining a TensorFlow benchmark requires extending the `tf.test.Benchmark`
+class and calling the `self.report_benchmark` method. Below, you'll find an example of benchmark code:
 
 ```python
 import time
@@ -54,20 +54,20 @@ Key points to note in the example above:
 
 ## Running with Python
 
-Use the `--benchmarks` flag to run the benchmark with python. A [BenchmarkEntries](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/util/test_log.proto) proto will be printed.
+Use the `--benchmarks` flag to run the benchmark with Python. A [BenchmarkEntries](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/util/test_log.proto) proto will be printed.
 
 ```
 python sample_benchmark.py --benchmarks=SampleBenchmark
 ```
 
-Setting the flag as `--benchmarks=.` or `--benchmarks=all` would work as well.
+Setting the flag as `--benchmarks=.` or `--benchmarks=all` works as well.
 
-(Please ensure that Tensorflow is installed to successfully import the package in the line `import tensorflow as tf`. For installation instructions, see [Installing TensorFlow](https://www.tensorflow.org/install/). This step is not necessary when running with bazel.)
+(Please ensure that Tensorflow is installed to successfully import the package in the line `import tensorflow as tf`. For installation instructions, see [Installing TensorFlow](https://www.tensorflow.org/install/). This step is not necessary when running with Bazel.)
 
 
 ## Adding a `bazel` Target
 
-We have a special target called `tf_py_logged_benchmark` for benchmarks defined under TensorFlow github repo. `tf_py_logged_benchmark` should wrap around a regular `py_test` target. Running a `tf_py_logged_benchmark` would print a [TestResults](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/util/test_log.proto) proto. Defining a `tf_py_logged_benchmark` also lets us run it with TensorFlow continuous build.
+We have a special target called `tf_py_logged_benchmark` for benchmarks defined under the TensorFlow github repo. `tf_py_logged_benchmark` should wrap around a regular `py_test` target. Running a `tf_py_logged_benchmark` would print a [TestResults](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/util/test_log.proto) proto. Defining a `tf_py_logged_benchmark` also lets us run it with TensorFlow continuous build.
 
 First, define a regular `py_test` target. See example below:
 
@@ -82,7 +82,7 @@ py_test(
 )
 ```
 
-You can run benchmarks in a `py_test` target by passing `--benchmarks` flag. The benchmark should just print out a [BenchmarkEntries](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/util/test_log.proto) proto.
+You can run benchmarks in a `py_test` target by passing the `--benchmarks` flag. The benchmark should just print out a [BenchmarkEntries](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/util/test_log.proto) proto.
 
 ```shell
 bazel test :sample_benchmark --test_arg=--benchmarks=all
@@ -90,7 +90,7 @@ bazel test :sample_benchmark --test_arg=--benchmarks=all
 
 
 Now, add the `tf_py_logged_benchmark` target (if available). This target would
-pass in `--benchmarks=all` to the wrapped `py_test` target and provide a way to store output for our TensorFlow continuous build. `tf_py_logged_benchmark` target should be available in TensorFlow repository.
+pass in `--benchmarks=all` to the wrapped `py_test` target and provide a way to store output for our TensorFlow continuous build. The target `tf_py_logged_benchmark` should be available in TensorFlow repository.
 
 ```build
 load("//tensorflow/tools/test:performance.bzl", "tf_py_logged_benchmark")
-- 
GitLab


From 9f2728bf9b5439fd5a286a1088d7543600974d4a Mon Sep 17 00:00:00 2001
From: Mark Daoust <markdaoust@google.com>
Date: Mon, 30 Apr 2018 11:01:54 -0700
Subject: [PATCH 047/395] Switch install get_started link

PiperOrigin-RevId: 194811871
---
 tensorflow/docs_src/install/install_linux.md   | 2 +-
 tensorflow/docs_src/install/install_mac.md     | 2 +-
 tensorflow/docs_src/install/install_sources.md | 2 +-
 tensorflow/docs_src/install/install_windows.md | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/docs_src/install/install_linux.md b/tensorflow/docs_src/install/install_linux.md
index 1a349f5412..e087b0c221 100644
--- a/tensorflow/docs_src/install/install_linux.md
+++ b/tensorflow/docs_src/install/install_linux.md
@@ -566,7 +566,7 @@ If you are new to machine learning, we recommend the following:
 *  @{$get_started/get_started_for_beginners$Getting Started for ML Beginners}
 
 If you are experienced with machine learning but new to TensorFlow, see
-@{$get_started/premade_estimators$Getting Started with TensorFlow}.
+@{$get_started/eager}.
 
 
 ## Common installation problems
diff --git a/tensorflow/docs_src/install/install_mac.md b/tensorflow/docs_src/install/install_mac.md
index a237d1af54..af24aaaca8 100644
--- a/tensorflow/docs_src/install/install_mac.md
+++ b/tensorflow/docs_src/install/install_mac.md
@@ -409,7 +409,7 @@ If you are new to machine learning, we recommend the following:
 *  @{$get_started/get_started_for_beginners$Getting Started for ML Beginners}
 
 If you are experienced with machine learning but new to TensorFlow, see
-@{$get_started/premade_estimators$Getting Started with TensorFlow}.
+@{$get_started/eager}.
 
 
 ## Common installation problems
diff --git a/tensorflow/docs_src/install/install_sources.md b/tensorflow/docs_src/install/install_sources.md
index 71f066e4cb..649c5b4751 100644
--- a/tensorflow/docs_src/install/install_sources.md
+++ b/tensorflow/docs_src/install/install_sources.md
@@ -388,7 +388,7 @@ TensorFlow programs:
 
 <pre>Hello, TensorFlow!</pre>
 
-If you are new to TensorFlow, see @{$get_started/premade_estimators$Getting Started with TensorFlow}.
+If you are new to TensorFlow, see @{$get_started/eager}.
 
 If the system outputs an error message instead of a greeting, see [Common
 installation problems](#common_installation_problems).
diff --git a/tensorflow/docs_src/install/install_windows.md b/tensorflow/docs_src/install/install_windows.md
index 86add74da1..a139a49661 100644
--- a/tensorflow/docs_src/install/install_windows.md
+++ b/tensorflow/docs_src/install/install_windows.md
@@ -163,7 +163,7 @@ If you are new to machine learning, we recommend the following:
 *  @{$get_started/get_started_for_beginners$Getting Started for ML Beginners}
 
 If you are experienced with machine learning but new to TensorFlow, see
-@{$get_started/premade_estimators$Getting Started with TensorFlow}.
+@{$get_started/eager}.
 
 
 ## Common installation problems
-- 
GitLab


From d6da4aa946e1f0763b9c3c2e6713c058eda0fdd4 Mon Sep 17 00:00:00 2001
From: Dustin Tran <trandustin@google.com>
Date: Mon, 30 Apr 2018 11:14:51 -0700
Subject: [PATCH 048/395] Add snippet illustrating discretized logistic mixture
 for WaveNet.

Currently, the example manually centers the bins in order to capture ?rounding? intervals and not ?ceiling? intervals. In the future, we may simplify the example by expanding QuantizedDistribution with a binning argument.

PiperOrigin-RevId: 194814662
---
 .../python/ops/quantized_distribution.py      | 64 +++++++++++++++++--
 1 file changed, 57 insertions(+), 7 deletions(-)

diff --git a/tensorflow/contrib/distributions/python/ops/quantized_distribution.py b/tensorflow/contrib/distributions/python/ops/quantized_distribution.py
index 1ef7651d03..eb94760ad7 100644
--- a/tensorflow/contrib/distributions/python/ops/quantized_distribution.py
+++ b/tensorflow/contrib/distributions/python/ops/quantized_distribution.py
@@ -128,7 +128,7 @@ The base distribution's `log_cdf` method must be defined on `y - 1`.
 class QuantizedDistribution(distributions.Distribution):
   """Distribution representing the quantization `Y = ceiling(X)`.
 
-  #### Definition in terms of sampling.
+  #### Definition in Terms of Sampling
 
   ```
   1. Draw X
@@ -138,7 +138,7 @@ class QuantizedDistribution(distributions.Distribution):
   5. Return Y
   ```
 
-  #### Definition in terms of the probability mass function.
+  #### Definition in Terms of the Probability Mass Function
 
   Given scalar random variable `X`, we define a discrete random variable `Y`
   supported on the integers as follows:
@@ -170,12 +170,62 @@ class QuantizedDistribution(distributions.Distribution):
 
   `P[Y = j]` is still the mass of `X` within the `jth` interval.
 
-  #### Caveats
+  #### Examples
+
+  We illustrate a mixture of discretized logistic distributions
+  [(Salimans et al., 2017)][1]. This is used, for example, for capturing 16-bit
+  audio in WaveNet [(van den Oord et al., 2017)][2]. The values range in
+  a 1-D integer domain of `[0, 2**16-1]`, and the discretization captures
+  `P(x - 0.5 < X <= x + 0.5)` for all `x` in the domain excluding the endpoints.
+  The lowest value has probability `P(X <= 0.5)` and the highest value has
+  probability `P(2**16 - 1.5 < X)`.
+
+  Below we assume a `wavenet` function. It takes as `input` right-shifted audio
+  samples of shape `[..., sequence_length]`. It returns a real-valued tensor of
+  shape `[..., num_mixtures * 3]`, i.e., each mixture component has a `loc` and
+  `scale` parameter belonging to the logistic distribution, and a `logits`
+  parameter determining the unnormalized probability of that component.
+
+  ```python
+  tfd = tf.contrib.distributions
+  tfb = tfd.bijectors
+
+  net = wavenet(inputs)
+  loc, unconstrained_scale, logits = tf.split(net,
+                                              num_or_size_splits=3,
+                                              axis=-1)
+  scale = tf.nn.softplus(unconstrained_scale)
+
+  # Form mixture of discretized logistic distributions. Note we shift the
+  # logistic distribution by -0.5. This lets the quantization capture "rounding"
+  # intervals, `(x-0.5, x+0.5]`, and not "ceiling" intervals, `(x-1, x]`.
+  discretized_logistic_dist = tfd.QuantizedDistribution(
+      distribution=tfd.TransformedDistribution(
+          distribution=tfd.Logistic(loc=loc, scale=scale),
+          bijector=tfb.AffineScalar(shift=-0.5)),
+      low=0.,
+      high=2**16 - 1.)
+  mixture_dist = tfd.MixtureSameFamily(
+      mixture_distribution=tfd.Categorical(logits=logits),
+      components_distribution=discretized_logistic_dist)
+
+  neg_log_likelihood = -tf.reduce_sum(mixture_dist.log_prob(targets))
+  train_op = tf.train.AdamOptimizer().minimize(neg_log_likelihood)
+  ```
+
+  After instantiating `mixture_dist`, we illustrate maximum likelihood by
+  calculating its log-probability of audio samples as `target` and optimizing.
+
+  #### References
 
-  Since evaluation of each `P[Y = j]` involves a cdf evaluation (rather than
-  a closed form function such as for a Poisson), computations such as mean and
-  entropy are better done with samples or approximations, and are not
-  implemented by this class.
+  [1]: Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P. Kingma.
+       PixelCNN++: Improving the PixelCNN with discretized logistic mixture
+       likelihood and other modifications.
+       _International Conference on Learning Representations_, 2017.
+       https://arxiv.org/abs/1701.05517
+  [2]: Aaron van den Oord et al. Parallel WaveNet: Fast High-Fidelity Speech
+       Synthesis. _arXiv preprint arXiv:1711.10433_, 2017.
+       https://arxiv.org/abs/1711.10433
   """
 
   def __init__(self,
-- 
GitLab


From 113cc4dbfc0c598d139bcf4118cde698c2e74589 Mon Sep 17 00:00:00 2001
From: Yifei Feng <yifeif@google.com>
Date: Mon, 30 Apr 2018 11:26:52 -0700
Subject: [PATCH 049/395] Add --keep_going flag to bazel query in
 pip_smoke_test to bypass bazel query cannot handle select statement.

PiperOrigin-RevId: 194816816
---
 tensorflow/tools/pip_package/pip_smoke_test.py | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/tensorflow/tools/pip_package/pip_smoke_test.py b/tensorflow/tools/pip_package/pip_smoke_test.py
index e2518f6cbf..1b692104f1 100644
--- a/tensorflow/tools/pip_package/pip_smoke_test.py
+++ b/tensorflow/tools/pip_package/pip_smoke_test.py
@@ -79,6 +79,16 @@ BLACKLIST = [
 ]
 
 
+def bazel_query(query_target):
+  """Run bazel query on target."""
+  try:
+    output = subprocess.check_output(
+        ["bazel", "query", "--keep_going", query_target])
+  except subprocess.CalledProcessError as e:
+    output = e.output
+  return output
+
+
 def main():
   """This script runs the pip smoke test.
 
@@ -93,15 +103,13 @@ def main():
   """
 
   # pip_package_dependencies_list is the list of included files in pip packages
-  pip_package_dependencies = subprocess.check_output(
-      ["bazel", "query", PIP_PACKAGE_QUERY_EXPRESSION])
+  pip_package_dependencies = bazel_query(PIP_PACKAGE_QUERY_EXPRESSION)
   pip_package_dependencies_list = pip_package_dependencies.strip().split("\n")
   print("Pip package superset size: %d" % len(pip_package_dependencies_list))
 
   # tf_py_test_dependencies is the list of dependencies for all python
   # tests in tensorflow
-  tf_py_test_dependencies = subprocess.check_output(
-      ["bazel", "query", PY_TEST_QUERY_EXPRESSION])
+  tf_py_test_dependencies = bazel_query(PY_TEST_QUERY_EXPRESSION)
   tf_py_test_dependencies_list = tf_py_test_dependencies.strip().split("\n")
   print("Pytest dependency subset size: %d" % len(tf_py_test_dependencies_list))
 
@@ -135,7 +143,7 @@ def main():
       print("Affected Tests:")
       rdep_query = ("rdeps(kind(py_test, //tensorflow/python/...), %s)" %
                     missing_dependency)
-      affected_tests = subprocess.check_output(["bazel", "query", rdep_query])
+      affected_tests = bazel_query(rdep_query)
       affected_tests_list = affected_tests.split("\n")[:-2]
       print("\n".join(affected_tests_list))
 
-- 
GitLab


From bdaa70c9e4b4215d68fd50ff120c8945ce53c18c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 11:51:03 -0700
Subject: [PATCH 050/395] -Miscellaneous code clean-up

PiperOrigin-RevId: 194821201
---
 .../graph_transformations/identify_relu1.cc   |  5 ++--
 .../graph_transformations/remove_unused_op.cc | 27 +++++--------------
 .../resolve_constant_stack.cc                 | 12 ++++++---
 .../contrib/lite/toco/import_tensorflow.cc    | 14 ++++++++--
 4 files changed, 29 insertions(+), 29 deletions(-)

diff --git a/tensorflow/contrib/lite/toco/graph_transformations/identify_relu1.cc b/tensorflow/contrib/lite/toco/graph_transformations/identify_relu1.cc
index de6d8889fb..bddb563206 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/identify_relu1.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/identify_relu1.cc
@@ -79,8 +79,9 @@ bool IdentifyRelu1::Run(Model* model, std::size_t op_index) {
   const auto* max_op =
       op_0->type == OperatorType::kTensorFlowMaximum ? op_0 : op_1;
 
-  CHECK_EQ(min_op->inputs.size(), 2);
-  CHECK_EQ(max_op->inputs.size(), 2);
+  if (min_op->inputs.size() != 2 || max_op->inputs.size() != 2) {
+    return false;
+  }
   if (min_op->outputs.size() != 1 || max_op->outputs.size() != 1) {
     return false;
   }
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc
index 8e6aaf544a..1956ab2d20 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc
@@ -88,13 +88,11 @@ bool RemoveUnusedOp::Run(Model* model, std::size_t op_index) {
   // At that point we know that none of the outputs is used, so we will
   // definitely remove the node and all its outputs.
 
-  // Remove any input array that is not used by anything else,
-  // and that is not the output of some other operator.
+  // Remove any input array that not the output of another op, and only used by
+  // this op.
   for (const auto& input : op->inputs) {
-    if (IsDiscardableArray(*model, input) &&
-        CountOpsWithInput(*model, input) == 1 &&
-        !GetOpWithOutput(*model, input)) {
-      model->EraseArray(input);
+    if (!GetOpWithOutput(*model, input)) {
+      DeleteArrayIfUsedOnce(input, model);
     }
   }
 
@@ -102,22 +100,9 @@ bool RemoveUnusedOp::Run(Model* model, std::size_t op_index) {
   for (const auto& output : op->outputs) {
     // If the output array is the model's input array, don't remove that.
     // That's the case when cropping a model at a given --input_array.
-    if (!IsDiscardableArray(*model, output)) {
-      continue;
-    }
-    // Likewise, if the output array is a RNN state array, don't remove that.
-    bool found_output_as_rnn_state_array = false;
-    for (const auto& rnn_state : model->flags.rnn_states()) {
-      if (output == rnn_state.state_array()) {
-        found_output_as_rnn_state_array = true;
-        break;
-      }
-    }
-    if (found_output_as_rnn_state_array) {
-      continue;
+    if (IsDiscardableArray(*model, output)) {
+      model->EraseArray(output);
     }
-    // Generic case: do delete this output array.
-    model->EraseArray(output);
   }
   model->operators.erase(it);
   return true;
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_stack.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_stack.cc
index ea0d6dc820..69db1942cd 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_stack.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_stack.cc
@@ -77,6 +77,13 @@ bool ResolveConstantStack::Run(Model* model, std::size_t op_index) {
     }
   }
 
+  int axis = op->axis;
+  if (axis < 0) {
+    // Handle negative axis
+    axis += model->GetArray(op->inputs[0]).shape().dims().size();
+  }
+  CHECK_EQ(axis, 0) << "Stacking only supported along 0th axis";
+
   CHECK(!output_array.buffer);
   switch (output_array.data_type) {
     case ArrayDataType::kFloat:
@@ -99,10 +106,7 @@ bool ResolveConstantStack::Run(Model* model, std::size_t op_index) {
 
   // Erase input arrays if no longer used
   for (const auto& input : op->inputs) {
-    if (IsDiscardableArray(*model, input) &&
-        CountOpsWithInput(*model, input) == 1) {
-      model->EraseArray(input);
-    }
+    toco::DeleteArrayIfUsedOnce(input, model);
   }
 
   // Erase the operator
diff --git a/tensorflow/contrib/lite/toco/import_tensorflow.cc b/tensorflow/contrib/lite/toco/import_tensorflow.cc
index 2ed05cb372..61e4c9d542 100644
--- a/tensorflow/contrib/lite/toco/import_tensorflow.cc
+++ b/tensorflow/contrib/lite/toco/import_tensorflow.cc
@@ -451,8 +451,18 @@ void ConvertConvOperator(const NodeDef& node,
   if (HasAttr(node, "dilations")) {
     const auto& dilations = GetListAttr(node, "dilations");
     CHECK_EQ(dilations.i_size(), 4);
-    CHECK_EQ(dilations.i(0), 1);
-    CHECK_EQ(dilations.i(3), 1);
+    CHECK_EQ(dilations.i(0), 1)
+        << "Can only import Conv ops with dilation along the height (1st) or "
+           "width (2nd) axis. TensorFlow op \""
+        << node.name() << "\" had dilations:[ " << dilations.i(0) << ", "
+        << dilations.i(1) << ", " << dilations.i(2) << ", " << dilations.i(3)
+        << "].";
+    CHECK_EQ(dilations.i(3), 1)
+        << "Can only import Conv ops with dilation along the height (1st) or "
+           "width (2nd) axis. TensorFlow op \""
+        << node.name() << "\" had dilations:[ " << dilations.i(0) << ", "
+        << dilations.i(1) << ", " << dilations.i(2) << ", " << dilations.i(3)
+        << "].";
     conv->dilation_height_factor = dilations.i(1);
     conv->dilation_width_factor = dilations.i(2);
   } else {
-- 
GitLab


From c3e9ca763cbacee961e247df02ec91b52cc59326 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 12:01:35 -0700
Subject: [PATCH 051/395] Fix bugs in AssignOp:  1. Releasing the unique_ptr
 would "leak" a TensorBuffer refcount.  2. The output shape is defined by rhs,
 not lhs.

PiperOrigin-RevId: 194822802
---
 tensorflow/core/kernels/assign_op.h | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/kernels/assign_op.h b/tensorflow/core/kernels/assign_op.h
index 2ed1628bf1..19b38f9e68 100644
--- a/tensorflow/core/kernels/assign_op.h
+++ b/tensorflow/core/kernels/assign_op.h
@@ -78,11 +78,10 @@ class AssignOp : public OpKernel {
       // 1. Try to reuse the rhs.
       std::unique_ptr<Tensor> input_alias = context->forward_input(
           1, OpKernelContext::Params::kNoReservation /*output_index*/,
-          old_lhs.dtype(), old_lhs.shape(), DEVICE_MEMORY, attr);
+          rhs.dtype(), rhs.shape(), DEVICE_MEMORY, attr);
       if (input_alias != nullptr) {
         // Transfer ownership to the ref.
-        context->replace_ref_input(0, *input_alias.release(),
-                                   /* lock_held */ true);
+        context->replace_ref_input(0, *input_alias, /* lock_held */ true);
         return;
       }
 
-- 
GitLab


From 8d5e87b157772c6ee131be7748245557e0df2c38 Mon Sep 17 00:00:00 2001
From: Benoit Steiner <bsteiner@google.com>
Date: Mon, 30 Apr 2018 12:13:00 -0700
Subject: [PATCH 052/395] Use the default rewriter config instead of a custom
 one

PiperOrigin-RevId: 194824761
---
 tensorflow/python/grappler/graph_placer.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/tensorflow/python/grappler/graph_placer.py b/tensorflow/python/grappler/graph_placer.py
index 1cd51df4d9..654013b23c 100644
--- a/tensorflow/python/grappler/graph_placer.py
+++ b/tensorflow/python/grappler/graph_placer.py
@@ -55,11 +55,6 @@ def PlaceGraph(metagraph,
 
   # Optimize the metagraph to speedup the placement
   rewriter_config = rewriter_config_pb2.RewriterConfig()
-  rewriter_config.optimizers.append("pruning")
-  rewriter_config.optimizers.append("constfold")
-  rewriter_config.optimizers.append("arithmetic")
-  rewriter_config.optimizers.append("dependency")
-  rewriter_config.optimizers.append("pruning")
   optimized_graph = tf_optimizer.OptimizeGraph(
       rewriter_config, metagraph, verbose=verbose, cluster=cluster)
   optimized_metagraph = meta_graph_pb2.MetaGraphDef()
-- 
GitLab


From 9d79acc6aae306e0444c193e945f0c87fe5bb509 Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Mon, 30 Apr 2018 12:33:21 -0700
Subject: [PATCH 053/395] [TF:XLA] Bump open source llvm revision to r331173

PiperOrigin-RevId: 194827639
---
 tensorflow/workspace.bzl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 5f57485d74..152da547c1 100644
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -452,11 +452,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""):
   tf_http_archive(
       name = "llvm",
       urls = [
-          "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/3b2f0b2c7e66d226a9342be5163da4240e2951a8.tar.gz",
-          "https://github.com/llvm-mirror/llvm/archive/3b2f0b2c7e66d226a9342be5163da4240e2951a8.tar.gz",
+          "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/068c967842b83d22007eee4515b57e8d9aaccb82.tar.gz",
+          "https://github.com/llvm-mirror/llvm/archive/068c967842b83d22007eee4515b57e8d9aaccb82.tar.gz",
       ],
-      sha256 = "49bb3cbb7c8e9af091c5a743fa7ae749656994408438f38c9b6ac6a052fdce56",
-      strip_prefix = "llvm-3b2f0b2c7e66d226a9342be5163da4240e2951a8",
+      sha256 = "4950432fb5cc68e5bf1f87a30b17dfdc69a5b93dac1e89d5274242d3ce7dae7c",
+      strip_prefix = "llvm-068c967842b83d22007eee4515b57e8d9aaccb82",
       build_file = clean_dep("//third_party/llvm:llvm.BUILD"),
   )
 
-- 
GitLab


From 8609ef4db1a2af0da0c2c20b26756031637de3ff Mon Sep 17 00:00:00 2001
From: Igor Saprykin <isaprykin@google.com>
Date: Mon, 30 Apr 2018 12:41:12 -0700
Subject: [PATCH 054/395] When a mirrored variable is fetched in cross-tower
 mode, fetch its primary variable.

This prevents errors like
ValueError: Fetch argument MirroredVariable({'/job:localhost/replica:0/task:0/device:GPU:0': <tf.Variable 'global_step:0' shape=() dtype=int64>, '/job:localhost/replica:0/task:0/device:GPU:1': <tf.Variable 'global_step/replica_1:0' shape=() dtype=int64>}) cannot be interpreted as a Tensor. (Device /job:localhost/replica:0/task:0/device:CPU:0 not found in ['/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1'] (current device ))

I ran distribute/examples/resnet with and without the change and it fixed the problem.

PiperOrigin-RevId: 194828672
---
 tensorflow/contrib/distribute/python/values.py   |  6 ++++++
 .../contrib/distribute/python/values_test.py     | 16 ++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/tensorflow/contrib/distribute/python/values.py b/tensorflow/contrib/distribute/python/values.py
index 8cb5276579..466678ef2e 100644
--- a/tensorflow/contrib/distribute/python/values.py
+++ b/tensorflow/contrib/distribute/python/values.py
@@ -229,6 +229,12 @@ class DistributedVariable(DistributedDelegate):
                               self._primary_var.op.type)
     return self.get().op
 
+  def _as_graph_element(self):
+    # pylint: disable=protected-access
+    if distribute_lib.get_cross_tower_context():
+      return self._primary_var._as_graph_element()
+    return self.get()._as_graph_element()
+
   def _should_act_as_resource_variable(self):
     """Pass resource_variable_ops.is_resource_variable check."""
     pass
diff --git a/tensorflow/contrib/distribute/python/values_test.py b/tensorflow/contrib/distribute/python/values_test.py
index e96ce54741..1d4e801cd8 100644
--- a/tensorflow/contrib/distribute/python/values_test.py
+++ b/tensorflow/contrib/distribute/python/values_test.py
@@ -34,6 +34,7 @@ from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import random_ops
 from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables as variables_lib
 from tensorflow.python.training import device_util
 from tensorflow.python.training import saver as saver_lib
 
@@ -582,6 +583,21 @@ class MirroredVariableTest(test.TestCase):
     save_path = self._save_normal()
     self._restore_mirrored(save_path)
 
+  @test_util.run_in_graph_and_eager_modes(config=config)
+  def testFetchAMirroredVariable(self):
+    if context.num_gpus() < 1 or context.executing_eagerly():
+      self.skipTest("A GPU is not available for this test or it's eager mode.")
+
+    with self.test_session(
+        graph=ops.Graph()) as sess, mirrored_strategy.MirroredStrategy(
+            ["/device:GPU:0"]).scope():
+      with ops.device("/device:GPU:0"):
+        v = variable_scope.get_variable(
+            name="v", initializer=1., use_resource=True)
+      mirrored = values.MirroredVariable({"/device:GPU:0": v}, v)
+      sess.run(variables_lib.global_variables_initializer())
+      sess.run({"complicated": mirrored})
+
 
 _devices = ["/device:GPU:0", "/device:CPU:0"]
 
-- 
GitLab


From 6e9d8abcdc44552a53475405f6cf0fdbffb40613 Mon Sep 17 00:00:00 2001
From: Tom Hennigan <tomhennigan@google.com>
Date: Mon, 30 Apr 2018 12:47:35 -0700
Subject: [PATCH 055/395] Fix typos in tf.GradientTape documentation.

PiperOrigin-RevId: 194829506
---
 tensorflow/python/eager/backprop.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/eager/backprop.py b/tensorflow/python/eager/backprop.py
index 07aec59cc8..d04b004451 100644
--- a/tensorflow/python/eager/backprop.py
+++ b/tensorflow/python/eager/backprop.py
@@ -681,8 +681,8 @@ class GradientTape(object):
     with tfe.GradientTape() as gg:
       gg.watch(x)
       y = x * x
-    dy_dx = gg.gradient(y, [x])[0]     # Will compute to 6.0
-  d2y_dx2 = g.gradient(dy_dx, [x])[0]  # Will compute to 2.0
+    dy_dx = gg.gradient(y, x)     # Will compute to 6.0
+  d2y_dx2 = g.gradient(dy_dx, x)  # Will compute to 2.0
   ```
 
   By default, the resources held by a GradientTape are released as soon as
@@ -697,8 +697,8 @@ class GradientTape(object):
     g.watch(x)
     y = x * x
     z = y * y
-  dy_dx = g.gradient(z, [x])[0]  # 6.0
-  dz_dx = g.gradient(y, [x])[0]  # 108.0 (4*x^3 at x = 3)
+  dz_dx = g.gradient(z, x)  # 108.0 (4*x^3 at x = 3)
+  dy_dx = g.gradient(y, x)  # 6.0
   del g  # Drop the reference to the tape
   """
 
-- 
GitLab


From 5cdcb47361e9923c418c16fee6510a472a928427 Mon Sep 17 00:00:00 2001
From: HyoukJoong Lee <hyouklee@google.com>
Date: Mon, 30 Apr 2018 12:49:33 -0700
Subject: [PATCH 056/395] Fix device assignment in xla/service/service.cc to
 build the assignment based on the provided device handles rather than using
 the default assignment.

PiperOrigin-RevId: 194829761
---
 .../xla/service/hlo_module_group_metadata.cc        |  7 +++++++
 .../xla/service/hlo_module_group_metadata.h         |  3 +++
 tensorflow/compiler/xla/service/service.cc          | 13 ++++++++++---
 3 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc
index 54c34ce116..3367d76ded 100644
--- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc
+++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc
@@ -194,6 +194,13 @@ int64 HloModuleGroupMetadata::GetModuleId(const HloModule* module) const {
   LOG(FATAL) << "unknown module";
 }
 
+int64 HloModuleGroupMetadata::GetDeviceModulesCount() const {
+  return std::count_if(modules_.begin(), modules_.end(),
+                       [](const HloModule* module) {
+                         return !module->config().is_host_module();
+                       });
+}
+
 Status HloModuleGroupMetadata::RecordInstructions() {
   const auto visitor = [this](HloInstruction* hlo) -> Status {
     if (hlo->opcode() == HloOpcode::kWhile) {
diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.h b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h
index c48a7ab0b5..d619082616 100644
--- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.h
+++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h
@@ -147,6 +147,9 @@ class HloModuleGroupMetadata {
   // the module in the module vector.
   int64 GetModuleId(const HloModule* module) const;
 
+  // Returns the number of modules for devices (excluding the host module).
+  int64 GetDeviceModulesCount() const;
+
   // Returns the companion instructions for the given instruction.
   //
   // Precondition: IsCompanionWhile(instruction) is true.
diff --git a/tensorflow/compiler/xla/service/service.cc b/tensorflow/compiler/xla/service/service.cc
index 6e0d07a12f..849488f4f9 100644
--- a/tensorflow/compiler/xla/service/service.cc
+++ b/tensorflow/compiler/xla/service/service.cc
@@ -542,9 +542,16 @@ Service::ExecuteParallelAndRegisterResult(
   // profiled.
   std::map<int64, se::Stream*> index_to_profiled_streams;
 
-  TF_ASSIGN_OR_RETURN(DeviceAssignment device_assignment,
-                      backend->computation_placer()->AssignDevices(
-                          options_.number_of_replicas(), executables.size()));
+  // Build DeviceAssignment for all cores based on the provided device handles.
+  DeviceAssignment device_assignment(options_.number_of_replicas(),
+                                     executables.size());
+  for (int64 i = 0; i < executables.size(); i++) {
+    TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*backend, device_handles[i]));
+    CHECK_EQ(replicas.size(), arguments[i].size());
+    for (int64 replica = 0; replica < replicas.size(); ++replica) {
+      device_assignment(replica, i) = replicas[replica]->device_ordinal();
+    }
+  }
 
   for (int64 i = 0; i < executables.size(); i++) {
     // Stream executors for the replicas of the current computation.
-- 
GitLab


From 1986f009218a5aa1653f91ed1f40e6321a91c922 Mon Sep 17 00:00:00 2001
From: "Joshua V. Dillon" <jvdillon@google.com>
Date: Mon, 30 Apr 2018 13:34:46 -0700
Subject: [PATCH 057/395] Cleanup handling of non-Tensor valued event_ndims in
 Bijector.

PiperOrigin-RevId: 194836408
---
 .../distributions/bijector_test.py            |  12 +
 .../python/ops/distributions/bijector_impl.py | 222 ++++++++++--------
 2 files changed, 137 insertions(+), 97 deletions(-)

diff --git a/tensorflow/python/kernel_tests/distributions/bijector_test.py b/tensorflow/python/kernel_tests/distributions/bijector_test.py
index 18582241e2..33db014279 100644
--- a/tensorflow/python/kernel_tests/distributions/bijector_test.py
+++ b/tensorflow/python/kernel_tests/distributions/bijector_test.py
@@ -24,6 +24,7 @@ import numpy as np
 import six
 
 from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops.distributions import bijector
 from tensorflow.python.platform import test
@@ -275,6 +276,17 @@ class BijectorReduceEventDimsTest(test.TestCase):
         8.,
         self.evaluate(bij.inverse_log_det_jacobian(x, event_ndims=2)))
 
+  def testHandlesNonStaticEventNdims(self):
+    x_ = [[[1., 2.], [3., 4.]]]
+    x = array_ops.placeholder_with_default(x_, shape=None)
+    event_ndims = array_ops.placeholder(dtype=np.int32, shape=[])
+    bij = ExpOnlyJacobian(forward_min_event_ndims=1)
+    bij.inverse_log_det_jacobian(x, event_ndims=event_ndims)
+    with self.test_session() as sess:
+      ildj = sess.run(bij.inverse_log_det_jacobian(x, event_ndims=event_ndims),
+                      feed_dict={event_ndims: 1})
+    self.assertAllClose(-np.log(x_), ildj)
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/python/ops/distributions/bijector_impl.py b/tensorflow/python/ops/distributions/bijector_impl.py
index 4ebc600d03..36eee5ce78 100644
--- a/tensorflow/python/ops/distributions/bijector_impl.py
+++ b/tensorflow/python/ops/distributions/bijector_impl.py
@@ -23,6 +23,7 @@ import collections
 import contextlib
 import re
 
+import numpy as np
 import six
 
 from tensorflow.python.framework import dtypes
@@ -146,15 +147,21 @@ class Bijector(object):
   for transforming a `Distribution` generated `Tensor`. A `Bijector` is
   characterized by three operations:
 
-  1. Forward\
+  1. Forward
+
      Useful for turning one random outcome into another random outcome from a
      different distribution.
-  2. Inverse\
+
+  2. Inverse
+
      Useful for "reversing" a transformation to compute one probability in
      terms of another.
-  3. `log_det_jacobian(x)`\
+
+  3. `log_det_jacobian(x)`
+
      "The log of the determinant of the matrix of all first-order partial
-     derivatives of the inverse function."\
+     derivatives of the inverse function."
+
      Useful for inverting a transformation to compute one probability in terms
      of another. Geometrically, the Jacobian determinant is the volume of the
      transformation and is used to scale the probability.
@@ -520,6 +527,8 @@ class Bijector(object):
       ValueError:  If a member of `graph_parents` is not a `Tensor`.
     """
     self._graph_parents = graph_parents or []
+    forward_min_event_ndims = get_static_value(forward_min_event_ndims)
+    inverse_min_event_ndims = get_static_value(inverse_min_event_ndims)
 
     if forward_min_event_ndims is None and inverse_min_event_ndims is None:
       raise ValueError("Must specify at least one of `forward_min_event_ndims` "
@@ -795,33 +804,37 @@ class Bijector(object):
         return self._constant_ildj_map[event_ndims]
       y = ops.convert_to_tensor(y, name="y")
       self._maybe_assert_dtype(y)
-      if not self._is_injective:  # No caching for non-injective
-        ildjs = self._inverse_log_det_jacobian(y, **kwargs)
-        return tuple(self._reduce_jacobian_det_over_event(
-            y, ildj, self.inverse_min_event_ndims, event_ndims)
-                     for ildj in ildjs)
-      mapping = self._lookup(y=y, kwargs=kwargs)
-      if mapping.ildj_map is not None and event_ndims in mapping.ildj_map:
-        return mapping.ildj_map[event_ndims]
-      try:
-        x = None  # Not needed; leave cache as is.
-        ildj = self._inverse_log_det_jacobian(y, **kwargs)
-        ildj = self._reduce_jacobian_det_over_event(
-            y, ildj, self.inverse_min_event_ndims, event_ndims)
-      except NotImplementedError as original_exception:
+      with ops.control_dependencies(self._check_valid_event_ndims(
+          min_event_ndims=self.inverse_min_event_ndims,
+          event_ndims=event_ndims)):
+        if not self._is_injective:  # No caching for non-injective
+          ildjs = self._inverse_log_det_jacobian(y, **kwargs)
+          return tuple(self._reduce_jacobian_det_over_event(
+              y, ildj, self.inverse_min_event_ndims, event_ndims)
+                       for ildj in ildjs)
+        mapping = self._lookup(y=y, kwargs=kwargs)
+        if mapping.ildj_map is not None and event_ndims in mapping.ildj_map:
+          return mapping.ildj_map[event_ndims]
         try:
-          x = mapping.x if mapping.x is not None else self._inverse(y, **kwargs)
-          ildj = -self._forward_log_det_jacobian(x, **kwargs)
+          x = None  # Not needed; leave cache as is.
+          ildj = self._inverse_log_det_jacobian(y, **kwargs)
           ildj = self._reduce_jacobian_det_over_event(
-              x, ildj, self.forward_min_event_ndims, event_ndims)
-        except NotImplementedError:
-          raise original_exception
-
-      mapping = mapping.merge(x=x, ildj_map={event_ndims: ildj})
-      self._cache(mapping)
-      if self.is_constant_jacobian:
-        self._constant_ildj_map[event_ndims] = ildj
-      return ildj
+              y, ildj, self.inverse_min_event_ndims, event_ndims)
+        except NotImplementedError as original_exception:
+          try:
+            x = (mapping.x if mapping.x is not None
+                 else self._inverse(y, **kwargs))
+            ildj = -self._forward_log_det_jacobian(x, **kwargs)
+            ildj = self._reduce_jacobian_det_over_event(
+                x, ildj, self.forward_min_event_ndims, event_ndims)
+          except NotImplementedError:
+            raise original_exception
+
+        mapping = mapping.merge(x=x, ildj_map={event_ndims: ildj})
+        self._cache(mapping)
+        if self.is_constant_jacobian:
+          self._constant_ildj_map[event_ndims] = ildj
+        return ildj
 
   def inverse_log_det_jacobian(
       self, y, event_ndims, name="inverse_log_det_jacobian"):
@@ -852,9 +865,7 @@ class Bijector(object):
         `self.dtype`.
       NotImplementedError: if `_inverse_log_det_jacobian` is not implemented.
     """
-    with ops.control_dependencies(self._check_valid_event_ndims(
-        min_event_ndims=self.inverse_min_event_ndims, event_ndims=event_ndims)):
-      return self._call_inverse_log_det_jacobian(y, event_ndims, name)
+    return self._call_inverse_log_det_jacobian(y, event_ndims, name)
 
   def _forward_log_det_jacobian(self, x):
     """Subclass implementation of `forward_log_det_jacobian` public function.
@@ -876,38 +887,46 @@ class Bijector(object):
         "forward_log_det_jacobian not implemented.")
 
   def _call_forward_log_det_jacobian(self, x, event_ndims, name, **kwargs):
+    if not self._is_injective:
+      raise NotImplementedError(
+          "forward_log_det_jacobian cannot be implemented for non-injective "
+          "transforms.")
     with self._name_scope(name, [x]):
-      if event_ndims in self._constant_ildj_map:
-        # Need "-1. *" to avoid invalid-unary-operand-type linter warning.
-        return -1. * self._constant_ildj_map[event_ndims]
-      x = ops.convert_to_tensor(x, name="x")
-      self._maybe_assert_dtype(x)
-      if not self._is_injective:
-        fldjs = self._forward_log_det_jacobian(x, **kwargs)  # No caching.
-        return tuple(self._reduce_jacobian_det_over_event(
-            x, fldj, self.forward_min_event_ndims, event_ndims)
-                     for fldj in fldjs)
-      mapping = self._lookup(x=x, kwargs=kwargs)
-      if mapping.ildj_map is not None and event_ndims in mapping.ildj_map:
-        return -mapping.ildj_map[event_ndims]
-      try:
-        y = None  # Not needed; leave cache as is.
-        ildj = -self._forward_log_det_jacobian(x, **kwargs)
-        ildj = self._reduce_jacobian_det_over_event(
-            x, ildj, self.forward_min_event_ndims, event_ndims)
-      except NotImplementedError as original_exception:
+      with ops.control_dependencies(self._check_valid_event_ndims(
+          min_event_ndims=self.forward_min_event_ndims,
+          event_ndims=event_ndims)):
+        if event_ndims in self._constant_ildj_map:
+          # Need "-1. *" to avoid invalid-unary-operand-type linter warning.
+          return -1. * self._constant_ildj_map[event_ndims]
+        x = ops.convert_to_tensor(x, name="x")
+        self._maybe_assert_dtype(x)
+        if not self._is_injective:
+          fldjs = self._forward_log_det_jacobian(x, **kwargs)  # No caching.
+          return tuple(self._reduce_jacobian_det_over_event(
+              x, fldj, self.forward_min_event_ndims, event_ndims)
+                       for fldj in fldjs)
+        mapping = self._lookup(x=x, kwargs=kwargs)
+        if mapping.ildj_map is not None and event_ndims in mapping.ildj_map:
+          return -mapping.ildj_map[event_ndims]
         try:
-          y = mapping.y if mapping.y is not None else self._forward(x, **kwargs)
-          ildj = self._inverse_log_det_jacobian(y, **kwargs)
+          y = None  # Not needed; leave cache as is.
+          ildj = -self._forward_log_det_jacobian(x, **kwargs)
           ildj = self._reduce_jacobian_det_over_event(
-              y, ildj, self.inverse_min_event_ndims, event_ndims)
-        except NotImplementedError:
-          raise original_exception
-      mapping = mapping.merge(y=y, ildj_map={event_ndims: ildj})
-      self._cache(mapping)
-      if self.is_constant_jacobian:
-        self._constant_ildj_map[event_ndims] = ildj
-      return -ildj
+              x, ildj, self.forward_min_event_ndims, event_ndims)
+        except NotImplementedError as original_exception:
+          try:
+            y = (mapping.y if mapping.y is not None
+                 else self._forward(x, **kwargs))
+            ildj = self._inverse_log_det_jacobian(y, **kwargs)
+            ildj = self._reduce_jacobian_det_over_event(
+                y, ildj, self.inverse_min_event_ndims, event_ndims)
+          except NotImplementedError:
+            raise original_exception
+        mapping = mapping.merge(y=y, ildj_map={event_ndims: ildj})
+        self._cache(mapping)
+        if self.is_constant_jacobian:
+          self._constant_ildj_map[event_ndims] = ildj
+        return -ildj
 
   def forward_log_det_jacobian(
       self, x, event_ndims, name="forward_log_det_jacobian"):
@@ -933,13 +952,7 @@ class Bijector(object):
         nor {`_inverse`, `_inverse_log_det_jacobian`} are implemented, or
         this is a non-injective bijector.
     """
-    if not self._is_injective:
-      raise NotImplementedError(
-          "forward_log_det_jacobian cannot be implemented for non-injective "
-          "transforms.")
-    with ops.control_dependencies(self._check_valid_event_ndims(
-        min_event_ndims=self.forward_min_event_ndims, event_ndims=event_ndims)):
-      return self._call_forward_log_det_jacobian(x, event_ndims, name)
+    return self._call_forward_log_det_jacobian(x, event_ndims, name)
 
   @contextlib.contextmanager
   def _name_scope(self, name=None, values=None):
@@ -981,12 +994,14 @@ class Bijector(object):
   def _reduce_jacobian_det_over_event(
       self, y, ildj, min_event_ndims, event_ndims):
     """Reduce jacobian over event_ndims - min_event_ndims."""
+    assert_static(min_event_ndims)
+
     if not self.is_constant_jacobian:
       return math_ops.reduce_sum(
           ildj,
           self._get_event_reduce_dims(min_event_ndims, event_ndims))
 
-    # In this case, we need to tile the jacobian over the event and reduce.
+    # In this case, we need to tile the Jacobian over the event and reduce.
     y_rank = array_ops.rank(y)
     y_shape = array_ops.shape(y)[
         y_rank - event_ndims : y_rank - min_event_ndims]
@@ -997,47 +1012,60 @@ class Bijector(object):
         axis=self._get_event_reduce_dims(min_event_ndims, event_ndims))
     # The multiplication by ones can change the inferred static shape so we try
     # to recover as much as possible.
-    if (isinstance(event_ndims, int) and
-        y.get_shape().ndims and ildj.get_shape().ndims):
-      y_shape = y.get_shape()
-      y_shape = y_shape[y_shape.ndims - event_ndims :
-                        y_shape.ndims - min_event_ndims]
-      ildj_shape = ildj.get_shape()
-      broadcast_shape = array_ops.broadcast_static_shape(
-          ildj_shape, y_shape)
+    event_ndims_ = get_static_value(event_ndims)
+    if (event_ndims_ is not None and
+        y.shape.ndims is not None and
+        ildj.shape.ndims is not None):
+      y_shape = y.shape[y.shape.ndims - event_ndims_ :
+                        y.shape.ndims - min_event_ndims]
+      broadcast_shape = array_ops.broadcast_static_shape(ildj.shape, y_shape)
       reduced_ildj.set_shape(
           broadcast_shape[: broadcast_shape.ndims - (
-              event_ndims - min_event_ndims)])
+              event_ndims_ - min_event_ndims)])
 
     return reduced_ildj
 
   def _get_event_reduce_dims(self, min_event_ndims, event_ndims):
     """Compute the reduction dimensions given event_ndims."""
-    min_event_ndims_ = (min_event_ndims if isinstance(min_event_ndims, int)
-                        else tensor_util.constant_value(min_event_ndims))
-    event_ndims_ = (event_ndims if isinstance(event_ndims, int)
-                    else tensor_util.constant_value(event_ndims))
+    assert_static(min_event_ndims)
+    event_ndims_ = get_static_value(event_ndims, np.int32)
 
-    if min_event_ndims_ is not None and event_ndims_ is not None:
-      return [-index for index in range(1, event_ndims_ - min_event_ndims_ + 1)]
+    if event_ndims_ is not None:
+      return [-index for index in range(1, event_ndims_ - min_event_ndims + 1)]
     else:
       reduce_ndims = event_ndims - min_event_ndims
       return math_ops.range(-reduce_ndims, 0)
 
   def _check_valid_event_ndims(self, min_event_ndims, event_ndims):
     """Check whether event_ndims is atleast min_event_ndims."""
-    min_event_ndims_ = (min_event_ndims if isinstance(min_event_ndims, int)
-                        else tensor_util.constant_value(min_event_ndims))
-    event_ndims_ = (event_ndims if isinstance(event_ndims, int)
-                    else tensor_util.constant_value(event_ndims))
-
-    if min_event_ndims_ is not None and event_ndims_ is not None:
-      if min_event_ndims_ > event_ndims_:
+    assert_static(min_event_ndims)
+    event_ndims_ = get_static_value(event_ndims, np.int32)
+    assertions = []
+    if event_ndims_ is not None:
+      if min_event_ndims > event_ndims_:
         raise ValueError("event_ndims ({}) must be larger than "
                          "min_event_ndims ({})".format(
-                             event_ndims_, min_event_ndims_))
-      return []
-
-    if self.validate_args:
-      return [check_ops.assert_greater_equal(event_ndims, min_event_ndims)]
-    return []
+                             event_ndims_, min_event_ndims))
+    elif self.validate_args:
+      assertions += [
+          check_ops.assert_greater_equal(event_ndims, min_event_ndims)]
+    return assertions
+
+
+def get_static_value(x, dtype=None):
+  """Helper which returns static value; casting when dtype is preferred."""
+  if x is None:
+    return x
+  try:
+    x_ = tensor_util.constant_value(x)
+  except TypeError:
+    x_ = x
+  if x_ is None or dtype is None:
+    return x_
+  return np.array(x_, dtype)
+
+
+def assert_static(x):
+  """Helper which asserts that input arg is known statically."""
+  if x is None or type(x) != type(get_static_value(x)):  # pylint: disable=unidiomatic-typecheck
+    raise TypeError("Input must be known statically.")
-- 
GitLab


From 57b7c7befa52ee4a205536c0552422a750cbcd21 Mon Sep 17 00:00:00 2001
From: Shashi Shekhar <shashishekhar@google.com>
Date: Mon, 30 Apr 2018 13:50:17 -0700
Subject: [PATCH 058/395] Fix a bug in profiler.

PiperOrigin-RevId: 194838948
---
 tensorflow/contrib/lite/interpreter.h           |  4 +---
 tensorflow/contrib/lite/profiling/profiler.h    | 17 +++++++++++------
 .../contrib/lite/profiling/profiler_test.cc     | 14 ++++++++++++++
 3 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/tensorflow/contrib/lite/interpreter.h b/tensorflow/contrib/lite/interpreter.h
index 6f3433abcf..1074f64263 100644
--- a/tensorflow/contrib/lite/interpreter.h
+++ b/tensorflow/contrib/lite/interpreter.h
@@ -325,9 +325,7 @@ class Interpreter {
 
   void SetProfiler(profiling::Profiler* profiler) { profiler_ = profiler; }
 
-  profiling::Profiler* GetProfiler(profiling::Profiler* profiler) {
-    return profiler_;
-  }
+  profiling::Profiler* GetProfiler() { return profiler_; }
 
   // The default capacity of `tensors_` vector.
   static constexpr int kTensorsReservedCapacity = 128;
diff --git a/tensorflow/contrib/lite/profiling/profiler.h b/tensorflow/contrib/lite/profiling/profiler.h
index dfa98a6708..8c3e4dc76d 100644
--- a/tensorflow/contrib/lite/profiling/profiler.h
+++ b/tensorflow/contrib/lite/profiling/profiler.h
@@ -85,7 +85,7 @@ class Profiler {
   std::vector<const ProfileEvent*> GetProfileEvents() {
     std::vector<const ProfileEvent*> profile_events;
     profile_events.reserve(buffer_.Size());
-    for (int i = 0; i < buffer_.Size(); i++) {
+    for (size_t i = 0; i < buffer_.Size(); i++) {
       profile_events.push_back(buffer_.At(i));
     }
     return profile_events;
@@ -103,7 +103,9 @@ class ScopedProfile {
   // Adds a profile event to profile that begins with the construction
   // of object and ends when the object goes out of scope.
   // The lifetime of tag should be at least the lifetime of profiler.
-  ScopedProfile(Profiler* profiler, const char* tag) {
+
+  ScopedProfile(Profiler* profiler, const char* tag)
+      : buffer_(nullptr), event_handle_(0) {
     if (profiler) {
       buffer_ = profiler->GetProfileBuffer();
       event_handle_ =
@@ -126,7 +128,8 @@ class ScopedOperatorProfile {
   // Adds a profile event to profile that begins with the construction
   // of object and ends when the object goes out of scope.
   // The lifetime of tag should be at least the lifetime of profiler.
-  ScopedOperatorProfile(Profiler* profiler, const char* tag, int node_index) {
+  ScopedOperatorProfile(Profiler* profiler, const char* tag, int node_index)
+      : buffer_(nullptr), event_handle_(0) {
     if (profiler) {
       buffer_ = profiler->GetProfileBuffer();
       event_handle_ = buffer_->BeginEvent(
@@ -148,9 +151,11 @@ class ScopedOperatorProfile {
 }  // namespace profiling
 }  // namespace tflite
 
-#define SCOPED_OPERATOR_PROFILE(profiler, node_index)                       \
-  tflite::profiling::ScopedOperatorProfile _profile((profiler), "OpInvoke", \
-                                                    (node_index))
+#define VARNAME_UNIQ(name, ctr) name##ctr
+
+#define SCOPED_OPERATOR_PROFILE(profiler, node_index)    \
+  tflite::profiling::ScopedOperatorProfile VARNAME_UNIQ( \
+      _profile_, __COUNTER__)((profiler), "OpInvoke", (node_index))
 #else
 
 namespace tflite {
diff --git a/tensorflow/contrib/lite/profiling/profiler_test.cc b/tensorflow/contrib/lite/profiling/profiler_test.cc
index 7ea1d8f7d3..0fba0450a0 100644
--- a/tensorflow/contrib/lite/profiling/profiler_test.cc
+++ b/tensorflow/contrib/lite/profiling/profiler_test.cc
@@ -93,6 +93,20 @@ TEST(ProfilingTest, ProfilesAreCollected) {
 #endif
 }
 
+TEST(ProfilingTest, NullProfiler) {
+  Profiler* profiler = nullptr;
+  { SCOPED_OPERATOR_PROFILE(profiler, 1); }
+}
+
+TEST(ProfilingTest, ScopedProfile) {
+  Profiler profiler;
+  profiler.StartProfiling();
+  { SCOPED_OPERATOR_PROFILE(&profiler, 1); }
+  profiler.StopProfiling();
+  auto profile_events = profiler.GetProfileEvents();
+  EXPECT_EQ(1, profile_events.size());
+}
+
 }  // namespace
 }  // namespace profiling
 }  // namespace tflite
-- 
GitLab


From b7cb5fa0059b2ef6a40aa15a4d97a01ba2e57d85 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 13:51:34 -0700
Subject: [PATCH 059/395] Extend SDCAOptimizer functionality to prune negative
 indices (the default value for OOV with tf.feature_column.FeatureColumn,
 sparse / categorical).

PiperOrigin-RevId: 194839178
---
 .../python/learn/estimators/linear_test.py    | 32 +++++++++++++++++++
 .../linear_optimizer/python/sdca_optimizer.py |  8 +++++
 2 files changed, 40 insertions(+)

diff --git a/tensorflow/contrib/learn/python/learn/estimators/linear_test.py b/tensorflow/contrib/learn/python/learn/estimators/linear_test.py
index d3bb0fda57..0a863f0e20 100644
--- a/tensorflow/contrib/learn/python/learn/estimators/linear_test.py
+++ b/tensorflow/contrib/learn/python/learn/estimators/linear_test.py
@@ -863,6 +863,38 @@ class LinearClassifierTest(test.TestCase):
     scores = classifier.evaluate(input_fn=input_fn, steps=1)
     self.assertGreater(scores['accuracy'], 0.9)
 
+  def testSdcaOptimizerWeightedSparseFeaturesOOVWithNoOOVBuckets(self):
+    """LinearClassifier with SDCAOptimizer with OOV features (-1 IDs)."""
+
+    def input_fn():
+      return {
+          'example_id':
+              constant_op.constant(['1', '2', '3']),
+          'price':
+              sparse_tensor.SparseTensor(
+                  values=[2., 3., 1.],
+                  indices=[[0, 0], [1, 0], [2, 0]],
+                  dense_shape=[3, 5]),
+          'country':
+              sparse_tensor.SparseTensor(
+                  # 'GB' is out of the vocabulary.
+                  values=['IT', 'US', 'GB'],
+                  indices=[[0, 0], [1, 0], [2, 0]],
+                  dense_shape=[3, 5])
+      }, constant_op.constant([[1], [0], [1]])
+
+    country = feature_column_lib.sparse_column_with_keys(
+        'country', keys=['US', 'CA', 'MK', 'IT', 'CN'])
+    country_weighted_by_price = feature_column_lib.weighted_sparse_column(
+        country, 'price')
+    sdca_optimizer = sdca_optimizer_lib.SDCAOptimizer(
+        example_id_column='example_id')
+    classifier = linear.LinearClassifier(
+        feature_columns=[country_weighted_by_price], optimizer=sdca_optimizer)
+    classifier.fit(input_fn=input_fn, steps=50)
+    scores = classifier.evaluate(input_fn=input_fn, steps=1)
+    self.assertGreater(scores['accuracy'], 0.9)
+
   def testSdcaOptimizerCrossedFeatures(self):
     """Tests LinearClassifier with SDCAOptimizer and crossed features."""
 
diff --git a/tensorflow/contrib/linear_optimizer/python/sdca_optimizer.py b/tensorflow/contrib/linear_optimizer/python/sdca_optimizer.py
index 213c2eced5..12039ecc6f 100644
--- a/tensorflow/contrib/linear_optimizer/python/sdca_optimizer.py
+++ b/tensorflow/contrib/linear_optimizer/python/sdca_optimizer.py
@@ -198,6 +198,14 @@ class SDCAOptimizer(object):
           example_ids = array_ops.reshape(id_tensor.indices[:, 0], [-1])
 
           flat_ids = array_ops.reshape(id_tensor.values, [-1])
+          # Prune invalid IDs (< 0) from the flat_ids, example_ids, and
+          # weight_tensor.  These can come from looking up an OOV entry in the
+          # vocabulary (default value being -1).
+          is_id_valid = math_ops.greater_equal(flat_ids, 0)
+          flat_ids = array_ops.boolean_mask(flat_ids, is_id_valid)
+          example_ids = array_ops.boolean_mask(example_ids, is_id_valid)
+          weight_tensor = array_ops.boolean_mask(weight_tensor, is_id_valid)
+
           projection_length = math_ops.reduce_max(flat_ids) + 1
           # project ids based on example ids so that we can dedup ids that
           # occur multiple times for a single example.
-- 
GitLab


From 0ed712e87742a455b56ece6fd828945f42765c52 Mon Sep 17 00:00:00 2001
From: Jiri Simsa <jsimsa@google.com>
Date: Mon, 30 Apr 2018 14:08:29 -0700
Subject: [PATCH 060/395] [tf.data] Adding support for `tf.SparseTensor` into
 `tf.contrib.data.scan()`

PiperOrigin-RevId: 194842266
---
 .../kernel_tests/scan_dataset_op_test.py      |  44 ++++++-
 .../contrib/data/python/ops/scan_ops.py       | 122 ++++++++++++------
 2 files changed, 121 insertions(+), 45 deletions(-)

diff --git a/tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py
index 1a97a84b2c..f544b1caa6 100644
--- a/tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py
@@ -28,6 +28,7 @@ from tensorflow.python.eager import context
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.platform import test
@@ -35,15 +36,19 @@ from tensorflow.python.platform import test
 
 class ScanDatasetTest(test.TestCase):
 
-  def _count(self, start, step):
-    return dataset_ops.Dataset.from_tensors(0).repeat(None).apply(
-        scan_ops.scan(start, lambda state, _: (state + step, state)))
+  def _counting_dataset(self, start, scan_fn):
+    return dataset_ops.Dataset.from_tensors(0).repeat().apply(
+        scan_ops.scan(start, scan_fn))
 
   def testCount(self):
+    def make_scan_fn(step):
+      return lambda state, _: (state + step, state)
+
     start = array_ops.placeholder(dtypes.int32, shape=[])
     step = array_ops.placeholder(dtypes.int32, shape=[])
     take = array_ops.placeholder(dtypes.int64, shape=[])
-    iterator = self._count(start, step).take(take).make_initializable_iterator()
+    iterator = self._counting_dataset(
+        start, make_scan_fn(step)).take(take).make_initializable_iterator()
     next_element = iterator.get_next()
 
     with self.test_session() as sess:
@@ -78,6 +83,37 @@ class ScanDatasetTest(test.TestCase):
     self.assertEqual(5, self.evaluate(next_element()))
     self.assertEqual(8, self.evaluate(next_element()))
 
+  def testSparseCount(self):
+    def _sparse(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=np.array([[0, 0]]),
+          values=(i * np.array([1])),
+          dense_shape=np.array([1, 1]))
+
+    def make_scan_fn(step):
+      return lambda state, _: (_sparse(state.values[0] + step), state)
+
+    start = array_ops.placeholder(dtypes.int32, shape=[])
+    step = array_ops.placeholder(dtypes.int32, shape=[])
+    take = array_ops.placeholder(dtypes.int64, shape=[])
+    iterator = self._counting_dataset(
+        _sparse(start),
+        make_scan_fn(step)).take(take).make_initializable_iterator()
+    next_element = iterator.get_next()
+
+    with self.test_session() as sess:
+
+      for start_val, step_val, take_val in [(0, 1, 10), (0, 1, 0), (10, 1, 10),
+                                            (10, 2, 10), (10, -1, 10),
+                                            (10, -2, 10)]:
+        sess.run(iterator.initializer,
+                 feed_dict={start: start_val, step: step_val, take: take_val})
+        for expected, _ in zip(
+            itertools.count(start_val, step_val), range(take_val)):
+          self.assertEqual(expected, sess.run(next_element).values[0])
+        with self.assertRaises(errors.OutOfRangeError):
+          sess.run(next_element)
+
   def testChangingStateShape(self):
     # Test the fixed-point shape invariant calculations: start with
     # initial values with known shapes, and use a scan function that
diff --git a/tensorflow/contrib/data/python/ops/scan_ops.py b/tensorflow/contrib/data/python/ops/scan_ops.py
index 60ef7efba4..e911ad0fa0 100644
--- a/tensorflow/contrib/data/python/ops/scan_ops.py
+++ b/tensorflow/contrib/data/python/ops/scan_ops.py
@@ -24,6 +24,7 @@ from tensorflow.python.data.util import nest
 from tensorflow.python.data.util import sparse
 from tensorflow.python.framework import function
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.ops import gen_dataset_ops
 
 
@@ -36,18 +37,22 @@ class _ScanDataset(dataset_ops.Dataset):
     self._input_dataset = input_dataset
 
     with ops.name_scope("initial_state"):
+      # Convert any `SparseTensorValue`s to `SparseTensor`s and all other
+      # values to tensors.
       self._initial_state = nest.pack_sequence_as(initial_state, [
-          ops.convert_to_tensor(t, name="component_%d" % i)
+          sparse_tensor.SparseTensor.from_value(t)
+          if sparse_tensor.is_sparse(t) else ops.convert_to_tensor(
+              t, name="component_%d" % i)
           for i, t in enumerate(nest.flatten(initial_state))
       ])
 
-    # Compute initial values for the state shapes and types based on
-    # the initial state. These will be refined by running
-    # `tf_scan_func` one or more times below.
-    # TODO(b/68937811): Allow the initial state to be a tf.SparseTensor.
+    # Compute initial values for the state classes, shapes and types based on
+    # the initial state. The shapes may be refined by running `tf_scan_func` one
+    # or more times below.
+    self._state_classes = sparse.get_classes(self._initial_state)
     self._state_shapes = nest.pack_sequence_as(
         self._initial_state,
-        [t.shape for t in nest.flatten(self._initial_state)])
+        [t.get_shape() for t in nest.flatten(self._initial_state)])
     self._state_types = nest.pack_sequence_as(
         self._initial_state,
         [t.dtype for t in nest.flatten(self._initial_state)])
@@ -62,67 +67,102 @@ class _ScanDataset(dataset_ops.Dataset):
     need_to_rerun = True
     while need_to_rerun:
 
-      flat_state_shapes = nest.flatten(self._state_shapes)
-      flat_state_types = nest.flatten(self._state_types)
-
-      # Create a list in which `tf_scan_func` will store the s
+      # Create a list in which `tf_scan_func` will store the new shapes.
       flat_new_state_shapes = []
 
-      @function.Defun(*(flat_state_types + nest.flatten(
-          sparse.as_dense_types(input_dataset.output_types,
-                                input_dataset.output_classes))))
+      @function.Defun(*(nest.flatten(
+          sparse.as_dense_types(
+              self._state_types, self._state_classes)) + nest.flatten(
+                  sparse.as_dense_types(input_dataset.output_types,
+                                        input_dataset.output_classes))))
       def tf_scan_func(*args):
         """A wrapper for Defun that facilitates shape inference."""
         # Pass in shape information from the state and input_dataset.
-        # TODO(b/69424092): Check that neither inputs nor outputs are sparse.
-        dense_shapes = sparse.as_dense_shapes(input_dataset.output_shapes,
-                                              input_dataset.output_classes)
-        for arg, shape in zip(args,
-                              flat_state_shapes + nest.flatten(dense_shapes)):
+        for arg, shape in zip(
+            args,
+            nest.flatten(
+                sparse.as_dense_shapes(self._state_shapes, self._state_classes))
+            + nest.flatten(
+                sparse.as_dense_shapes(input_dataset.output_shapes,
+                                       input_dataset.output_classes))):
           arg.set_shape(shape)
 
-        pivot = len(flat_state_shapes)
-        old_state = nest.pack_sequence_as(self._initial_state, args[:pivot])
-        input_value = nest.pack_sequence_as(input_dataset.output_types,
-                                            args[pivot:])
-
-        ret = scan_func(old_state, input_value)
+        pivot = len(nest.flatten(self._state_shapes))
+        print(self._state_classes)
+        nested_state_args = nest.pack_sequence_as(self._state_types,
+                                                  args[:pivot])
+        nested_state_args = sparse.deserialize_sparse_tensors(
+            nested_state_args, self._state_types, self._state_shapes,
+            self._state_classes)
+        print(input_dataset.output_classes)
+        nested_input_args = nest.pack_sequence_as(input_dataset.output_types,
+                                                  args[pivot:])
+        nested_input_args = sparse.deserialize_sparse_tensors(
+            nested_input_args, input_dataset.output_types,
+            input_dataset.output_shapes, input_dataset.output_classes)
+
+        ret = scan_func(nested_state_args, nested_input_args)
         if not isinstance(ret, collections.Sequence) or len(ret) != 2:
           raise TypeError("The scan function must return a pair comprising the "
                           "new state and the output value.")
+
+        # Convert any `SparseTensorValue`s to `SparseTensor`s and all other
+        # values to tensors.
+        ret = nest.pack_sequence_as(ret, [
+            sparse_tensor.SparseTensor.from_value(t)
+            if sparse_tensor.is_sparse(t) else ops.convert_to_tensor(t)
+            for t in nest.flatten(ret)
+        ])
         new_state, output_value = ret
 
-        flat_new_state = [
-            ops.convert_to_tensor(t) for t in nest.flatten(new_state)
-        ]
-        flat_output_value = [
-            ops.convert_to_tensor(t) for t in nest.flatten(output_value)
-        ]
+        # Extract and validate class information from the returned values.
+        for t, clazz in zip(
+            nest.flatten(new_state), nest.flatten(self._state_classes)):
+          if not isinstance(t, clazz):
+            raise TypeError(
+                "The element classes for the new state must match the initial "
+                "state. Expected %s; got %s." %
+                (self._state_classes,
+                 nest.pack_sequence_as(
+                     self._state_types,
+                     [type(t) for t in nest.flatten(new_state)])))
+        self._output_classes = sparse.get_classes(output_value)
 
         # Extract shape information from the returned values.
-        flat_new_state_shapes.extend([t.shape for t in flat_new_state])
+        flat_new_state_shapes.extend(
+            [t.get_shape() for t in nest.flatten(new_state)])
         self._output_shapes = nest.pack_sequence_as(
-            output_value, [t.shape for t in flat_output_value])
+            output_value, [t.get_shape() for t in nest.flatten(output_value)])
 
         # Extract and validate type information from the returned values.
-        for t, dtype in zip(flat_new_state, flat_state_types):
+        for t, dtype in zip(
+            nest.flatten(new_state), nest.flatten(self._state_types)):
           if t.dtype != dtype:
             raise TypeError(
                 "The element types for the new state must match the initial "
                 "state. Expected %s; got %s." %
-                (self._state_types, nest.pack_sequence_as(
-                    self._state_types, [t.dtype for t in flat_new_state])))
-        self._output_classes = nest.pack_sequence_as(
-            output_value, [ops.Tensor for _ in flat_output_value])
+                (self._state_types,
+                 nest.pack_sequence_as(
+                     self._state_types,
+                     [t.dtype for t in nest.flatten(new_state)])))
         self._output_types = nest.pack_sequence_as(
-            output_value, [t.dtype for t in flat_output_value])
-
-        return flat_new_state + flat_output_value
+            output_value, [t.dtype for t in nest.flatten(output_value)])
+
+        # Serialize any sparse tensors.
+        new_state = nest.pack_sequence_as(new_state, [
+            t for t in nest.flatten(sparse.serialize_sparse_tensors(new_state))
+        ])
+        output_value = nest.pack_sequence_as(output_value, [
+            t for t in nest.flatten(
+                sparse.serialize_sparse_tensors(output_value))
+        ])
+        return nest.flatten(new_state) + nest.flatten(output_value)
 
       # Use the private method that will execute `tf_scan_func` but delay
       # adding it to the graph in case we need to rerun the function.
       tf_scan_func._create_definition_if_needed()  # pylint: disable=protected-access
 
+      flat_state_shapes = nest.flatten(self._state_shapes)
       weakened_state_shapes = [
           original.most_specific_compatible_shape(new)
           for original, new in zip(flat_state_shapes, flat_new_state_shapes)
@@ -150,7 +190,7 @@ class _ScanDataset(dataset_ops.Dataset):
     input_t = self._input_dataset._as_variant_tensor()  # pylint: disable=protected-access
     return gen_dataset_ops.scan_dataset(
         input_t,
-        nest.flatten(self._initial_state),
+        nest.flatten(sparse.serialize_sparse_tensors(self._initial_state)),
         self._scan_func.captured_inputs,
         f=self._scan_func,
         output_types=nest.flatten(
-- 
GitLab


From 4041ae0fac83060e6d17d26fa3a46ee7b69f9919 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 14:14:43 -0700
Subject: [PATCH 061/395] Push down const inputs into the function of
 specialized functions.

PiperOrigin-RevId: 194843380
---
 .../grappler/optimizers/function_optimizer.cc | 132 ++++++++++++++++--
 .../optimizers/function_optimizer_test.cc     |  61 ++++++++
 2 files changed, 183 insertions(+), 10 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/function_optimizer.cc b/tensorflow/core/grappler/optimizers/function_optimizer.cc
index 3a6de9e3b2..1bec9086f7 100644
--- a/tensorflow/core/grappler/optimizers/function_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/function_optimizer.cc
@@ -79,6 +79,7 @@ class FunctionOptimizerContext {
   explicit FunctionOptimizerContext(RewriterConfig::Toggle opt_level,
                                     const GrapplerItem& item)
       : function_library_(OpRegistry::Global(), item.graph.library()) {
+    InitializeTrulyConstNodes(item);
     InitializeInlinedFunctions(opt_level, item);
   }
 
@@ -86,20 +87,41 @@ class FunctionOptimizerContext {
     return function_library_;
   }
 
-  FunctionLibraryDefinition& mutable_function_library() {
-    return function_library_;
+  FunctionLibraryDefinition* mutable_function_library() {
+    return &function_library_;
   }
 
   bool IsInlinedFunction(const string& name) const {
     return inlined_functions_.count(name) > 0;
   }
 
+  bool IsTrulyConst(const string& name) const {
+    return TrulyConstNode(name) != nullptr;
+  }
+
+  const NodeDef* TrulyConstNode(const string& name) const {
+    return gtl::FindWithDefault(truly_const_nodes_, name, nullptr);
+  }
+
   // Find inlining candidate by name. Return nullptr if not found.
   const FunctionDef* FindInlinedFunction(const string& name) const {
     return gtl::FindWithDefault(inlined_functions_, name, nullptr);
   }
 
  private:
+  void InitializeTrulyConstNodes(const GrapplerItem& item) {
+    std::unordered_set<string> feed_nodes;
+    for (const auto& feed : item.feed) {
+      feed_nodes.insert(NodeName(feed.first));
+    }
+
+    for (const NodeDef& node : item.graph.node()) {
+      if (IsConstant(node) && feed_nodes.count(node.name()) == 0) {
+        truly_const_nodes_[node.name()] = &node;
+      }
+    }
+  }
+
   void InitializeInlinedFunctions(RewriterConfig::Toggle opt_level,
                                   const GrapplerItem& item) {
     bool aggressive = opt_level == RewriterConfig::AGGRESSIVE;
@@ -123,10 +145,20 @@ class FunctionOptimizerContext {
   FunctionLibraryDefinition function_library_;
   // Functions that can be inlined into optimized graph.
   std::unordered_map<string, const FunctionDef*> inlined_functions_;
+  // Nodes that are Const and not in feed.
+  std::unordered_map<string, const NodeDef*> truly_const_nodes_;
 
   TF_DISALLOW_COPY_AND_ASSIGN(FunctionOptimizerContext);
 };
 
+bool HasTrulyConstInputs(const NodeDef& node,
+                         const FunctionOptimizerContext& ctx) {
+  const auto is_truly_const = [&ctx](const string& input) {
+    return ctx.IsTrulyConst(NodeName(input));
+  };
+  return std::any_of(node.input().begin(), node.input().end(), is_truly_const);
+}
+
 // Return trimmed FunctionDefLibrary with functions that are reachable from
 // the optimized graph.
 FunctionDefLibrary TrimFunctionLibrary(const FunctionLibraryDefinition& flib,
@@ -208,6 +240,77 @@ FunctionDefLibrary TrimFunctionLibrary(const FunctionLibraryDefinition& flib,
   return lib;
 }
 
+// Push all constant inputs of an instantiating node into the function body.
+Status PushDownConstInputs(const NodeDef& func_node,
+                           const FunctionOptimizerContext& ctx,
+                           GrapplerFunctionItem* item,
+                           std::unordered_set<string>* const_inputs,
+                           std::unordered_set<string>* control_deps) {
+  // Record node control dependencies in the control_deps set.
+  const auto record_control_deps = [&](const NodeDef* const_input) {
+    for (int i = const_input->input_size() - 1; i >= 0; --i) {
+      const string& input = const_input->input(i);
+      if (IsControlInput(input))
+        control_deps->insert(input);
+      else
+        break;
+    }
+  };
+
+  for (int i = func_node.input_size() - 1; i >= 0; --i) {
+    const string& input = func_node.input(i);
+    if (IsControlInput(input)) continue;
+
+    const string node_name = NodeName(input);
+    if (ctx.IsTrulyConst(node_name)) {
+      VLOG(3) << "Push const into function body: input=" << input;
+      const auto* const_input = CHECK_NOTNULL(ctx.TrulyConstNode(node_name));
+      const_inputs->insert(input);
+      record_control_deps(const_input);
+      TF_RETURN_IF_ERROR(ReplaceInputWithConst(*const_input, i, item));
+    }
+  }
+
+  return Status::OK();
+}
+
+// Remove inputs that were pushed into the function body, and attach their
+// control dependencies to the function caller node.
+void RemovePushedDownConstInputs(const std::unordered_set<string>& const_inputs,
+                                 const std::unordered_set<string>& control_deps,
+                                 NodeDef* specialized_func_node) {
+  // Nothing to do if it was no const inputs to the function node.
+  if (const_inputs.empty()) return;
+
+  // Keep only non-const inputs.
+  std::vector<string> keep_inputs;
+  const auto& inputs = specialized_func_node->input();
+  std::copy_if(inputs.begin(), inputs.end(), std::back_inserter(keep_inputs),
+               [&](const string& input) {
+                 return const_inputs.find(input) == const_inputs.end();
+               });
+
+  specialized_func_node->clear_input();
+  for (const auto& keep : keep_inputs) specialized_func_node->add_input(keep);
+
+  // Attach control dependencies of pushed down const input to the caller node.
+  if (!control_deps.empty()) {
+    std::unordered_set<string> existing_control_deps;
+
+    for (const string& input : keep_inputs) {
+      existing_control_deps.insert(AsControlDependency(NodeName(input)));
+    }
+
+    for (const string& ctrl : control_deps) {
+      if (existing_control_deps.find(ctrl) == existing_control_deps.end()) {
+        VLOG(3) << "Forward control dependency to function caller node: input="
+                << ctrl;
+        specialized_func_node->add_input(ctrl);
+      }
+    }
+  }
+}
+
 Status SpecializeFunction(const NodeDef& func_node, const FunctionDef& func,
                           FunctionOptimizerContext* ctx,
                           GraphDef* optimized_graph) {
@@ -219,11 +322,19 @@ Status SpecializeFunction(const NodeDef& func_node, const FunctionDef& func,
 
   const auto& flib = ctx->function_library();
 
-  // Make a GrapplerFunctionItem and immediately convert it back to FunctionDef.
+  // Make a GrapplerFunctionItem and convert it back to FunctionDef after
+  // pushing all constant inputs into the function body.
   GrapplerFunctionItem item;
   TF_RETURN_IF_ERROR(MakeGrapplerFunctionItem(func, func_attr, flib, &item));
 
-  // TODO(ezhulenev): Push down const inputs and known input shapes.
+  // Push const inputs into the function body, and keep track of their control
+  // dependencies.
+  std::unordered_set<string> const_inputs;
+  std::unordered_set<string> control_deps;
+  TF_RETURN_IF_ERROR(PushDownConstInputs(func_node, *ctx, &item, &const_inputs,
+                                         &control_deps));
+
+  // TODO(ezhulenev): Push down known input shapes.
   FunctionDef specialized_func;
   TF_RETURN_IF_ERROR(MakeFunctionDef(item, flib, &specialized_func));
 
@@ -237,13 +348,16 @@ Status SpecializeFunction(const NodeDef& func_node, const FunctionDef& func,
 
   // Add specialized function to the library.
   TF_RETURN_IF_ERROR(
-      ctx->mutable_function_library().AddFunctionDef(specialized_func));
+      ctx->mutable_function_library()->AddFunctionDef(specialized_func));
 
   // Add a function call node for the specialized function.
   NodeDef* specialized_func_node = optimized_graph->add_node();
   *specialized_func_node = func_node;
   specialized_func_node->set_op(specialized_func_name);
 
+  // Update specialized node to remove inputs for pushed down consts.
+  RemovePushedDownConstInputs(const_inputs, control_deps,
+                              specialized_func_node);
   return Status::OK();
 }
 
@@ -582,11 +696,9 @@ Status FunctionOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
       // Do not specialize if function has custom gradient.
       const string grad_func = ctx.function_library().FindGradient(func_name);
 
-      if (specialize_func && grad_func.empty() && IsParametrized(*func)) {
-        // TODO(ezhulenev): Specialize function call if input is a Const or has
-        // a known shape. Const input tensors can be pushed into the function
-        // body and removed from function inputs.
-
+      if (specialize_func && grad_func.empty() &&
+          (IsParametrized(*func) || HasTrulyConstInputs(node, ctx))) {
+        // TODO(ezhulenev): Specialize function call if input has a known shape.
         // Specialize function body for its instantiation attributes and inputs.
         TF_RETURN_IF_ERROR(
             SpecializeFunction(node, *func, &ctx, optimized_graph));
diff --git a/tensorflow/core/grappler/optimizers/function_optimizer_test.cc b/tensorflow/core/grappler/optimizers/function_optimizer_test.cc
index 6147e8a27c..147a264421 100644
--- a/tensorflow/core/grappler/optimizers/function_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/function_optimizer_test.cc
@@ -657,5 +657,66 @@ TEST_F(FunctionOptimizerTest, SpecializeFunction_XTimesTwo) {
   test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
 }
 
+TEST_F(FunctionOptimizerTest, SpecializeFunction_PushDownConstInput) {
+  using test::function::NDef;
+
+  FunctionOptimizer optimizer(RewriterConfig::DEFAULT);
+
+  FunctionDef mul_func = FunctionDefHelper::Create(
+      "MyMul", {"x:T", "y:T"}, {"z:T"}, {"T: {float, double}"},
+      {{{"output"}, "Mul", {"x", "y"}, {{"T", "$T"}}}},
+      /* Mapping between function returns and function node outputs. */
+      {{"z", "output:z:0"}});
+
+  // Mark MyMul as noinline.
+  (*mul_func.mutable_attr())["_noinline"].set_b(true);
+  std::vector<FunctionDef> function_library = {mul_func};
+
+  // Build a graph to compute y = MyMul(x, 2.0).
+  const Tensor kTwo = test::AsScalar<float>(2.0);
+
+  GrapplerItem item;
+  item.graph = test::function::GDef(
+      {NDef("x", "Placeholder", {}, {{"dtype", DT_FLOAT}}, kDevice),
+       NDef("init", "NoOp", {}, {}, kDevice),
+       NDef("two", "Const", {"^init", "^x"},
+            {{"dtype", DT_FLOAT}, {"value", kTwo}}, kDevice),
+       NDef("y", "MyMul", {"x", "two"}, {{"T", DT_FLOAT}}, kDevice),
+       NDef("z", "Identity", {"y"}, {{"T", DT_FLOAT}}, kDevice)},
+      function_library);
+
+  GraphDef output;
+  TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
+
+  // Make sure that specialized function was added to the library and original
+  // function was removed.
+  ASSERT_EQ(1, output.library().function_size());
+
+  const FunctionDef& specialized = output.library().function(0);
+  EXPECT_EQ("MyMul_specialized_for_y", specialized.signature().name());
+  EXPECT_EQ(1, specialized.signature().input_arg_size());
+
+  // And 'y' node has control dependencies of a pushed down const node.
+  int count = 0;
+  for (const NodeDef& node : output.node()) {
+    if (node.name() == "y" && count++) {
+      ASSERT_EQ(2, node.input_size());
+      EXPECT_EQ("x", node.input(0));
+      EXPECT_EQ("^init", node.input(1));
+    }
+  }
+  EXPECT_EQ(1, count);
+
+  // And that graph evaluation yields the same result.
+  Tensor pi = test::AsScalar<float>(3.14f);
+  item.fetch = {"z"};
+  item.feed.emplace_back("x", pi);
+
+  auto tensors_expected = EvaluateFetchNodes(item);
+  GrapplerItem optimized(item, std::move(output));
+  auto tensors = EvaluateFetchNodes(optimized);
+  test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
+}
+
 }  // namespace grappler
 }  // namespace tensorflow
-- 
GitLab


From fac11a7fbeed495938f2d1eafb75f77c88ebd068 Mon Sep 17 00:00:00 2001
From: Petros Mol <pmol@google.com>
Date: Mon, 30 Apr 2018 14:26:08 -0700
Subject: [PATCH 062/395] Removing an obsolete TODO

PiperOrigin-RevId: 194845376
---
 .../python/mappers/random_fourier_features_test.py          | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/tensorflow/contrib/kernel_methods/python/mappers/random_fourier_features_test.py b/tensorflow/contrib/kernel_methods/python/mappers/random_fourier_features_test.py
index 91929184a2..2ff4d41d75 100644
--- a/tensorflow/contrib/kernel_methods/python/mappers/random_fourier_features_test.py
+++ b/tensorflow/contrib/kernel_methods/python/mappers/random_fourier_features_test.py
@@ -31,7 +31,7 @@ from tensorflow.python.platform import googletest
 
 
 def _inner_product(x, y):
-  """Inner product between tensors x and y.
+  r"""Inner product between tensors x and y.
 
   The input tensors are assumed to be in ROW representation, that is, the method
   returns \\(x * y^T\\).
@@ -131,10 +131,6 @@ class RandomFourierFeatureMapperTest(TensorFlowTestCase):
     mapped_dim = 5000
     stddev = 5.0
 
-    # TODO(sibyl-vie3Poto): Reduce test's running time before moving to third_party. One
-    # possible way to speed the test up is to compute both the approximate and
-    # the exact kernel matrix directly using matrix operations instead of
-    # computing the values for each pair of points separately.
     points_shape = [1, input_dim]
     points = [
         random_ops.random_uniform(shape=points_shape, maxval=1.0)
-- 
GitLab


From c0f0720445226375ac8a176d9d3de9c5e647fa4a Mon Sep 17 00:00:00 2001
From: Dimitris Vardoulakis <dimvar@google.com>
Date: Mon, 30 Apr 2018 14:28:46 -0700
Subject: [PATCH 063/395] [TF:XLA] Fix some unexpected memory leak in
 hlo_graph_dumper_test.

PiperOrigin-RevId: 194845792
---
 tensorflow/compiler/xla/service/BUILD         |  1 -
 .../xla/service/hlo_graph_dumper_test.cc      | 19 +++++++++----------
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD
index ed0da47681..6e2510aa10 100644
--- a/tensorflow/compiler/xla/service/BUILD
+++ b/tensorflow/compiler/xla/service/BUILD
@@ -2421,7 +2421,6 @@ tf_cc_test(
         ":hlo_graph_dumper",
         "//tensorflow/compiler/xla:test",
         "//tensorflow/compiler/xla:xla_proto",
-        "//tensorflow/compiler/xla/tests:hlo_test_base",
         "//tensorflow/compiler/xla/tests:test_utils",
         "//tensorflow/compiler/xla/tests:xla_internal_test_main",  # fixdeps: keep
         "//tensorflow/core:lib",
diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper_test.cc b/tensorflow/compiler/xla/service/hlo_graph_dumper_test.cc
index b589cd573d..4843963243 100644
--- a/tensorflow/compiler/xla/service/hlo_graph_dumper_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_graph_dumper_test.cc
@@ -20,7 +20,6 @@ limitations under the License.
 #include "tensorflow/compiler/xla/service/hlo_module.h"
 #include "tensorflow/compiler/xla/service/hlo_opcode.h"
 #include "tensorflow/compiler/xla/test.h"
-#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
 #include "tensorflow/compiler/xla/tests/test_utils.h"
 #include "tensorflow/compiler/xla/xla.pb.h"
 #include "tensorflow/core/lib/strings/strcat.h"
@@ -48,9 +47,7 @@ class DotRenderer : public hlo_graph_dumper::GraphRendererInterface {
 
 XLA_REGISTER_GRAPH_RENDERER(DotRenderer);
 
-class HloGraphDumperTest : public HloTestBase {};
-
-TEST_F(HloGraphDumperTest, NestedFusion) {
+TEST(HloGraphDumperTest, NestedFusion) {
   HloComputation::Builder b("b");
 
   // Build param0 + param1 + param2 + param3 + param4.
@@ -67,9 +64,10 @@ TEST_F(HloGraphDumperTest, NestedFusion) {
     sums.push_back(b.AddInstruction(HloInstruction::CreateBinary(
         shape, HloOpcode::kAdd, sums[i], params[i + 2])));
   }
-  auto m = CreateNewModule();
-  m->AddEntryComputation(b.Build());
-  HloComputation* root_computation = m->entry_computation();
+  HloModuleConfig config;
+  HloModule m(TestName(), config);
+  m.AddEntryComputation(b.Build());
+  HloComputation* root_computation = m.entry_computation();
 
   // Fuse into fusion(param0 + param1 + param2 + param3 + param4).
   auto* outer_fusion = root_computation->CreateFusionInstruction(
@@ -119,13 +117,14 @@ TEST_F(HloGraphDumperTest, NestedFusion) {
       HasSubstr(inner_sum->name()));
 }
 
-TEST_F(HloGraphDumperTest, Constant) {
+TEST(HloGraphDumperTest, Constant) {
   HloComputation::Builder b("b");
   auto instruction = b.AddInstruction(
       HloInstruction::CreateConstant(Literal::CreateR0<float>(-42)));
   instruction->set_name("i_am_a_constant_root_instruction");
-  auto m = CreateNewModule();
-  HloComputation* root_computation = m->AddEntryComputation(b.Build());
+  HloModuleConfig config;
+  HloModule m(TestName(), config);
+  HloComputation* root_computation = m.AddEntryComputation(b.Build());
   string graph = hlo_graph_dumper::DumpGraph(
       *root_computation, /*label=*/"an_empty_graph", DebugOptions());
   EXPECT_THAT(graph, HasSubstr("an_empty_graph"));
-- 
GitLab


From ab02bce13e49fbd001c6db241d213dc2886a5792 Mon Sep 17 00:00:00 2001
From: Alexandre Passos <apassos@google.com>
Date: Mon, 30 Apr 2018 14:34:01 -0700
Subject: [PATCH 064/395] Do not cast int64 to int32 in keras embedding
 lookups.

Often when working on the GPU with tf int64s are more efficient as int32s will
be copied back and forth to the host quite a bit.

PiperOrigin-RevId: 194846629
---
 tensorflow/python/keras/_impl/keras/layers/embeddings.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/keras/_impl/keras/layers/embeddings.py b/tensorflow/python/keras/_impl/keras/layers/embeddings.py
index 2b353ac007..f7398845d4 100644
--- a/tensorflow/python/keras/_impl/keras/layers/embeddings.py
+++ b/tensorflow/python/keras/_impl/keras/layers/embeddings.py
@@ -153,7 +153,8 @@ class Embedding(Layer):
       return (input_shape[0],) + tuple(in_lens) + (self.output_dim,)
 
   def call(self, inputs):
-    if K.dtype(inputs) != 'int32':
+    dtype = K.dtype(inputs)
+    if dtype != 'int32' and dtype != 'int64':
       inputs = math_ops.cast(inputs, 'int32')
     out = embedding_ops.embedding_lookup(self.embeddings, inputs)
     return out
-- 
GitLab


From d2a4227636955958f9acbc7c60c72eb8cd9f6480 Mon Sep 17 00:00:00 2001
From: Eli Bendersky <eliben@google.com>
Date: Mon, 30 Apr 2018 14:39:25 -0700
Subject: [PATCH 065/395] Add XLA logo to its documentation page

PiperOrigin-RevId: 194847599
---
 tensorflow/docs_src/performance/xla/index.md | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensorflow/docs_src/performance/xla/index.md b/tensorflow/docs_src/performance/xla/index.md
index a884783074..8f5de83ea6 100644
--- a/tensorflow/docs_src/performance/xla/index.md
+++ b/tensorflow/docs_src/performance/xla/index.md
@@ -1,5 +1,9 @@
 # XLA Overview
 
+<div style="width:50%; margin:auto; margin-bottom:10px; margin-top:20px;">
+<img style="width:50%" src="/images/xlalogo.png">
+</div>
+
 > Note: XLA is experimental and considered alpha.  Most use cases will not
 > see improvements in performance (speed or decreased memory usage). We have
 > released XLA early so the Open Source Community can contribute to its
-- 
GitLab


From b8197b2190c185a138b18716100621192ee02b79 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 14:56:13 -0700
Subject: [PATCH 066/395] Implement unary chain hoisting optimization for
 Concat, Split, and SplitV.

For Concat, hoist prefix chains of unary ops before concatenation, e.g.
  // Rewrites
  //       Concat({Cos(Exp(a)), Cos(Exp(b)), Cos(Exp(c))})
  // into
  //       Cos(Exp(Concat({a, b, c}))).

For Split/SplitV hoist unary postfix chains before the split, e.g.
  // Rewrites
  //          [Cos(Exp(y)) for y in Split(x)]
  // into
  //          [y for y in Split(Cos(Exp(x)))].

The new optimization is off by default.

PiperOrigin-RevId: 194850318
---
 .../optimizers/arithmetic_optimizer.cc        | 420 ++++++++++++------
 .../optimizers/arithmetic_optimizer.h         |   5 +-
 .../optimizers/arithmetic_optimizer_test.cc   | 179 +++++++-
 3 files changed, 459 insertions(+), 145 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
index 18076eee96..bf59b25449 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
@@ -302,6 +302,11 @@ class ArithmeticOptimizerStage : public GraphOptimizerStage<string> {
     }
   }
 
+  bool IsInPreserveSet(const NodeDef& node) const {
+    return ctx().nodes_to_preserve->find(node.name()) !=
+           ctx().nodes_to_preserve->end();
+  }
+
  private:
   // Extended context required for ArithmeticOptimizer.
   const ArithmeticOptimizerContext ctx_ext_;
@@ -474,11 +479,6 @@ class ArithmeticNodesGroupOptimizerStage : public ArithmeticOptimizerStage {
     return group.root_node->device() == node.device();
   }
 
-  bool IsInPreserveSet(const NodeDef& node) const {
-    return ctx().nodes_to_preserve->find(node.name()) !=
-           ctx().nodes_to_preserve->end();
-  }
-
   bool IsAlreadyOptimized(const NodeDef& node) const {
     return optimized_nodes_.find(node.name()) != optimized_nodes_.end();
   }
@@ -1340,65 +1340,143 @@ class RemoveNegationStage : public ArithmeticOptimizerStage {
 };
 
 // This optimization hoists the common prefix of unary ops of the inputs to
-// concat out of the concat.
-// For example: Concat([Exp(Sin(x)), Exp(Sin(y)), Exp(Sin(z))]) ->
-// Exp(Sin(Concat([x, y, z]))).
+// concat out of the concat, for example:
+//    Concat([Exp(Sin(x)), Exp(Sin(y)), Exp(Sin(z))])
+// becomes
+//    Exp(Sin(Concat([x, y, z]))).
+// Similarly, it will hoist the common postfix of unary ops into Split or
+// SplitV nodes, for example:
+//    [Exp(Sin(y)) for y in Split(x)]
+// becomes
+//    [y for y in Split(Exp(Sin(x))]
+//
 // TODO(rmlarsen): Support casting. We would have to change the type attribute
-// on the concat node.
-class HoistCWiseUnaryFromConcatStage : public ArithmeticOptimizerStage {
+// on the concat/split node.
+// TODO(rmlarsen): Handle Enter/Exit.
+class HoistCWiseUnaryChainsStage : public ArithmeticOptimizerStage {
  public:
-  explicit HoistCWiseUnaryFromConcatStage(
-      const GraphOptimizerContext& ctx,
-      const ArithmeticOptimizerContext& ctx_ext)
+  explicit HoistCWiseUnaryChainsStage(const GraphOptimizerContext& ctx,
+                                      const ArithmeticOptimizerContext& ctx_ext)
       : ArithmeticOptimizerStage("", ctx, ctx_ext) {}
 
-  ~HoistCWiseUnaryFromConcatStage() override = default;
+  ~HoistCWiseUnaryChainsStage() override = default;
+
+  struct ChainLink {
+    ChainLink() = default;
+    ChainLink(NodeDef* _node, int _port_origin)
+        : node(_node), port_origin(_port_origin) {}
+    NodeDef* node;    // Node in a chain.
+    int port_origin;  // Port on concat/split node from which this chain
+                      // originates.
+
+    bool operator<(const ChainLink& other) const {
+      if (port_origin < other.port_origin) {
+        return true;
+      } else if (port_origin > other.port_origin) {
+        return false;
+      } else {
+        return node->name() < other.node->name();
+      }
+    }
+  };
+
+  // We use an ordinary set sorted on port and node name, so the order, and
+  // hence the node name used for the hoisted chain, will be deterministic.
+  using ChainLinkSet = std::set<ChainLink>;
 
   bool IsSupported(const NodeDef* node) const override {
-    if (!IsConcat(*node)) return false;
-    const int n = node->attr().at("N").i();
-    return n > 1;
+    if (IsInPreserveSet(*node)) return false;
+    if (IsConcat(*node)) {
+      const int n = node->attr().at("N").i();
+      return n > 1;
+    } else if (IsSplit(*node) || IsSplitV(*node)) {
+      const int num_split = node->attr().at("num_split").i();
+      return num_split > 1 && !IsAlreadyOptimized(*node);
+    }
+    return false;
   }
 
-  Status TrySimplify(NodeDef* concat_node,
-                     string* simplified_node_name) override {
+  Status TrySimplify(NodeDef* node, string* simplified_node_name) override {
+    node_is_concat_ = IsConcat(*node);
     int prefix_length;
     std::set<string> ctrl_inputs;
+    ChainLinkSet tails;
     TF_RETURN_IF_ERROR(
-        FindCommonUnaryOpPrefix(*concat_node, &prefix_length, &ctrl_inputs));
-    if (prefix_length > 0) {
+        FindCommonUnaryOpChain(*node, &prefix_length, &tails, &ctrl_inputs));
+    if (prefix_length > 0 && !tails.empty()) {
       TF_RETURN_IF_ERROR(
-          HoistUnaryOpPrefix(prefix_length, &ctrl_inputs, concat_node));
-      AddToOptimizationQueue(concat_node);
+          HoistUnaryOpChain(prefix_length, tails, &ctrl_inputs, node));
     }
     return Status::OK();
   }
 
  private:
-  void RemoveControlInputs(std::set<string>* removed_ctrl_inputs,
-                           NodeDef* node) const {
-    const int num_inputs = node->input_size();
-    for (int idx = num_inputs - 1; idx >= 0; --idx) {
-      const string& input = node->input(idx);
-      if (IsControlInput(input)) {
-        removed_ctrl_inputs->insert(input);
-        ctx().node_map->RemoveOutput(NodeName(input), node->name());
-        node->mutable_input()->RemoveLast();
-      } else {
-        break;
+  // Returns the length of the common unary chain of ops that can be
+  // hoisted to the other side of concat or split.
+  Status FindCommonUnaryOpChain(const NodeDef& root_node, int* prefix_length,
+                                ChainLinkSet* tails,
+                                std::set<string>* ctrl_inputs) const {
+    *prefix_length = 0;
+    // Follow the chains starting at each concat input or split output as long
+    // as all the following conditions hold:
+    //   1. The ops in all chains are the same.
+    //   2. The ops are unary elemenwise op.
+    //   3. The op output has only a single consumer (concat only).
+    ChainLinkSet cur_tails;
+    TF_RETURN_IF_ERROR(InitializeChains(root_node, &cur_tails));
+    if (cur_tails.size() < 2) {
+      return Status::OK();
+    }
+    ctrl_inputs->clear();
+    bool stop = false;
+    while (!stop && !cur_tails.empty() &&
+           OpsAreSafeToHoist(root_node, cur_tails)) {
+      // We found one more link that can be hoisted.
+      ++(*prefix_length);
+      tails->swap(cur_tails);
+      GatherControlInputs(ctrl_inputs, *tails);
+
+      // Advance tail pointers to the next level.
+      TF_RETURN_IF_ERROR(AdvanceTails(*tails, &cur_tails, &stop));
+    }
+    return Status::OK();
+  }
+
+  // Hoists the chains to the other side of concat or split and attaches the
+  // control inputs gathered from them to the concat or split node.
+  Status HoistUnaryOpChain(const int prefix_length, const ChainLinkSet& tails,
+                           std::set<string>* ctrl_inputs, NodeDef* root_node) {
+    if (tails.empty()) {
+      return Status::OK();
+    }
+    AddControlInputs(ctrl_inputs, root_node);
+    AddToOptimizationQueue(root_node);
+    optimized_nodes_.insert(root_node->name());
+    if (node_is_concat_) {
+      return HoistChainForConcat(prefix_length, tails, root_node);
+    } else {
+      return HoistChainForSplit(prefix_length, tails, root_node);
+    }
+  }
+
+  void GatherControlInputs(std::set<string>* ctrl_inputs,
+                           const ChainLinkSet& ops) const {
+    for (const auto& link : ops) {
+      const NodeDef* node = link.node;
+      for (int i = node->input_size() - 1; i >= 0; --i) {
+        const string& input = node->input(i);
+        if (!IsControlInput(input)) break;
+        ctrl_inputs->insert(input);
       }
     }
   }
 
   void AddControlInputs(std::set<string>* new_ctrl_inputs,
                         NodeDef* node) const {
-    for (int idx = node->input_size() - 1; idx >= 0; --idx) {
-      const string& existing_input = node->input(idx);
-      if (IsControlInput(existing_input)) {
-        new_ctrl_inputs->erase(existing_input);
-      } else {
-        break;
-      }
+    for (int i = node->input_size() - 1; i >= 0; --i) {
+      const string& existing_input = node->input(i);
+      if (!IsControlInput(existing_input)) break;
+      new_ctrl_inputs->erase(existing_input);
     }
     for (const string& new_input : *new_ctrl_inputs) {
       ctx().node_map->AddOutput(NodeName(new_input), node->name());
@@ -1406,113 +1484,193 @@ class HoistCWiseUnaryFromConcatStage : public ArithmeticOptimizerStage {
     }
   }
 
-  // Returns the length of the common unary prefix chain of ops that can be
-  // hoisted out of concat.
-  Status FindCommonUnaryOpPrefix(const NodeDef& concat_node, int* prefix_length,
-                                 std::set<string>* ctrl_inputs) const {
-    *prefix_length = 0;
-    const int n = concat_node.attr().at("N").i();
-    // Follow the chains backwards from each concat input as long as all the
-    // following conditions hold:
-    //   1. The ops in all chains are the same.
-    //   2. The op is a unary elemenwise op.
-    //   3. The op output has only a single consumer.
-    std::vector<NodeDef*> tail(n, nullptr);
-    const int start = concat_node.op() == "Concat" ? 1 : 0;
-    const int end = start + n;
-    // Set up tail pointers to point to the immediate inputs to Concat.
-    for (int i = start; i < end; ++i) {
-      if (IsControlInput(concat_node.input(i))) {
-        return errors::FailedPrecondition("Got control input ",
-                                          concat_node.input(i),
-                                          " where normal input was expected.");
-      }
-      TF_RETURN_IF_ERROR(GetInputNode(concat_node.input(i), &tail[i - start]));
-    }
-
-    bool stop = false;
-    ctrl_inputs->clear();
-    while (!stop) {
-      const NodeDef* tail0 = tail[0];
-      if (!IsUnaryElementWise(*tail0)) break;
-      for (int chain = 0; chain < n; ++chain) {
-        // TODO(rmlarsen): Allow and hoist outgoing control edges.
-        if (tail[chain]->op() != tail0->op() ||
-            ctx().node_map->GetOutputs(tail[chain]->name()).size() > 1) {
-          stop = true;
-          break;
+  Status InitializeChains(const NodeDef& node, ChainLinkSet* tails) const {
+    if (node_is_concat_) {
+      // Handle concat nodes by looking backwards in the graph.
+      const int n = node.attr().at("N").i();
+      const int start = node.op() == "Concat" ? 1 : 0;
+      const int end = start + n;
+      // Set up tail pointers to point to the immediate inputs to Concat.
+      for (int input_port = start; input_port < end; ++input_port) {
+        if (IsControlInput(node.input(input_port))) {
+          return errors::FailedPrecondition(
+              "Got control input ", node.input(input_port),
+              " where normal input was expected.");
         }
+        NodeDef* tail;
+        TF_RETURN_IF_ERROR(GetInputNode(node.input(input_port), &tail));
+        tails->insert(ChainLink(tail, input_port));
       }
-      if (stop) break;
-      // We found one more op that can be hoisted.
-      ++(*prefix_length);
-      for (int chain = 0; chain < n; ++chain) {
-        RemoveControlInputs(ctrl_inputs, tail[chain]);
-      }
-      // Advance tail pointers to the next level.
-      for (int chain = 0; chain < n; ++chain) {
-        if (tail[chain]->input_size() == 0 ||
-            IsControlInput(tail[chain]->input(0))) {
-          stop = true;
-          break;
+      return Status::OK();
+    } else {
+      // Handle split nodes by looking forwards in the graph.
+      const auto& outputs = ctx().node_map->GetOutputs(node.name());
+      for (NodeDef* output : outputs) {
+        if (IsControlInput(output->input(0))) continue;
+        int port;
+        const string node_name = ParseNodeName(output->input(0), &port);
+        if (node_name == node.name()) {
+          tails->insert(ChainLink(output, port));
         } else {
-          NodeDef* new_tail = nullptr;
-          TF_RETURN_IF_ERROR(GetInputNode(tail[chain]->input(0), &new_tail));
-          tail[chain] = new_tail;
+          // This output node has a non-control input other than the split node,
+          // abort.
+          tails->clear();
+          return Status::OK();
         }
       }
     }
     return Status::OK();
   }
 
-  Status HoistUnaryOpPrefix(const int prefix_length,
-                            std::set<string>* ctrl_inputs,
-                            NodeDef* concat_node) {
-    const int n = concat_node->attr().at("N").i();
-    const int start = concat_node->op() == "Concat" ? 1 : 0;
-    const int end = start + n;
-    const std::set<NodeDef*> consumers =
-        ctx().node_map->GetOutputs(concat_node->name());
-    AddControlInputs(ctrl_inputs, concat_node);
-    for (int chain = 0; chain < (end - start); ++chain) {
-      NodeDef* tail = nullptr;
-      const string concat_input = concat_node->input(chain + start);
-      for (int distance = 0; distance < prefix_length; ++distance) {
-        if (distance == 0) {
-          TF_RETURN_IF_ERROR(GetInputNode(concat_input, &tail));
-        } else {
-          TF_RETURN_IF_ERROR(GetInputNode(tail->input(0), &tail));
+  bool OpsAreSafeToHoist(const NodeDef& root_node,
+                         const ChainLinkSet& ops) const {
+    if (ops.empty()) return true;
+    const NodeDef* op0 = ops.begin()->node;
+    if (!IsUnaryElementWise(*op0)) return false;
+    for (const auto& link : ops) {
+      const NodeDef* op = link.node;
+      if (op->device() != root_node.device() || op->op() != op0->op() ||
+          IsInPreserveSet(*op)) {
+        return false;
+      }
+      if (node_is_concat_ &&
+          ctx().node_map->GetOutputs(op->name()).size() > 1) {
+        // TODO(rmlarsen): Allow and hoist outgoing control edges.
+        return false;
+      }
+    }
+    return true;
+  }
+
+  Status AdvanceTails(const ChainLinkSet& tails, ChainLinkSet* new_tails,
+                      bool* stop) const {
+    *stop = true;
+    new_tails->clear();
+    for (const auto& link : tails) {
+      const NodeDef* tail = link.node;
+      if (node_is_concat_) {
+        if (tail->input_size() == 0 || IsControlInput(tail->input(0))) {
+          return Status::OK();
+        }
+        NodeDef* new_tail;
+        TF_RETURN_IF_ERROR(GetInputNode(tail->input(0), &new_tail));
+        // Remember original port.
+        new_tails->insert(ChainLink(new_tail, link.port_origin));
+      } else {
+        for (NodeDef* new_tail : ctx().node_map->GetOutputs(tail->name())) {
+          int port;
+          const string node_name = ParseNodeName(new_tail->input(0), &port);
+          if (node_name != tail->name()) {
+            return Status::OK();
+          }
+          // Skip control outputs.
+          if (port >= 0) {
+            // Remember original port.
+            new_tails->insert(ChainLink(new_tail, link.port_origin));
+          }
         }
       }
+    }
+    *stop = false;
+    return Status::OK();
+  }
 
+  Status HoistChainForConcat(const int prefix_length, const ChainLinkSet& tails,
+                             NodeDef* concat_node) {
+    const string& concat_name = concat_node->name();
+    const int first_input = concat_node->op() == "Concat" ? 1 : 0;
+    for (const auto& link : tails) {
+      NodeDef* tail = CHECK_NOTNULL(link.node);
+      const int concat_port = link.port_origin;
+      CHECK_GE(concat_port, 0);
+      CHECK_LT(concat_port, concat_node->input_size());
+      const string concat_input = concat_node->input(concat_port);
       // Hook the node following tail directly into the concat node.
       const string tail_input = tail->input(0);
-      concat_node->set_input(chain + start, tail_input);
-      ctx().node_map->UpdateInput(concat_node->name(), concat_input,
-                                  tail_input);
-
-      if (chain == 0) {
-        // Reuse nodes in the first chain to process output of concat.
-        tail->set_input(0, concat_node->name());
-        ctx().node_map->UpdateInput(tail->name(), tail_input,
-                                    concat_node->name());
+      concat_node->set_input(concat_port, tail_input);
+      ctx().node_map->UpdateInput(concat_name, concat_input, tail_input);
 
+      if (concat_port == first_input) {
         // Update the consumers of concat to consume the end of the chain
         // instead.
-        for (NodeDef* consumer : consumers) {
-          for (int idx = 0; idx < consumer->input_size(); ++idx) {
-            if (consumer->input(idx) == concat_node->name()) {
-              consumer->set_input(idx, concat_input);
-              ctx().node_map->UpdateInput(consumer->name(), concat_node->name(),
-                                          concat_input);
-            }
-          }
-          AddToOptimizationQueue(consumer);
-        }
+        UpdateConsumers(concat_node, concat_input);
+        // Reuse nodes in the first chain to process output of concat.
+        tail->set_input(0, concat_name);
+        ctx().node_map->UpdateInput(tail->name(), tail_input, concat_name);
       }
     }
     return Status::OK();
   }
+
+  Status HoistChainForSplit(const int prefix_length, const ChainLinkSet& tails,
+                            NodeDef* split_node) {
+    // Create a new chain before the split node to process the input tensor.
+    const string& split_name = split_node->name();
+    auto root_scope_and_name = ParseNodeScopeAndName(split_name);
+
+    // We use the first tail node in the set as a template to get the list of
+    // ops to apply (starting from the end).
+    NodeDef* cur_tail = tails.begin()->node;
+    NodeDef* cur_copy = AddCopyNode(
+        OptimizedNodeName(root_scope_and_name, cur_tail->name()), cur_tail);
+    cur_copy->clear_input();
+
+    // Update the split to take its input from the tail of the new chain.
+    const int value_slot = split_node->op() == "SplitV" ? 0 : 1;
+    const string orig_input = split_node->input(value_slot);
+    split_node->set_input(value_slot, cur_copy->name());
+    ctx().node_map->UpdateInput(split_node->name(), orig_input,
+                                cur_copy->name());
+    TF_RETURN_IF_ERROR(GetInputNode(cur_tail->input(0), &cur_tail));
+
+    // Now walk backwards creating the rest of the chain.
+    while (cur_tail != split_node) {
+      NodeDef* new_copy = AddCopyNode(
+          OptimizedNodeName(root_scope_and_name, cur_tail->name()), cur_tail);
+      new_copy->clear_input();
+      cur_copy->add_input(new_copy->name());
+      ctx().node_map->AddOutput(new_copy->name(), cur_copy->name());
+      cur_copy = new_copy;
+      TF_RETURN_IF_ERROR(GetInputNode(cur_tail->input(0), &cur_tail));
+    }
+    // Connect the original input to the head of the new chain.
+    cur_copy->add_input(orig_input);
+    ctx().node_map->UpdateOutput(NodeName(orig_input), split_name,
+                                 cur_copy->name());
+
+    // Connect all consumers of the tail nodes directly to the
+    // output port of Split from which the chain started.
+    for (const auto& link : tails) {
+      UpdateConsumers(link.node,
+                      link.port_origin == 0
+                          ? split_name
+                          : strings::StrCat(split_name, ":", link.port_origin));
+    }
+    return Status::OK();
+  }
+
+  // Update consumers of node to take new_input as input instead.
+  void UpdateConsumers(NodeDef* node, const string& new_input) {
+    const string& node_name = node->name();
+    const std::set<NodeDef*> consumers = ctx().node_map->GetOutputs(node_name);
+    for (NodeDef* consumer : consumers) {
+      for (int i = 0; i < consumer->input_size(); ++i) {
+        if (consumer->input(i) == node_name) {
+          consumer->set_input(i, new_input);
+          ctx().node_map->UpdateInput(consumer->name(), node_name, new_input);
+        }
+      }
+      AddToOptimizationQueue(consumer);
+    }
+  }
+
+  bool IsAlreadyOptimized(const NodeDef& node) const {
+    return optimized_nodes_.find(node.name()) != optimized_nodes_.end();
+  }
+
+ private:
+  bool node_is_concat_;
+  std::unordered_set<string> optimized_nodes_;
 };
 
 // Performs the conversion:
@@ -2200,8 +2358,8 @@ Status ArithmeticOptimizer::SimplifyArithmeticOps(bool can_use_shapes) {
     pipeline.AddStage<RemoveRedundantCastStage>(ctx, ctx_ext);
   if (options_.remove_negation)
     pipeline.AddStage<RemoveNegationStage>(ctx, ctx_ext);
-  if (options_.hoist_unary_out_of_concat)
-    pipeline.AddStage<HoistCWiseUnaryFromConcatStage>(ctx, ctx_ext);
+  if (options_.hoist_cwise_unary_chains)
+    pipeline.AddStage<HoistCWiseUnaryChainsStage>(ctx, ctx_ext);
   if (options_.convert_sqrt_div_to_rsqrt_mul)
     pipeline.AddStage<SqrtDivToRsqrtMulStage>(ctx, ctx_ext);
 
@@ -2304,5 +2462,5 @@ void ArithmeticOptimizer::Feedback(Cluster* /*cluster*/,
   // Nothing to do for ArithmeticOptimizer.
 }
 
-}  // end namespace grappler
-}  // end namespace tensorflow
+}  // namespace grappler
+}  // namespace tensorflow
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
index 24a2a50719..3b297ec0aa 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
@@ -65,7 +65,7 @@ class ArithmeticOptimizer : public GraphOptimizer {
     bool remove_redundant_bitcast = true;
     bool remove_redundant_cast = true;
     bool remove_negation = true;
-    bool hoist_unary_out_of_concat = false;
+    bool hoist_cwise_unary_chains = false;
     bool convert_sqrt_div_to_rsqrt_mul = false;
 
     // Choose which arithmetic optimizer stages will be enabled for a given
@@ -73,9 +73,6 @@ class ArithmeticOptimizer : public GraphOptimizer {
     static ArithmeticOptimizerOptions Default(
         RewriterConfig::Toggle opt_level) {
       ArithmeticOptimizerOptions options;
-      if (opt_level == RewriterConfig::AGGRESSIVE) {
-        options.hoist_unary_out_of_concat = true;
-      }
       return options;
     }
   };
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
index 7485d99c3b..f903f53a35 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
@@ -94,6 +94,16 @@ class ArithmeticOptimizerTest : public GrapplerTest {
     TF_EXPECT_OK(optimizer->Optimize(nullptr, *item, output));
   }
 
+  // Run ArithmeticOptimizer twice to make sure the rewrite is idempotent.
+  void OptimizeTwiceAndPrune(ArithmeticOptimizer* optimizer, GrapplerItem* item,
+                             GraphDef* output) {
+    TF_EXPECT_OK(optimizer->Optimize(nullptr, *item, output));
+    item->graph.Swap(output);
+    TF_EXPECT_OK(optimizer->Optimize(nullptr, *item, output));
+    item->graph.Swap(output);
+    TF_EXPECT_OK(ModelPruner().Optimize(nullptr, *item, output));
+  }
+
   // TODO(ezhulenev): Make private. After migration to stages each test
   // should explicitly enable required optimization for tests isolation
   void DisableAllStages(ArithmeticOptimizer* optimizer) {
@@ -149,9 +159,9 @@ class ArithmeticOptimizerTest : public GrapplerTest {
     optimizer->options_.remove_negation = true;
   }
 
-  void EnableOnlyHoistCWiseUnaryFromConcat(ArithmeticOptimizer* optimizer) {
+  void EnableOnlyHoistCWiseUnaryChains(ArithmeticOptimizer* optimizer) {
     DisableAllStages(optimizer);
-    optimizer->options_.hoist_unary_out_of_concat = true;
+    optimizer->options_.hoist_cwise_unary_chains = true;
   }
 
   void EnableOnlySqrtDivToRsqrtMul(ArithmeticOptimizer* optimizer) {
@@ -2136,14 +2146,18 @@ TEST_F(ArithmeticOptimizerTest, MinimizeBroadcasts_BuildTreeUp) {
 
 TEST_F(ArithmeticOptimizerTest, HoistCWiseUnaryFromConcat) {
   tensorflow::Scope s = tensorflow::Scope::NewRootScope();
-  Output a = ops::Variable(s.WithOpName("a"), {32}, DT_FLOAT);
-  Output b = ops::Variable(s.WithOpName("b"), {32}, DT_FLOAT);
-  Output c = ops::Variable(s.WithOpName("c"), {32}, DT_FLOAT);
+  Output a = ops::Const(s.WithOpName("a"), 3.14f, {32});
+  Output b = ops::Const(s.WithOpName("b"), 1.0f, {32});
+  Output c = ops::Const(s.WithOpName("c"), 42.0f, {32});
   Output axis = ops::Const(s.WithOpName("axis"), 0, {});
   Output ctrl1 = ops::Const(s.WithOpName("ctrl1"), 1, {});
   Output ctrl2 = ops::Const(s.WithOpName("ctrl2"), 2, {});
   Output ctrl3 = ops::Const(s.WithOpName("ctrl3"), 3, {});
   // Test case with chains of length 1.
+  // Rewrites
+  //       Concat({Exp(a), Exp(b), Exp(c)})
+  // into
+  //       Exp(Concat({a, b, c})).
   Output sin_a =
       ops::Sin(s.WithOpName("sin_a").WithControlDependencies(ctrl3), a);
   Output exp_a =
@@ -2156,6 +2170,10 @@ TEST_F(ArithmeticOptimizerTest, HoistCWiseUnaryFromConcat) {
   Output id = ops::Identity(s.WithOpName("id"), concat);
 
   // Test case with chains of length 2.
+  // Rewrites
+  //       Concat({Cos(Exp(a)), Cos(Exp(b)), Cos(Exp(c))})
+  // into
+  //       Cos(Exp(Concat({a, b, c}))).
   Output exp_a2 =
       ops::Exp(s.WithOpName("exp_a2").WithControlDependencies(ctrl1), sin_a);
   Output exp_b2 = ops::Exp(s.WithOpName("exp_b2"), b);
@@ -2173,11 +2191,13 @@ TEST_F(ArithmeticOptimizerTest, HoistCWiseUnaryFromConcat) {
   item.fetch = {"id", "id2"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch);
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
-  EnableOnlyHoistCWiseUnaryFromConcat(&optimizer);
+  EnableOnlyHoistCWiseUnaryChains(&optimizer);
+  OptimizeTwiceAndPrune(&optimizer, &item, &output);
 
-  OptimizeAndPrune(&optimizer, &item, &output);
   int found = 0;
   for (const NodeDef& node : output.node()) {
     if (node.name() == "concat") {
@@ -2191,8 +2211,9 @@ TEST_F(ArithmeticOptimizerTest, HoistCWiseUnaryFromConcat) {
       found++;
     }
     if (node.name() == "exp_a") {
-      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ(2, node.input_size());
       EXPECT_EQ("concat", node.input(0));
+      EXPECT_EQ("^ctrl1", node.input(1));
       found++;
     }
     if (node.name() == "id") {
@@ -2213,13 +2234,15 @@ TEST_F(ArithmeticOptimizerTest, HoistCWiseUnaryFromConcat) {
       found++;
     }
     if (node.name() == "exp_a2") {
-      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ(2, node.input_size());
       EXPECT_EQ("concat2", node.input(0));
+      EXPECT_EQ("^ctrl1", node.input(1));
       found++;
     }
     if (node.name() == "cos_exp_a2") {
-      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ(2, node.input_size());
       EXPECT_EQ("exp_a2", node.input(0));
+      EXPECT_EQ("^ctrl1", node.input(1));
       found++;
     }
     if (node.name() == "id2") {
@@ -2229,6 +2252,142 @@ TEST_F(ArithmeticOptimizerTest, HoistCWiseUnaryFromConcat) {
     }
   }
   EXPECT_EQ(7, found);
+
+  auto tensors = EvaluateNodes(output, item.fetch);
+  EXPECT_EQ(tensors.size(), tensors_expected.size());
+  EXPECT_EQ(tensors.size(), item.fetch.size());
+  for (int i = 0; i < item.fetch.size(); ++i) {
+    test::ExpectTensorNear<float>(tensors_expected[i], tensors[i], 1e-6);
+  }
+}
+
+TEST_F(ArithmeticOptimizerTest, HoistCWiseUnaryIntoSplit) {
+  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+  Output x = ops::Const(s.WithOpName("x"), 3.1415f, {32});
+  Output axis = ops::Const(s.WithOpName("axis"), 0, {});
+  Output ctrl1 = ops::Const(s.WithOpName("ctrl1"), 1, {});
+  Output ctrl2 = ops::Const(s.WithOpName("ctrl2"), 2, {});
+  Output ctrl3 = ops::Const(s.WithOpName("ctrl3"), 3, {});
+  // Test case with chains of length 1.
+  // Rewrites
+  //          [Sin(y) for y in Split(x)]
+  // into
+  //          [y for y in Split(Sin(x))].
+  ops::Split split1(s.WithOpName("split1"), axis, x, 2);
+  Output sin_a =
+      ops::Sin(s.WithOpName("sin_a").WithControlDependencies(ctrl1), split1[0]);
+  Output id_a = ops::Identity(s.WithOpName("id_a"), sin_a);
+  Output sin_b = ops::Sin(s.WithOpName("sin_b"), split1[1]);
+  Output exp_b = ops::Exp(s.WithOpName("exp_b"), sin_b);
+  Output id_b = ops::Identity(s.WithOpName("id_b"), exp_b);
+
+  // Test case with SplitV and chains of length 2.
+  // Rewrites
+  //          [Cos(Exp(y)) for y in Split(x)]
+  // into
+  //          [y for y in Split(Cos(Exp(x)))].
+  Output size_splits2 = ops::Const(s.WithOpName("size_splits2"), {20, 12}, {2});
+  ops::SplitV split2(s.WithOpName("split2"), x, size_splits2, axis, 2);
+  Output exp_a2 = ops::Exp(
+      s.WithOpName("exp_a2").WithControlDependencies(ctrl1), split2[0]);
+  Output exp_b2 = ops::Exp(s.WithOpName("exp_b2"), split2[1]);
+  Output cos_exp_a2 = ops::Cos(
+      s.WithOpName("cos_exp_a2").WithControlDependencies(ctrl2), exp_a2);
+  Output cos_exp_b2 = ops::Cos(
+      s.WithOpName("cos_exp_b2").WithControlDependencies(ctrl3), exp_b2);
+  Output id_a2 = ops::Identity(s.WithOpName("id_a2"), cos_exp_a2);
+  Output id_b2 = ops::Identity(s.WithOpName("id_b2"), cos_exp_b2);
+
+  GrapplerItem item;
+  item.fetch = {"id_a", "id_b", "id_a2", "id_b2"};
+  TF_CHECK_OK(s.ToGraphDef(&item.graph));
+
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch);
+
+  GraphDef output;
+  ArithmeticOptimizer optimizer;
+  EnableOnlyHoistCWiseUnaryChains(&optimizer);
+  OptimizeTwiceAndPrune(&optimizer, &item, &output);
+
+  int found = 0;
+  for (const NodeDef& node : output.node()) {
+    // The following 6 nodes should be pruned.
+    EXPECT_NE(node.name(), "sin_a");
+    EXPECT_NE(node.name(), "sin_b");
+    EXPECT_NE(node.name(), "exp_a2");
+    EXPECT_NE(node.name(), "exp_b2");
+    EXPECT_NE(node.name(), "cos_exp_a2");
+    EXPECT_NE(node.name(), "cos_exp_b2");
+
+    if (node.name() == "split1") {
+      EXPECT_EQ(3, node.input_size());
+      EXPECT_EQ("axis", node.input(0));
+      EXPECT_EQ("ArithmeticOptimizer/_sin_a_split1", node.input(1));
+      EXPECT_EQ("^ctrl1", node.input(2));
+      found++;
+    }
+    if (node.name() == "ArithmeticOptimizer/_sin_a_split1") {
+      EXPECT_EQ("Sin", node.op());
+      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ("x", node.input(0));
+      found++;
+    }
+    if (node.name() == "id_a") {
+      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ("split1", node.input(0));
+      found++;
+    }
+    if (node.name() == "exp_b") {
+      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ("split1:1", node.input(0));
+      found++;
+    }
+    if (node.name() == "id_b") {
+      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ("exp_b", node.input(0));
+      found++;
+    }
+    if (node.name() == "ArithmeticOptimizer/_exp_a2_split2") {
+      EXPECT_EQ("Exp", node.op());
+      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ("x", node.input(0));
+      found++;
+    }
+    if (node.name() == "ArithmeticOptimizer/_cos_exp_a2_split2") {
+      EXPECT_EQ("Cos", node.op());
+      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ("ArithmeticOptimizer/_exp_a2_split2", node.input(0));
+      found++;
+    }
+    if (node.name() == "split2") {
+      EXPECT_EQ(6, node.input_size());
+      EXPECT_EQ("ArithmeticOptimizer/_cos_exp_a2_split2", node.input(0));
+      EXPECT_EQ("size_splits2", node.input(1));
+      EXPECT_EQ("axis", node.input(2));
+      EXPECT_EQ("^ctrl1", node.input(3));
+      EXPECT_EQ("^ctrl2", node.input(4));
+      EXPECT_EQ("^ctrl3", node.input(5));
+      found++;
+    }
+    if (node.name() == "id_a2") {
+      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ("split2", node.input(0));
+      found++;
+    }
+    if (node.name() == "id_b2") {
+      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ("split2:1", node.input(0));
+      found++;
+    }
+  }
+  EXPECT_EQ(10, found);
+
+  auto tensors = EvaluateNodes(output, item.fetch);
+  EXPECT_EQ(tensors.size(), tensors_expected.size());
+  EXPECT_EQ(tensors.size(), item.fetch.size());
+  for (int i = 0; i < item.fetch.size(); ++i) {
+    test::ExpectTensorNear<float>(tensors_expected[i], tensors[i], 1e-6);
+  }
 }
 
 }  // namespace grappler
-- 
GitLab


From 9c961e80a6be0136fc43821f1ad01ea00f83acb3 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 15:19:00 -0700
Subject: [PATCH 067/395] Enhancements to GRAPHVIZ_DOT output:     -edge
 weights added to encourage straighter main data-flow     -line thickness
 proportional to log(data_size)     -set global parameter "nslimit" to prevent
 excessive layout time for difficult graphs

PiperOrigin-RevId: 194854051
---
 tensorflow/contrib/lite/toco/dump_graphviz.cc | 45 +++++++++++++++++--
 1 file changed, 41 insertions(+), 4 deletions(-)

diff --git a/tensorflow/contrib/lite/toco/dump_graphviz.cc b/tensorflow/contrib/lite/toco/dump_graphviz.cc
index 5bb0e3ba4d..166ead9184 100644
--- a/tensorflow/contrib/lite/toco/dump_graphviz.cc
+++ b/tensorflow/contrib/lite/toco/dump_graphviz.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/contrib/lite/toco/dump_graphviz.h"
 
+#include <cmath>
 #include <memory>
 #include <set>
 #include <unordered_set>
@@ -63,6 +64,7 @@ struct NodeProperties {
   // color will be chosen for the 'fontcolor' for the inside text
   // label, see Color::TextColorString.
   Color color;
+  float log2_buffer_size;
 };
 
 // All colors in this file are from:
@@ -162,9 +164,12 @@ NodeProperties GetPropertiesForArray(const Model& model,
     }
     node_properties.label += "]";
 
+    int buffer_size = RequiredBufferSizeForShape(array.shape());
+    node_properties.log2_buffer_size =
+        std::log2(static_cast<float>(buffer_size));
+
     if (array.buffer) {
       const auto& array = model.GetArray(array_name);
-      int buffer_size = RequiredBufferSizeForShape(array.shape());
       if (buffer_size <= 4) {
         AppendF(&node_properties.label, " = ");
         if (array.shape().dimensions_count() > 0) {
@@ -194,6 +199,8 @@ NodeProperties GetPropertiesForArray(const Model& model,
         AppendF(&node_properties.label, "}");
       }
     }
+  } else {
+    node_properties.log2_buffer_size = 0.0f;
   }
 
   if (array.minmax) {
@@ -325,12 +332,18 @@ std::vector<const Operator*> OperatorsToDump(const Model& model) {
 
 void DumpGraphviz(const Model& model, string* output_file_contents) {
   AppendF(output_file_contents, "digraph Computegraph {\n");
+  // 'nslimit' is a graphviz (dot) paramater that limits the iterations during
+  // the layout phase. Omitting it allows infinite iterations, causing some
+  // complex graphs to never finish. A value of 125 produces good graphs
+  // while allowing complex graphs to finish.
+  AppendF(output_file_contents, "\t nslimit=125;\n");
 
   constexpr char kNodeFormat[] =
       "\t \"%s\" [label=\"%s\", shape=%s, style=filled, fillcolor=\"#%s\", "
       "fontcolor = \"#%sDD\"];\n";
 
-  constexpr char kEdgeFormat[] = "\t \"%s\" -> \"%s\";\n";
+  constexpr char kEdgeFormat[] =
+      "\t \"%s\" -> \"%s\" [penwidth=%f, weight=%f];\n";
 
   constexpr char kRNNBackEdgeFormat[] =
       "\t \"%s\" -> \"%s\" [color=\"#0F9D58\"];\n";
@@ -358,7 +371,22 @@ void DumpGraphviz(const Model& model, string* output_file_contents) {
                 array_properties.color.FillColorString().c_str(),
                 array_properties.color.TextColorString().c_str());
       }
-      AppendF(output_file_contents, kEdgeFormat, input, operator_id);
+
+      // Draw lines that transport more data thicker (Otherwise, where would the
+      // data fit? right?).
+      float line_width =
+          std::max(0.5f, array_properties.log2_buffer_size / 3.0f);
+      // Keep edges that transport more data shorter than those with less.
+      float weight = std::max(1.0f, array_properties.log2_buffer_size);
+      if (!IsInputArray(model, input) &&
+          GetOpWithOutput(model, input) == nullptr) {
+        // Give the main line of data flow a straighter path by penalizing edges
+        // to standalone buffers. Weights are generally very large buffers that
+        // otherwise skew the layout without this.
+        weight = 1.0f;
+      }
+      AppendF(output_file_contents, kEdgeFormat, input, operator_id, line_width,
+              weight);
       already_added_arrays.insert(input);
     }
     // Add nodes and edges for all outputs of the operator.
@@ -374,7 +402,16 @@ void DumpGraphviz(const Model& model, string* output_file_contents) {
                 array_properties.color.FillColorString().c_str(),
                 array_properties.color.TextColorString().c_str());
       }
-      AppendF(output_file_contents, kEdgeFormat, operator_id, output);
+
+      // See comments above regarding weight and line_width calculations.
+      float line_width =
+          std::max(0.5f, array_properties.log2_buffer_size / 3.0f);
+      float weight = std::max(1.0f, array_properties.log2_buffer_size);
+      if (!IsArrayConsumed(model, output)) {
+        weight = 1.0f;
+      }
+      AppendF(output_file_contents, kEdgeFormat, operator_id, output,
+              line_width, weight);
       already_added_arrays.insert(output);
     }
   }
-- 
GitLab


From 286d61b246280b3a8dea39ac2f7d48b7cdbd48dc Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 15:41:28 -0700
Subject: [PATCH 068/395] Do not allocate memory for literal as it will be
 allocated later.

PiperOrigin-RevId: 194857422
---
 tensorflow/compiler/xla/literal_util.h | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensorflow/compiler/xla/literal_util.h b/tensorflow/compiler/xla/literal_util.h
index 8aa19222dc..956ff7d21c 100644
--- a/tensorflow/compiler/xla/literal_util.h
+++ b/tensorflow/compiler/xla/literal_util.h
@@ -74,6 +74,10 @@ class Literal {
   Literal(const Literal& other) = delete;
   Literal& operator=(const Literal& other) = delete;
   Literal(Literal&& other);
+  // 'allocate_arrays' indicates whether to allocate memory for the arrays in
+  // the shape. If false, buffer pointers inside of the Literal::Pieces are set
+  // to nullptr.
+  Literal(const Shape& shape, bool allocate_arrays);
   Literal& operator=(Literal&& other);
 
   // Literals are equal if they have compatible shapes and the same data
@@ -659,11 +663,6 @@ class Literal {
   int64 sparse_element_count() const;
 
  protected:
-  // 'allocate_arrays' indicates whether to allocate memory for the arrays in
-  // the shape. If false, buffer pointers inside of the Literal::Pieces are set
-  // to nullptr.
-  Literal(const Shape& shape, bool allocate_arrays);
-
   // Internal template helper for the Literal::CopySliceFrom(), matching its
   // arguments one by one.
   template <typename NativeT>
-- 
GitLab


From 30fcdecc05e6b25ab8d451997904e40b2a76acd4 Mon Sep 17 00:00:00 2001
From: Yifei Feng <yifeif@google.com>
Date: Mon, 30 Apr 2018 15:56:26 -0700
Subject: [PATCH 069/395] Improve error message for pip_smoke_test.

PiperOrigin-RevId: 194859591
---
 tensorflow/tools/pip_package/pip_smoke_test.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensorflow/tools/pip_package/pip_smoke_test.py b/tensorflow/tools/pip_package/pip_smoke_test.py
index 1b692104f1..b23dde2019 100644
--- a/tensorflow/tools/pip_package/pip_smoke_test.py
+++ b/tensorflow/tools/pip_package/pip_smoke_test.py
@@ -147,10 +147,11 @@ def main():
       affected_tests_list = affected_tests.split("\n")[:-2]
       print("\n".join(affected_tests_list))
 
-    raise RuntimeError("""One or more dependencies are not in the pip package.
-Please either blacklist the dependencies in
-//tensorflow/tools/pip_package/pip_smoke_test.py
-or add them to //tensorflow/tools/pip_package/BUILD.""")
+    raise RuntimeError("""
+    One or more added test dependencies are not in the pip package.
+If these test dependencies need to be in TensorFlow pip package, please add them to //tensorflow/tools/pip_package/BUILD.
+Else either blacklist the dependencies in //tensorflow/tools/pip_package/pip_smoke_test.py
+or add no_pip tag to the test.""")
 
   else:
     print("TEST PASSED")
-- 
GitLab


From 18343616da47a9c3eab79b5028ac3d8bf786f2ff Mon Sep 17 00:00:00 2001
From: Bixia Zheng <bixia@google.com>
Date: Mon, 30 Apr 2018 16:11:38 -0700
Subject: [PATCH 070/395] [XLA] Change the TF2XLA bridge to perform F16
 reduction using F32 data type.

Add test cases to test that reduce sum for bfloat16 and float16 doesn't lose too
much precision.

PiperOrigin-RevId: 194862078
---
 tensorflow/compiler/tests/reduce_ops_test.py | 64 ++++++++++++++++++++
 tensorflow/compiler/tf2xla/xla_helpers.cc    |  2 +-
 2 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/tensorflow/compiler/tests/reduce_ops_test.py b/tensorflow/compiler/tests/reduce_ops_test.py
index 2c084b04fa..7420724bdb 100644
--- a/tensorflow/compiler/tests/reduce_ops_test.py
+++ b/tensorflow/compiler/tests/reduce_ops_test.py
@@ -19,6 +19,7 @@ from __future__ import division
 from __future__ import print_function
 
 import functools
+import itertools
 import numpy as np
 
 from tensorflow.compiler.tests.xla_test import XLATestCase
@@ -155,5 +156,68 @@ class ReduceOpsTest(XLATestCase):
     self._testReduction(math_ops.reduce_any, np.any, np.bool, self.BOOL_DATA)
 
 
+class ReduceOpPrecisionTest(XLATestCase):
+
+  def _testReduceSum(self,
+                     expected_result,
+                     dtype,
+                     test_inputs,
+                     rtol=1e-3,
+                     atol=1e-4):
+    """Tests reduce sum on a list of input arrays.
+
+    For each array in test_inputs, check that performing reduce sum on the array
+    produces a value that is close to the expected result.
+
+    Args:
+      expected_result: the expected result.
+      dtype: the data type of the reduce sum operation.
+      test_inputs: a list of input arrays for the reduce sum operation.
+      rtol: the relative error.
+      atol: the absolute error.
+    """
+
+    for test_input in test_inputs:
+      with self.test_session() as sess:
+        with self.test_scope():
+          a = array_ops.placeholder(dtype)
+          index = array_ops.placeholder(dtypes.int32)
+          out = math_ops.reduce_sum(a, index)
+        result = sess.run(out, {
+            a: np.array(test_input, dtype=dtype),
+            index: [0]
+        })
+        # Compare the results using float32 type.
+        self.assertAllClose(
+            np.float32(result),
+            np.float32(expected_result),
+            rtol=rtol,
+            atol=atol)
+
+  def testReduceSumF16(self):
+    """Tests the reduce sum of float16 doesn't lose too much precision."""
+
+    if np.float16 not in self.all_types:
+      return
+
+    f16_max = np.finfo(np.float16).max
+    self._testReduceSum(
+        f16_max, np.float16,
+        itertools.permutations([f16_max, f16_max, f16_max * (-1.0)], 3))
+
+  def testReduceSumBF16(self):
+    """Tests the reduce sum of bfloat16 doesn't lose too much precision."""
+
+    if dtypes.bfloat16.as_numpy_dtype not in self.all_types:
+      return
+
+    bf16_max = np.float32(dtypes.bfloat16.max)
+    f32_max = dtypes.float32.max
+    value = min(bf16_max, f32_max - bf16_max)
+    self._testReduceSum(
+        dtypes.bfloat16.as_numpy_dtype(value), dtypes.bfloat16.as_numpy_dtype,
+        itertools.permutations([bf16_max, value, bf16_max * (-1.0)], 3))
+
+
 if __name__ == '__main__':
   googletest.main()
diff --git a/tensorflow/compiler/tf2xla/xla_helpers.cc b/tensorflow/compiler/tf2xla/xla_helpers.cc
index 62a5114837..a3deb02a1f 100644
--- a/tensorflow/compiler/tf2xla/xla_helpers.cc
+++ b/tensorflow/compiler/tf2xla/xla_helpers.cc
@@ -278,7 +278,7 @@ Status XlaHelpers::OneHot(xla::ComputationBuilder* builder, int64 depth,
 }
 
 DataType XlaHelpers::SumAccumulationType(const DataType& dtype) {
-  if (dtype == DT_BFLOAT16) {
+  if (dtype == DT_BFLOAT16 || dtype == DT_HALF) {
     return DT_FLOAT;
   }
   return dtype;
-- 
GitLab


From 7141ed55dd0f36f698143812b44aeffc6129257b Mon Sep 17 00:00:00 2001
From: Yuefeng Zhou <yuefengz@google.com>
Date: Mon, 30 Apr 2018 16:12:33 -0700
Subject: [PATCH 071/395] Add MultiNodeDataset and MultiNodeIterator which are
 intended to work for multi-node distribution strategy.

PiperOrigin-RevId: 194862215
---
 tensorflow/contrib/distribute/python/BUILD    |  22 +++
 .../python/multi_worker_test_base.py          |  90 +++++++++++++
 .../contrib/distribute/python/values.py       |  95 +++++++++++++
 .../contrib/distribute/python/values_test.py  | 127 ++++++++++++++++++
 4 files changed, 334 insertions(+)
 create mode 100644 tensorflow/contrib/distribute/python/multi_worker_test_base.py

diff --git a/tensorflow/contrib/distribute/python/BUILD b/tensorflow/contrib/distribute/python/BUILD
index c2834d8226..aa1a956a2d 100644
--- a/tensorflow/contrib/distribute/python/BUILD
+++ b/tensorflow/contrib/distribute/python/BUILD
@@ -42,6 +42,7 @@ cuda_py_test(
     srcs = ["values_test.py"],
     additional_deps = [
         ":mirrored_strategy",
+        ":multi_worker_test_base",
         ":values",
         "//tensorflow/core:protos_all_py",
         "//tensorflow/python/data/ops:dataset_ops",
@@ -57,6 +58,9 @@ cuda_py_test(
         "//tensorflow/python/eager:test",
         "//tensorflow/python/estimator:model_fn",
     ],
+    tags = [
+        "no_pip",
+    ],
 )
 
 py_library(
@@ -216,6 +220,24 @@ cuda_py_test(
     ],
 )
 
+py_library(
+    name = "multi_worker_test_base",
+    testonly = 1,
+    srcs = ["multi_worker_test_base.py"],
+    srcs_version = "PY2AND3",
+    tags = [
+        "no_pip",
+    ],
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python:distributed_framework_test_lib",
+        "//tensorflow/python:platform",
+        "//tensorflow/python:session",
+        "//tensorflow/python:training",
+        "//tensorflow/python/eager:test",
+    ],
+)
+
 py_library(
     name = "step_fn",
     srcs = ["step_fn.py"],
diff --git a/tensorflow/contrib/distribute/python/multi_worker_test_base.py b/tensorflow/contrib/distribute/python/multi_worker_test_base.py
new file mode 100644
index 0000000000..f659be5f42
--- /dev/null
+++ b/tensorflow/contrib/distribute/python/multi_worker_test_base.py
@@ -0,0 +1,90 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Base testing class for strategies that require multiple nodes."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import contextlib
+import copy
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import rewriter_config_pb2
+from tensorflow.python.client import session
+from tensorflow.python.eager import test
+from tensorflow.python.framework import test_util
+
+
+class MultiWorkerTestBase(test.TestCase):
+  """Base class for testing multi node strategy and dataset."""
+
+  @classmethod
+  def setUpClass(cls):
+    """Create a local cluster with 2 workers."""
+    num_workers = 2
+    # Leave some memory for cuda runtime.
+    gpu_mem_frac = 0.7 / num_workers
+    default_config = config_pb2.ConfigProto()
+    default_config.gpu_options.per_process_gpu_memory_fraction = gpu_mem_frac
+
+    # The local cluster takes some portion of the local GPUs and there is no way
+    # for the cluster to terminate unless using multiple processes. Therefore,
+    # we have to only create only one cluster throughout a test process.
+    workers, _ = test_util.create_local_cluster(
+        num_workers, num_ps=0, worker_config=default_config)
+    cls._master_target = workers[0].target
+
+  @contextlib.contextmanager
+  def test_session(self, graph=None, config=None):
+    """Create a test session with master target set to the testing cluster.
+
+    This overrides the base class' method, removes arguments that are not needed
+    by the multi-node case and creates a test session that connects to the local
+    testing cluster.
+
+    Args:
+      graph: Optional graph to use during the returned session.
+      config: An optional config_pb2.ConfigProto to use to configure the
+        session.
+
+    Yields:
+      A Session object that should be used as a context manager to surround
+      the graph building and execution code in a test case.
+    """
+    if self.id().endswith('.test_session'):
+      self.skipTest('Not a test.')
+
+    if config is None:
+      config = config_pb2.ConfigProto(allow_soft_placement=True)
+    else:
+      config = copy.deepcopy(config)
+    # Don't perform optimizations for tests so we don't inadvertently run
+    # gpu ops on cpu
+    config.graph_options.optimizer_options.opt_level = -1
+    config.graph_options.rewrite_options.constant_folding = (
+        rewriter_config_pb2.RewriterConfig.OFF)
+
+    if graph is None:
+      if self._cached_session is None:  # pylint: disable=access-member-before-definition
+        self._cached_session = session.Session(
+            graph=None, config=config, target=self._master_target)
+      sess = self._cached_session
+      with sess.graph.as_default(), sess.as_default():
+        yield sess
+    else:
+      with session.Session(
+          graph=graph, config=config, target=self._master_target) as sess:
+        yield sess
diff --git a/tensorflow/contrib/distribute/python/values.py b/tensorflow/contrib/distribute/python/values.py
index 466678ef2e..18afdaa7b0 100644
--- a/tensorflow/contrib/distribute/python/values.py
+++ b/tensorflow/contrib/distribute/python/values.py
@@ -29,6 +29,7 @@ import six
 from tensorflow.contrib.data.python.ops import batching
 from tensorflow.contrib.distribute.python import prefetching_ops_v2
 from tensorflow.python.eager import context
+from tensorflow.python.framework import device as tf_device
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
@@ -576,6 +577,100 @@ class PerDeviceDataset(object):
         dataset_iterator, self._devices, self._prefetch_on_device)
 
 
+class MultiWorkerDataIterator(object):
+  """An iterator (like `tf.data.Iterator`) into a `MultiWorkerDataset`."""
+
+  def __init__(self, iterators, worker_device_map):
+    """Initialize the MultiWorkerDataIterator object.
+
+    Args:
+      iterators: a dict mapping from each worker to an iterator for
+        that worker.
+      worker_device_map: a dict mapping from each worker's devices to a list of
+        devices that belong to this worker.
+
+    Raises:
+      ValueError: if iterators and worker_device_map are not compatible.
+    """
+    self._iterators = iterators
+    self._worker_device_map = worker_device_map
+    if set(self._iterators) != set(self._worker_device_map):
+      raise ValueError("iterators and worker_device_map are not compatible.")
+
+  @property
+  def initializer(self):
+    return control_flow_ops.group(
+        [iterator.initializer for iterator in self._iterators.values()])
+
+  def get_next(self, name=None):
+    """Scatter the input across hosts and devices."""
+    index = {}
+    for worker, iterator in six.iteritems(self._iterators):
+      if name is not None:
+        d = tf_device.DeviceSpec.from_string(worker)
+        new_name = "%s_%s_%d" % (name, d.job, d.task)
+      else:
+        new_name = None
+      with ops.device(worker):
+        data_per_worker = iterator.get_next(name=new_name)
+
+      worker_devices = self._worker_device_map[worker]
+      # Ungroup these per-device value so as to get a flat map from devices to
+      # values.
+      for d in worker_devices:
+        v = select_device(d, data_per_worker)
+        if d in index:
+          raise ValueError("Duplicated devices in worker_device_map: %r" % v)
+        index[d] = v
+
+    return regroup(index)
+
+
+class MultiWorkerDataset(object):
+  """Like a `tf.data.Dataset` that distributes data to different workers.
+
+  Each worker gets one shard of the input dataset. It is currently not working
+  in
+  eager mode.
+  """
+
+  def __init__(self, dataset_fn, worker_device_map, prefetch_on_device=None):
+    """Initialize the MultiWorkerDataset object.
+
+    Args:
+      dataset_fn: a function that returns a `tf.data.Dataset`.
+      worker_device_map: a dict mapping from each worker to a list of devices
+        that belong to this worker.
+      prefetch_on_device: whether to prefetch to devices.
+    """
+    self._worker_device_map = worker_device_map
+    self._datasets = {}
+    # TODO(yuefengz, priyag): support different set of jobs for input
+    # processing.
+    for i, (worker, worker_devices) in enumerate(
+        six.iteritems(worker_device_map)):
+      with ops.device(worker):
+        worker_input = dataset_fn()
+        # TODO(yuefengz, priyag): support efficient sharding.
+        worker_input = worker_input.shard(len(worker_device_map), i)
+        self._datasets[worker] = PerDeviceDataset(
+            worker_input, worker_devices, prefetch_on_device=prefetch_on_device)
+
+  def make_one_shot_iterator(self):
+    iterators = {}
+    for worker, dataset in six.iteritems(self._datasets):
+      with ops.device(worker):
+        iterators[worker] = dataset.make_one_shot_iterator()
+    return MultiWorkerDataIterator(iterators, self._worker_device_map)
+
+  def make_initializable_iterator(self):
+    iterators = {}
+    for worker, dataset in six.iteritems(self._datasets):
+      with ops.device(worker):
+        iterators[worker] = dataset.make_initializable_iterator()
+    return MultiWorkerDataIterator(iterators, self._worker_device_map)
+
+
 class PerIteration(object):
   """Holds input for multiple iterations at once."""
 
diff --git a/tensorflow/contrib/distribute/python/values_test.py b/tensorflow/contrib/distribute/python/values_test.py
index 1d4e801cd8..9aeef9fa3e 100644
--- a/tensorflow/contrib/distribute/python/values_test.py
+++ b/tensorflow/contrib/distribute/python/values_test.py
@@ -18,9 +18,11 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import os
 
 from tensorflow.contrib.distribute.python import mirrored_strategy
+from tensorflow.contrib.distribute.python import multi_worker_test_base
 from tensorflow.contrib.distribute.python import values
 from tensorflow.core.protobuf import config_pb2
 from tensorflow.python.data.ops import dataset_ops
@@ -37,6 +39,7 @@ from tensorflow.python.ops import variable_scope
 from tensorflow.python.ops import variables as variables_lib
 from tensorflow.python.training import device_util
 from tensorflow.python.training import saver as saver_lib
+from tensorflow.python.util import nest
 
 
 @test_util.with_c_api
@@ -437,6 +440,130 @@ class PerDeviceDatasetTest(test.TestCase):
         self.evaluate(next_element)
 
 
+class MultiWorkerDatasetTest(multi_worker_test_base.MultiWorkerTestBase):
+
+  def _test_iterator(self, iterator, devices, expected_values):
+    next_element = iterator.get_next()
+    for device in devices:
+      v = values.select_device(device, next_element)
+      # The `v` here can be a tuple.
+      for element in nest.flatten(v):
+        self.assertTrue(element.device in device)
+
+    for expected_value in expected_values:
+      actual = self.evaluate(
+          [values.select_device(d, next_element) for d in devices])
+      self.assertEqual(expected_value, actual)
+
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate([values.select_device(d, next_element) for d in devices])
+
+  def _test_dataset(self, dataset_fn, worker_device_map, devices,
+                    expected_values):
+    multi_worker_dataset = values.MultiWorkerDataset(
+        dataset_fn, worker_device_map, prefetch_on_device=False)
+    multi_worker_iterator = multi_worker_dataset.make_one_shot_iterator()
+    self._test_iterator(multi_worker_iterator, devices, expected_values)
+
+  def _cpu_devices(self):
+    worker_device_map = collections.OrderedDict(
+        [("/job:worker/replica:0/task:0",
+          ["/job:worker/replica:0/task:0/device:CPU:0"]),
+         ("/job:worker/replica:0/task:1",
+          ["/job:worker/replica:0/task:1/device:CPU:0"])])
+    devices = [
+        "/job:worker/replica:0/task:0/device:CPU:0",
+        "/job:worker/replica:0/task:1/device:CPU:0"
+    ]
+    return worker_device_map, devices
+
+  def _cpu_and_one_gpu_devices(self):
+    # The worker_device_map doesn't have to be a OrderDict object, this is just
+    # to simplify the testing so that we can pass expected values as a list
+    # instead of a dict.
+    worker_device_map = collections.OrderedDict(
+        [("/job:worker/replica:0/task:0", [
+            "/job:worker/replica:0/task:0/device:GPU:0",
+            "/job:worker/replica:0/task:0/device:CPU:0"
+        ]), ("/job:worker/replica:0/task:1", [
+            "/job:worker/replica:0/task:1/device:GPU:0",
+            "/job:worker/replica:0/task:1/device:CPU:0"
+        ])])
+    devices = [
+        "/job:worker/replica:0/task:0/device:GPU:0",
+        "/job:worker/replica:0/task:0/device:CPU:0",
+        "/job:worker/replica:0/task:1/device:GPU:0",
+        "/job:worker/replica:0/task:1/device:CPU:0"
+    ]
+    return worker_device_map, devices
+
+  def testDataDistributionOneDevicePerWorker(self):
+    worker_device_map, devices = self._cpu_devices()
+    with context.graph_mode():
+      dataset_fn = lambda: dataset_ops.Dataset.range(8)
+      self._test_dataset(dataset_fn, worker_device_map, devices,
+                         [[0, 1], [2, 3], [4, 5], [6, 7]])
+
+  def testDataDistributionTwoDevicePerWorker(self):
+    if context.num_gpus() < 1:
+      self.skipTest("A GPU is not available for this test.")
+    worker_device_map, devices = self._cpu_and_one_gpu_devices()
+    with context.graph_mode():
+      dataset_fn = lambda: dataset_ops.Dataset.range(8)
+      self._test_dataset(dataset_fn, worker_device_map, devices,
+                         [[0, 2, 1, 3], [4, 6, 5, 7]])
+
+  def testTupleDataset(self):
+    worker_device_map, devices = self._cpu_devices()
+
+    with context.graph_mode():
+
+      def dataset_fn():
+        dataset1 = dataset_ops.Dataset.range(8)
+        dataset2 = dataset_ops.Dataset.range(8).map(lambda x: x**2)
+        return dataset_ops.Dataset.zip((dataset1, dataset2))
+
+      expected_values = [
+          [(i, i**2), (i + 1, (i + 1)**2)] for i in range(0, 8, 2)
+      ]
+      self._test_dataset(dataset_fn, worker_device_map, devices,
+                         expected_values)
+
+  def testInitializableIterator(self):
+    worker_device_map, devices = self._cpu_devices()
+    with context.graph_mode():
+      dataset_fn = lambda: dataset_ops.Dataset.range(8)
+      multi_worker_dataset = values.MultiWorkerDataset(
+          dataset_fn, worker_device_map, prefetch_on_device=False)
+      multi_worker_iterator = multi_worker_dataset.make_initializable_iterator()
+
+      self.evaluate(multi_worker_iterator.initializer)
+      self._test_iterator(multi_worker_iterator, devices,
+                          [[0, 1], [2, 3], [4, 5], [6, 7]])
+
+      # After re-initializing the iterator, should be able to iterate again.
+      self.evaluate(multi_worker_iterator.initializer)
+      self._test_iterator(multi_worker_iterator, devices,
+                          [[0, 1], [2, 3], [4, 5], [6, 7]])
+
+  def testValueErrorForIterator(self):
+    # Incompatiable arguments.
+    with self.assertRaises(ValueError):
+      values.MultiWorkerDataIterator({"w1": None}, {"w1": "d1", "w2": "d2"})
+
+    # Test duplicated devices under same worker.
+    worker_device_map, _ = self._cpu_devices()
+    worker_device_map["/job:worker/replica:0/task:0"].append(
+        "/job:worker/replica:0/task:0/device:CPU:0")
+    with context.graph_mode():
+      dataset_fn = lambda: dataset_ops.Dataset.range(8)
+      multi_worker_dataset = values.MultiWorkerDataset(
+          dataset_fn, worker_device_map, prefetch_on_device=False)
+      multi_worker_iterator = multi_worker_dataset.make_initializable_iterator()
+      with self.assertRaises(ValueError):
+        multi_worker_iterator.get_next()
+
+
 @test_util.with_c_api
 class MirroredVariableTest(test.TestCase):
 
-- 
GitLab


From 1ff23a314f355a9ebaaf207dbeae56ebc1634d63 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 16:43:14 -0700
Subject: [PATCH 072/395] Small fix to prevent a crash if the delegate has not
 implemented FreeBufferHandle.

PiperOrigin-RevId: 194866595
---
 tensorflow/contrib/lite/interpreter.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/contrib/lite/interpreter.cc b/tensorflow/contrib/lite/interpreter.cc
index 9d8ea55fd1..ebb0aedc20 100644
--- a/tensorflow/contrib/lite/interpreter.cc
+++ b/tensorflow/contrib/lite/interpreter.cc
@@ -125,7 +125,8 @@ Interpreter::~Interpreter() {
 
   for (int i = 0; i < context_.tensors_size; i++) {
     TfLiteTensor* tensor = &context_.tensors[i];
-    if (tensor->buffer_handle != kTfLiteNullBufferHandle) {
+    if (tensor->buffer_handle != kTfLiteNullBufferHandle &&
+        tensor->delegate->FreeBufferHandle != nullptr) {
       tensor->delegate->FreeBufferHandle(tensor->delegate,
                                          &tensor->buffer_handle);
     }
-- 
GitLab


From 64bb1de61377f12859a719448b65b452b03047a7 Mon Sep 17 00:00:00 2001
From: Alexandre Passos <apassos@google.com>
Date: Mon, 30 Apr 2018 17:11:40 -0700
Subject: [PATCH 073/395] Faster reduce_logsoftmax (specially in eager) and
 bugfixes in broadcast_to

PiperOrigin-RevId: 194870645
---
 tensorflow/core/kernels/broadcast_to_op.h     |  34 +++++-
 tensorflow/core/ops/array_ops.cc              |   2 +-
 tensorflow/python/kernel_tests/BUILD          |  16 +++
 .../kernel_tests/reduce_benchmark_test.py     | 107 ++++++++++++++++++
 tensorflow/python/ops/math_ops.py             |  11 +-
 5 files changed, 161 insertions(+), 9 deletions(-)
 create mode 100644 tensorflow/python/kernel_tests/reduce_benchmark_test.py

diff --git a/tensorflow/core/kernels/broadcast_to_op.h b/tensorflow/core/kernels/broadcast_to_op.h
index 608e9b6ac9..73fdd5d28e 100644
--- a/tensorflow/core/kernels/broadcast_to_op.h
+++ b/tensorflow/core/kernels/broadcast_to_op.h
@@ -34,14 +34,37 @@ struct BroadcastTo {
                   const TensorShape &input_shape) {
 #define BROADCAST_SHAPE(broadcast, reshape, NDIMS, input_shape, output_shape) \
   for (int i = 0; i < NDIMS; i++) {                                           \
-    OP_REQUIRES(ctx, (broadcast[i] % reshape[i] == 0),                        \
-                errors::InvalidArgument("invalid shape to broadcast from ",   \
-                                        input_shape.DebugString(), " to ",    \
-                                        output_shape.DebugString()));         \
-    broadcast[i] = broadcast[i] / reshape[i];                                 \
+    if (reshape[i] != broadcast[i]) {                                         \
+      OP_REQUIRES(ctx,                                                        \
+                  ((reshape[i] != 0) && (broadcast[i] % reshape[i] == 0)),    \
+                  errors::InvalidArgument("invalid shape to broadcast from ", \
+                                          input_shape.DebugString(), " to ",  \
+                                          output_shape.DebugString()));       \
+      broadcast[i] = broadcast[i] / reshape[i];                               \
+    } else {                                                                  \
+      broadcast[i] = 1;                                                       \
+    }                                                                         \
   }
 
+    if (output_shape.num_elements() == 0) {
+      return;
+    }
+    if (output_shape == input_shape) {
+      output_tensor.flat<T>().device(d) = input_tensor.flat<T>();
+      return;
+    }
+
     switch (output_shape.dims()) {
+      case 0: {
+        if (input_shape.dims() > 0) {
+          ctx->CtxFailure(errors::InvalidArgument(
+              "invalid shape to broadcast from ", input_shape.DebugString(),
+              " to ", output_shape.DebugString()));
+          break;
+        }
+        output_tensor.scalar<T>().device(d) = input_tensor.scalar<T>();
+        break;
+      }
       case 1: {
         auto reshape = AsEigenDSizesWithPrefix<1>(input_shape);
         auto broadcast = output_shape.AsEigenDSizes<1>();
@@ -125,7 +148,6 @@ struct BroadcastTo {
         auto broadcast = output_shape.AsEigenDSizes<4>();
 
         BROADCAST_SHAPE(broadcast, reshape, 4, input_shape, output_shape);
-
         auto output = output_tensor.tensor<T, 4>();
         switch (input_shape.dims()) {
           case 0: {
diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc
index 88fc03826a..fce0b93cd7 100644
--- a/tensorflow/core/ops/array_ops.cc
+++ b/tensorflow/core/ops/array_ops.cc
@@ -466,7 +466,7 @@ REGISTER_OP("BroadcastTo")
           // so no check needed.
           if (i >= in_offset) {
             DimensionHandle in_dim = c->Dim(in, i - in_offset);
-            if (c->ValueKnown(in_dim)) {
+            if (c->ValueKnown(in_dim) && c->Value(in_dim) != 0) {
               if (c->Value(dim) % c->Value(in_dim) != 0) {
                 return errors::InvalidArgument(
                     "Cannot broadcast a tensor with shape ", c->DebugString(in),
diff --git a/tensorflow/python/kernel_tests/BUILD b/tensorflow/python/kernel_tests/BUILD
index b4ff094cdf..c892b6ee9a 100644
--- a/tensorflow/python/kernel_tests/BUILD
+++ b/tensorflow/python/kernel_tests/BUILD
@@ -112,6 +112,22 @@ cuda_py_test(
     tags = ["no_windows"],
 )
 
+cuda_py_test(
+    name = "reduce_benchmark_test",
+    srcs = ["reduce_benchmark_test.py"],
+    additional_deps = [
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python:framework",
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:gradients",
+        "//tensorflow/python:math_ops",
+        "//tensorflow/python:platform",
+        "//tensorflow/python:platform_benchmark",
+    ],
+)
+
 tf_py_test(
     name = "bincount_op_test",
     size = "small",
diff --git a/tensorflow/python/kernel_tests/reduce_benchmark_test.py b/tensorflow/python/kernel_tests/reduce_benchmark_test.py
new file mode 100644
index 0000000000..3a2fb81157
--- /dev/null
+++ b/tensorflow/python/kernel_tests/reduce_benchmark_test.py
@@ -0,0 +1,107 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""Simple benchmarks for reductions and their gradients."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import time
+
+import numpy as np
+from six.moves import range  # pylint: disable=redefined-builtin
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.client import session
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class ReduceBenchmarks(test.Benchmark):
+  """Benchmarks for reductions."""
+
+  def _run(self, func, num_iters):
+    # call func to maybe warm up the GPU
+    func()
+    start = time.time()
+    for _ in range(num_iters):
+      func()
+    end = time.time()
+    mean_us = (end - start) * 1e6 / num_iters
+    self.report_benchmark(
+        iters=num_iters,
+        wall_time=mean_us,
+        extras={"examples_per_sec": num_iters / (end - start)})
+
+  def benchmark_reduce_sum_grad_eager(self):
+    with context.eager_mode():
+      tensor = array_ops.zeros([100, 1000])
+
+      def fn():
+        backprop.gradients_function(math_ops.reduce_sum, [0])(tensor)
+
+      self._run(fn, 10000)
+
+  def benchmark_reduce_sum_grad_eager_cpu(self):
+    with context.eager_mode(), ops.device("/cpu:0"):
+      tensor = array_ops.zeros([100, 1000])
+
+      def fn():
+        backprop.gradients_function(math_ops.reduce_sum, [0])(tensor)
+
+      self._run(fn, 10000)
+
+  def benchmark_reduce_sum_grad_graph(self):
+    config = config_pb2.ConfigProto(
+        graph_options=config_pb2.GraphOptions(
+            optimizer_options=config_pb2.OptimizerOptions(
+                opt_level=config_pb2.OptimizerOptions.L0)))
+    with ops.Graph().as_default(), session.Session(config=config) as sess:
+
+      tensor = constant_op.constant(np.zeros([100, 1000], dtype=np.float32))
+      reduction = math_ops.reduce_sum(tensor)
+      grad, = gradients_impl.gradients(reduction, tensor)
+
+      def fn():
+        sess.run(grad.op)
+
+      self._run(fn, 10000)
+
+  def benchmark_reduce_sum_grad_graph_cpu(self):
+    config = config_pb2.ConfigProto(
+        graph_options=config_pb2.GraphOptions(
+            optimizer_options=config_pb2.OptimizerOptions(
+                opt_level=config_pb2.OptimizerOptions.L0)))
+    with ops.Graph().as_default(), session.Session(config=config) as sess:
+
+      with ops.device("/cpu:0"):
+        tensor = constant_op.constant(np.zeros([100, 1000], dtype=np.float32))
+        reduction = math_ops.reduce_sum(tensor)
+        grad, = gradients_impl.gradients(reduction, tensor)
+
+      def fn():
+        sess.run(grad.op)
+
+      self._run(fn, 10000)
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py
index b937273137..57660578aa 100644
--- a/tensorflow/python/ops/math_ops.py
+++ b/tensorflow/python/ops/math_ops.py
@@ -1757,6 +1757,7 @@ def reduce_logsumexp(input_tensor,
                                                     "keep_dims", keep_dims)
   if keepdims is None:
     keepdims = False
+  input_tensor = ops.convert_to_tensor(input_tensor)
   with ops.name_scope(name, "ReduceLogSumExp", [input_tensor]) as name:
     raw_max = reduce_max(
         input_tensor,
@@ -1769,13 +1770,13 @@ def reduce_logsumexp(input_tensor,
             array_ops.zeros_like(raw_max)))
     result = gen_math_ops.log(
         reduce_sum(
-            gen_math_ops.exp(input_tensor - my_max),
+            gen_math_ops.exp(gen_math_ops.sub(input_tensor, my_max)),
             axis,
             keepdims=keepdims,
             reduction_indices=reduction_indices))
     if not keepdims:
       my_max = array_ops.reshape(my_max, array_ops.shape(result))
-    result += my_max
+    result = gen_math_ops.add(result, my_max)
     return _may_reduce_to_scalar(keepdims, axis, reduction_indices, result)
 
 
@@ -2475,6 +2476,12 @@ def reduced_shape(input_shape, axes):
   """
   # Example:
   # cast needed for SparseTensor reductions
+  if context.executing_eagerly():
+    input_shape = input_shape.numpy()
+    axes = axes.numpy()
+    input_shape[axes] = 1
+    return input_shape
+
   input_shape = to_int32(input_shape)  # [2, 3, 5, 7]
   axes = to_int32(axes)  # [1, 2]
 
-- 
GitLab


From b7978d48f4588feb717157a9dbfd2e1df678628b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 17:14:50 -0700
Subject: [PATCH 074/395] Internal cleanup.

PiperOrigin-RevId: 194871141
---
 .../org/tensorflow/lite/NativeInterpreterWrapper.java     | 8 ++++----
 .../test/java/org/tensorflow/lite/InterpreterTest.java    | 8 ++++----
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
index 2fc803715b..a43251cad1 100644
--- a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
+++ b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
@@ -173,8 +173,8 @@ final class NativeInterpreterWrapper implements AutoCloseable {
     } else {
       throw new IllegalArgumentException(
           String.format(
-              "Input error: %s is not a valid name for any input. "
-                  + "The indexes of the inputs are %s",
+              "Input error: '%s' is not a valid name for any input. Names of inputs and their "
+                  + "indexes are %s",
               name, inputsIndexes.toString()));
     }
   }
@@ -195,8 +195,8 @@ final class NativeInterpreterWrapper implements AutoCloseable {
     } else {
       throw new IllegalArgumentException(
           String.format(
-              "Input error: %s is not a valid name for any output. "
-                  + "The indexes of the outputs are %s",
+              "Input error: '%s' is not a valid name for any output. Names of outputs and their "
+                  + "indexes are %s",
               name, outputsIndexes.toString()));
     }
   }
diff --git a/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java b/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java
index 61d6c35ec8..210d943724 100644
--- a/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java
+++ b/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java
@@ -195,8 +195,8 @@ public final class InterpreterTest {
       assertThat(e)
           .hasMessageThat()
           .contains(
-              "WrongInputName is not a valid name for any input. The indexes of the inputs"
-                  + " are {input=0}");
+              "'WrongInputName' is not a valid name for any input. Names of inputs and their "
+                  + "indexes are {input=0}");
     }
     int index = interpreter.getInputIndex("input");
     assertThat(index).isEqualTo(0);
@@ -212,8 +212,8 @@ public final class InterpreterTest {
       assertThat(e)
           .hasMessageThat()
           .contains(
-              "WrongOutputName is not a valid name for any output. The indexes of the outputs"
-                  + " are {MobilenetV1/Predictions/Softmax=0}");
+              "'WrongOutputName' is not a valid name for any output. Names of outputs and their"
+                  + " indexes are {MobilenetV1/Predictions/Softmax=0}");
     }
     int index = interpreter.getOutputIndex("MobilenetV1/Predictions/Softmax");
     assertThat(index).isEqualTo(0);
-- 
GitLab


From c89a1d9605427d74079774af7da37933f9ca153c Mon Sep 17 00:00:00 2001
From: Jiri Simsa <jsimsa@google.com>
Date: Mon, 30 Apr 2018 17:38:38 -0700
Subject: [PATCH 075/395] [tf.data] Adding an experimental `group_by_reducer`
 transformation which groups elements of an input pipeline by a key, applies a
 reduce function to elements of each group "on-the-fly", and outputs the
 results once all input elements have been processed.

PiperOrigin-RevId: 194874087
---
 .../python/kernel_tests/bucketing_test.py     | 174 ++++++++
 .../kernel_tests/scan_dataset_op_test.py      |   2 +-
 .../contrib/data/python/ops/grouping.py       | 301 +++++++++++++
 .../api_def_GroupByReducerDataset.pbtxt       |  69 +++
 tensorflow/core/kernels/data/BUILD            |  15 +
 .../core/kernels/data/captured_function.cc    |  14 +
 .../core/kernels/data/captured_function.h     |  11 +
 .../data/group_by_reducer_dataset_op.cc       | 422 ++++++++++++++++++
 .../data/group_by_window_dataset_op.cc        |   2 +-
 tensorflow/core/ops/dataset_ops.cc            |  20 +
 10 files changed, 1028 insertions(+), 2 deletions(-)
 create mode 100644 tensorflow/core/api_def/base_api/api_def_GroupByReducerDataset.pbtxt
 create mode 100644 tensorflow/core/kernels/data/group_by_reducer_dataset_op.cc

diff --git a/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py b/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py
index 55a56b83a8..bd3e034211 100644
--- a/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/bucketing_test.py
@@ -28,6 +28,7 @@ from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.framework import tensor_shape
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
@@ -35,6 +36,179 @@ from tensorflow.python.ops import string_ops
 from tensorflow.python.platform import test
 
 
+class GroupByReducerTest(test.TestCase):
+
+  def checkResults(self, dataset, shapes, values):
+    self.assertEqual(shapes, dataset.output_shapes)
+    get_next = dataset.make_one_shot_iterator().get_next()
+    with self.test_session() as sess:
+      for expected in values:
+        got = sess.run(get_next)
+        self.assertEqual(got, expected)
+      with self.assertRaises(errors.OutOfRangeError):
+        sess.run(get_next)
+
+  def testSum(self):
+    reducer = grouping.Reducer(
+        init_func=lambda _: np.int64(0),
+        reduce_func=lambda x, y: x + y,
+        finalize_func=lambda x: x)
+    for i in range(1, 11):
+      dataset = dataset_ops.Dataset.range(2 * i).apply(
+          grouping.group_by_reducer(lambda x: x % 2, reducer))
+      self.checkResults(
+          dataset, shapes=tensor_shape.scalar(), values=[(i - 1) * i, i * i])
+
+  def testAverage(self):
+
+    def reduce_fn(x, y):
+      return (x[0] * x[1] + math_ops.cast(y, dtypes.float32)) / (
+          x[1] + 1), x[1] + 1
+
+    reducer = grouping.Reducer(
+        init_func=lambda _: (0.0, 0.0),
+        reduce_func=reduce_fn,
+        finalize_func=lambda x: x[0])
+    for i in range(1, 11):
+      dataset = dataset_ops.Dataset.range(2 * i).apply(
+          grouping.group_by_reducer(
+              lambda x: math_ops.cast(x, dtypes.int64) % 2, reducer))
+      self.checkResults(
+          dataset, shapes=tensor_shape.scalar(), values=[i - 1, i])
+
+  def testConcat(self):
+    components = np.array(list("abcdefghijklmnopqrst")).view(np.chararray)
+    reducer = grouping.Reducer(
+        init_func=lambda x: "",
+        reduce_func=lambda x, y: x + y[0],
+        finalize_func=lambda x: x)
+    for i in range(1, 11):
+      dataset = dataset_ops.Dataset.zip(
+          (dataset_ops.Dataset.from_tensor_slices(components),
+           dataset_ops.Dataset.range(2 * i))).apply(
+               grouping.group_by_reducer(lambda x, y: y % 2, reducer))
+      self.checkResults(
+          dataset,
+          shapes=tensor_shape.scalar(),
+          values=[b"acegikmoqs" [:i], b"bdfhjlnprt" [:i]])
+
+  def testSparseSum(self):
+    def _sparse(i):
+      return sparse_tensor.SparseTensorValue(
+          indices=np.array([[0, 0]]),
+          values=(i * np.array([1], dtype=np.int64)),
+          dense_shape=np.array([1, 1]))
+
+    reducer = grouping.Reducer(
+        init_func=lambda _: _sparse(np.int64(0)),
+        reduce_func=lambda x, y: _sparse(x.values[0] + y.values[0]),
+        finalize_func=lambda x: x.values[0])
+    for i in range(1, 11):
+      dataset = dataset_ops.Dataset.range(2 * i).map(_sparse).apply(
+          grouping.group_by_reducer(lambda x: x.values[0] % 2, reducer))
+      self.checkResults(
+          dataset, shapes=tensor_shape.scalar(), values=[(i - 1) * i, i * i])
+
+  def testChangingStateShape(self):
+
+    def reduce_fn(x, _):
+      # Statically known rank, but dynamic length.
+      larger_dim = array_ops.concat([x[0], x[0]], 0)
+      # Statically unknown rank.
+      larger_rank = array_ops.expand_dims(x[1], 0)
+      return larger_dim, larger_rank
+
+    reducer = grouping.Reducer(
+        init_func=lambda x: ([0], 1),
+        reduce_func=reduce_fn,
+        finalize_func=lambda x: x)
+
+    for i in range(1, 11):
+      dataset = dataset_ops.Dataset.from_tensors(np.int64(0)).repeat(i).apply(
+          grouping.group_by_reducer(lambda x: x, reducer))
+      self.assertEqual([None], dataset.output_shapes[0].as_list())
+      self.assertIs(None, dataset.output_shapes[1].ndims)
+      iterator = dataset.make_one_shot_iterator()
+      get_next = iterator.get_next()
+      with self.test_session() as sess:
+        x, y = sess.run(get_next)
+        self.assertAllEqual([0] * (2**i), x)
+        self.assertAllEqual(np.array(1, ndmin=i), y)
+        with self.assertRaises(errors.OutOfRangeError):
+          sess.run(get_next)
+
+  def testTypeMismatch(self):
+    reducer = grouping.Reducer(
+        init_func=lambda x: constant_op.constant(1, dtype=dtypes.int32),
+        reduce_func=lambda x, y: constant_op.constant(1, dtype=dtypes.int64),
+        finalize_func=lambda x: x)
+
+    dataset = dataset_ops.Dataset.range(10)
+    with self.assertRaisesRegexp(
+        TypeError,
+        "The element types for the new state must match the initial state."):
+      dataset.apply(
+          grouping.group_by_reducer(lambda _: np.int64(0), reducer))
+
+  # TODO(b/78665031): Remove once non-scalar keys are supported.
+  def testInvalidKeyShape(self):
+    reducer = grouping.Reducer(
+        init_func=lambda x: np.int64(0),
+        reduce_func=lambda x, y: x + y,
+        finalize_func=lambda x: x)
+
+    dataset = dataset_ops.Dataset.range(10)
+    with self.assertRaisesRegexp(
+        ValueError, "`key_func` must return a single tf.int64 tensor."):
+      dataset.apply(
+          grouping.group_by_reducer(lambda _: np.int64((0, 0)), reducer))
+
+  # TODO(b/78665031): Remove once non-int64 keys are supported.
+  def testInvalidKeyType(self):
+    reducer = grouping.Reducer(
+        init_func=lambda x: np.int64(0),
+        reduce_func=lambda x, y: x + y,
+        finalize_func=lambda x: x)
+
+    dataset = dataset_ops.Dataset.range(10)
+    with self.assertRaisesRegexp(
+        ValueError, "`key_func` must return a single tf.int64 tensor."):
+      dataset.apply(
+          grouping.group_by_reducer(lambda _: "wrong", reducer))
+
+
+class GroupByReducerSerializationTest(
+    dataset_serialization_test_base.DatasetSerializationTestBase):
+
+  def _build_dataset(self, components):
+    reducer = grouping.Reducer(
+        init_func=lambda _: np.int64(0),
+        reduce_func=lambda x, y: x + y,
+        finalize_func=lambda x: x)
+
+    return dataset_ops.Dataset.from_tensor_slices(components).apply(
+        grouping.group_by_reducer(lambda x: x % 5, reducer))
+
+  def testCoreGroupByReducer(self):
+    components = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=np.int64)
+    self.verify_unused_iterator(
+        lambda: self._build_dataset(components), 5, verify_exhausted=True)
+    self.verify_init_before_restore(
+        lambda: self._build_dataset(components), 5, verify_exhausted=True)
+    self.verify_multiple_breaks(
+        lambda: self._build_dataset(components), 5, verify_exhausted=True)
+    self.verify_reset_restored_iterator(
+        lambda: self._build_dataset(components), 5, verify_exhausted=True)
+    self.verify_restore_in_empty_graph(
+        lambda: self._build_dataset(components), 5, verify_exhausted=True)
+    diff_components = np.array([5, 4, 3, 2, 1, 0], dtype=np.int64)
+    self.verify_restore_in_modified_graph(
+        lambda: self._build_dataset(components),
+        lambda: self._build_dataset(diff_components),
+        5,
+        verify_exhausted=True)
+
+
 class GroupByWindowTest(test.TestCase):
 
   def testSimple(self):
diff --git a/tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py
index f544b1caa6..eb2ceff893 100644
--- a/tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/scan_dataset_op_test.py
@@ -168,7 +168,7 @@ class ScanDatasetTest(test.TestCase):
           scan_ops.scan(constant_op.constant(1, dtype=dtypes.int32), _scan_fn))
 
 
-class ScanDatasetSerialzationTest(
+class ScanDatasetSerializationTest(
     dataset_serialization_test_base.DatasetSerializationTestBase):
 
   def _build_dataset(self, num_elements):
diff --git a/tensorflow/contrib/data/python/ops/grouping.py b/tensorflow/contrib/data/python/ops/grouping.py
index 0531f9cbb9..ea229b5b27 100644
--- a/tensorflow/contrib/data/python/ops/grouping.py
+++ b/tensorflow/contrib/data/python/ops/grouping.py
@@ -26,6 +26,7 @@ from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import function
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.framework import tensor_shape
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import check_ops
@@ -33,6 +34,35 @@ from tensorflow.python.ops import gen_dataset_ops
 from tensorflow.python.ops import math_ops
 
 
+def group_by_reducer(key_func, reducer):
+  """A transformation that groups elements and performs a reduction.
+
+  This transformation maps element of a dataset to a key using `key_func` and
+  groups the elements by key. The `reducer` is used to process each group; its
+  `init_func` is used to initialize state for each group when it is created, the
+  `reduce_func` is used to update the state every time an element is mapped to
+  the matching group, and the `finalize_func` is used to map the final state to
+  an output value.
+
+  Args:
+    key_func: A function mapping a nested structure of tensors
+      (having shapes and types defined by `self.output_shapes` and
+      `self.output_types`) to a scalar `tf.int64` tensor.
+    reducer: An instance of `Reducer`, which captures the reduction logic using
+      the `init_func`, `reduce_func`, and `finalize_func` functions.
+
+  Returns:
+    A `Dataset` transformation function, which can be passed to
+    @{tf.data.Dataset.apply}.
+  """
+
+  def _apply_fn(dataset):
+    """Function from `Dataset` to `Dataset` that applies the transformation."""
+    return GroupByReducerDataset(dataset, key_func, reducer)
+
+  return _apply_fn
+
+
 def group_by_window(key_func,
                     reduce_func,
                     window_size=None,
@@ -227,6 +257,250 @@ class _VariantDataset(dataset_ops.Dataset):
     return self._output_types
 
 
+class GroupByReducerDataset(dataset_ops.Dataset):
+  """A `Dataset` that groups its input and performs a reduction."""
+
+  def __init__(self, input_dataset, key_func, reducer):
+    """See `group_by_reducer()` for details."""
+    super(GroupByReducerDataset, self).__init__()
+
+    self._input_dataset = input_dataset
+
+    self._make_key_func(key_func, input_dataset)
+    self._make_init_func(reducer.init_func)
+    self._make_reduce_func(reducer.reduce_func, input_dataset)
+    self._make_finalize_func(reducer.finalize_func)
+
+  def _make_key_func(self, key_func, input_dataset):
+    """Make wrapping Defun for key_func."""
+
+    @function.Defun(*nest.flatten(
+        sparse.as_dense_types(input_dataset.output_types,
+                              input_dataset.output_classes)))
+    def tf_key_func(*args):
+      """A wrapper for Defun that facilitates shape inference."""
+      # Pass in shape information from the input_dataset.
+      dense_shapes = sparse.as_dense_shapes(input_dataset.output_shapes,
+                                            input_dataset.output_classes)
+      for arg, shape in zip(args, nest.flatten(dense_shapes)):
+        arg.set_shape(shape)
+
+      nested_args = nest.pack_sequence_as(input_dataset.output_types, args)
+      nested_args = sparse.deserialize_sparse_tensors(
+          nested_args, input_dataset.output_types, input_dataset.output_shapes,
+          input_dataset.output_classes)
+      # pylint: disable=protected-access
+      if dataset_ops._should_unpack_args(nested_args):
+        ret = key_func(*nested_args)
+      # pylint: enable=protected-access
+      else:
+        ret = key_func(nested_args)
+      ret = ops.convert_to_tensor(ret)
+      if ret.dtype != dtypes.int64 or ret.get_shape() != tensor_shape.scalar():
+        raise ValueError(
+            "`key_func` must return a single tf.int64 tensor. "
+            "Got type=%s and shape=%s" % (ret.dtype, ret.get_shape()))
+      return ret
+
+    self._key_func = tf_key_func
+    self._key_func.add_to_graph(ops.get_default_graph())
+
+  def _make_init_func(self, init_func):
+    """Make wrapping Defun for init_func."""
+
+    @function.Defun(dtypes.int64)
+    def tf_init_func(key):
+      """A wrapper for Defun that facilitates shape inference."""
+      key.set_shape([])
+      ret = init_func(key)
+      # Convert any `SparseTensorValue`s to `SparseTensor`s and all other
+      # values to tensors.
+      ret = nest.pack_sequence_as(ret, [
+          sparse_tensor.SparseTensor.from_value(t)
+          if sparse_tensor.is_sparse(t) else ops.convert_to_tensor(t)
+          for t in nest.flatten(ret)
+      ])
+
+      self._state_classes = sparse.get_classes(ret)
+      self._state_shapes = nest.pack_sequence_as(
+          ret, [t.get_shape() for t in nest.flatten(ret)])
+      self._state_types = nest.pack_sequence_as(
+          ret, [t.dtype for t in nest.flatten(ret)])
+
+      # Serialize any sparse tensors.
+      ret = nest.pack_sequence_as(
+          ret, [t for t in nest.flatten(sparse.serialize_sparse_tensors(ret))])
+      return nest.flatten(ret)
+
+    self._init_func = tf_init_func
+    self._init_func.add_to_graph(ops.get_default_graph())
+
+  def _make_reduce_func(self, reduce_func, input_dataset):
+    """Make wrapping Defun for reduce_func."""
+
+    # Iteratively rerun the reduce function until reaching a fixed point on
+    # `self._state_shapes`.
+    need_to_rerun = True
+    while need_to_rerun:
+
+      # Create a list in which `tf_reduce_func` will store the new shapes.
+      flat_new_state_shapes = []
+
+      @function.Defun(*(nest.flatten(
+          sparse.as_dense_types(
+              self._state_types, self._state_classes)) + nest.flatten(
+                  sparse.as_dense_types(input_dataset.output_types,
+                                        input_dataset.output_classes))))
+      def tf_reduce_func(*args):
+        """A wrapper for Defun that facilitates shape inference."""
+        for arg, shape in zip(
+            args,
+            nest.flatten(
+                sparse.as_dense_shapes(self._state_shapes, self._state_classes))
+            + nest.flatten(
+                sparse.as_dense_shapes(input_dataset.output_shapes,
+                                       input_dataset.output_classes))):
+          arg.set_shape(shape)
+
+        pivot = len(nest.flatten(self._state_shapes))
+        nested_state_args = nest.pack_sequence_as(self._state_types,
+                                                  args[:pivot])
+        nested_state_args = sparse.deserialize_sparse_tensors(
+            nested_state_args, self._state_types, self._state_shapes,
+            self._state_classes)
+        nested_input_args = nest.pack_sequence_as(input_dataset.output_types,
+                                                  args[pivot:])
+        nested_input_args = sparse.deserialize_sparse_tensors(
+            nested_input_args, input_dataset.output_types,
+            input_dataset.output_shapes, input_dataset.output_classes)
+
+        ret = reduce_func(nested_state_args, nested_input_args)
+
+        # Convert any `SparseTensorValue`s to `SparseTensor`s and all other
+        # values to tensors.
+        ret = nest.pack_sequence_as(ret, [
+            sparse_tensor.SparseTensor.from_value(t)
+            if sparse_tensor.is_sparse(t) else ops.convert_to_tensor(t)
+            for t in nest.flatten(ret)
+        ])
+
+        # Extract shape information from the returned values.
+        flat_new_state = nest.flatten(ret)
+        flat_new_state_shapes.extend([t.get_shape() for t in flat_new_state])
+
+        # Extract and validate type information from the returned values.
+        for t, dtype in zip(flat_new_state, nest.flatten(self._state_types)):
+          if t.dtype != dtype:
+            raise TypeError(
+                "The element types for the new state must match the initial "
+                "state. Expected %s; got %s." %
+                (self._state_types,
+                 nest.pack_sequence_as(self._state_types,
+                                       [t.dtype for t in flat_new_state])))
+
+        # Serialize any sparse tensors.
+        ret = nest.pack_sequence_as(
+            ret,
+            [t for t in nest.flatten(sparse.serialize_sparse_tensors(ret))])
+        return nest.flatten(ret)
+
+      # Use the private method that will execute `tf_reduce_func` but delay
+      # adding it to the graph in case we need to rerun the function.
+      tf_reduce_func._create_definition_if_needed()  # pylint: disable=protected-access
+
+      flat_state_shapes = nest.flatten(self._state_shapes)
+      weakened_state_shapes = [
+          old.most_specific_compatible_shape(new)
+          for old, new in zip(flat_state_shapes, flat_new_state_shapes)
+      ]
+
+      need_to_rerun = False
+      for old_shape, weakened_shape in zip(flat_state_shapes,
+                                           weakened_state_shapes):
+        if old_shape.ndims is not None and (
+            weakened_shape.ndims is None or
+            old_shape.as_list() != weakened_shape.as_list()):
+          need_to_rerun = True
+          break
+
+      if need_to_rerun:
+        self._state_shapes = nest.pack_sequence_as(self._state_shapes,
+                                                   weakened_state_shapes)
+
+    self._reduce_func = tf_reduce_func
+    self._reduce_func.add_to_graph(ops.get_default_graph())
+
+  def _make_finalize_func(self, finalize_func):
+    """Make wrapping Defun for finalize_func."""
+
+    @function.Defun(*(nest.flatten(
+        sparse.as_dense_types(self._state_types, self._state_classes))))
+    def tf_finalize_func(*args):
+      """A wrapper for Defun that facilitates shape inference."""
+      for arg, shape in zip(
+          args,
+          nest.flatten(
+              sparse.as_dense_shapes(self._state_shapes, self._state_classes))):
+        arg.set_shape(shape)
+
+      nested_args = nest.pack_sequence_as(self._state_types, args)
+      nested_args = sparse.deserialize_sparse_tensors(
+          nested_args, self._state_types, self._state_shapes,
+          self._state_classes)
+
+      ret = finalize_func(nested_args)
+
+      # Convert any `SparseTensorValue`s to `SparseTensor`s and all other
+      # values to tensors.
+      ret = nest.pack_sequence_as(ret, [
+          sparse_tensor.SparseTensor.from_value(t)
+          if sparse_tensor.is_sparse(t) else ops.convert_to_tensor(t)
+          for t in nest.flatten(ret)
+      ])
+
+      self._output_classes = sparse.get_classes(ret)
+      self._output_shapes = nest.pack_sequence_as(
+          ret, [t.get_shape() for t in nest.flatten(ret)])
+      self._output_types = nest.pack_sequence_as(
+          ret, [t.dtype for t in nest.flatten(ret)])
+
+      # Serialize any sparse tensors.
+      ret = nest.pack_sequence_as(
+          ret, [t for t in nest.flatten(sparse.serialize_sparse_tensors(ret))])
+      return nest.flatten(ret)
+
+    self._finalize_func = tf_finalize_func
+    self._finalize_func.add_to_graph(ops.get_default_graph())
+
+  @property
+  def output_classes(self):
+    return self._output_classes
+
+  @property
+  def output_shapes(self):
+    return self._output_shapes
+
+  @property
+  def output_types(self):
+    return self._output_types
+
+  def _as_variant_tensor(self):
+    return gen_dataset_ops.group_by_reducer_dataset(
+        self._input_dataset._as_variant_tensor(),  # pylint: disable=protected-access
+        self._key_func.captured_inputs,
+        self._init_func.captured_inputs,
+        self._reduce_func.captured_inputs,
+        self._finalize_func.captured_inputs,
+        key_func=self._key_func,
+        init_func=self._init_func,
+        reduce_func=self._reduce_func,
+        finalize_func=self._finalize_func,
+        output_types=nest.flatten(
+            sparse.as_dense_types(self.output_types, self.output_classes)),
+        output_shapes=nest.flatten(
+            sparse.as_dense_shapes(self.output_shapes, self.output_classes)))
+
+
 class GroupByWindowDataset(dataset_ops.Dataset):
   """A `Dataset` that groups its input and performs a windowed reduction."""
 
@@ -336,3 +610,30 @@ class GroupByWindowDataset(dataset_ops.Dataset):
             sparse.as_dense_types(self.output_types, self.output_classes)),
         output_shapes=nest.flatten(
             sparse.as_dense_shapes(self.output_shapes, self.output_classes)))
+
+
+class Reducer(object):
+  """A reducer is used for reducing a set of elements.
+
+  A reducer is represented as a tuple of the three functions:
+    1) initialization function: key => initial state
+    2) reduce function: (old state, input) => new state
+    3) finalization function: state => result
+  """
+
+  def __init__(self, init_func, reduce_func, finalize_func):
+    self._init_func = init_func
+    self._reduce_func = reduce_func
+    self._finalize_func = finalize_func
+
+  @property
+  def init_func(self):
+    return self._init_func
+
+  @property
+  def reduce_func(self):
+    return self._reduce_func
+
+  @property
+  def finalize_func(self):
+    return self._finalize_func
diff --git a/tensorflow/core/api_def/base_api/api_def_GroupByReducerDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_GroupByReducerDataset.pbtxt
new file mode 100644
index 0000000000..067ad4018b
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_GroupByReducerDataset.pbtxt
@@ -0,0 +1,69 @@
+op {
+  graph_op_name: "GroupByReducerDataset"
+  visibility: HIDDEN
+  in_arg {
+    name: "input_dataset"
+    description: <<END
+A variant tensor representing the input dataset.
+END
+  }
+  in_arg {
+    name: "key_func_other_arguments"
+    description: <<END
+A list of tensors, typically values that were captured when
+building a closure for `key_func`.
+END
+  }
+  attr {
+    name: "key_func"
+    description: <<END
+A function mapping an element of `input_dataset`, concatenated
+with `key_func_other_arguments` to a scalar value of type DT_INT64.
+END
+  }
+  in_arg {
+    name: "init_func_other_arguments"
+    description: <<END
+A list of tensors, typically values that were captured when
+building a closure for `init_func`.
+END
+  }
+  attr {
+    name: "init_func"
+    description: <<END
+A function mapping a key of type DT_INT64, concatenated with
+`init_func_other_arguments` to the initial reducer state.
+END
+  }
+  in_arg {
+    name: "reduce_func_other_arguments"
+    description: <<END
+A list of tensors, typically values that were captured when
+building a closure for `reduce_func`.
+END
+  }
+  attr {
+    name: "reduce_func"
+    description: <<END
+A function mapping the current reducer state and an element of `input_dataset`,
+concatenated with `reduce_func_other_arguments` to a new reducer state.
+END
+  }
+  in_arg {
+    name: "finalize_func_other_arguments"
+    description: <<END
+A list of tensors, typically values that were captured when
+building a closure for `finalize_func`.
+END
+  }
+  attr {
+    name: "finalize_func"
+    description: <<END
+A function mapping the final reducer state to an output element.
+END
+  }
+  summary: "Creates a dataset that computes a group-by on `input_dataset`."
+  description: <<END
+Creates a dataset that computes a group-by on `input_dataset`.
+END
+}
diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index c78e0aff83..9ded2667eb 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -123,6 +123,20 @@ tf_kernel_library(
     ],
 )
 
+tf_kernel_library(
+    name = "group_by_reducer_dataset_op",
+    srcs = ["group_by_reducer_dataset_op.cc"],
+    deps = [
+        ":captured_function",
+        ":dataset",
+        "//tensorflow/core:core_cpu_internal",
+        "//tensorflow/core:dataset_ops_op_lib",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+    ],
+)
+
 tf_kernel_library(
     name = "group_by_window_dataset_op",
     srcs = ["group_by_window_dataset_op.cc"],
@@ -550,6 +564,7 @@ tf_kernel_library(
         ":filter_dataset_op",
         ":flat_map_dataset_op",
         ":generator_dataset_op",
+        ":group_by_reducer_dataset_op",
         ":group_by_window_dataset_op",
         ":interleave_dataset_op",
         ":iterator_ops",
diff --git a/tensorflow/core/kernels/data/captured_function.cc b/tensorflow/core/kernels/data/captured_function.cc
index dd61b7daee..ee58341cfd 100644
--- a/tensorflow/core/kernels/data/captured_function.cc
+++ b/tensorflow/core/kernels/data/captured_function.cc
@@ -32,6 +32,20 @@ Status CapturedFunction::Create(
   return Status::OK();
 }
 
+/* static */
+Status CapturedFunction::Create(
+    const NameAttrList& func, OpKernelContext* ctx, const string& argument,
+    std::unique_ptr<CapturedFunction>* out_function) {
+  OpInputList argument_inputs;
+  TF_RETURN_IF_ERROR(ctx->input_list(argument, &argument_inputs));
+  std::vector<Tensor> arguments_t;
+  arguments_t.reserve(argument_inputs.size());
+  for (const Tensor& t : argument_inputs) {
+    arguments_t.push_back(t);
+  }
+  return CapturedFunction::Create(func, std::move(arguments_t), out_function);
+}
+
 CapturedFunction::~CapturedFunction() {
   if (lib_ != nullptr && f_handle_ != kInvalidHandle) {
     lib_->ReleaseHandle(f_handle_).IgnoreError();
diff --git a/tensorflow/core/kernels/data/captured_function.h b/tensorflow/core/kernels/data/captured_function.h
index 490f5cd1e3..e9ad3e381d 100644
--- a/tensorflow/core/kernels/data/captured_function.h
+++ b/tensorflow/core/kernels/data/captured_function.h
@@ -40,12 +40,20 @@ class ResourceMgr;
 // context.
 class CapturedFunction {
  public:
+  // Creates a new instance from a list of named attributes and captured inputs.
+  //
   // NOTE(mrry): The `captured_inputs` are passed by value. For
   // efficiency, you are recommended to move this argument into the call.
   static Status Create(const NameAttrList& func,
                        std::vector<Tensor> captured_inputs,
                        std::unique_ptr<CapturedFunction>* out_function);
 
+  // Creates a new instance using a list of named attributes, fetching captured
+  // inputs from a context argument.
+  static Status Create(const NameAttrList& func, OpKernelContext* ctx,
+                       const string& argument,
+                       std::unique_ptr<CapturedFunction>* out_function);
+
   ~CapturedFunction();
 
   // Runs the "Captured function" using the given FLR and caches the lib and
@@ -87,6 +95,9 @@ class CapturedFunction {
                 std::vector<Tensor>* rets,
                 FunctionLibraryRuntime::DoneCallback done);
 
+  // Returns the named list of function arguments.
+  const NameAttrList& func() { return func_; }
+
   // Returns that additional captured inputs that will be passed to the function
   // when `Run*()` is called.
   const std::vector<Tensor>& captured_inputs() { return captured_inputs_; }
diff --git a/tensorflow/core/kernels/data/group_by_reducer_dataset_op.cc b/tensorflow/core/kernels/data/group_by_reducer_dataset_op.cc
new file mode 100644
index 0000000000..c8aeaab9cb
--- /dev/null
+++ b/tensorflow/core/kernels/data/group_by_reducer_dataset_op.cc
@@ -0,0 +1,422 @@
+/* Copyright 2018 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.
+==============================================================================*/
+#include <map>
+
+#include "tensorflow/core/common_runtime/function.h"
+#include "tensorflow/core/framework/partial_tensor_shape.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/kernels/data/captured_function.h"
+#include "tensorflow/core/kernels/data/dataset.h"
+#include "tensorflow/core/lib/random/random.h"
+
+namespace tensorflow {
+namespace {
+
+// See documentation in ../ops/dataset_ops.cc for a high-level
+// description of the following op.
+class GroupByReducerDatasetOp : public UnaryDatasetOpKernel {
+ public:
+  explicit GroupByReducerDatasetOp(OpKernelConstruction* ctx)
+      : UnaryDatasetOpKernel(ctx),
+        graph_def_version_(ctx->graph_def_version()) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("key_func", &key_func_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("init_func", &init_func_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("reduce_func", &reduce_func_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("finalize_func", &finalize_func_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("output_types", &output_types_));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("output_shapes", &output_shapes_));
+  }
+
+  void MakeDataset(OpKernelContext* ctx, DatasetBase* input,
+                   DatasetBase** output) override {
+    std::unique_ptr<CapturedFunction> captured_key_func;
+    OP_REQUIRES_OK(ctx, CapturedFunction::Create(key_func_, ctx,
+                                                 "key_func_other_arguments",
+                                                 &captured_key_func));
+    std::unique_ptr<CapturedFunction> captured_init_func;
+    OP_REQUIRES_OK(ctx, CapturedFunction::Create(init_func_, ctx,
+                                                 "init_func_other_arguments",
+                                                 &captured_init_func));
+    std::unique_ptr<CapturedFunction> captured_reduce_func;
+    OP_REQUIRES_OK(ctx, CapturedFunction::Create(reduce_func_, ctx,
+                                                 "reduce_func_other_arguments",
+                                                 &captured_reduce_func));
+    std::unique_ptr<CapturedFunction> captured_finalize_func;
+    OP_REQUIRES_OK(ctx,
+                   CapturedFunction::Create(finalize_func_, ctx,
+                                            "finalize_func_other_arguments",
+                                            &captured_finalize_func));
+
+    *output = new Dataset(
+        ctx, input, std::move(captured_key_func), std::move(captured_init_func),
+        std::move(captured_reduce_func), std::move(captured_finalize_func),
+        output_types_, output_shapes_);
+  }
+
+ private:
+  class Dataset : public GraphDatasetBase {
+   public:
+    Dataset(OpKernelContext* ctx, const DatasetBase* input,
+            std::unique_ptr<CapturedFunction> captured_key_func,
+            std::unique_ptr<CapturedFunction> captured_init_func,
+            std::unique_ptr<CapturedFunction> captured_reduce_func,
+            std::unique_ptr<CapturedFunction> captured_finalize_func,
+            const DataTypeVector& output_types,
+            const std::vector<PartialTensorShape>& output_shapes)
+        : GraphDatasetBase(ctx),
+          input_(input),
+          captured_key_func_(std::move(captured_key_func)),
+          captured_init_func_(std::move(captured_init_func)),
+          captured_reduce_func_(std::move(captured_reduce_func)),
+          captured_finalize_func_(std::move(captured_finalize_func)),
+          output_types_(output_types),
+          output_shapes_(output_shapes) {
+      input_->Ref();
+    }
+
+    ~Dataset() override { input_->Unref(); }
+
+    std::unique_ptr<IteratorBase> MakeIterator(
+        const string& prefix) const override {
+      return std::unique_ptr<IteratorBase>(
+          new Iterator({this, strings::StrCat(prefix, "::GroupByReducer")}));
+    }
+
+    const DataTypeVector& output_dtypes() const override {
+      return output_types_;
+    }
+    const std::vector<PartialTensorShape>& output_shapes() const override {
+      return output_shapes_;
+    }
+
+    string DebugString() override { return "GroupByReducerDatasetOp::Dataset"; }
+
+   protected:
+    Status AsGraphDefInternal(OpKernelContext* ctx, DatasetGraphDefBuilder* b,
+                              Node** output) const override {
+      TF_RETURN_IF_ERROR(b->AddFunction(ctx, key_func().name()));
+      TF_RETURN_IF_ERROR(b->AddFunction(ctx, init_func().name()));
+      TF_RETURN_IF_ERROR(b->AddFunction(ctx, reduce_func().name()));
+      TF_RETURN_IF_ERROR(b->AddFunction(ctx, finalize_func().name()));
+      Node* input_graph_node = nullptr;
+      TF_RETURN_IF_ERROR(b->AddParentDataset(ctx, input_, &input_graph_node));
+
+      std::vector<Node*> key_func_other_arguments_node;
+      DataTypeVector key_func_other_arguments_types;
+      TF_RETURN_IF_ERROR(OtherArgumentsNodeAndType(
+          b, captured_key_func_, &key_func_other_arguments_node,
+          &key_func_other_arguments_types));
+
+      std::vector<Node*> init_func_other_arguments_node;
+      DataTypeVector init_func_other_arguments_types;
+      TF_RETURN_IF_ERROR(OtherArgumentsNodeAndType(
+          b, captured_init_func_, &init_func_other_arguments_node,
+          &init_func_other_arguments_types));
+
+      std::vector<Node*> reduce_func_other_arguments_node;
+      DataTypeVector reduce_func_other_arguments_types;
+      TF_RETURN_IF_ERROR(OtherArgumentsNodeAndType(
+          b, captured_reduce_func_, &reduce_func_other_arguments_node,
+          &reduce_func_other_arguments_types));
+
+      std::vector<Node*> finalize_func_other_arguments_node;
+      DataTypeVector finalize_func_other_arguments_types;
+      TF_RETURN_IF_ERROR(OtherArgumentsNodeAndType(
+          b, captured_finalize_func_, &finalize_func_other_arguments_node,
+          &finalize_func_other_arguments_types));
+
+      AttrValue key_func;
+      b->BuildAttrValue(this->key_func(), &key_func);
+      AttrValue init_func;
+      b->BuildAttrValue(this->init_func(), &init_func);
+      AttrValue reduce_func;
+      b->BuildAttrValue(this->reduce_func(), &reduce_func);
+      AttrValue finalize_func;
+      b->BuildAttrValue(this->finalize_func(), &finalize_func);
+
+      AttrValue key_func_other_arguments_types_attr;
+      b->BuildAttrValue(key_func_other_arguments_types,
+                        &key_func_other_arguments_types_attr);
+      AttrValue init_func_other_arguments_types_attr;
+      b->BuildAttrValue(init_func_other_arguments_types,
+                        &init_func_other_arguments_types_attr);
+      AttrValue reduce_func_other_arguments_types_attr;
+      b->BuildAttrValue(reduce_func_other_arguments_types,
+                        &reduce_func_other_arguments_types_attr);
+      AttrValue finalize_func_other_arguments_types_attr;
+      b->BuildAttrValue(finalize_func_other_arguments_types,
+                        &finalize_func_other_arguments_types_attr);
+
+      TF_RETURN_IF_ERROR(b->AddDataset(
+          this, {{0, input_graph_node}},
+          {{1, key_func_other_arguments_node},
+           {2, init_func_other_arguments_node},
+           {3, reduce_func_other_arguments_node},
+           {4, finalize_func_other_arguments_node}},
+          {{"key_func", key_func},
+           {"init_func", init_func},
+           {"reduce_func", reduce_func},
+           {"finalize_func", finalize_func},
+           {"Tkey_func_other_arguments", key_func_other_arguments_types_attr},
+           {"Tinit_func_other_arguments", init_func_other_arguments_types_attr},
+           {"Treduce_func_other_arguments",
+            reduce_func_other_arguments_types_attr},
+           {"Tfinalize_func_other_arguments",
+            finalize_func_other_arguments_types_attr}},
+          output));
+      return Status::OK();
+    }
+
+   private:
+    class Iterator : public DatasetIterator<Dataset> {
+     public:
+      explicit Iterator(const Params& params)
+          : DatasetIterator<Dataset>(params),
+            input_impl_(params.dataset->input_->MakeIterator(params.prefix)) {}
+
+      Status GetNextInternal(IteratorContext* ctx,
+                             std::vector<Tensor>* out_tensors,
+                             bool* end_of_sequence) override {
+        mutex_lock l(mu_);
+
+        // Iterate through the input dataset, keying input elements to reducers.
+        while (!end_of_input_) {
+          std::vector<Tensor> next_input_element;
+          TF_RETURN_IF_ERROR(
+              input_impl_->GetNext(ctx, &next_input_element, &end_of_input_));
+
+          if (!end_of_input_) {
+            // Run the key function on the input element.
+            std::vector<Tensor> key_func_output;
+            TF_RETURN_IF_ERROR(
+                dataset()->captured_key_func_->RunWithBorrowedArgs(
+                    ctx, next_input_element, &key_func_output));
+
+            if (key_func_output.size() != 1 ||
+                key_func_output[0].dtype() != DT_INT64 ||
+                key_func_output[0].NumElements() != 1) {
+              // TODO(b/78665031): Support non-int64 keys.
+              return errors::InvalidArgument(
+                  "`key_func` must return a scalar int64.");
+            }
+            const int64 key = key_func_output[0].scalar<int64>()();
+
+            if (states_.find(key) == states_.end()) {
+              // Run the init function to create the initial state.
+              std::vector<Tensor> init_func_output;
+              TF_RETURN_IF_ERROR(dataset()->captured_init_func_->Run(
+                  ctx, std::move(key_func_output), &init_func_output));
+              states_[key] = init_func_output;
+            }
+
+            // Run the reduce function to update the current state.
+            std::vector<Tensor> args;
+            args.reserve(states_[key].size() + next_input_element.size());
+            std::copy(states_[key].begin(), states_[key].end(),
+                      std::back_inserter(args));
+            std::copy(next_input_element.begin(), next_input_element.end(),
+                      std::back_inserter(args));
+
+            std::vector<Tensor> reduce_func_output;
+            TF_RETURN_IF_ERROR(dataset()->captured_reduce_func_->Run(
+                ctx, std::move(args), &reduce_func_output));
+            states_[key] = reduce_func_output;
+          } else {
+            keys_.resize(states_.size());
+            int idx = 0;
+            for (auto it = states_.begin(); it != states_.end(); ++idx, ++it) {
+              keys_[idx] = it->first;
+            }
+          }
+        }
+
+        if (keys_index_ == keys_.size()) {
+          *end_of_sequence = true;
+          return Status::OK();
+        }
+        TF_RETURN_IF_ERROR(
+            dataset()->captured_finalize_func_->RunWithBorrowedArgs(
+                ctx, states_[keys_[keys_index_++]], out_tensors));
+        return Status::OK();
+      }
+
+     protected:
+      Status SaveInternal(IteratorStateWriter* writer) override {
+        mutex_lock l(mu_);
+        TF_RETURN_IF_ERROR(SaveParent(writer, input_impl_));
+
+        if (end_of_input_) {
+          TF_RETURN_IF_ERROR(
+              writer->WriteScalar(full_name("end_of_input"), ""));
+        }
+
+        // Saving states_.
+        if (!states_.empty()) {
+          TF_RETURN_IF_ERROR(
+              writer->WriteScalar(full_name("states_size"), states_.size()));
+          int idx = 0;
+          for (auto it = states_.begin(); it != states_.end(); ++idx, ++it) {
+            int64 key = it->first;
+            TF_RETURN_IF_ERROR(writer->WriteScalar(
+                full_name(strings::StrCat("states[", idx, "]->key")), key));
+            if (!it->second.empty()) {
+              TF_RETURN_IF_ERROR(writer->WriteScalar(
+                  full_name(strings::StrCat("states[", idx, "]->state_size")),
+                  it->second.size()));
+              for (int j = 0; j < it->second.size(); ++j) {
+                TF_RETURN_IF_ERROR(writer->WriteTensor(
+                    full_name(
+                        strings::StrCat("states[", idx, "]->state[", j, "]")),
+                    it->second[j]));
+              }
+            }
+          }
+        }
+
+        // Saving keys_index_ and keys_.
+        if (end_of_input_) {
+          TF_RETURN_IF_ERROR(
+              writer->WriteScalar(full_name("keys_index"), keys_index_));
+          if (!keys_.empty()) {
+            TF_RETURN_IF_ERROR(
+                writer->WriteScalar(full_name("keys_size"), keys_.size()));
+            for (int idx = 0; idx < keys_.size(); ++idx) {
+              TF_RETURN_IF_ERROR(writer->WriteScalar(
+                  full_name(strings::StrCat("keys[", idx, "]")), keys_[idx]));
+            }
+          }
+        }
+
+        return Status::OK();
+      }
+
+      Status RestoreInternal(IteratorContext* ctx,
+                             IteratorStateReader* reader) override {
+        mutex_lock l(mu_);
+        TF_RETURN_IF_ERROR(RestoreParent(ctx, reader, input_impl_));
+
+        if (reader->Contains(full_name("end_of_input"))) end_of_input_ = true;
+
+        // Restoring states_.
+        if (reader->Contains(full_name("states_size"))) {
+          int64 size;
+          TF_RETURN_IF_ERROR(
+              reader->ReadScalar(full_name("states_size"), &size));
+          for (int idx = 0; idx < size; ++idx) {
+            int64 key;
+            TF_RETURN_IF_ERROR(reader->ReadScalar(
+                full_name(strings::StrCat("states[", idx, "]->key")), &key));
+            std::vector<Tensor> state;
+            if (reader->Contains(full_name(
+                    strings::StrCat("states[", idx, "]->state_size")))) {
+              int64 state_size;
+              TF_RETURN_IF_ERROR(reader->ReadScalar(
+                  full_name(strings::StrCat("states[", idx, "]->state_size")),
+                  &state_size));
+              state.resize(state_size);
+              for (int j = 0; j < state_size; ++j) {
+                TF_RETURN_IF_ERROR(reader->ReadTensor(
+                    full_name(
+                        strings::StrCat("states[", idx, "]->state[", j, "]")),
+                    &state[j]));
+              }
+            }
+            states_[key] = state;
+          }
+        }
+
+        // Restoring keys_index_ and keys_.
+        if (end_of_input_) {
+          TF_RETURN_IF_ERROR(
+              reader->ReadScalar(full_name("keys_index"), &keys_index_));
+          if (reader->Contains(full_name("keys_size"))) {
+            int64 size;
+            TF_RETURN_IF_ERROR(
+                reader->ReadScalar(full_name("keys_size"), &size));
+            keys_.resize(size);
+            for (int idx = 0; idx < size; ++idx) {
+              int64 key;
+              TF_RETURN_IF_ERROR(reader->ReadScalar(
+                  full_name(strings::StrCat("keys[", idx, "]")), &key));
+              keys_[idx] = key;
+            }
+          }
+        }
+
+        return Status::OK();
+      }
+
+     private:
+      mutex mu_;
+      std::unique_ptr<IteratorBase> input_impl_ GUARDED_BY(mu_);
+      bool end_of_input_ GUARDED_BY(mu_) = false;
+      std::map<int64, std::vector<Tensor>> states_ GUARDED_BY(mu_);
+      std::vector<int64> keys_ GUARDED_BY(mu_);
+      int64 keys_index_ GUARDED_BY(mu_) = 0;
+    };
+
+    const NameAttrList& key_func() const { return captured_key_func_->func(); }
+
+    const NameAttrList& init_func() const {
+      return captured_init_func_->func();
+    }
+
+    const NameAttrList& reduce_func() const {
+      return captured_reduce_func_->func();
+    }
+
+    const NameAttrList& finalize_func() const {
+      return captured_finalize_func_->func();
+    }
+
+    Status OtherArgumentsNodeAndType(
+        DatasetGraphDefBuilder* b,
+        const std::unique_ptr<CapturedFunction>& captured_func,
+        std::vector<Node*>* other_arguments_node,
+        DataTypeVector* other_arguments_types) const {
+      other_arguments_node->reserve(captured_func->captured_inputs().size());
+      other_arguments_types->reserve(captured_func->captured_inputs().size());
+      for (const Tensor& t : captured_func->captured_inputs()) {
+        Node* node;
+        TF_RETURN_IF_ERROR(b->AddTensor(t, &node));
+        other_arguments_node->emplace_back(node);
+        other_arguments_types->emplace_back(t.dtype());
+      }
+      return Status::OK();
+    }
+
+    const DatasetBase* const input_;
+    const std::unique_ptr<CapturedFunction> captured_key_func_;
+    const std::unique_ptr<CapturedFunction> captured_init_func_;
+    const std::unique_ptr<CapturedFunction> captured_reduce_func_;
+    const std::unique_ptr<CapturedFunction> captured_finalize_func_;
+    const DataTypeVector output_types_;
+    const std::vector<PartialTensorShape> output_shapes_;
+  };
+
+  const int graph_def_version_;
+  DataTypeVector output_types_;
+  std::vector<PartialTensorShape> output_shapes_;
+  NameAttrList key_func_;
+  NameAttrList init_func_;
+  NameAttrList reduce_func_;
+  NameAttrList finalize_func_;
+};
+
+REGISTER_KERNEL_BUILDER(Name("GroupByReducerDataset").Device(DEVICE_CPU),
+                        GroupByReducerDatasetOp);
+
+}  // namespace
+}  // namespace tensorflow
diff --git a/tensorflow/core/kernels/data/group_by_window_dataset_op.cc b/tensorflow/core/kernels/data/group_by_window_dataset_op.cc
index 46f43dd1b1..03f847ce9c 100644
--- a/tensorflow/core/kernels/data/group_by_window_dataset_op.cc
+++ b/tensorflow/core/kernels/data/group_by_window_dataset_op.cc
@@ -241,7 +241,7 @@ class GroupByWindowDatasetOp : public UnaryDatasetOpKernel {
               if (key_func_output.size() != 1 ||
                   key_func_output[0].dtype() != DT_INT64 ||
                   key_func_output[0].NumElements() != 1) {
-                // TODO(mrry): Support non-int64 keys.
+                // TODO(b/78665031): Support non-int64 keys.
                 return errors::InvalidArgument(
                     "`key_func` must return a scalar int64.");
               }
diff --git a/tensorflow/core/ops/dataset_ops.cc b/tensorflow/core/ops/dataset_ops.cc
index 4ba3f15ef0..5f10ad24b6 100644
--- a/tensorflow/core/ops/dataset_ops.cc
+++ b/tensorflow/core/ops/dataset_ops.cc
@@ -270,6 +270,26 @@ REGISTER_OP("ParallelInterleaveDataset")
     .Attr("output_shapes: list(shape) >= 1")
     .SetShapeFn(shape_inference::ScalarShape);
 
+REGISTER_OP("GroupByReducerDataset")
+    .Input("input_dataset: variant")
+    .Input("key_func_other_arguments: Tkey_func_other_arguments")
+    .Input("init_func_other_arguments: Tinit_func_other_arguments")
+    .Input("reduce_func_other_arguments: Treduce_func_other_arguments")
+    .Input("finalize_func_other_arguments: Tfinalize_func_other_arguments")
+    .Output("handle: variant")
+    .Attr("key_func: func")
+    .Attr("init_func: func")
+    .Attr("reduce_func: func")
+    .Attr("finalize_func: func")
+    .Attr("Tkey_func_other_arguments: list(type) >= 0")
+    .Attr("Tinit_func_other_arguments: list(type) >= 0")
+    .Attr("Treduce_func_other_arguments: list(type) >= 0")
+    .Attr("Tfinalize_func_other_arguments: list(type) >= 0")
+    .Attr("output_types: list(type) >= 1")
+    .Attr("output_shapes: list(shape) >= 1")
+    .SetIsStateful()
+    .SetShapeFn(shape_inference::ScalarShape);
+
 REGISTER_OP("GroupByWindowDataset")
     .Input("input_dataset: variant")
     .Input("key_func_other_arguments: Tkey_func_other_arguments")
-- 
GitLab


From 45bafe9a3589fc735c22c3c703f8689ea9c1e71e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 17:41:33 -0700
Subject: [PATCH 076/395] [XLA] Redesign: migrate tensorflow/compiler/tf2xla,
 tensorflow/compiler/aot: - xla::ComputationBuilder -> xla::XlaBuilder -
 xla::ComputationDataHandle -> xla::XlaOp - xla::Computation ->
 xla::XlaComputation - xla::CompileOnlyClient::AotComputationInstance ->
 xla::CompileOnlyClient::AotXlaComputationInstance - xla::SessionModule ->
 xla::HloSnapshot

PiperOrigin-RevId: 194874462
---
 tensorflow/compiler/aot/compile.cc            |  12 +-
 .../compiler/aot/tests/tfcompile_test.cc      |  14 +-
 tensorflow/compiler/tf2xla/BUILD              |   7 +-
 tensorflow/compiler/tf2xla/graph_compiler.cc  |   7 +-
 tensorflow/compiler/tf2xla/kernels/BUILD      |   8 +-
 .../compiler/tf2xla/kernels/aggregate_ops.cc  |   2 +-
 .../compiler/tf2xla/kernels/batch_norm_op.cc  |  18 +--
 .../tf2xla/kernels/batchtospace_op.cc         |  16 +--
 .../compiler/tf2xla/kernels/bias_ops.cc       |   4 +-
 .../compiler/tf2xla/kernels/binary_ops.cc     |  32 ++---
 tensorflow/compiler/tf2xla/kernels/cast_op.cc |  12 +-
 .../compiler/tf2xla/kernels/categorical_op.cc |   6 +-
 .../tf2xla/kernels/clip_by_value_op.cc        |   2 +-
 .../compiler/tf2xla/kernels/concat_op.cc      |  10 +-
 .../compiler/tf2xla/kernels/const_op.cc       |   2 +-
 .../compiler/tf2xla/kernels/conv_ops.cc       |  49 ++++---
 .../compiler/tf2xla/kernels/cross_op.cc       |   2 +-
 .../compiler/tf2xla/kernels/cwise_ops.cc      |  12 +-
 .../compiler/tf2xla/kernels/cwise_ops.h       |  19 ++-
 .../tf2xla/kernels/depthtospace_op.cc         |  12 +-
 tensorflow/compiler/tf2xla/kernels/diag_op.cc |  36 +++--
 .../tf2xla/kernels/dynamic_slice_ops.cc       |   4 +-
 .../tf2xla/kernels/dynamic_stitch_op.cc       |   6 +-
 tensorflow/compiler/tf2xla/kernels/elu_op.cc  |  11 +-
 .../kernels/extract_image_patches_op.cc       |   6 +-
 .../tf2xla/kernels/fake_quantize_ops.cc       | 134 ++++++++----------
 tensorflow/compiler/tf2xla/kernels/fft_ops.cc |   5 +-
 tensorflow/compiler/tf2xla/kernels/fill_op.cc |   4 +-
 .../compiler/tf2xla/kernels/gather_op.cc      |  20 ++-
 .../tf2xla/kernels/gather_op_helpers.h        |  14 +-
 tensorflow/compiler/tf2xla/kernels/if_op.cc   |  10 +-
 .../compiler/tf2xla/kernels/image_ops.cc      |  69 +++++----
 .../tf2xla/kernels/image_resize_ops.cc        |  48 ++++---
 .../compiler/tf2xla/kernels/index_ops.cc      |   6 +-
 .../compiler/tf2xla/kernels/index_ops_cpu.cc  |   6 +-
 .../compiler/tf2xla/kernels/l2loss_op.cc      |   4 +-
 tensorflow/compiler/tf2xla/kernels/lrn_ops.cc |  14 +-
 .../compiler/tf2xla/kernels/matmul_op.cc      |   4 +-
 .../tf2xla/kernels/matrix_band_part_op.cc     |  12 +-
 .../tf2xla/kernels/matrix_set_diag_op.cc      |  10 +-
 .../compiler/tf2xla/kernels/mirror_pad_op.cc  |  17 +--
 .../compiler/tf2xla/kernels/one_hot_op.cc     |   2 +-
 tensorflow/compiler/tf2xla/kernels/pack_op.cc |   4 +-
 tensorflow/compiler/tf2xla/kernels/pad_op.cc  |   2 +-
 .../compiler/tf2xla/kernels/pooling_ops.cc    |  45 +++---
 .../kernels/quantize_and_dequantize_op.cc     |  16 +--
 .../compiler/tf2xla/kernels/random_ops.cc     |  58 ++++----
 .../tf2xla/kernels/reduce_window_op.cc        |  11 +-
 .../compiler/tf2xla/kernels/reduction_ops.cc  |  63 ++++----
 .../compiler/tf2xla/kernels/reduction_ops.h   |  22 ++-
 .../tf2xla/kernels/reduction_ops_common.cc    |  13 +-
 tensorflow/compiler/tf2xla/kernels/relu_op.cc |  10 +-
 .../compiler/tf2xla/kernels/retval_op.cc      |   4 +-
 .../compiler/tf2xla/kernels/reverse_op.cc     |   4 +-
 .../tf2xla/kernels/reverse_sequence_op.cc     |   4 +-
 .../compiler/tf2xla/kernels/scan_ops.cc       |   6 +-
 .../compiler/tf2xla/kernels/scatter_nd_op.cc  |   2 +-
 .../tf2xla/kernels/segment_reduction_ops.cc   |  10 +-
 .../compiler/tf2xla/kernels/select_op.cc      |   2 +-
 .../compiler/tf2xla/kernels/sendrecv_ops.cc   |   2 +-
 .../compiler/tf2xla/kernels/softmax_op.cc     |  29 ++--
 .../tf2xla/kernels/spacetobatch_op.cc         |  17 +--
 .../tf2xla/kernels/spacetodepth_op.cc         |  12 +-
 .../compiler/tf2xla/kernels/split_op.cc       |   2 +-
 .../compiler/tf2xla/kernels/stack_ops.cc      |  32 ++---
 .../tf2xla/kernels/stateless_random_ops.cc    |  36 +++--
 .../tf2xla/kernels/strided_slice_op.cc        |   8 +-
 .../tf2xla/kernels/tensor_array_ops.cc        |  82 +++++------
 .../compiler/tf2xla/kernels/tile_ops.cc       |   2 +-
 .../compiler/tf2xla/kernels/training_ops.cc   | 103 +++++++-------
 .../compiler/tf2xla/kernels/unary_ops.cc      |  36 +++--
 .../compiler/tf2xla/kernels/variable_ops.cc   |  14 +-
 .../compiler/tf2xla/kernels/while_op.cc       |  16 +--
 tensorflow/compiler/tf2xla/lib/BUILD          |  26 ++--
 tensorflow/compiler/tf2xla/lib/batch_dot.cc   |  50 ++++---
 tensorflow/compiler/tf2xla/lib/batch_dot.h    |  12 +-
 tensorflow/compiler/tf2xla/lib/cholesky.cc    |  50 +++----
 tensorflow/compiler/tf2xla/lib/cholesky.h     |   9 +-
 tensorflow/compiler/tf2xla/lib/scatter.cc     |  58 ++++----
 tensorflow/compiler/tf2xla/lib/scatter.h      |  18 ++-
 .../compiler/tf2xla/lib/triangular_solve.cc   | 131 +++++++++--------
 .../compiler/tf2xla/lib/triangular_solve.h    |  21 +--
 .../tf2xla/lib/triangular_solve_test.cc       |  50 +++----
 tensorflow/compiler/tf2xla/lib/util.cc        |  92 ++++++------
 tensorflow/compiler/tf2xla/lib/util.h         |  67 +++++----
 tensorflow/compiler/tf2xla/lib/util_test.cc   |  17 ++-
 tensorflow/compiler/tf2xla/lib/while_loop.cc  |  52 +++----
 tensorflow/compiler/tf2xla/lib/while_loop.h   |  29 ++--
 tensorflow/compiler/tf2xla/tf2xla.cc          |   6 +-
 tensorflow/compiler/tf2xla/tf2xla.h           |  12 +-
 tensorflow/compiler/tf2xla/tf2xla_test.cc     |   2 +-
 .../compiler/tf2xla/xla_compilation_device.cc |   7 +-
 .../compiler/tf2xla/xla_compilation_device.h  |  10 +-
 tensorflow/compiler/tf2xla/xla_compiler.cc    |  47 +++---
 tensorflow/compiler/tf2xla/xla_compiler.h     |  18 ++-
 .../compiler/tf2xla/xla_compiler_test.cc      |  36 ++---
 tensorflow/compiler/tf2xla/xla_context.cc     |  33 +++--
 tensorflow/compiler/tf2xla/xla_context.h      |  36 +++--
 tensorflow/compiler/tf2xla/xla_helpers.cc     |  95 ++++++-------
 tensorflow/compiler/tf2xla/xla_helpers.h      |  66 ++++-----
 .../tf2xla/xla_jit_compiled_cpu_function.cc   |   4 +-
 tensorflow/compiler/tf2xla/xla_op_kernel.cc   |  71 +++++-----
 tensorflow/compiler/tf2xla/xla_op_kernel.h    |  30 ++--
 tensorflow/compiler/tf2xla/xla_resource.cc    |  33 ++---
 tensorflow/compiler/tf2xla/xla_resource.h     |  29 ++--
 tensorflow/compiler/xla/client/BUILD          |   1 +
 tensorflow/compiler/xla/client/local_client.h |   1 +
 107 files changed, 1218 insertions(+), 1356 deletions(-)

diff --git a/tensorflow/compiler/aot/compile.cc b/tensorflow/compiler/aot/compile.cc
index 31044ff85d..bbc35da2ef 100644
--- a/tensorflow/compiler/aot/compile.cc
+++ b/tensorflow/compiler/aot/compile.cc
@@ -44,7 +44,7 @@ namespace {
 
 // Compiles the XLA computation into executable code.
 Status CompileXla(xla::CompileOnlyClient* client,
-                  const xla::Computation& computation,
+                  const xla::XlaComputation& computation,
                   const xla::cpu::CpuAotCompilationOptions& aot_opts,
                   CompileResult* compile_result) {
   // Retrieves arg and result layouts from the computation.
@@ -62,7 +62,7 @@ Status CompileXla(xla::CompileOnlyClient* client,
   for (int i = 0; i < pshape->parameters_size(); ++i) {
     arg_layouts.push_back(pshape->mutable_parameters(i));
   }
-  xla::CompileOnlyClient::AotComputationInstance instance;
+  xla::CompileOnlyClient::AotXlaComputationInstance instance;
   instance.computation = &computation;
   instance.argument_layouts = std::move(arg_layouts);
   instance.result_layout = &pshape->result();
@@ -93,14 +93,14 @@ Status CompileGraph(const GraphDef& graph_def, const tf2xla::Config& config,
   xla::CompileOnlyClient* client =
       xla::ClientLibrary::GetOrCreateCompileOnlyClient(cpu_platform)
           .ValueOrDie();
-  xla::Computation computation;
+  xla::XlaComputation computation;
   TF_RETURN_IF_ERROR(
       ConvertGraphDefToXla(graph_def, config, client, &computation));
   if (!flags.out_session_module.empty()) {
-    TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::SessionModule> module,
+    TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::HloSnapshot> module,
                         computation.Snapshot());
-    // Serialize the SessionModule deterministically so that all the outputs of
-    // a tf_library genrule are deterministic.
+    // Serialize the HloSnapshot deterministically so that all the outputs of a
+    // tf_library genrule are deterministic.
     string proto;
     TF_RET_CHECK(SerializeToStringDeterministic(*module, &proto));
     TF_RETURN_IF_ERROR(
diff --git a/tensorflow/compiler/aot/tests/tfcompile_test.cc b/tensorflow/compiler/aot/tests/tfcompile_test.cc
index aa9d968265..27ba42b31f 100644
--- a/tensorflow/compiler/aot/tests/tfcompile_test.cc
+++ b/tensorflow/compiler/aot/tests/tfcompile_test.cc
@@ -525,14 +525,16 @@ TEST(TFCompileTest, HloProfiling) {
   auto header = HasSubstr("Execution profile for");
   auto total_cycles_profile_line = HasSubstr("[total]");
   auto dot_profile_line = HasSubstr(
-      "%dot = f32[2,2]{1,0} dot(f32[2,2]{1,0} %arg0, f32[2,2]{1,0} %arg1)");
+      "%dot.0.2 = f32[2,2]{1,0} dot(f32[2,2]{1,0} %arg0.0.0, f32[2,2]{1,0} "
+      "%arg1.0.1)");
   auto add_profile_line = HasSubstr(
-      "%add = f32[2,2]{1,0} add(f32[2,2]{1,0} %arg0, f32[2,2]{1,0} %arg1)");
+      "%add.0.5 = f32[2,2]{1,0} add(f32[2,2]{1,0} %arg0.0.0, f32[2,2]{1,0} "
+      "%arg1.0.1)");
   auto tuple_profile_line = HasSubstr(
-      "%tuple.2 = (f32[2,2]{1,0}, f32[2,2]{1,0}) tuple(f32[2,2]{1,0} %dot, "
-      "f32[2,2]{1,0} %add)");
-  auto arg0_profile_line = HasSubstr("%arg0 = f32[2,2]{1,0} parameter(0)");
-  auto arg1_profile_line = HasSubstr("%arg1 = f32[2,2]{1,0} parameter(1)");
+      "%tuple.0.8 = (f32[2,2]{1,0}, f32[2,2]{1,0}) tuple(f32[2,2]{1,0} "
+      "%dot.0.2, f32[2,2]{1,0} %add.0.5)");
+  auto arg0_profile_line = HasSubstr("%arg0.0.0 = f32[2,2]{1,0} parameter(0)");
+  auto arg1_profile_line = HasSubstr("%arg1.0.1 = f32[2,2]{1,0} parameter(1)");
 
   hlo_profile_lines.erase(hlo_profile_lines.begin() + 7,
                           hlo_profile_lines.end());
diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD
index 942504e6bd..4fca51f54d 100644
--- a/tensorflow/compiler/tf2xla/BUILD
+++ b/tensorflow/compiler/tf2xla/BUILD
@@ -81,7 +81,7 @@ cc_library(
         "//tensorflow/compiler/tf2xla/kernels:xla_cpu_only_ops",
         "//tensorflow/compiler/tf2xla/kernels:xla_ops",
         "//tensorflow/compiler/xla/client",
-        "//tensorflow/compiler/xla/client:computation",
+        "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/core:core_cpu",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
@@ -168,9 +168,9 @@ cc_library(
         "//tensorflow/compiler/xla:types",
         "//tensorflow/compiler/xla:xla_data_proto",
         "//tensorflow/compiler/xla/client:client_library",
-        "//tensorflow/compiler/xla/client:computation",
-        "//tensorflow/compiler/xla/client:computation_builder",
         "//tensorflow/compiler/xla/client:local_client",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/core:core_cpu",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
@@ -215,7 +215,6 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         "//tensorflow/compiler/xla:status_macros",
-        "//tensorflow/compiler/xla/client:computation_builder",
         "//tensorflow/compiler/xla/client:sharding_builder",
         "//tensorflow/core:core_cpu",
         "//tensorflow/core:core_cpu_internal",
diff --git a/tensorflow/compiler/tf2xla/graph_compiler.cc b/tensorflow/compiler/tf2xla/graph_compiler.cc
index b20c1ffc7d..8115a26210 100644
--- a/tensorflow/compiler/tf2xla/graph_compiler.cc
+++ b/tensorflow/compiler/tf2xla/graph_compiler.cc
@@ -51,6 +51,7 @@ Status PrepareArguments(XlaOpKernelContext* ctx, Graph* graph,
                         const std::vector<const XlaExpression*>& expressions,
                         std::vector<XlaCompiler::Argument>* args) {
   auto builder = ctx->builder();
+  auto client = ctx->compiler()->client();
   std::vector<bool> compile_time_constant_flags(expressions.size());
 
   TF_RETURN_IF_ERROR(
@@ -72,8 +73,10 @@ Status PrepareArguments(XlaOpKernelContext* ctx, Graph* graph,
       arg.kind = XlaCompiler::Argument::kConstant;
       TF_RET_CHECK(expressions[i]->resource() == nullptr)
           << "Input with resource is not yet implemented.";
+      TF_ASSIGN_OR_RETURN(auto constant_graph, builder->BuildConstantSubGraph(
+                                                   expressions[i]->handle()));
       TF_ASSIGN_OR_RETURN(auto literal,
-                          builder->ComputeConstant(expressions[i]->handle()));
+                          client->ComputeConstant(constant_graph));
       TF_RETURN_IF_ERROR(
           LiteralToHostTensor(*literal, arg.type, &arg.constant_value));
     } else {
@@ -212,7 +215,7 @@ Status GraphCompiler::CompileFunctionalNode(Node* n,
 
   TF_RET_CHECK(arguments.size() == expressions.size());
 
-  std::vector<xla::ComputationDataHandle> handles;
+  std::vector<xla::XlaOp> handles;
   for (int64 i = 0; i < expressions.size(); ++i) {
     if (arguments[i].kind == XlaCompiler::Argument::kConstant) {
       continue;
diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD
index 00fd08b1a0..85ab4c41bf 100644
--- a/tensorflow/compiler/tf2xla/kernels/BUILD
+++ b/tensorflow/compiler/tf2xla/kernels/BUILD
@@ -114,8 +114,8 @@ tf_kernel_library(
         "//tensorflow/compiler/xla:util",
         "//tensorflow/compiler/xla:xla_data_proto",
         "//tensorflow/compiler/xla/client:client_library",
-        "//tensorflow/compiler/xla/client:computation_builder",
         "//tensorflow/compiler/xla/client/lib:arithmetic",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
         "//tensorflow/core:framework",
         "//tensorflow/core:image_ops_op_lib",
         "//tensorflow/core:lib",
@@ -151,7 +151,7 @@ tf_kernel_library(
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/compiler/tf2xla/ops:xla_ops",
         "//tensorflow/compiler/xla:literal_util",
-        "//tensorflow/compiler/xla/client:computation_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
@@ -167,7 +167,7 @@ tf_kernel_library(
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/compiler/tf2xla/ops:xla_ops",
         "//tensorflow/compiler/xla:literal_util",
-        "//tensorflow/compiler/xla/client:computation_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
@@ -203,8 +203,8 @@ tf_kernel_library(
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/compiler/xla:literal_util",
         "//tensorflow/compiler/xla/client:client_library",
-        "//tensorflow/compiler/xla/client:computation_builder",
         "//tensorflow/compiler/xla/client/lib:arithmetic",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core/kernels:argmax_op",
diff --git a/tensorflow/compiler/tf2xla/kernels/aggregate_ops.cc b/tensorflow/compiler/tf2xla/kernels/aggregate_ops.cc
index 5c9f66df10..1e59868621 100644
--- a/tensorflow/compiler/tf2xla/kernels/aggregate_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/aggregate_ops.cc
@@ -29,7 +29,7 @@ class AddNOp : public XlaOpKernel {
     OP_REQUIRES(ctx, ctx->num_inputs() >= 1,
                 errors::InvalidArgument("AddN requires at least one argument"));
 
-    xla::ComputationDataHandle sum = ctx->Input(0);
+    xla::XlaOp sum = ctx->Input(0);
     for (int i = 1; i < ctx->num_inputs(); ++i) {
       sum = ctx->builder()->Add(sum, ctx->Input(i));
     }
diff --git a/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc b/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc
index 931175be11..15e1815a4c 100644
--- a/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc
@@ -48,9 +48,9 @@ class FusedBatchNormOp : public XlaOpKernel {
     OP_REQUIRES_OK(ctx,
                    DataTypeToPrimitiveType(ctx->input_type(1), &scale_type));
 
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
 
-    xla::ComputationDataHandle input = ctx->Input(0);
+    xla::XlaOp input = ctx->Input(0);
     TensorShape input_shape = ctx->InputShape(0);
 
     int feature_index =
@@ -62,7 +62,7 @@ class FusedBatchNormOp : public XlaOpKernel {
     input = builder->ConvertElementType(input, scale_type);
 
     if (is_training_) {
-      xla::ComputationDataHandle output = builder->BatchNormTraining(
+      xla::XlaOp output = builder->BatchNormTraining(
           input, ctx->Input(1), ctx->Input(2), epsilon_, feature_index);
 
       // In training mode, outputs the normalized value as well as the
@@ -79,7 +79,7 @@ class FusedBatchNormOp : public XlaOpKernel {
       ctx->SetOutput(3, builder->GetTupleElement(output, 1));
       ctx->SetOutput(4, builder->GetTupleElement(output, 2));
     } else {
-      xla::ComputationDataHandle output = builder->BatchNormInference(
+      xla::XlaOp output = builder->BatchNormInference(
           input, ctx->Input(1), ctx->Input(2), ctx->Input(3), ctx->Input(4),
           epsilon_, feature_index);
       ctx->SetOutput(0, builder->ConvertElementType(output, input_type));
@@ -118,7 +118,7 @@ class FusedBatchNormGradOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* const b = ctx->builder();
+    xla::XlaBuilder* const b = ctx->builder();
     DataType input_dtype = ctx->input_type(0);
     DataType scale_dtype = ctx->input_type(2);
 
@@ -137,11 +137,11 @@ class FusedBatchNormGradOp : public XlaOpKernel {
     const int feature_index =
         GetTensorFeatureDimIndex(input_dims, data_format_);
 
-    xla::ComputationDataHandle x_backprop;
-    xla::ComputationDataHandle scale_backprop;
-    xla::ComputationDataHandle offset_backprop;
+    xla::XlaOp x_backprop;
+    xla::XlaOp scale_backprop;
+    xla::XlaOp offset_backprop;
     if (is_training_) {
-      xla::ComputationDataHandle output =
+      xla::XlaOp output =
           b->BatchNormGrad(activations, scale, mean, var, grad_backprop,
                            epsilon_, feature_index);
 
diff --git a/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc b/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc
index 569950c2df..642278ab99 100644
--- a/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc
@@ -20,9 +20,8 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
-void BatchToSpace(XlaOpKernelContext* ctx,
-                  const xla::ComputationDataHandle& input, DataType input_dtype,
-                  const TensorShape& input_tensor_shape,
+void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp& input,
+                  DataType input_dtype, const TensorShape& input_tensor_shape,
                   gtl::ArraySlice<int64> block_shape,
                   const xla::Literal& crops) {
   const int input_rank = input_tensor_shape.dims();
@@ -46,7 +45,7 @@ void BatchToSpace(XlaOpKernelContext* ctx,
                               ", 2] instead of ",
                               xla::ShapeUtil::HumanString(crops.shape())));
 
-  xla::ComputationBuilder* b = ctx->builder();
+  xla::XlaBuilder* b = ctx->builder();
   const int64 batch_size = input_shape[0];
 
   // Compute the product of the block_shape values.
@@ -73,7 +72,7 @@ void BatchToSpace(XlaOpKernelContext* ctx,
   reshaped_shape[block_rank] = batch_size / block_num_elems;
   std::copy(input_shape.begin() + 1, input_shape.end(),
             reshaped_shape.begin() + block_rank + 1);
-  xla::ComputationDataHandle reshaped = b->Reshape(input, reshaped_shape);
+  xla::XlaOp reshaped = b->Reshape(input, reshaped_shape);
 
   // 2. Permute dimensions of `reshaped` to produce `permuted` of shape
   //      [batch / prod(block_shape),
@@ -91,7 +90,7 @@ void BatchToSpace(XlaOpKernelContext* ctx,
   }
   std::iota(permutation.begin() + 1 + block_rank * 2, permutation.end(),
             1 + block_rank * 2);
-  xla::ComputationDataHandle permuted = b->Transpose(reshaped, permutation);
+  xla::XlaOp permuted = b->Transpose(reshaped, permutation);
 
   // 3. Reshape `permuted` to produce `reshaped_permuted` of shape
   //      [batch / prod(block_shape),
@@ -111,8 +110,7 @@ void BatchToSpace(XlaOpKernelContext* ctx,
   std::copy(remainder_shape.begin(), remainder_shape.end(),
             reshaped_permuted_shape.begin() + 1 + block_rank);
 
-  xla::ComputationDataHandle reshaped_permuted =
-      b->Reshape(permuted, reshaped_permuted_shape);
+  xla::XlaOp reshaped_permuted = b->Reshape(permuted, reshaped_permuted_shape);
 
   // 4. Crop the start and end of dimensions `[1, ..., M]` of
   //    `reshaped_permuted` according to `crops` to produce the output of shape:
@@ -139,7 +137,7 @@ void BatchToSpace(XlaOpKernelContext* ctx,
             "Cropped size must be non-negative: start: ", crop_start,
             " end: ", crop_end, " size ", reshaped_permuted_shape[1 + i]));
   }
-  xla::ComputationDataHandle output =
+  xla::XlaOp output =
       b->Slice(reshaped_permuted, start_indices, end_indices, strides);
   ctx->SetOutput(0, output);
 }
diff --git a/tensorflow/compiler/tf2xla/kernels/bias_ops.cc b/tensorflow/compiler/tf2xla/kernels/bias_ops.cc
index ed33b8ed2e..9d677f4266 100644
--- a/tensorflow/compiler/tf2xla/kernels/bias_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/bias_ops.cc
@@ -60,7 +60,7 @@ class BiasOp : public XlaOpKernel {
             "of the input tensor: ",
             bias_shape.DebugString(), " vs. ", input_shape.DebugString()));
 
-    xla::ComputationDataHandle result =
+    xla::XlaOp result =
         ctx->builder()->Add(ctx->Input(0), ctx->Input(1), {feature_dim});
     ctx->SetOutput(0, result);
   }
@@ -103,7 +103,7 @@ class BiasAddGradOp : public XlaOpKernel {
     std::iota(reduce_dims.begin(), reduce_dims.begin() + feature_dim, 0);
     std::iota(reduce_dims.begin() + feature_dim, reduce_dims.end(),
               feature_dim + 1);
-    xla::ComputationBuilder* const b = ctx->builder();
+    xla::XlaBuilder* const b = ctx->builder();
     const DataType accumulation_type =
         XlaHelpers::SumAccumulationType(input_type(0));
     auto converted =
diff --git a/tensorflow/compiler/tf2xla/kernels/binary_ops.cc b/tensorflow/compiler/tf2xla/kernels/binary_ops.cc
index 2436a6074a..f04cde878e 100644
--- a/tensorflow/compiler/tf2xla/kernels/binary_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/binary_ops.cc
@@ -19,7 +19,7 @@ limitations under the License.
 #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/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/core/framework/kernel_def_builder.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/types.h"
@@ -34,14 +34,13 @@ namespace {
   class NAME##Op : public XlaBinaryOp {                                 \
    public:                                                              \
     explicit NAME##Op(OpKernelConstruction* ctx) : XlaBinaryOp(ctx) {}  \
-    xla::ComputationDataHandle Computation(                             \
-        XlaOpKernelContext* ctx, const xla::ComputationDataHandle& lhs, \
-        const gtl::ArraySlice<int64>& lhs_shape,                        \
-        const xla::ComputationDataHandle& rhs,                          \
+    xla::XlaOp Computation(                                             \
+        XlaOpKernelContext* ctx, const xla::XlaOp& lhs,                 \
+        const gtl::ArraySlice<int64>& lhs_shape, const xla::XlaOp& rhs, \
         const gtl::ArraySlice<int64>& rhs_shape,                        \
         const BCast& broadcast_helper,                                  \
         const std::vector<int64>& extend_dimensions) override {         \
-      xla::ComputationBuilder* b = ctx->builder();                      \
+      xla::XlaBuilder* b = ctx->builder();                              \
       return HLO;                                                       \
     }                                                                   \
   };                                                                    \
@@ -63,11 +62,8 @@ XLA_MAKE_BINARY(Complex, b->Complex(lhs, rhs, extend_dimensions));
 // } else {
 //   return x / y;
 // }
-static xla::ComputationDataHandle FloorDivImpl(xla::ComputationBuilder* b,
-                                               DataType dtype,
-                                               xla::ComputationDataHandle x,
-                                               xla::ComputationDataHandle y,
-                                               const BCast& broadcast_helper) {
+static xla::XlaOp FloorDivImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
+                               xla::XlaOp y, const BCast& broadcast_helper) {
   std::tie(x, y) = XlaBinaryOp::Broadcast(b, x, y, broadcast_helper);
   auto zero = XlaHelpers::Zero(b, dtype);
   auto one = XlaHelpers::One(b, dtype);
@@ -87,11 +83,8 @@ XLA_MAKE_BINARY(FloorDiv,
 // Implementation of FloorMod. Pseudo-code:
 // T trunc_mod = std::fmod(x, y);
 // return (x < T(0)) == (y < T(0)) ? trunc_mod : std::fmod(trunc_mod + y, y);
-static xla::ComputationDataHandle FloorModImpl(xla::ComputationBuilder* b,
-                                               DataType dtype,
-                                               xla::ComputationDataHandle x,
-                                               xla::ComputationDataHandle y,
-                                               const BCast& broadcast_helper) {
+static xla::XlaOp FloorModImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
+                               xla::XlaOp y, const BCast& broadcast_helper) {
   std::tie(x, y) = XlaBinaryOp::Broadcast(b, x, y, broadcast_helper);
   auto zero = XlaHelpers::Zero(b, dtype);
   auto same_sign = b->Eq(b->Lt(x, zero), b->Lt(y, zero));
@@ -127,8 +120,7 @@ XLA_MAKE_BINARY(SqrtGrad,
                               XlaHelpers::FloatLiteral(b, input_type(0), 0.5)),
                        lhs, extend_dimensions));
 
-static xla::ComputationDataHandle Square(xla::ComputationBuilder* builder,
-                                         const xla::ComputationDataHandle& x) {
+static xla::XlaOp Square(xla::XlaBuilder* builder, const xla::XlaOp& x) {
   return builder->Mul(x, x);
 }
 
@@ -175,11 +167,11 @@ class ApproximateEqualOp : public XlaOpKernel {
 
   // Computes the max of the scalar input x and 0.
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
     auto abs = b->Abs(b->Sub(ctx->Input(0), ctx->Input(1)));
     auto abs_shape = b->GetShape(abs);
     OP_REQUIRES_OK(ctx, abs_shape.status());
-    auto abs_type = abs_shape.ValueOrDie()->element_type();
+    auto abs_type = abs_shape.ValueOrDie().element_type();
     auto result = b->Lt(
         abs, b->ConvertElementType(b->ConstantR0<float>(tolerance_), abs_type));
     ctx->SetOutput(0, result);
diff --git a/tensorflow/compiler/tf2xla/kernels/cast_op.cc b/tensorflow/compiler/tf2xla/kernels/cast_op.cc
index c52b2dcb7e..e9d98c7685 100644
--- a/tensorflow/compiler/tf2xla/kernels/cast_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/cast_op.cc
@@ -33,9 +33,9 @@ class CastOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* builder = ctx->builder();
-    xla::ComputationDataHandle input = ctx->Input(0);
-    xla::ComputationDataHandle output;
+    xla::XlaBuilder* builder = ctx->builder();
+    xla::XlaOp input = ctx->Input(0);
+    xla::XlaOp output;
 
     if (src_dtype_ == dst_dtype_) {
       output = input;
@@ -72,9 +72,9 @@ class BitcastOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* builder = ctx->builder();
-    xla::ComputationDataHandle input = ctx->Input(0);
-    xla::ComputationDataHandle output;
+    xla::XlaBuilder* builder = ctx->builder();
+    xla::XlaOp input = ctx->Input(0);
+    xla::XlaOp output;
 
     if (src_dtype_ == dst_dtype_) {
       output = input;
diff --git a/tensorflow/compiler/tf2xla/kernels/categorical_op.cc b/tensorflow/compiler/tf2xla/kernels/categorical_op.cc
index 545aa364f9..835a7f5689 100644
--- a/tensorflow/compiler/tf2xla/kernels/categorical_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/categorical_op.cc
@@ -34,7 +34,7 @@ class CategoricalOp : public XlaOpKernel {
 
   void Compile(XlaOpKernelContext* ctx) override {
     // Get the logits
-    const xla::ComputationDataHandle& logits = ctx->Input(0);
+    const xla::XlaOp& logits = ctx->Input(0);
     TensorShape logits_shape = ctx->InputShape(0);
     int64 num_samples;
     OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(1, &num_samples));
@@ -56,7 +56,7 @@ class CategoricalOp : public XlaOpKernel {
     const int64 batch_size = logits_shape.dim_size(0);
     const int64 num_classes = logits_shape.dim_size(1);
 
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
 
     std::array<int64, 3> uniform_shape_array = {
         {batch_size, num_samples, num_classes}};
@@ -78,7 +78,7 @@ class CategoricalOp : public XlaOpKernel {
                      /*broadcast_dimensions=*/{0, 2});
 
     TensorShape softmax_shape(uniform_shape_array);
-    xla::ComputationDataHandle argmax;
+    xla::XlaOp argmax;
     OP_REQUIRES_OK(
         ctx,
         XlaHelpers::ArgMax(builder, ctx, softmax_entries, softmax_shape,
diff --git a/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc b/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc
index fdf75be7b1..a00bc912f9 100644
--- a/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc
@@ -29,7 +29,7 @@ class ClipByValueOp : public XlaOpKernel {
     const TensorShape min_shape = ctx->InputShape(1);
     const TensorShape max_shape = ctx->InputShape(2);
 
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
     auto input = ctx->Input(0);
     auto min = ctx->Input(1);
     auto max = ctx->Input(2);
diff --git a/tensorflow/compiler/tf2xla/kernels/concat_op.cc b/tensorflow/compiler/tf2xla/kernels/concat_op.cc
index 1a246e8df9..78285affa1 100644
--- a/tensorflow/compiler/tf2xla/kernels/concat_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/concat_op.cc
@@ -54,7 +54,7 @@ class ConcatBaseOp : public XlaOpKernel {
     // TODO(annarev): add a helper to support int64 input.
     const int32 concat_dim = literal.Get<int>({});
 
-    std::vector<xla::ComputationDataHandle> values;
+    std::vector<xla::XlaOp> values;
     std::vector<TensorShape> shapes;
     OP_REQUIRES_OK(ctx, ctx->InputList("values", &values, &shapes));
     const int N = values.size();
@@ -70,13 +70,13 @@ class ConcatBaseOp : public XlaOpKernel {
                     "[",
                     -input_dims, ", ", input_dims, "), but got ", concat_dim));
 
-    // Make a vector holding the ComputationDataHandles for each of
-    // the inputs that has non-zero elements.
-    std::vector<xla::ComputationDataHandle> input_data;
+    // Make a vector holding the XlaOp for each of the inputs that has non-zero
+    // elements.
+    std::vector<xla::XlaOp> input_data;
     int output_concat_dim = 0;
     const bool input_is_scalar = IsLegacyScalar(input_shape);
     for (int i = 0; i < N; ++i) {
-      xla::ComputationDataHandle handle = values[i];
+      xla::XlaOp handle = values[i];
       const TensorShape& in_shape = shapes[i];
       const bool in_is_scalar = IsLegacyScalar(in_shape);
       OP_REQUIRES(
diff --git a/tensorflow/compiler/tf2xla/kernels/const_op.cc b/tensorflow/compiler/tf2xla/kernels/const_op.cc
index 8f78b4c8f9..59d06c654d 100644
--- a/tensorflow/compiler/tf2xla/kernels/const_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/const_op.cc
@@ -45,7 +45,7 @@ class ConstOp : public XlaOpKernel {
       ctx->SetInvalidOutput(0);
       return;
     }
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     // To avoid blowups for large constants filled with the same value,
     // recognize that case and emit a scalar broadcast instead.
diff --git a/tensorflow/compiler/tf2xla/kernels/conv_ops.cc b/tensorflow/compiler/tf2xla/kernels/conv_ops.cc
index c0ee0c9c2e..627bad12f3 100644
--- a/tensorflow/compiler/tf2xla/kernels/conv_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/conv_ops.cc
@@ -47,9 +47,8 @@ TensorShape ExpandedFilterShapeForDepthwiseConvolution(
 }
 
 // Broadcast zeros to ExpandedFilterShapeForDepthwiseConvolution.
-xla::ComputationDataHandle CreateExpandedZero(
-    const TensorShape& filter_shape, DataType dtype,
-    xla::ComputationBuilder* builder) {
+xla::XlaOp CreateExpandedZero(const TensorShape& filter_shape, DataType dtype,
+                              xla::XlaBuilder* builder) {
   TensorShape expanded_filter_shape =
       ExpandedFilterShapeForDepthwiseConvolution(filter_shape);
   return builder->Broadcast(XlaHelpers::Zero(builder, dtype),
@@ -87,8 +86,8 @@ xla::ComputationDataHandle CreateExpandedZero(
 //
 // Finally compare A and broadcasted B in dimension 2 amd return the result at
 // the beginning of the comment.
-xla::ComputationDataHandle CreateExpandedFilterMask(
-    const TensorShape& filter_shape, xla::ComputationBuilder* builder) {
+xla::XlaOp CreateExpandedFilterMask(const TensorShape& filter_shape,
+                                    xla::XlaBuilder* builder) {
   TensorShape expanded_filter_shape =
       ExpandedFilterShapeForDepthwiseConvolution(filter_shape);
   int64 depthwise_multiplier = filter_shape.dim_size(filter_shape.dims() - 1);
@@ -96,11 +95,11 @@ xla::ComputationDataHandle CreateExpandedFilterMask(
 
   // Create a M sized linspace and an M*N sized linspace that will be
   // broadcasted into perpendicular dimensions and compared.
-  xla::ComputationDataHandle input_feature_iota;
+  xla::XlaOp input_feature_iota;
   // DT_INT32 Iota will always return status::OK().
   TF_CHECK_OK(XlaHelpers::Iota(builder, DataType::DT_INT32, input_feature,
                                &input_feature_iota));
-  xla::ComputationDataHandle expanded_feature_iota;
+  xla::XlaOp expanded_feature_iota;
   TF_CHECK_OK(XlaHelpers::Iota(builder, DataType::DT_INT32,
                                input_feature * depthwise_multiplier,
                                &expanded_feature_iota));
@@ -126,10 +125,10 @@ xla::ComputationDataHandle CreateExpandedFilterMask(
 
 // Expands a filter of shape [H, W, ..., M, N] to [H, W, ..., M, M*N] by adding
 // zeros for the cross-depth filters. Used to build a depthwise convolution.
-xla::ComputationDataHandle ExpandFilterForDepthwiseConvolution(
-    const TensorShape& filter_shape, DataType dtype,
-    const xla::ComputationDataHandle& filter,
-    xla::ComputationBuilder* builder) {
+xla::XlaOp ExpandFilterForDepthwiseConvolution(const TensorShape& filter_shape,
+                                               DataType dtype,
+                                               const xla::XlaOp& filter,
+                                               xla::XlaBuilder* builder) {
   int64 depthwise_multiplier = filter_shape.dim_size(filter_shape.dims() - 1);
   int64 input_feature = filter_shape.dim_size(filter_shape.dims() - 2);
   TensorShape expanded_filter_shape =
@@ -156,10 +155,11 @@ xla::ComputationDataHandle ExpandFilterForDepthwiseConvolution(
 }
 
 // Inverse of ExpandFilterForDepthwiseConvolution.
-xla::ComputationDataHandle ContractFilterForDepthwiseBackprop(
-    XlaOpKernelContext* ctx, const TensorShape& filter_shape, DataType dtype,
-    const xla::ComputationDataHandle& filter_backprop,
-    xla::ComputationBuilder* builder) {
+xla::XlaOp ContractFilterForDepthwiseBackprop(XlaOpKernelContext* ctx,
+                                              const TensorShape& filter_shape,
+                                              DataType dtype,
+                                              const xla::XlaOp& filter_backprop,
+                                              xla::XlaBuilder* builder) {
   TensorShape expanded_filter_shape =
       ExpandedFilterShapeForDepthwiseConvolution(filter_shape);
   auto masked_expanded_filter = builder->Select(
@@ -248,9 +248,9 @@ class ConvOp : public XlaOpKernel {
                     "input and filter must have the same depth: ", in_depth,
                     " vs ", input_shape.dim_size(feature_dim)));
 
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
-    xla::ComputationDataHandle filter = ctx->Input(1);
+    xla::XlaOp filter = ctx->Input(1);
     TensorShape expanded_filter_shape = filter_shape;
     if (depthwise_) {
       filter = ExpandFilterForDepthwiseConvolution(
@@ -288,7 +288,7 @@ class ConvOp : public XlaOpKernel {
                    &unused_output_size, &padding[i].first, &padding[i].second));
     }
 
-    xla::ComputationDataHandle conv =
+    xla::XlaOp conv =
         b->ConvGeneralDilated(ctx->Input(0), filter, window_strides, padding,
                               lhs_dilation, rhs_dilation, dims);
     ctx->SetOutput(0, conv);
@@ -391,7 +391,7 @@ class ConvBackpropInputOp : public XlaOpKernel {
                        expanded_filter_shape, out_backprop_shape, dilations_,
                        strides_, padding_, data_format_, &dims));
 
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
     auto filter = ctx->Input(1);
     auto out_backprop = ctx->Input(2);
 
@@ -435,12 +435,11 @@ class ConvBackpropInputOp : public XlaOpKernel {
     }
 
     // Mirror the filter in the spatial dimensions.
-    xla::ComputationDataHandle mirrored_weights =
-        b->Rev(filter, kernel_spatial_dims);
+    xla::XlaOp mirrored_weights = b->Rev(filter, kernel_spatial_dims);
 
     // activation gradients
     //   = gradients (with padding and dilation) <conv> mirrored_weights
-    xla::ComputationDataHandle in_backprop = b->ConvGeneralDilated(
+    xla::XlaOp in_backprop = b->ConvGeneralDilated(
         out_backprop, mirrored_weights, /*window_strides=*/ones, padding,
         lhs_dilation, rhs_dilation, dnums);
 
@@ -546,9 +545,9 @@ class ConvBackpropFilterOp : public XlaOpKernel {
                        expanded_filter_shape, out_backprop_shape, dilations_,
                        strides_, padding_, data_format_, &dims));
 
-    xla::ComputationBuilder* b = ctx->builder();
-    xla::ComputationDataHandle activations = ctx->Input(0);
-    xla::ComputationDataHandle gradients = ctx->Input(2);
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp activations = ctx->Input(0);
+    xla::XlaOp gradients = ctx->Input(2);
 
     // The filter gradients are computed by a convolution of the input
     // activations and the output gradients, with some appropriate padding.
diff --git a/tensorflow/compiler/tf2xla/kernels/cross_op.cc b/tensorflow/compiler/tf2xla/kernels/cross_op.cc
index 3df8c00f1b..7fcd4170fb 100644
--- a/tensorflow/compiler/tf2xla/kernels/cross_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/cross_op.cc
@@ -53,7 +53,7 @@ class CrossOp : public XlaOpKernel {
     }
     std::vector<int64> strides(in0_shape.dims(), 1);
 
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
     auto in0 = ctx->Input(0);
     auto in1 = ctx->Input(1);
     starts.back() = 0;
diff --git a/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc b/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc
index 0cf03ceb94..01aa1a83e7 100644
--- a/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc
@@ -22,7 +22,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "tensorflow/compiler/xla/client/client_library.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/util/bcast.h"
@@ -75,7 +75,7 @@ void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) {
   }
 
   // Call virtual method to emit the computation.
-  xla::ComputationDataHandle output =
+  xla::XlaOp output =
       Computation(ctx, lhs_handle, lhs_shape.dim_sizes(), rhs_handle,
                   rhs_shape.dim_sizes(), bcast, extend_dimension);
 
@@ -85,11 +85,9 @@ void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) {
   ctx->SetOutput(0, output);
 }
 
-/* static */ std::pair<xla::ComputationDataHandle, xla::ComputationDataHandle>
-XlaBinaryOp::Broadcast(xla::ComputationBuilder* builder,
-                       const xla::ComputationDataHandle& lhs,
-                       const xla::ComputationDataHandle& rhs,
-                       const BCast& broadcast_helper) {
+/* static */ std::pair<xla::XlaOp, xla::XlaOp> XlaBinaryOp::Broadcast(
+    xla::XlaBuilder* builder, const xla::XlaOp& lhs, const xla::XlaOp& rhs,
+    const BCast& broadcast_helper) {
   // Manually construct the broadcasting since MapN does not do
   // automatic broadcasting. The bcast helper ensures that
   // lhs.reshape(bcast.x_reshape()).broadcast(bcast.x_bcast()) and
diff --git a/tensorflow/compiler/tf2xla/kernels/cwise_ops.h b/tensorflow/compiler/tf2xla/kernels/cwise_ops.h
index 5bc1d5fb1f..4f92dbc874 100644
--- a/tensorflow/compiler/tf2xla/kernels/cwise_ops.h
+++ b/tensorflow/compiler/tf2xla/kernels/cwise_ops.h
@@ -20,7 +20,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/xla/client/client_library.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/util/bcast.h"
 
@@ -30,7 +30,7 @@ namespace tensorflow {
 // inputs that can be broadcast to the same shape. The base class
 // contains pure virtual methods to override: description is a textual
 // description of the operation; and Computation adds the
-// implementation of the operation to a xla::ComputationBuilder. For most
+// implementation of the operation to a xla::XlaBuilder. For most
 // arithmetic Ops XLA handles the broadcasting automatically given the input
 // tensors.
 class XlaBinaryOp : public XlaOpKernel {
@@ -55,10 +55,9 @@ class XlaBinaryOp : public XlaOpKernel {
   // higher-rank input should be matched when broadcasting the
   // lower-rank input. See comment below and the documentation on broadcasting
   // in the XLA documentation.
-  virtual xla::ComputationDataHandle Computation(
-      XlaOpKernelContext* ctx, const xla::ComputationDataHandle& lhs,
-      const gtl::ArraySlice<int64>& lhs_shape,
-      const xla::ComputationDataHandle& rhs,
+  virtual xla::XlaOp Computation(
+      XlaOpKernelContext* ctx, const xla::XlaOp& lhs,
+      const gtl::ArraySlice<int64>& lhs_shape, const xla::XlaOp& rhs,
       const gtl::ArraySlice<int64>& rhs_shape, const BCast& broadcast_helper,
       const std::vector<int64>& extend_dimensions) = 0;
 
@@ -67,11 +66,9 @@ class XlaBinaryOp : public XlaOpKernel {
   // Helper function that performs the broadcasting described by
   // 'broadcast_helper', yielding arguments 'lhs' and 'rhs' that have the same
   // shape.
-  static std::pair<xla::ComputationDataHandle, xla::ComputationDataHandle>
-  Broadcast(xla::ComputationBuilder* builder,
-            const xla::ComputationDataHandle& lhs,
-            const xla::ComputationDataHandle& rhs,
-            const BCast& broadcast_helper);
+  static std::pair<xla::XlaOp, xla::XlaOp> Broadcast(
+      xla::XlaBuilder* builder, const xla::XlaOp& lhs, const xla::XlaOp& rhs,
+      const BCast& broadcast_helper);
 };
 
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/tf2xla/kernels/depthtospace_op.cc b/tensorflow/compiler/tf2xla/kernels/depthtospace_op.cc
index 96d7809f79..23243f6246 100644
--- a/tensorflow/compiler/tf2xla/kernels/depthtospace_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/depthtospace_op.cc
@@ -50,8 +50,8 @@ class DepthToSpaceOp : public XlaOpKernel {
     const gtl::InlinedVector<int64, 4> input_shape =
         input_tensor_shape.dim_sizes();
 
-    xla::ComputationBuilder* b = ctx->builder();
-    xla::ComputationDataHandle input = ctx->Input(0);
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp input = ctx->Input(0);
 
     int feature_dim = GetTensorFeatureDimIndex(input_rank, data_format_);
     int num_spatial_dims = GetTensorSpatialDims(input_rank, data_format_);
@@ -130,7 +130,7 @@ class DepthToSpaceOp : public XlaOpKernel {
                     ") is not divisible by square of the block size (",
                     block_size_, ")"));
 
-    xla::ComputationDataHandle reshaped = b->Reshape(input, reshaped_shape);
+    xla::XlaOp reshaped = b->Reshape(input, reshaped_shape);
 
     // 2. Permute dimensions of `reshaped` to produce
     //    `permuted_reshaped` of shape:
@@ -141,8 +141,7 @@ class DepthToSpaceOp : public XlaOpKernel {
     //       input_shape[2],
     //       block_size_,
     //       depth / (block_size_ * block_size_)]
-    xla::ComputationDataHandle permuted_reshaped =
-        b->Transpose(reshaped, transpose_order);
+    xla::XlaOp permuted_reshaped = b->Transpose(reshaped, transpose_order);
 
     // 3. Reshape `permuted_reshaped` to flatten `block_shape` into the
     //    batch dimension, producing an output tensor of shape:
@@ -152,8 +151,7 @@ class DepthToSpaceOp : public XlaOpKernel {
     //       input_shape[2] * block_size_,
     //       depth / (block_size_ * block_size_)]
     //
-    xla::ComputationDataHandle output =
-        b->Reshape(permuted_reshaped, output_shape);
+    xla::XlaOp output = b->Reshape(permuted_reshaped, output_shape);
 
     ctx->SetOutput(0, output);
   }
diff --git a/tensorflow/compiler/tf2xla/kernels/diag_op.cc b/tensorflow/compiler/tf2xla/kernels/diag_op.cc
index 765ea922a5..931705ba83 100644
--- a/tensorflow/compiler/tf2xla/kernels/diag_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/diag_op.cc
@@ -25,10 +25,10 @@ namespace tensorflow {
 namespace {
 
 // Create a diagonal / batch diagonal matrix with 'input' on the diagonal.
-xla::StatusOr<xla::ComputationDataHandle> CreateDiagonal(
-    const xla::ComputationDataHandle& input, int64 last_dim_size,
+xla::StatusOr<xla::XlaOp> CreateDiagonal(
+    const xla::XlaOp& input, int64 last_dim_size,
     tensorflow::gtl::ArraySlice<int64> other_dims, XlaOpKernelContext* ctx,
-    xla::ComputationBuilder* builder) {
+    xla::XlaBuilder* builder) {
   // Create two matrices that have the following forms, and compare them:
   //
   // [[0, 0, 0, 0]            [[0, 1, 2, 3]
@@ -38,12 +38,11 @@ xla::StatusOr<xla::ComputationDataHandle> CreateDiagonal(
   //
   // This produces a predicate matrix of the right size, with "true" on the
   // diagonal.
-  xla::ComputationDataHandle iota;
+  xla::XlaOp iota;
   TF_RETURN_IF_ERROR(
       XlaHelpers::Iota(builder, DataType::DT_INT32, last_dim_size, &iota));
-  xla::ComputationDataHandle iota_broadcast =
-      builder->Broadcast(iota, {last_dim_size});
-  xla::ComputationDataHandle mask = builder->Eq(iota_broadcast, iota, {0});
+  xla::XlaOp iota_broadcast = builder->Broadcast(iota, {last_dim_size});
+  xla::XlaOp mask = builder->Eq(iota_broadcast, iota, {0});
 
   // If this is a batched diagonal, broadcast the mask across the other
   // dimensions.
@@ -65,8 +64,7 @@ xla::StatusOr<xla::ComputationDataHandle> CreateDiagonal(
   std::vector<int64> broadcast_dims(other_dims.begin(), other_dims.end());
   broadcast_dims.push_back(1LL);
   broadcast_dims.push_back(last_dim_size);
-  xla::ComputationDataHandle input_broadcast =
-      builder->Reshape(input, broadcast_dims);
+  xla::XlaOp input_broadcast = builder->Reshape(input, broadcast_dims);
 
   broadcast_dims[broadcast_dims.size() - 2] = last_dim_size;
   xla::PrimitiveType element_type;
@@ -74,7 +72,7 @@ xla::StatusOr<xla::ComputationDataHandle> CreateDiagonal(
       DataTypeToPrimitiveType(ctx->input_type(0), &element_type));
   auto broadcast_shape =
       xla::ShapeUtil::MakeShape(element_type, broadcast_dims);
-  xla::ComputationDataHandle zeros = Zeros(builder, broadcast_shape);
+  xla::XlaOp zeros = Zeros(builder, broadcast_shape);
 
   input_broadcast = builder->Add(input_broadcast, zeros);
   return builder->Select(mask, input_broadcast, zeros);
@@ -85,7 +83,7 @@ class DiagOp : public XlaOpKernel {
   explicit DiagOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
 
     OP_REQUIRES(ctx, ctx->num_inputs() >= 1,
                 errors::InvalidArgument("Diag op must have at an input"));
@@ -96,7 +94,7 @@ class DiagOp : public XlaOpKernel {
                 errors::InvalidArgument("Expected 1 <= dims, got shape ",
                                         input_shape.DebugString()));
 
-    xla::ComputationDataHandle input = ctx->Input(0);
+    xla::XlaOp input = ctx->Input(0);
 
     // Picture:
     // tf.diag([1, 2, 3, 4]) ==> [[1, 0, 0, 0]
@@ -112,7 +110,7 @@ class DiagOp : public XlaOpKernel {
     auto diag_or_status =
         CreateDiagonal(input, size, /*other_dims=*/{}, ctx, builder);
     OP_REQUIRES_OK(ctx, diag_or_status.status());
-    xla::ComputationDataHandle diag = diag_or_status.ValueOrDie();
+    xla::XlaOp diag = diag_or_status.ValueOrDie();
 
     // Reshapes to the final shape.
     std::vector<int64> new_dims(dims.size() * 2);
@@ -131,7 +129,7 @@ class DiagPartOp : public XlaOpKernel {
   explicit DiagPartOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
 
     const TensorShape input_shape = ctx->InputShape(0);
     auto dims = input_shape.dim_sizes();
@@ -158,7 +156,7 @@ class DiagPartOp : public XlaOpKernel {
       new_dims.push_back(dims[i]);
     }
 
-    xla::ComputationDataHandle diag = ctx->Input(0);
+    xla::XlaOp diag = ctx->Input(0);
 
     // TODO(b/30878775): use Slice with strides when supported, in place of
     // the Pad -> Reshape -> Slice.
@@ -199,7 +197,7 @@ class MatrixDiagOp : public XlaOpKernel {
   explicit MatrixDiagOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
 
     OP_REQUIRES(ctx, ctx->num_inputs() >= 1,
                 errors::InvalidArgument("MatrixDiag op must have at an input"));
@@ -210,7 +208,7 @@ class MatrixDiagOp : public XlaOpKernel {
                 errors::InvalidArgument("Expected 1 <= dims, got shape ",
                                         input_shape.DebugString()));
 
-    xla::ComputationDataHandle diag = ctx->Input(0);
+    xla::XlaOp diag = ctx->Input(0);
 
     int last_dim = dims.size() - 1;
     int64 last_dim_size = input_shape.dim_size(last_dim);
@@ -232,7 +230,7 @@ class MatrixDiagPartOp : public XlaOpKernel {
   explicit MatrixDiagPartOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
 
     const TensorShape input_shape = ctx->InputShape(0);
     auto dims = input_shape.dim_sizes();
@@ -241,7 +239,7 @@ class MatrixDiagPartOp : public XlaOpKernel {
                 errors::InvalidArgument("Expected 2 <= dims, got shape ",
                                         input_shape.DebugString()));
 
-    xla::ComputationDataHandle diag = ctx->Input(0);
+    xla::XlaOp diag = ctx->Input(0);
 
     int last_dim = dims.size() - 1;
     int64 last_dim_size = dims[last_dim];
diff --git a/tensorflow/compiler/tf2xla/kernels/dynamic_slice_ops.cc b/tensorflow/compiler/tf2xla/kernels/dynamic_slice_ops.cc
index 800ef5ab98..0419de78b2 100644
--- a/tensorflow/compiler/tf2xla/kernels/dynamic_slice_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/dynamic_slice_ops.cc
@@ -18,7 +18,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/shape_util.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/core/framework/op_kernel.h"
 
 #include "tensorflow/compiler/tf2xla/type_util.h"
@@ -57,7 +57,7 @@ class DynamicUpdateSliceOp : public XlaOpKernel {
                                 input_shape.DebugString(), "; update shape is ",
                                 update_shape.DebugString()));
 
-    xla::ComputationDataHandle result = ctx->builder()->DynamicUpdateSlice(
+    xla::XlaOp result = ctx->builder()->DynamicUpdateSlice(
         ctx->Input(0), ctx->Input(1), ctx->Input(2));
     ctx->SetOutput(0, result);
   }
diff --git a/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc b/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc
index f2cd21ffb9..dd4a169087 100644
--- a/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc
@@ -56,7 +56,7 @@ class DynamicStitchOp : public XlaOpKernel {
     std::vector<xla::Literal> indices_input;
     OP_REQUIRES_OK(ctx, ctx->ConstantInputList("indices", &indices_input));
 
-    std::vector<xla::ComputationDataHandle> data;
+    std::vector<xla::XlaOp> data;
     std::vector<TensorShape> data_shapes;
     OP_REQUIRES_OK(ctx, ctx->InputList("data", &data, &data_shapes));
 
@@ -136,7 +136,7 @@ class DynamicStitchOp : public XlaOpKernel {
 
     // Look up all the children expressions that represent the data
     // inputs.
-    std::vector<xla::ComputationDataHandle> input(indices.size());
+    std::vector<xla::XlaOp> input(indices.size());
     for (int input_num = 0; input_num < indices.size(); input_num++) {
       TensorShape new_shape;
       // first reshaped dimension is the number of indices for this input.
@@ -166,7 +166,7 @@ class DynamicStitchOp : public XlaOpKernel {
     for (int d = indices0_shape.dims(); d < data0_shape.dims(); d++) {
       slice_limit[1 + d - indices0_shape.dims()] = data0_shape.dim_size(d);
     }
-    std::vector<xla::ComputationDataHandle> to_concat(number_of_indices);
+    std::vector<xla::XlaOp> to_concat(number_of_indices);
     for (int index_num = 0; index_num < number_of_indices; index_num++) {
       const auto& expression = input[src_input_vector[index_num]];
       // Take the appropriate slice of data.
diff --git a/tensorflow/compiler/tf2xla/kernels/elu_op.cc b/tensorflow/compiler/tf2xla/kernels/elu_op.cc
index 2fd27c5ca7..ed7462c166 100644
--- a/tensorflow/compiler/tf2xla/kernels/elu_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/elu_op.cc
@@ -18,7 +18,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/kernels/cwise_ops.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/core/framework/kernel_def_builder.h"
 #include "tensorflow/core/framework/types.h"
@@ -32,7 +32,7 @@ class EluOp : public XlaOpKernel {
   explicit EluOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
   // Computes the max of the scalar input x and 0.
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
     const auto zero = XlaHelpers::Zero(b, input_type(0));
     const auto one = XlaHelpers::One(b, input_type(0));
     const auto pred = b->Gt(ctx->Input(0), zero);
@@ -47,7 +47,7 @@ class EluGradOp : public XlaOpKernel {
   // Return the lhs (incoming gradient) if the rhs (input feature) > 0,
   // otherwise return lhs * (1 + rhs).
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
     const auto zero = XlaHelpers::Zero(b, input_type(0));
     const auto one = XlaHelpers::One(b, input_type(0));
     const auto grad = ctx->Input(0);
@@ -66,7 +66,7 @@ class SeluOp : public XlaOpKernel {
   explicit SeluOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
   // Computes the max of the scalar input x and 0.
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
     const auto zero = XlaHelpers::Zero(b, input_type(0));
     const auto one = XlaHelpers::One(b, input_type(0));
     const auto scale = XlaHelpers::FloatLiteral(b, input_type(0),
@@ -86,9 +86,8 @@ class SeluGradOp : public XlaOpKernel {
   // Return the lhs (incoming gradient) if the rhs (input feature) > 0,
   // otherwise return lhs * (1 + rhs).
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
     const auto zero = XlaHelpers::Zero(b, input_type(0));
-    const auto one = XlaHelpers::One(b, input_type(0));
     const auto scale = XlaHelpers::FloatLiteral(b, input_type(0),
             1.0507009873554804934193349852946);
     const auto scale_alpha = XlaHelpers::FloatLiteral(b, input_type(0),
diff --git a/tensorflow/compiler/tf2xla/kernels/extract_image_patches_op.cc b/tensorflow/compiler/tf2xla/kernels/extract_image_patches_op.cc
index b2970eae20..6df01cabbf 100644
--- a/tensorflow/compiler/tf2xla/kernels/extract_image_patches_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/extract_image_patches_op.cc
@@ -93,7 +93,7 @@ class ExtractImagePatchesOp : public XlaOpKernel {
                                 input_shape.DebugString()));
     const int64 depth = input_shape.dim_size(feature_dim);
 
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
 
     // The following code is equivalent to:
     // eye = np.eye(kH * kW * D).reshape([kH, kW, D, kH * kW * kD])
@@ -110,7 +110,7 @@ class ExtractImagePatchesOp : public XlaOpKernel {
     // Builds an identity matrix as a broadcast equality of iotas.
     // iota = np.arange(np.prod(ksize), depth)
     // filter = np.equal(np.reshape(iota, [-1, 1]), iota).astype(np.float32)
-    xla::ComputationDataHandle iota;
+    xla::XlaOp iota;
     TF_CHECK_OK(XlaHelpers::Iota(builder, DataType::DT_INT32,
                                  kernel_size * depth, &iota));
 
@@ -147,7 +147,7 @@ class ExtractImagePatchesOp : public XlaOpKernel {
                    &padding[i].first, &padding[i].second));
     }
 
-    xla::ComputationDataHandle conv =
+    xla::XlaOp conv =
         builder->ConvGeneralDilated(ctx->Input(0), filter, window_strides,
                                     padding, lhs_dilation, rhs_dilation, dims);
     ctx->SetOutput(0, conv);
diff --git a/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc b/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc
index 99470d70e7..8f0de0a524 100644
--- a/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc
@@ -44,23 +44,20 @@ void CpuNudge(const float min, const float max, const float quant_min,
 }
 
 // An XLA version of CpuNudge().
-void XlaNudge(xla::ComputationBuilder* b, const DataType data_type,
-              const xla::ComputationDataHandle& min,
-              const xla::ComputationDataHandle& max,
+void XlaNudge(xla::XlaBuilder* b, const DataType data_type,
+              const xla::XlaOp& min, const xla::XlaOp& max,
               const float quant_min_value, const float quant_max_value,
-              xla::ComputationDataHandle* nudged_min,
-              xla::ComputationDataHandle* nudged_max,
-              xla::ComputationDataHandle* scale) {
+              xla::XlaOp* nudged_min, xla::XlaOp* nudged_max,
+              xla::XlaOp* scale) {
   *scale = b->Div(b->Sub(max, min),
                   XlaHelpers::FloatLiteral(b, data_type,
                                            quant_max_value - quant_min_value));
-  xla::ComputationDataHandle quant_min =
+  xla::XlaOp quant_min =
       XlaHelpers::FloatLiteral(b, data_type, quant_min_value);
-  xla::ComputationDataHandle zero_point_from_min =
-      b->Sub(quant_min, b->Div(min, *scale));
-  xla::ComputationDataHandle quant_max =
+  xla::XlaOp zero_point_from_min = b->Sub(quant_min, b->Div(min, *scale));
+  xla::XlaOp quant_max =
       XlaHelpers::FloatLiteral(b, data_type, quant_max_value);
-  xla::ComputationDataHandle nudged_zero_point =
+  xla::XlaOp nudged_zero_point =
       b->Select(b->Le(zero_point_from_min, quant_min), quant_min,
                 b->Select(b->Ge(zero_point_from_min, quant_max), quant_max,
                           b->Round(zero_point_from_min)));
@@ -68,22 +65,18 @@ void XlaNudge(xla::ComputationBuilder* b, const DataType data_type,
   *nudged_max = b->Mul(b->Sub(quant_max, nudged_zero_point), *scale);
 }
 
-xla::ComputationDataHandle Quantize(
-    xla::ComputationBuilder* b, const xla::ComputationDataHandle& input,
-    const DataType data_type,
-    const xla::ComputationDataHandle& nudged_input_min,
-    const xla::ComputationDataHandle& nudged_input_max,
-    const xla::ComputationDataHandle& input_scale) {
-  xla::ComputationDataHandle one = XlaHelpers::FloatLiteral(b, data_type, 1.0f);
-  xla::ComputationDataHandle inv_scale = b->Div(one, input_scale);
-  xla::ComputationDataHandle half =
-      XlaHelpers::FloatLiteral(b, data_type, 0.5f);
-
-  xla::ComputationDataHandle clamped =
-      b->Clamp(nudged_input_min, input, nudged_input_max);
-  xla::ComputationDataHandle clamped_shifted =
-      b->Sub(clamped, nudged_input_min);
-  xla::ComputationDataHandle rounded =
+xla::XlaOp Quantize(xla::XlaBuilder* b, const xla::XlaOp& input,
+                    const DataType data_type,
+                    const xla::XlaOp& nudged_input_min,
+                    const xla::XlaOp& nudged_input_max,
+                    const xla::XlaOp& input_scale) {
+  xla::XlaOp one = XlaHelpers::FloatLiteral(b, data_type, 1.0f);
+  xla::XlaOp inv_scale = b->Div(one, input_scale);
+  xla::XlaOp half = XlaHelpers::FloatLiteral(b, data_type, 0.5f);
+
+  xla::XlaOp clamped = b->Clamp(nudged_input_min, input, nudged_input_max);
+  xla::XlaOp clamped_shifted = b->Sub(clamped, nudged_input_min);
+  xla::XlaOp rounded =
       b->Floor(b->Add(b->Mul(clamped_shifted, inv_scale), half));
   return b->Add(b->Mul(rounded, input_scale), nudged_input_min);
 }
@@ -111,18 +104,18 @@ class FakeQuantWithMinMaxArgsOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationDataHandle input = ctx->Input(0);
+    xla::XlaOp input = ctx->Input(0);
     const DataType data_type = ctx->input_type(0);
 
-    xla::ComputationBuilder* b = ctx->builder();
-    xla::ComputationDataHandle nudged_input_min =
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp nudged_input_min =
         XlaHelpers::FloatLiteral(b, data_type, nudged_input_min_);
-    xla::ComputationDataHandle nudged_input_max =
+    xla::XlaOp nudged_input_max =
         XlaHelpers::FloatLiteral(b, data_type, nudged_input_max_);
-    xla::ComputationDataHandle input_scale =
+    xla::XlaOp input_scale =
         XlaHelpers::FloatLiteral(b, data_type, input_scale_);
-    xla::ComputationDataHandle output = Quantize(
-        b, input, data_type, nudged_input_min, nudged_input_max, input_scale);
+    xla::XlaOp output = Quantize(b, input, data_type, nudged_input_min,
+                                 nudged_input_max, input_scale);
     ctx->SetOutput(0, output);
   }
 
@@ -159,23 +152,22 @@ class FakeQuantWithMinMaxArgsGradOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationDataHandle gradient = ctx->Input(0);
+    xla::XlaOp gradient = ctx->Input(0);
     const TensorShape gradient_shape = ctx->InputShape(0);
-    xla::ComputationDataHandle input = ctx->Input(1);
+    xla::XlaOp input = ctx->Input(1);
     const DataType data_type = ctx->input_type(1);
 
-    xla::ComputationBuilder* b = ctx->builder();
-    xla::ComputationDataHandle nudged_input_min =
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp nudged_input_min =
         XlaHelpers::FloatLiteral(b, data_type, nudged_input_min_);
-    xla::ComputationDataHandle nudged_input_max =
+    xla::XlaOp nudged_input_max =
         XlaHelpers::FloatLiteral(b, data_type, nudged_input_max_);
 
-    xla::ComputationDataHandle between_nudged_min_max =
+    xla::XlaOp between_nudged_min_max =
         b->And(b->Le(nudged_input_min, input), b->Le(input, nudged_input_max));
-    xla::ComputationDataHandle zeroes = b->Broadcast(
-        XlaHelpers::Zero(b, data_type), gradient_shape.dim_sizes());
-    xla::ComputationDataHandle output =
-        b->Select(between_nudged_min_max, gradient, zeroes);
+    xla::XlaOp zeroes = b->Broadcast(XlaHelpers::Zero(b, data_type),
+                                     gradient_shape.dim_sizes());
+    xla::XlaOp output = b->Select(between_nudged_min_max, gradient, zeroes);
     ctx->SetOutput(0, output);
   }
 
@@ -204,18 +196,18 @@ class FakeQuantWithMinMaxVarsOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationDataHandle input = ctx->Input(0);
+    xla::XlaOp input = ctx->Input(0);
     const DataType data_type = ctx->input_type(0);
-    xla::ComputationDataHandle input_min = ctx->Input(1);
-    xla::ComputationDataHandle input_max = ctx->Input(2);
+    xla::XlaOp input_min = ctx->Input(1);
+    xla::XlaOp input_max = ctx->Input(2);
 
-    xla::ComputationBuilder* b = ctx->builder();
-    xla::ComputationDataHandle nudged_input_min, nudged_input_max, input_scale;
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp nudged_input_min, nudged_input_max, input_scale;
     XlaNudge(b, data_type, input_min, input_max, quant_min_, quant_max_,
              &nudged_input_min, &nudged_input_max, &input_scale);
 
-    xla::ComputationDataHandle output = Quantize(
-        b, input, data_type, nudged_input_min, nudged_input_max, input_scale);
+    xla::XlaOp output = Quantize(b, input, data_type, nudged_input_min,
+                                 nudged_input_max, input_scale);
     ctx->SetOutput(0, output);
   }
 
@@ -243,47 +235,43 @@ class FakeQuantWithMinMaxVarsGradOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationDataHandle gradient = ctx->Input(0);
+    xla::XlaOp gradient = ctx->Input(0);
     const TensorShape gradient_shape = ctx->InputShape(0);
-    xla::ComputationDataHandle input = ctx->Input(1);
+    xla::XlaOp input = ctx->Input(1);
     const DataType data_type = ctx->input_type(1);
     const DataType accumulation_type =
         XlaHelpers::SumAccumulationType(data_type);
-    xla::ComputationDataHandle input_min = ctx->Input(2);
-    xla::ComputationDataHandle input_max = ctx->Input(3);
+    xla::XlaOp input_min = ctx->Input(2);
+    xla::XlaOp input_max = ctx->Input(3);
 
-    xla::ComputationBuilder* b = ctx->builder();
-    xla::ComputationDataHandle nudged_input_min, nudged_input_max, input_scale;
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp nudged_input_min, nudged_input_max, input_scale;
     XlaNudge(b, data_type, input_min, input_max, quant_min_, quant_max_,
              &nudged_input_min, &nudged_input_max, &input_scale);
 
-    xla::ComputationDataHandle between_nudged_min_max =
+    xla::XlaOp between_nudged_min_max =
         b->And(b->Le(nudged_input_min, input), b->Le(input, nudged_input_max));
-    xla::ComputationDataHandle zero = XlaHelpers::Zero(b, data_type);
-    xla::ComputationDataHandle zeroes =
-        b->Broadcast(zero, gradient_shape.dim_sizes());
-    xla::ComputationDataHandle output0 =
-        b->Select(between_nudged_min_max, gradient, zeroes);
+    xla::XlaOp zero = XlaHelpers::Zero(b, data_type);
+    xla::XlaOp zeroes = b->Broadcast(zero, gradient_shape.dim_sizes());
+    xla::XlaOp output0 = b->Select(between_nudged_min_max, gradient, zeroes);
     ctx->SetOutput(0, output0);
 
-    xla::ComputationDataHandle below_min = b->Lt(input, nudged_input_min);
-    xla::ComputationDataHandle select1 = b->Select(below_min, gradient, zeroes);
-    xla::ComputationDataHandle reduce1 = b->ReduceAll(
+    xla::XlaOp below_min = b->Lt(input, nudged_input_min);
+    xla::XlaOp select1 = b->Select(below_min, gradient, zeroes);
+    xla::XlaOp reduce1 = b->ReduceAll(
         XlaHelpers::ConvertElementType(b, select1, accumulation_type),
         XlaHelpers::Zero(b, accumulation_type),
         *ctx->GetOrCreateAdd(accumulation_type));
-    xla::ComputationDataHandle output1 =
-        XlaHelpers::ConvertElementType(b, reduce1, data_type);
+    xla::XlaOp output1 = XlaHelpers::ConvertElementType(b, reduce1, data_type);
     ctx->SetOutput(1, output1);
 
-    xla::ComputationDataHandle above_max = b->Gt(input, nudged_input_max);
-    xla::ComputationDataHandle select2 = b->Select(above_max, gradient, zeroes);
-    xla::ComputationDataHandle reduce2 = b->ReduceAll(
+    xla::XlaOp above_max = b->Gt(input, nudged_input_max);
+    xla::XlaOp select2 = b->Select(above_max, gradient, zeroes);
+    xla::XlaOp reduce2 = b->ReduceAll(
         XlaHelpers::ConvertElementType(b, select2, accumulation_type),
         XlaHelpers::Zero(b, accumulation_type),
         *ctx->GetOrCreateAdd(accumulation_type));
-    xla::ComputationDataHandle output2 =
-        XlaHelpers::ConvertElementType(b, reduce2, data_type);
+    xla::XlaOp output2 = XlaHelpers::ConvertElementType(b, reduce2, data_type);
     ctx->SetOutput(2, output2);
   }
 
diff --git a/tensorflow/compiler/tf2xla/kernels/fft_ops.cc b/tensorflow/compiler/tf2xla/kernels/fft_ops.cc
index a4f3c1c3ad..fcb927dab0 100644
--- a/tensorflow/compiler/tf2xla/kernels/fft_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/fft_ops.cc
@@ -62,9 +62,8 @@ class GenericFftOp : public XlaOpKernel {
       }
     }
 
-    xla::ComputationBuilder* b = ctx->builder();
-    xla::ComputationDataHandle fft =
-        b->Fft(ctx->Input(0), fft_type_, fft_length);
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp fft = b->Fft(ctx->Input(0), fft_type_, fft_length);
     ctx->SetOutput(0, fft);
   }
 
diff --git a/tensorflow/compiler/tf2xla/kernels/fill_op.cc b/tensorflow/compiler/tf2xla/kernels/fill_op.cc
index eaa13b8dfa..e4467a0fb1 100644
--- a/tensorflow/compiler/tf2xla/kernels/fill_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/fill_op.cc
@@ -48,7 +48,7 @@ class FillOp : public XlaOpKernel {
                             0, {dims_shape.num_elements()}, &dims_literal));
 
     // Convert the dims literal into a vector that we can pass to
-    // ComputationBuilder.
+    // XlaBuilder.
     std::vector<int64> broadcast;
     broadcast.reserve(dims_literal.shape().dimensions(0));
     for (int i = 0; i < dims_literal.shape().dimensions(0); ++i) {
@@ -56,7 +56,7 @@ class FillOp : public XlaOpKernel {
     }
     // Look up the value input, reshaping to a scalar if it was a
     // 'legacy' scalar (secretly a vector).
-    xla::ComputationDataHandle data = ctx->Input(1);
+    xla::XlaOp data = ctx->Input(1);
     if (value_shape.dims() > 0) {
       CHECK_EQ(value_shape.dims(), 1);
       data = ctx->builder()->Reshape(data, {});
diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op.cc b/tensorflow/compiler/tf2xla/kernels/gather_op.cc
index 0b79cb0916..d13e25bcdd 100644
--- a/tensorflow/compiler/tf2xla/kernels/gather_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/gather_op.cc
@@ -26,13 +26,11 @@ limitations under the License.
 
 namespace tensorflow {
 
-Status XlaGather(const xla::ComputationDataHandle& input,
-                 const TensorShape& input_shape,
-                 const xla::ComputationDataHandle& indices,
-                 const TensorShape& indices_shape, int64 axis,
-                 bool indices_are_nd, DataType dtype, DataType index_type,
-                 xla::ComputationBuilder* builder,
-                 xla::ComputationDataHandle* gather_output) {
+Status XlaGather(const xla::XlaOp& input, const TensorShape& input_shape,
+                 const xla::XlaOp& indices, const TensorShape& indices_shape,
+                 int64 axis, bool indices_are_nd, DataType dtype,
+                 DataType index_type, xla::XlaBuilder* builder,
+                 xla::XlaOp* gather_output) {
   // There is no deep reason why we need this precondition, but this is the only
   // combination that is used and tested today.
   CHECK(!indices_are_nd || axis == 0);
@@ -153,7 +151,7 @@ class GatherOp : public XlaOpKernel {
   explicit GatherOp(OpKernelConstruction* context) : XlaOpKernel(context) {}
 
   void Compile(XlaOpKernelContext* context) override {
-    xla::ComputationBuilder* builder = context->builder();
+    xla::XlaBuilder* builder = context->builder();
     auto input = context->Input(0);
     auto input_shape = context->InputShape(0);
     auto indices = context->Input(1);
@@ -182,7 +180,7 @@ class GatherOp : public XlaOpKernel {
     OP_REQUIRES(context, index_type == DT_INT32 || index_type == DT_INT64,
                 errors::InvalidArgument("indices must be int32 or int64"));
 
-    xla::ComputationDataHandle gather;
+    xla::XlaOp gather;
     OP_REQUIRES_OK(
         context, XlaGather(input, input_shape, indices, indices_shape, axis,
                            /*indices_are_nd=*/false, input_type(0), index_type,
@@ -220,10 +218,10 @@ class GatherNdOp : public XlaOpKernel {
             indices_shape.dim_size(indices_shape.dims() - 1), " vs. ",
             params_shape.dims()));
 
-    xla::ComputationBuilder* builder = context->builder();
+    xla::XlaBuilder* builder = context->builder();
     auto params = context->Input(0);
     auto indices = context->Input(1);
-    xla::ComputationDataHandle gather;
+    xla::XlaOp gather;
     OP_REQUIRES_OK(context, XlaGather(params, params_shape, indices,
                                       indices_shape, /*axis=*/0,
                                       /*indices_are_nd=*/true, params_type,
diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h b/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h
index f9376f0eab..d898e43b85 100644
--- a/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h
+++ b/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h
@@ -20,7 +20,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/xla/client/client_library.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/util/bcast.h"
 
@@ -33,13 +33,11 @@ namespace tensorflow {
 // If `indices_are_nd` is true, the last dimension of `indices` are treated as
 // a multidimensional index values. Otherwise, `indices` is treated as a tensor
 // of scalar indices.
-Status XlaGather(const xla::ComputationDataHandle& input,
-                 const TensorShape& input_shape,
-                 const xla::ComputationDataHandle& indices,
-                 const TensorShape& indices_shape, int64 axis,
-                 bool indices_are_nd, DataType dtype, DataType index_type,
-                 xla::ComputationBuilder* builder,
-                 xla::ComputationDataHandle* gather_output);
+Status XlaGather(const xla::XlaOp& input, const TensorShape& input_shape,
+                 const xla::XlaOp& indices, const TensorShape& indices_shape,
+                 int64 axis, bool indices_are_nd, DataType dtype,
+                 DataType index_type, xla::XlaBuilder* builder,
+                 xla::XlaOp* gather_output);
 
 }  // namespace tensorflow
 
diff --git a/tensorflow/compiler/tf2xla/kernels/if_op.cc b/tensorflow/compiler/tf2xla/kernels/if_op.cc
index eefbe55c81..8b9b026643 100644
--- a/tensorflow/compiler/tf2xla/kernels/if_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/if_op.cc
@@ -37,7 +37,7 @@ XlaIfOp::XlaIfOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
 // TODO(b/35949885): There is duplication here with the handling of the
 // while_op. Refactor the common code out/rework.
 void XlaIfOp::Compile(XlaOpKernelContext* ctx) {
-  xla::ComputationBuilder* b = ctx->builder();
+  xla::XlaBuilder* b = ctx->builder();
 
   OP_REQUIRES(ctx, cond_type_ == DT_BOOL,
               errors::InvalidArgument(
@@ -48,7 +48,7 @@ void XlaIfOp::Compile(XlaOpKernelContext* ctx) {
 
   VLOG(1) << "Building If: " << input_types_.size() << " inputs";
 
-  std::vector<xla::ComputationDataHandle> inputs(input_types_.size());
+  std::vector<xla::XlaOp> inputs(input_types_.size());
   std::vector<XlaCompiler::Argument> arguments(input_types_.size());
   for (int i = 0; i < input_types_.size(); ++i) {
     XlaCompiler::Argument& arg = arguments[i];
@@ -175,19 +175,19 @@ void XlaIfOp::Compile(XlaOpKernelContext* ctx) {
             "Mismatch in resource of then and else branch for resource ", i));
   }
 
-  xla::ComputationDataHandle outputs =
+  xla::XlaOp outputs =
       b->Conditional(ctx->Input(0), b->Tuple(inputs), *then_result.computation,
                      b->Tuple(inputs), *else_result.computation);
   // Sets non-variable outputs.
   for (int i = 0; i < output_types_.size(); ++i) {
     if (ctx->input_type(i) != DT_RESOURCE) {
-      xla::ComputationDataHandle output_handle = b->GetTupleElement(outputs, i);
+      xla::XlaOp output_handle = b->GetTupleElement(outputs, i);
       if (VLOG_IS_ON(2)) {
         LOG(INFO) << "Setting output " << i;
         auto shape_or = b->GetShape(output_handle);
         if (shape_or.ok()) {
           LOG(INFO) << "Shape for output " << i << ": "
-                    << xla::ShapeUtil::HumanString(*shape_or.ValueOrDie());
+                    << xla::ShapeUtil::HumanString(shape_or.ValueOrDie());
         } else {
           LOG(INFO) << "Shape unknown for output " << i;
         }
diff --git a/tensorflow/compiler/tf2xla/kernels/image_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
index 5eeda79a93..1568b33679 100644
--- a/tensorflow/compiler/tf2xla/kernels/image_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
@@ -23,10 +23,9 @@ namespace {
 
 // Converts 'input' from RGB format to HSV format.
 // 'shape' is the shape of the red/green/blue tensors.
-std::array<xla::ComputationDataHandle, 3> RGBToHSV(
-    XlaOpKernelContext* ctx, xla::ComputationBuilder* b,
-    const std::array<xla::ComputationDataHandle, 3>& rgb, DataType dtype,
-    const TensorShape& shape) {
+std::array<xla::XlaOp, 3> RGBToHSV(XlaOpKernelContext* ctx, xla::XlaBuilder* b,
+                                   const std::array<xla::XlaOp, 3>& rgb,
+                                   DataType dtype, const TensorShape& shape) {
   auto zero = XlaHelpers::Zero(b, dtype);
   auto one = XlaHelpers::One(b, dtype);
 
@@ -54,12 +53,12 @@ std::array<xla::ComputationDataHandle, 3> RGBToHSV(
 }
 
 // Converts 'input' from HSV format to RGB format.
-std::array<xla::ComputationDataHandle, 3> HSVToRGB(
-    xla::ComputationBuilder* b,
-    const std::array<xla::ComputationDataHandle, 3>& hsv, DataType dtype) {
-  xla::ComputationDataHandle hue = hsv[0];
-  xla::ComputationDataHandle saturation = hsv[1];
-  xla::ComputationDataHandle value = hsv[2];
+std::array<xla::XlaOp, 3> HSVToRGB(xla::XlaBuilder* b,
+                                   const std::array<xla::XlaOp, 3>& hsv,
+                                   DataType dtype) {
+  xla::XlaOp hue = hsv[0];
+  xla::XlaOp saturation = hsv[1];
+  xla::XlaOp value = hsv[2];
   auto zero = XlaHelpers::Zero(b, dtype);
   auto one = XlaHelpers::FloatLiteral(b, dtype, 1.0);
   auto two = XlaHelpers::FloatLiteral(b, dtype, 2.0);
@@ -95,16 +94,16 @@ class RGBToHSVOp : public XlaOpKernel {
         errors::FailedPrecondition("input must have 3 channels but input has ",
                                    channels, " channels."));
 
-    xla::ComputationBuilder* b = context->builder();
-    xla::ComputationDataHandle input = context->Input(0);
+    xla::XlaBuilder* b = context->builder();
+    xla::XlaOp input = context->Input(0);
 
-    xla::ComputationDataHandle red =
+    xla::XlaOp red =
         b->SliceInDim(input, /*start_index=*/0, /*limit_index=*/1, /*stride=*/1,
                       /*dimno=*/channel_dim);
-    xla::ComputationDataHandle green =
+    xla::XlaOp green =
         b->SliceInDim(input, /*start_index=*/1, /*limit_index=*/2, /*stride=*/1,
                       /*dimno=*/channel_dim);
-    xla::ComputationDataHandle blue =
+    xla::XlaOp blue =
         b->SliceInDim(input, /*start_index=*/2, /*limit_index=*/3, /*stride=*/1,
                       /*dimno=*/channel_dim);
     TensorShape channel_shape = input_shape;
@@ -133,15 +132,15 @@ class HSVToRGBOp : public XlaOpKernel {
         errors::FailedPrecondition("input must have 3 channels but input has ",
                                    channels, " channels."));
 
-    xla::ComputationBuilder* b = context->builder();
-    xla::ComputationDataHandle input = context->Input(0);
-    xla::ComputationDataHandle hue =
+    xla::XlaBuilder* b = context->builder();
+    xla::XlaOp input = context->Input(0);
+    xla::XlaOp hue =
         b->SliceInDim(input, /*start_index=*/0, /*limit_index=*/1, /*stride=*/1,
                       /*dimno=*/channel_dim);
-    xla::ComputationDataHandle saturation =
+    xla::XlaOp saturation =
         b->SliceInDim(input, /*start_index=*/1, /*limit_index=*/2, /*stride=*/1,
                       /*dimno=*/channel_dim);
-    xla::ComputationDataHandle value =
+    xla::XlaOp value =
         b->SliceInDim(input, /*start_index=*/2, /*limit_index=*/3, /*stride=*/1,
                       /*dimno=*/channel_dim);
 
@@ -174,9 +173,9 @@ class AdjustContrastOpV2 : public XlaOpKernel {
                 errors::InvalidArgument("contrast_factor must be scalar: ",
                                         factor_shape.DebugString()));
 
-    xla::ComputationBuilder* b = context->builder();
-    xla::ComputationDataHandle input = context->Input(0);
-    xla::ComputationDataHandle factor = context->Input(1);
+    xla::XlaBuilder* b = context->builder();
+    xla::XlaOp input = context->Input(0);
+    xla::XlaOp factor = context->Input(1);
 
     DataType type = context->input_type(0);
 
@@ -221,19 +220,19 @@ class AdjustSaturationOp : public XlaOpKernel {
         errors::InvalidArgument("input must have 3 channels but instead has ",
                                 channels, " channels."));
 
-    xla::ComputationBuilder* b = context->builder();
-    xla::ComputationDataHandle input = context->Input(0);
-    xla::ComputationDataHandle scale = context->Input(1);
+    xla::XlaBuilder* b = context->builder();
+    xla::XlaOp input = context->Input(0);
+    xla::XlaOp scale = context->Input(1);
 
     DataType type = context->input_type(0);
 
-    xla::ComputationDataHandle red =
+    xla::XlaOp red =
         b->SliceInDim(input, /*start_index=*/0, /*limit_index=*/1, /*stride=*/1,
                       /*dimno=*/channel_dim);
-    xla::ComputationDataHandle green =
+    xla::XlaOp green =
         b->SliceInDim(input, /*start_index=*/1, /*limit_index=*/2, /*stride=*/1,
                       /*dimno=*/channel_dim);
-    xla::ComputationDataHandle blue =
+    xla::XlaOp blue =
         b->SliceInDim(input, /*start_index=*/2, /*limit_index=*/3, /*stride=*/1,
                       /*dimno=*/channel_dim);
     TensorShape channel_shape = input_shape;
@@ -271,19 +270,19 @@ class AdjustHueOp : public XlaOpKernel {
         errors::InvalidArgument("input must have 3 channels but instead has ",
                                 channels, " channels."));
 
-    xla::ComputationBuilder* b = context->builder();
-    xla::ComputationDataHandle input = context->Input(0);
-    xla::ComputationDataHandle delta = context->Input(1);
+    xla::XlaBuilder* b = context->builder();
+    xla::XlaOp input = context->Input(0);
+    xla::XlaOp delta = context->Input(1);
 
     DataType type = context->input_type(0);
 
-    xla::ComputationDataHandle red =
+    xla::XlaOp red =
         b->SliceInDim(input, /*start_index=*/0, /*limit_index=*/1, /*stride=*/1,
                       /*dimno=*/channel_dim);
-    xla::ComputationDataHandle green =
+    xla::XlaOp green =
         b->SliceInDim(input, /*start_index=*/1, /*limit_index=*/2, /*stride=*/1,
                       /*dimno=*/channel_dim);
-    xla::ComputationDataHandle blue =
+    xla::XlaOp blue =
         b->SliceInDim(input, /*start_index=*/2, /*limit_index=*/3, /*stride=*/1,
                       /*dimno=*/channel_dim);
     TensorShape channel_shape = input_shape;
diff --git a/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc
index f36b3f5948..9058cbc747 100644
--- a/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc
@@ -99,9 +99,9 @@ ResizeConvolutionDims ComputeResizeConvolutionParameters(
   return dims;
 }
 
-xla::ComputationDataHandle MakeBilinearResizeKernel(
-    xla::ComputationBuilder* builder, gtl::ArraySlice<int64> kernel_size,
-    int64 channels) {
+xla::XlaOp MakeBilinearResizeKernel(xla::XlaBuilder* builder,
+                                    gtl::ArraySlice<int64> kernel_size,
+                                    int64 channels) {
   // Form a 2D convolution kernel like:
   //       1 2 3 2 1
   //       2 4 6 4 2
@@ -120,7 +120,7 @@ xla::ComputationDataHandle MakeBilinearResizeKernel(
     return kernel;
   };
 
-  xla::ComputationDataHandle channels_iota;
+  xla::XlaOp channels_iota;
   // DT_INT32 Iota will always return status::OK().
   TF_CHECK_OK(
       XlaHelpers::Iota(builder, DataType::DT_INT32, channels, &channels_iota));
@@ -139,10 +139,12 @@ xla::ComputationDataHandle MakeBilinearResizeKernel(
       /*broadcast_dimensions=*/{0});
 }
 
-xla::ComputationDataHandle ResizeUsingDilationAndConvolution(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& input,
-    const int num_spatial_dims, std::vector<int64> in_size,
-    std::vector<int64> out_size, const int64 channels) {
+xla::XlaOp ResizeUsingDilationAndConvolution(xla::XlaBuilder* builder,
+                                             const xla::XlaOp& input,
+                                             const int num_spatial_dims,
+                                             std::vector<int64> in_size,
+                                             std::vector<int64> out_size,
+                                             const int64 channels) {
   // Picture for a 1x3 to 1x4 resize:
   // stride = 2, kernel size = 3
   // Input:
@@ -168,9 +170,9 @@ xla::ComputationDataHandle ResizeUsingDilationAndConvolution(
 
   ResizeConvolutionDims dims =
       ComputeResizeConvolutionParameters(in_size, out_size);
-  xla::ComputationDataHandle kernel =
+  xla::XlaOp kernel =
       MakeBilinearResizeKernel(builder, dims.kernel_size, channels);
-  xla::ComputationDataHandle output = builder->ConvGeneralDilated(
+  xla::XlaOp output = builder->ConvGeneralDilated(
       input, kernel, dims.stride,
       /*padding=*/
       {{dims.kernel_size[0] - 1, dims.kernel_size[0] - 1},
@@ -189,10 +191,12 @@ xla::ComputationDataHandle ResizeUsingDilationAndConvolution(
   return output;
 }
 
-xla::ComputationDataHandle ResizeUsingDilationAndConvolutionGradOp(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& grad,
-    const int num_spatial_dims, std::vector<int64> in_size,
-    std::vector<int64> grad_size, const int64 channels) {
+xla::XlaOp ResizeUsingDilationAndConvolutionGradOp(xla::XlaBuilder* builder,
+                                                   const xla::XlaOp& grad,
+                                                   const int num_spatial_dims,
+                                                   std::vector<int64> in_size,
+                                                   std::vector<int64> grad_size,
+                                                   const int64 channels) {
   ResizeConvolutionDims dims =
       ComputeResizeConvolutionParameters(in_size, grad_size);
 
@@ -210,7 +214,7 @@ xla::ComputationDataHandle ResizeUsingDilationAndConvolutionGradOp(
   }
   dimension_numbers.set_kernel_input_feature_dimension(num_spatial_dims);
   dimension_numbers.set_kernel_output_feature_dimension(num_spatial_dims + 1);
-  xla::ComputationDataHandle kernel =
+  xla::XlaOp kernel =
       MakeBilinearResizeKernel(builder, dims.kernel_size, channels);
 
   // Broadcast the input kernel where the forward op expanded from a size == 1
@@ -223,7 +227,7 @@ xla::ComputationDataHandle ResizeUsingDilationAndConvolutionGradOp(
     }
   }
 
-  xla::ComputationDataHandle output = builder->ConvGeneralDilated(
+  xla::XlaOp output = builder->ConvGeneralDilated(
       grad, kernel, /*window_strides=*/dims.kernel_size,
       /*padding=*/
       {{dims.kernel_size[0] - 1, dims.kernel_size[0] - 1},
@@ -258,7 +262,7 @@ class ResizeBilinearOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     TensorShape input_shape = ctx->InputShape(0);
     OP_REQUIRES(ctx, input_shape.dims() == 4,
@@ -283,7 +287,7 @@ class ResizeBilinearOp : public XlaOpKernel {
 
     const int num_spatial_dims = 2;
 
-    xla::ComputationDataHandle input = ctx->Input(0);
+    xla::XlaOp input = ctx->Input(0);
 
     // If in_size[i] > 1 and out_size[i] == 1, slice out the first input in
     // dimension i.
@@ -318,7 +322,7 @@ class ResizeBilinearOp : public XlaOpKernel {
     // from image of size axb -> cxd is same as resizing axb -> exf -> cxd.
     //
     // This makes the convolutions kernels smaller and the operation faster.
-    xla::ComputationDataHandle output = input;
+    xla::XlaOp output = input;
     while (in_size != out_size) {
       if (in_size[0] != 1 && in_size[1] != 1) {
         std::vector<float> k = {
@@ -369,7 +373,7 @@ class ResizeBilinearGradOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     TensorShape input_shape = ctx->InputShape(1);
     OP_REQUIRES(ctx, input_shape.dims() == 4,
@@ -406,9 +410,9 @@ class ResizeBilinearGradOp : public XlaOpKernel {
 
     const int num_spatial_dims = 2;
 
-    xla::ComputationDataHandle grad = ctx->Input(0);
+    xla::XlaOp grad = ctx->Input(0);
 
-    xla::ComputationDataHandle output = grad;
+    xla::XlaOp output = grad;
     while (in_size != grad_size) {
       if (in_size[0] != 1 && in_size[1] != 1) {
         std::vector<float> k = {
diff --git a/tensorflow/compiler/tf2xla/kernels/index_ops.cc b/tensorflow/compiler/tf2xla/kernels/index_ops.cc
index 7bf4b435f5..36eb4c7545 100644
--- a/tensorflow/compiler/tf2xla/kernels/index_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/index_ops.cc
@@ -61,10 +61,10 @@ void XlaArgMinMaxOp::Compile(XlaOpKernelContext* ctx) {
 
   DataType index_type = output_type(0);
 
-  xla::ComputationBuilder* b = ctx->builder();
-  xla::ComputationDataHandle input = ctx->Input(0);
+  xla::XlaBuilder* b = ctx->builder();
+  xla::XlaOp input = ctx->Input(0);
 
-  xla::ComputationDataHandle output;
+  xla::XlaOp output;
   if (is_min_) {
     OP_REQUIRES_OK(ctx,
                    XlaHelpers::ArgMin(b, ctx, input, input_shape, input_type(0),
diff --git a/tensorflow/compiler/tf2xla/kernels/index_ops_cpu.cc b/tensorflow/compiler/tf2xla/kernels/index_ops_cpu.cc
index b1f3c3c298..2c2d88486f 100644
--- a/tensorflow/compiler/tf2xla/kernels/index_ops_cpu.cc
+++ b/tensorflow/compiler/tf2xla/kernels/index_ops_cpu.cc
@@ -71,10 +71,10 @@ class ArgMaxCustomCallOp : public XlaOpKernel {
     OP_REQUIRES(ctx, XlaContext::Get(ctx).allow_cpu_custom_calls(),
                 errors::InvalidArgument(
                     "ArgMax implementation requires a CustomCall on CPU"));
-    xla::ComputationBuilder& b = *ctx->builder();
+    xla::XlaBuilder& b = *ctx->builder();
 
     // XLA passes <out> to the function, so it is not included here.
-    std::vector<xla::ComputationDataHandle> args;
+    std::vector<xla::XlaOp> args;
     args.push_back(ctx->Input(0));
     args.push_back(b.ConstantLiteral(
         *xla::Literal::CreateR1<int64>(input_shape.dim_sizes())));
@@ -91,7 +91,7 @@ class ArgMaxCustomCallOp : public XlaOpKernel {
 
     // Tell XLA to call the custom code, defined in
     // index_ops_kernel_argmax_float_1d.cc.
-    xla::ComputationDataHandle output;
+    xla::XlaOp output;
     switch (input_shape.dims()) {
       case 1:
         output = b.CustomCall("argmax_float_1d_xla_impl", args, xla_shape);
diff --git a/tensorflow/compiler/tf2xla/kernels/l2loss_op.cc b/tensorflow/compiler/tf2xla/kernels/l2loss_op.cc
index c177f08d9c..1decf7d72d 100644
--- a/tensorflow/compiler/tf2xla/kernels/l2loss_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/l2loss_op.cc
@@ -16,7 +16,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/kernels/no_op.h"
@@ -33,7 +33,7 @@ class L2LossOp : public XlaOpKernel {
     std::iota(dims.begin(), dims.end(), 0);
 
     DataType dtype = ctx->input_type(0);
-    xla::ComputationBuilder* const b = ctx->builder();
+    xla::XlaBuilder* const b = ctx->builder();
 
     //  output = sum(t ** 2) / 2
     const DataType accumulation_type = XlaHelpers::SumAccumulationType(dtype);
diff --git a/tensorflow/compiler/tf2xla/kernels/lrn_ops.cc b/tensorflow/compiler/tf2xla/kernels/lrn_ops.cc
index 1cfee3070f..39fbf98a62 100644
--- a/tensorflow/compiler/tf2xla/kernels/lrn_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/lrn_ops.cc
@@ -38,8 +38,8 @@ class LRNOp : public XlaOpKernel {
     OP_REQUIRES(ctx, in_shape.dims() == 4,
                 errors::InvalidArgument("in must be 4-dimensional"));
 
-    xla::ComputationBuilder* builder = ctx->builder();
-    xla::ComputationDataHandle input = ctx->Input(0);
+    xla::XlaBuilder* builder = ctx->builder();
+    xla::XlaOp input = ctx->Input(0);
 
     // sqr_sum[a, b, c, d] =
     //    sum(input[a, b, c, d - depth_radius : d + depth_radius + 1] ** 2)
@@ -111,10 +111,10 @@ class LRNGradOp : public XlaOpKernel {
             "input_grads, input_image, and out_image should have the same "
             "shape"));
 
-    xla::ComputationBuilder* builder = ctx->builder();
-    xla::ComputationDataHandle in_grads = ctx->Input(0);
-    xla::ComputationDataHandle in_image = ctx->Input(1);
-    xla::ComputationDataHandle out_image = ctx->Input(2);
+    xla::XlaBuilder* builder = ctx->builder();
+    xla::XlaOp in_grads = ctx->Input(0);
+    xla::XlaOp in_image = ctx->Input(1);
+    xla::XlaOp out_image = ctx->Input(2);
 
     // This code is ported from tensorflow/core/kernels/lrn_op.cc. In Python
     // pseudo-code, the Eigen code does this for each spatial position:
@@ -166,7 +166,7 @@ class LRNGradOp : public XlaOpKernel {
     auto dy_reduced =
         XlaHelpers::ConvertElementType(builder, dy_reduce, input_type(0));
 
-    xla::ComputationDataHandle gradients = builder->Add(
+    xla::XlaOp gradients = builder->Add(
         builder->Mul(in_image, dy_reduced),
         builder->Mul(in_grads,
                      builder->Pow(norm, builder->ConstantR0<float>(-beta_))));
diff --git a/tensorflow/compiler/tf2xla/kernels/matmul_op.cc b/tensorflow/compiler/tf2xla/kernels/matmul_op.cc
index 886baf8115..6949b296f4 100644
--- a/tensorflow/compiler/tf2xla/kernels/matmul_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/matmul_op.cc
@@ -66,8 +66,8 @@ class MatMulOp : public XlaOpKernel {
                                         a_shape.DebugString(), ", In[1]: ",
                                         b_shape.DebugString()));
 
-    xla::ComputationDataHandle a = ctx->Input(0);
-    xla::ComputationDataHandle b = ctx->Input(1);
+    xla::XlaOp a = ctx->Input(0);
+    xla::XlaOp b = ctx->Input(1);
     if (is_sparse_) {
       if (a_type_ == DT_BFLOAT16) {
         a = ctx->builder()->ConvertElementType(a, xla::F32);
diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_band_part_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_band_part_op.cc
index faa415a97b..fbd5dc0fda 100644
--- a/tensorflow/compiler/tf2xla/kernels/matrix_band_part_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/matrix_band_part_op.cc
@@ -44,10 +44,10 @@ class MatrixBandPartOp : public XlaOpKernel {
                 errors::InvalidArgument("num_upper must be scalar, got shape ",
                                         num_upper_in_shape.DebugString()));
 
-    xla::ComputationBuilder* builder = context->builder();
-    xla::ComputationDataHandle input = context->Input(0);
-    xla::ComputationDataHandle num_lower = context->Input(1);
-    xla::ComputationDataHandle num_upper = context->Input(2);
+    xla::XlaBuilder* builder = context->builder();
+    xla::XlaOp input = context->Input(0);
+    xla::XlaOp num_lower = context->Input(1);
+    xla::XlaOp num_upper = context->Input(2);
     DataType input_type = context->input_type(0);
     DataType index_type = context->input_type(1);
 
@@ -58,10 +58,10 @@ class MatrixBandPartOp : public XlaOpKernel {
 
     // Compute 'offset', which is how many diagonals we are above/below the
     // diagonal.
-    xla::ComputationDataHandle iota_m;
+    xla::XlaOp iota_m;
     OP_REQUIRES_OK(context, XlaHelpers::Iota(builder, index_type, m, &iota_m));
 
-    xla::ComputationDataHandle iota_n;
+    xla::XlaOp iota_n;
     OP_REQUIRES_OK(context, XlaHelpers::Iota(builder, index_type, n, &iota_n));
 
     auto offset = builder->Sub(builder->Broadcast(iota_n, {m}), iota_m,
diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_set_diag_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_set_diag_op.cc
index b2940bdcff..db53f6fef8 100644
--- a/tensorflow/compiler/tf2xla/kernels/matrix_set_diag_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/matrix_set_diag_op.cc
@@ -54,16 +54,16 @@ class MatrixSetDiagOp : public XlaOpKernel {
                     input_shape.DebugString(),
                     " and diagonal shape: ", diag_shape.DebugString()));
 
-    xla::ComputationBuilder* builder = context->builder();
-    xla::ComputationDataHandle input = context->Input(0);
-    xla::ComputationDataHandle diag = context->Input(1);
+    xla::XlaBuilder* builder = context->builder();
+    xla::XlaOp input = context->Input(0);
+    xla::XlaOp diag = context->Input(1);
 
     auto zero = XlaHelpers::Zero(builder, context->input_type(0));
 
     // Create an indicator tensor that is true only on the diagonal.
-    xla::ComputationDataHandle iota_m;
+    xla::XlaOp iota_m;
     OP_REQUIRES_OK(context, XlaHelpers::Iota(builder, DT_INT32, m, &iota_m));
-    xla::ComputationDataHandle iota_n;
+    xla::XlaOp iota_n;
     OP_REQUIRES_OK(context, XlaHelpers::Iota(builder, DT_INT32, n, &iota_n));
     auto indicator = builder->Eq(iota_m,
                                  builder->Broadcast(iota_n, {m}),
diff --git a/tensorflow/compiler/tf2xla/kernels/mirror_pad_op.cc b/tensorflow/compiler/tf2xla/kernels/mirror_pad_op.cc
index 05a36a031a..7e9de3ef9b 100644
--- a/tensorflow/compiler/tf2xla/kernels/mirror_pad_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/mirror_pad_op.cc
@@ -25,10 +25,11 @@ class MirrorPadOp : public XlaOpKernel {
  public:
   explicit MirrorPadOp(OpKernelConstruction* context) : XlaOpKernel(context) {}
 
-  xla::StatusOr<xla::ComputationDataHandle> DoMirrorPad(
-      const xla::ComputationDataHandle& t, const xla::Shape& original_shape,
-      const xla::Literal& pad_literal, xla::ComputationBuilder* b) {
-    xla::ComputationDataHandle accum = t;
+  xla::StatusOr<xla::XlaOp> DoMirrorPad(const xla::XlaOp& t,
+                                        const xla::Shape& original_shape,
+                                        const xla::Literal& pad_literal,
+                                        xla::XlaBuilder* b) {
+    xla::XlaOp accum = t;
     for (int64 dimno = xla::ShapeUtil::Rank(original_shape) - 1; dimno >= 0;
          --dimno) {
       auto t_rev = b->Rev(accum, {dimno});
@@ -76,12 +77,12 @@ class MirrorPadOp : public XlaOpKernel {
     OP_REQUIRES_OK(
         ctx, ctx->ConstantInputReshaped(1, {fixed_dims, 2}, &pad_literal));
 
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
     auto in0 = ctx->Input(0);
-    xla::StatusOr<std::unique_ptr<xla::Shape>> in0_shape = b->GetShape(in0);
+    xla::StatusOr<xla::Shape> in0_shape = b->GetShape(in0);
     OP_REQUIRES(ctx, in0_shape.ok(), in0_shape.status());
-    xla::StatusOr<xla::ComputationDataHandle> accum_status =
-        DoMirrorPad(in0, *in0_shape.ValueOrDie(), pad_literal, b);
+    xla::StatusOr<xla::XlaOp> accum_status =
+        DoMirrorPad(in0, in0_shape.ValueOrDie(), pad_literal, b);
 
     OP_REQUIRES_OK(ctx, accum_status.status());
 
diff --git a/tensorflow/compiler/tf2xla/kernels/one_hot_op.cc b/tensorflow/compiler/tf2xla/kernels/one_hot_op.cc
index 9f7c991380..cac2eea96e 100644
--- a/tensorflow/compiler/tf2xla/kernels/one_hot_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/one_hot_op.cc
@@ -62,7 +62,7 @@ class OneHotOp : public XlaOpKernel {
         ctx, depth >= 0,
         errors::InvalidArgument("depth must be non-negative, got: ", depth));
 
-    xla::ComputationDataHandle one_hot;
+    xla::XlaOp one_hot;
     OP_REQUIRES_OK(
         ctx, XlaHelpers::OneHot(ctx->builder(), depth, axis, input_type(0),
                                 indices_shape, ctx->Input(0), ctx->Input(2),
diff --git a/tensorflow/compiler/tf2xla/kernels/pack_op.cc b/tensorflow/compiler/tf2xla/kernels/pack_op.cc
index a4318e29d2..aecaabb6dc 100644
--- a/tensorflow/compiler/tf2xla/kernels/pack_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/pack_op.cc
@@ -43,7 +43,7 @@ class PackOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    std::vector<xla::ComputationDataHandle> values;
+    std::vector<xla::XlaOp> values;
     std::vector<TensorShape> shapes;
     OP_REQUIRES_OK(ctx, ctx->InputList("values", &values, &shapes));
     const int num = values.size();
@@ -69,7 +69,7 @@ class PackOp : public XlaOpKernel {
                                         -expanded_num_dims, ", ",
                                         expanded_num_dims, ")"));
 
-    std::vector<xla::ComputationDataHandle> reshaped_inputs(num);
+    std::vector<xla::XlaOp> reshaped_inputs(num);
 
     TensorShape child_shape(shapes[0]);
     child_shape.InsertDim(axis, 1);
diff --git a/tensorflow/compiler/tf2xla/kernels/pad_op.cc b/tensorflow/compiler/tf2xla/kernels/pad_op.cc
index 791351637a..7c95475e7b 100644
--- a/tensorflow/compiler/tf2xla/kernels/pad_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/pad_op.cc
@@ -70,7 +70,7 @@ class PadOp : public XlaOpKernel {
     }
 
     // PadV2 added a "constant_values" input that indicates the pad value.
-    xla::ComputationDataHandle constant_values;
+    xla::XlaOp constant_values;
     if (ctx->num_inputs() == 3) {
       OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(ctx->InputShape(2)),
                   errors::InvalidArgument("constant_values must be a scalar."));
diff --git a/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc
index 5f635dd1bc..f8e7b48a0f 100644
--- a/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc
@@ -66,15 +66,15 @@ class PoolingOp : public XlaOpKernel {
   int num_dims() const { return num_spatial_dims_ + 2; }
 
   // Method that builds an initial value to use in reductions.
-  virtual xla::ComputationDataHandle InitValue(xla::ComputationBuilder* b) = 0;
+  virtual xla::XlaOp InitValue(xla::XlaBuilder* b) = 0;
 
   // The reduction operation to apply to each window.
-  virtual const xla::Computation* Reduction(XlaOpKernelContext* ctx) = 0;
+  virtual const xla::XlaComputation* Reduction(XlaOpKernelContext* ctx) = 0;
 
   // A post-processing operation to apply on the outputs of the ReduceWindow.
-  virtual xla::ComputationDataHandle PostProcessOutput(
-      XlaOpKernelContext* ctx, const xla::ComputationDataHandle& output,
-      DataType dtype, const TensorShape& input_shape) = 0;
+  virtual xla::XlaOp PostProcessOutput(XlaOpKernelContext* ctx,
+                                       const xla::XlaOp& output, DataType dtype,
+                                       const TensorShape& input_shape) = 0;
 
   void Compile(XlaOpKernelContext* ctx) override {
     std::vector<int64> ksize = ksize_;
@@ -110,7 +110,7 @@ class PoolingOp : public XlaOpKernel {
                                         " operator must have ", num_dims(),
                                         " dimensions"));
 
-    xla::ComputationBuilder* const b = ctx->builder();
+    xla::XlaBuilder* const b = ctx->builder();
     auto input =
         XlaHelpers::ConvertElementType(b, ctx->Input(0), reduction_type_);
     auto reduce = ctx->builder()->ReduceWindow(
@@ -135,17 +135,17 @@ class MaxPoolOp : public PoolingOp {
       : PoolingOp(ctx, /*num_spatial_dims=*/num_spatial_dims,
                   /*reduction_type=*/ctx->input_type(0)) {}
 
-  xla::ComputationDataHandle InitValue(xla::ComputationBuilder* b) override {
+  xla::XlaOp InitValue(xla::XlaBuilder* b) override {
     return XlaHelpers::MinValue(b, reduction_type_);
   }
 
-  const xla::Computation* Reduction(XlaOpKernelContext* ctx) override {
+  const xla::XlaComputation* Reduction(XlaOpKernelContext* ctx) override {
     return ctx->GetOrCreateMax(reduction_type_);
   }
 
-  xla::ComputationDataHandle PostProcessOutput(
-      XlaOpKernelContext* ctx, const xla::ComputationDataHandle& output,
-      DataType dtype, const TensorShape& input_shape) override {
+  xla::XlaOp PostProcessOutput(XlaOpKernelContext* ctx,
+                               const xla::XlaOp& output, DataType dtype,
+                               const TensorShape& input_shape) override {
     return output;
   }
 };
@@ -176,9 +176,9 @@ REGISTER_XLA_OP(Name("MaxPool3D"), MaxPool3DOp);
 // Common computation shared between AvgPool and AvgPoolGrad. Divide each
 // element of an image by the count of elements that contributed to that
 // element during pooling.
-static xla::ComputationDataHandle AvgPoolDivideByCount(
-    XlaOpKernelContext* ctx, const xla::ComputationDataHandle& output,
-    DataType dtype, const TensorShape& input_shape, xla::Padding padding,
+static xla::XlaOp AvgPoolDivideByCount(
+    XlaOpKernelContext* ctx, const xla::XlaOp& output, DataType dtype,
+    const TensorShape& input_shape, xla::Padding padding,
     const std::vector<int64>& ksize, const std::vector<int64>& stride,
     int num_spatial_dims, TensorFormat data_format) {
   if (padding == xla::Padding::kValid) {
@@ -234,17 +234,17 @@ class AvgPoolOp : public PoolingOp {
                   /*reduction_type=*/
                   XlaHelpers::SumAccumulationType(ctx->input_type(0))) {}
 
-  xla::ComputationDataHandle InitValue(xla::ComputationBuilder* b) override {
+  xla::XlaOp InitValue(xla::XlaBuilder* b) override {
     return XlaHelpers::Zero(b, reduction_type_);
   }
 
-  const xla::Computation* Reduction(XlaOpKernelContext* ctx) override {
+  const xla::XlaComputation* Reduction(XlaOpKernelContext* ctx) override {
     return ctx->GetOrCreateAdd(reduction_type_);
   }
 
-  xla::ComputationDataHandle PostProcessOutput(
-      XlaOpKernelContext* ctx, const xla::ComputationDataHandle& output,
-      DataType dtype, const TensorShape& input_shape) override {
+  xla::XlaOp PostProcessOutput(XlaOpKernelContext* ctx,
+                               const xla::XlaOp& output, DataType dtype,
+                               const TensorShape& input_shape) override {
     return AvgPoolDivideByCount(ctx, output, dtype, input_shape, padding_,
                                 ksize_, stride_, num_spatial_dims_,
                                 data_format_);
@@ -344,11 +344,10 @@ class MaxPoolGradOp : public XlaOpKernel {
 
     xla::PrimitiveType element_type;
     OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(input_type(2), &element_type));
-    xla::ComputationDataHandle init_value =
-        XlaHelpers::Zero(ctx->builder(), input_type(2));
+    xla::XlaOp init_value = XlaHelpers::Zero(ctx->builder(), input_type(2));
     auto select = CreateScalarGeComputation(element_type, ctx->builder());
     auto scatter = CreateScalarAddComputation(element_type, ctx->builder());
-    xla::ComputationDataHandle gradients = ctx->builder()->SelectAndScatter(
+    xla::XlaOp gradients = ctx->builder()->SelectAndScatter(
         input, select, ksize_, stride_, xla_padding, out_backprop, init_value,
         scatter);
 
@@ -462,7 +461,7 @@ class AvgPoolGradOp : public XlaOpKernel {
     // The input gradients are computed by a convolution of the output gradients
     // and the filter, with some appropriate padding. See the comment at the top
     // of conv_grad_ops.h for details.
-    xla::ComputationBuilder* const b = ctx->builder();
+    xla::XlaBuilder* const b = ctx->builder();
     auto out_backprop = ctx->Input(1);
     auto dtype = input_type(1);
     xla::Padding xla_padding =
diff --git a/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc b/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc
index 4171e076ff..661cd5923e 100644
--- a/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc
@@ -35,7 +35,7 @@ class QuantizeAndDequantizeOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationDataHandle input = ctx->Input(0);
+    xla::XlaOp input = ctx->Input(0);
     const DataType data_type = ctx->input_type(0);
 
     // Comments taken from semantics description at
@@ -46,8 +46,8 @@ class QuantizeAndDequantizeOp : public XlaOpKernel {
     // m = max(abs(input_min), abs(input_max)) if range_given is true,
     // m = max(abs(min_elem(input)),
     //         abs(max_elem(input))) otherwise.
-    xla::ComputationBuilder* b = ctx->builder();
-    xla::ComputationDataHandle input_min, input_max;
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp input_min, input_max;
     if (range_given_) {
       double input_min_value, input_max_value;
       OP_REQUIRES_OK(ctx, ctx->ConstantInputAsFloatScalar(1, &input_min_value));
@@ -55,14 +55,14 @@ class QuantizeAndDequantizeOp : public XlaOpKernel {
       input_min = XlaHelpers::FloatLiteral(b, data_type, input_min_value);
       input_max = XlaHelpers::FloatLiteral(b, data_type, input_max_value);
     } else {
-      const xla::Computation* fmax = ctx->GetOrCreateMax(data_type);
-      const xla::Computation* fmin = ctx->GetOrCreateMin(data_type);
+      const xla::XlaComputation* fmax = ctx->GetOrCreateMax(data_type);
+      const xla::XlaComputation* fmin = ctx->GetOrCreateMin(data_type);
       input_min =
           b->ReduceAll(input, XlaHelpers::MaxValue(b, data_type), *fmin);
       input_max =
           b->ReduceAll(input, XlaHelpers::MinValue(b, data_type), *fmax);
     }
-    xla::ComputationDataHandle m = b->Max(b->Abs(input_min), b->Abs(input_max));
+    xla::XlaOp m = b->Max(b->Abs(input_min), b->Abs(input_max));
 
     // Next, we choose our fixed-point quantization buckets, [min_fixed,
     // max_fixed]. If signed_input is true, this is
@@ -85,7 +85,7 @@ class QuantizeAndDequantizeOp : public XlaOpKernel {
     // From this we compute our scaling factor, s:
     //
     // s = (max_fixed - min_fixed) / (2 * m).
-    xla::ComputationDataHandle s =
+    xla::XlaOp s =
         b->Div(XlaHelpers::FloatLiteral(b, data_type, max_fixed - min_fixed),
                b->Mul(XlaHelpers::FloatLiteral(b, data_type, 2.0), m));
 
@@ -93,7 +93,7 @@ class QuantizeAndDequantizeOp : public XlaOpKernel {
     // e is transformed into e':
     //
     // e' = (e * s).round_to_nearest() / s.
-    xla::ComputationDataHandle result = b->Div(b->Round(b->Mul(input, s)), s);
+    xla::XlaOp result = b->Div(b->Round(b->Mul(input, s)), s);
 
     ctx->SetOutput(0, result);
   }
diff --git a/tensorflow/compiler/tf2xla/kernels/random_ops.cc b/tensorflow/compiler/tf2xla/kernels/random_ops.cc
index c0994c434b..5f5bd58637 100644
--- a/tensorflow/compiler/tf2xla/kernels/random_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/random_ops.cc
@@ -41,9 +41,9 @@ class RandomUniformOp : public XlaOpKernel {
     xla::Shape xla_shape;
     OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype, shape, &xla_shape));
 
-    xla::ComputationBuilder* b = ctx->builder();
-    xla::ComputationDataHandle result = b->RngUniform(
-        XlaHelpers::Zero(b, dtype), XlaHelpers::One(b, dtype), xla_shape);
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp result = b->RngUniform(XlaHelpers::Zero(b, dtype),
+                                      XlaHelpers::One(b, dtype), xla_shape);
 
     ctx->SetOutput(0, result);
   }
@@ -100,11 +100,11 @@ class RandomStandardNormalOp : public XlaOpKernel {
     xla::Shape xla_shape;
     OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype, shape, &xla_shape));
 
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     // Normal distribution with a mean of 0 and a standard deviation of 1:
-    xla::ComputationDataHandle result = b->RngNormal(
-        XlaHelpers::Zero(b, dtype), XlaHelpers::One(b, dtype), xla_shape);
+    xla::XlaOp result = b->RngNormal(XlaHelpers::Zero(b, dtype),
+                                     XlaHelpers::One(b, dtype), xla_shape);
 
     ctx->SetOutput(0, result);
   }
@@ -130,19 +130,18 @@ class TruncatedNormalOp : public XlaOpKernel {
     xla::Shape xla_element_shape =
         xla::ShapeUtil::MakeShape(xla_shape.element_type(), {});
 
-    xla::ComputationBuilder* b = ctx->builder();
-    xla::ComputationDataHandle mean = XlaHelpers::Zero(b, dtype);
-    xla::ComputationDataHandle stddev = XlaHelpers::One(b, dtype);
-    xla::ComputationDataHandle candidate =
-        b->RngNormal(mean, stddev, xla_shape);
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp mean = XlaHelpers::Zero(b, dtype);
+    xla::XlaOp stddev = XlaHelpers::One(b, dtype);
+    xla::XlaOp candidate = b->RngNormal(mean, stddev, xla_shape);
 
-    auto two_sd = [dtype](bool negate, xla::ComputationBuilder* b) {
+    auto two_sd = [dtype](bool negate, xla::XlaBuilder* b) {
       return XlaHelpers::FloatLiteral(b, dtype, negate ? -2.0 : 2.0);
     };
-    auto out_of_range_mask = [two_sd](xla::ComputationDataHandle candidate,
-                                      xla::ComputationBuilder* b) {
-      xla::ComputationDataHandle too_large = b->Gt(candidate, two_sd(false, b));
-      xla::ComputationDataHandle too_small = b->Lt(candidate, two_sd(true, b));
+    auto out_of_range_mask = [two_sd](xla::XlaOp candidate,
+                                      xla::XlaBuilder* b) {
+      xla::XlaOp too_large = b->Gt(candidate, two_sd(false, b));
+      xla::XlaOp too_small = b->Lt(candidate, two_sd(true, b));
       return b->Or(too_large, too_small);
     };
 
@@ -152,35 +151,32 @@ class TruncatedNormalOp : public XlaOpKernel {
     //   out_of_range_mask := candidate < mean-2*sd || candidate > mean+2*sd
     //   candidate = select(out_of_range_mask, rng_normal(), candidate)
     // }
-    std::unique_ptr<xla::ComputationBuilder> test_builder =
+    std::unique_ptr<xla::XlaBuilder> test_builder =
         b->CreateSubBuilder("truncated_normal_test");
     {
       auto* b = test_builder.get();
-      xla::ComputationDataHandle candidate =
-          b->Parameter(0, xla_shape, "candidate");
-      xla::ComputationDataHandle oor_mask = out_of_range_mask(candidate, b);
+      xla::XlaOp candidate = b->Parameter(0, xla_shape, "candidate");
+      out_of_range_mask(candidate, b);
       OP_REQUIRES_OK(ctx, Any(out_of_range_mask(candidate, b), b).status());
     }
 
-    std::unique_ptr<xla::ComputationBuilder> body_builder =
+    std::unique_ptr<xla::XlaBuilder> body_builder =
         b->CreateSubBuilder("truncated_normal_body");
     {
       auto* b = body_builder.get();
-      xla::ComputationDataHandle candidate =
-          b->Parameter(0, xla_shape, "candidate");
-      xla::ComputationDataHandle to_resample = out_of_range_mask(candidate, b);
-      xla::ComputationDataHandle mean = XlaHelpers::Zero(b, dtype);
-      xla::ComputationDataHandle stddev = XlaHelpers::One(b, dtype);
+      xla::XlaOp candidate = b->Parameter(0, xla_shape, "candidate");
+      xla::XlaOp to_resample = out_of_range_mask(candidate, b);
+      xla::XlaOp mean = XlaHelpers::Zero(b, dtype);
+      xla::XlaOp stddev = XlaHelpers::One(b, dtype);
       b->Select(to_resample, b->RngNormal(mean, stddev, xla_shape), candidate);
     }
 
-    xla::StatusOr<xla::Computation> test_computation = test_builder->Build();
+    xla::StatusOr<xla::XlaComputation> test_computation = test_builder->Build();
     OP_REQUIRES_OK(ctx, test_computation.status());
-    xla::StatusOr<xla::Computation> body_computation = body_builder->Build();
+    xla::StatusOr<xla::XlaComputation> body_computation = body_builder->Build();
     OP_REQUIRES_OK(ctx, body_computation.status());
-    xla::ComputationDataHandle result =
-        b->While(test_computation.ValueOrDie(), body_computation.ValueOrDie(),
-                 candidate);
+    xla::XlaOp result = b->While(test_computation.ValueOrDie(),
+                                 body_computation.ValueOrDie(), candidate);
 
     ctx->SetOutput(0, result);
   }
diff --git a/tensorflow/compiler/tf2xla/kernels/reduce_window_op.cc b/tensorflow/compiler/tf2xla/kernels/reduce_window_op.cc
index cb144bea9e..08894489ac 100644
--- a/tensorflow/compiler/tf2xla/kernels/reduce_window_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/reduce_window_op.cc
@@ -19,7 +19,6 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
 #include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/op_kernel.h"
 
@@ -65,7 +64,7 @@ class ReduceWindowOp : public XlaOpKernel {
                     "rank (",
                     padding_high_.size(), " vs. ", rank, ")"));
 
-    xla::ComputationBuilder* builder = context->builder();
+    xla::XlaBuilder* builder = context->builder();
 
     // Build the reducer function.
     XlaCompiler::Argument reducer_arg;
@@ -95,15 +94,15 @@ class ReduceWindowOp : public XlaOpKernel {
                     xla::ShapeUtil::HumanString(reducer.xla_output_shape)));
 
     // Wraps the reducer in a computation that unpacks the output tuple.
-    xla::Computation wrapper;
+    xla::XlaComputation wrapper;
     {
-      std::unique_ptr<xla::ComputationBuilder> cb =
+      std::unique_ptr<xla::XlaBuilder> cb =
           builder->CreateSubBuilder("wrapper");
       auto x = cb->Parameter(0, scalar_shape, "x");
       auto y = cb->Parameter(1, scalar_shape, "y");
       auto outputs = cb->Call(*reducer.computation, {x, y});
       cb->GetTupleElement(outputs, 0);
-      xla::StatusOr<xla::Computation> result = cb->Build();
+      xla::StatusOr<xla::XlaComputation> result = cb->Build();
       OP_REQUIRES_OK(context, result.status());
       wrapper = std::move(result.ValueOrDie());
     }
@@ -113,7 +112,7 @@ class ReduceWindowOp : public XlaOpKernel {
       padding[i] = {padding_low_[i], padding_high_[i]};
     }
 
-    xla::ComputationDataHandle output = builder->ReduceWindowWithGeneralPadding(
+    xla::XlaOp output = builder->ReduceWindowWithGeneralPadding(
         context->Input(0), context->Input(1), wrapper, window_dimensions_,
         window_strides_, padding);
     context->SetOutput(0, output);
diff --git a/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc b/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc
index 812d258cd1..0f42563779 100644
--- a/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc
@@ -30,13 +30,11 @@ class SumOp : public XlaReductionOp {
   explicit SumOp(OpKernelConstruction* ctx)
       : XlaReductionOp(ctx,
                        XlaHelpers::SumAccumulationType(ctx->input_type(0))) {}
-  xla::ComputationDataHandle InitialValue(
-      xla::ComputationBuilder* builder) override {
+  xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
     return XlaHelpers::Zero(builder, reduction_type_);
   }
-  void BuildReducer(xla::ComputationBuilder* builder,
-                    const xla::ComputationDataHandle& scalar_lhs,
-                    const xla::ComputationDataHandle& scalar_rhs) override {
+  void BuildReducer(xla::XlaBuilder* builder, const xla::XlaOp& scalar_lhs,
+                    const xla::XlaOp& scalar_rhs) override {
     builder->Add(scalar_lhs, scalar_rhs);
   }
 };
@@ -49,14 +47,12 @@ class ProdOp : public XlaReductionOp {
       : XlaReductionOp(ctx,
                        XlaHelpers::SumAccumulationType(ctx->input_type(0))) {}
 
-  xla::ComputationDataHandle InitialValue(
-      xla::ComputationBuilder* builder) override {
+  xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
     return XlaHelpers::One(builder, reduction_type_);
   }
 
-  void BuildReducer(xla::ComputationBuilder* builder,
-                    const xla::ComputationDataHandle& scalar_lhs,
-                    const xla::ComputationDataHandle& scalar_rhs) override {
+  void BuildReducer(xla::XlaBuilder* builder, const xla::XlaOp& scalar_lhs,
+                    const xla::XlaOp& scalar_rhs) override {
     builder->Mul(scalar_lhs, scalar_rhs);
   }
 };
@@ -69,14 +65,12 @@ class MinOp : public XlaReductionOp {
   explicit MinOp(OpKernelConstruction* ctx)
       : XlaReductionOp(ctx, ctx->input_type(0)) {}
 
-  xla::ComputationDataHandle InitialValue(
-      xla::ComputationBuilder* builder) override {
+  xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
     return XlaHelpers::MaxValue(builder, reduction_type_);
   }
 
-  void BuildReducer(xla::ComputationBuilder* builder,
-                    const xla::ComputationDataHandle& scalar_lhs,
-                    const xla::ComputationDataHandle& scalar_rhs) override {
+  void BuildReducer(xla::XlaBuilder* builder, const xla::XlaOp& scalar_lhs,
+                    const xla::XlaOp& scalar_rhs) override {
     builder->Min(scalar_lhs, scalar_rhs);
   }
 };
@@ -88,14 +82,12 @@ class MaxOp : public XlaReductionOp {
   explicit MaxOp(OpKernelConstruction* ctx)
       : XlaReductionOp(ctx, ctx->input_type(0)) {}
 
-  xla::ComputationDataHandle InitialValue(
-      xla::ComputationBuilder* builder) override {
+  xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
     return XlaHelpers::MinValue(builder, reduction_type_);
   }
 
-  void BuildReducer(xla::ComputationBuilder* builder,
-                    const xla::ComputationDataHandle& scalar_lhs,
-                    const xla::ComputationDataHandle& scalar_rhs) override {
+  void BuildReducer(xla::XlaBuilder* builder, const xla::XlaOp& scalar_lhs,
+                    const xla::XlaOp& scalar_rhs) override {
     builder->Max(scalar_lhs, scalar_rhs);
   }
 };
@@ -108,20 +100,17 @@ class MeanOp : public XlaReductionOp {
       : XlaReductionOp(ctx,
                        XlaHelpers::SumAccumulationType(ctx->input_type(0))) {}
 
-  xla::ComputationDataHandle InitialValue(
-      xla::ComputationBuilder* builder) override {
+  xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
     return XlaHelpers::Zero(builder, reduction_type_);
   }
-  void BuildReducer(xla::ComputationBuilder* builder,
-                    const xla::ComputationDataHandle& scalar_lhs,
-                    const xla::ComputationDataHandle& scalar_rhs) override {
+  void BuildReducer(xla::XlaBuilder* builder, const xla::XlaOp& scalar_lhs,
+                    const xla::XlaOp& scalar_rhs) override {
     builder->Add(scalar_lhs, scalar_rhs);
   }
 
-  xla::ComputationDataHandle BuildFinalizer(
-      xla::ComputationBuilder* builder,
-      const xla::ComputationDataHandle& reduce_output,
-      int64 num_elements_reduced) override {
+  xla::XlaOp BuildFinalizer(xla::XlaBuilder* builder,
+                            const xla::XlaOp& reduce_output,
+                            int64 num_elements_reduced) override {
     auto divisor = XlaHelpers::IntegerLiteral(builder, input_type(0),
                                               num_elements_reduced);
     return builder->Div(reduce_output, divisor);
@@ -136,14 +125,12 @@ class AllOp : public XlaReductionOp {
   explicit AllOp(OpKernelConstruction* ctx)
       : XlaReductionOp(ctx, ctx->input_type(0)) {}
 
-  xla::ComputationDataHandle InitialValue(
-      xla::ComputationBuilder* builder) override {
+  xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
     return builder->ConstantR0<bool>(true);
   }
 
-  void BuildReducer(xla::ComputationBuilder* builder,
-                    const xla::ComputationDataHandle& scalar_lhs,
-                    const xla::ComputationDataHandle& scalar_rhs) override {
+  void BuildReducer(xla::XlaBuilder* builder, const xla::XlaOp& scalar_lhs,
+                    const xla::XlaOp& scalar_rhs) override {
     builder->And(scalar_lhs, scalar_rhs);
   }
 };
@@ -155,14 +142,12 @@ class AnyOp : public XlaReductionOp {
   explicit AnyOp(OpKernelConstruction* ctx)
       : XlaReductionOp(ctx, ctx->input_type(0)) {}
 
-  xla::ComputationDataHandle InitialValue(
-      xla::ComputationBuilder* builder) override {
+  xla::XlaOp InitialValue(xla::XlaBuilder* builder) override {
     return builder->ConstantR0<bool>(false);
   }
 
-  void BuildReducer(xla::ComputationBuilder* builder,
-                    const xla::ComputationDataHandle& scalar_lhs,
-                    const xla::ComputationDataHandle& scalar_rhs) override {
+  void BuildReducer(xla::XlaBuilder* builder, const xla::XlaOp& scalar_lhs,
+                    const xla::XlaOp& scalar_rhs) override {
     builder->Or(scalar_lhs, scalar_rhs);
   }
 };
diff --git a/tensorflow/compiler/tf2xla/kernels/reduction_ops.h b/tensorflow/compiler/tf2xla/kernels/reduction_ops.h
index f3181f0dad..2ecfb854a1 100644
--- a/tensorflow/compiler/tf2xla/kernels/reduction_ops.h
+++ b/tensorflow/compiler/tf2xla/kernels/reduction_ops.h
@@ -19,7 +19,7 @@ limitations under the License.
 #define TENSORFLOW_COMPILER_TF2XLA_KERNELS_REDUCTION_OPS_H_
 
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/core/framework/op_kernel.h"
 
 namespace tensorflow {
@@ -28,35 +28,33 @@ namespace tensorflow {
 // to override: description is a textual description of the mapped
 // function; InitialValue constructs the base case for the reduction;
 // BuildReducer adds the implementation of the reduction lambda to a
-// xla::ComputationBuilder and BuildFinalizer adds the
+// xla::XlaBuilder and BuildFinalizer adds the
 // implementation of the finalizer lambda (if there is one) to a
-// xla::ComputationBuilder.
+// xla::XlaBuilder.
 class XlaReductionOp : public XlaOpKernel {
  public:
   XlaReductionOp(OpKernelConstruction* ctx, DataType reduction_type);
   ~XlaReductionOp() override {}
 
   // Return the base case for the reduction.
-  virtual xla::ComputationDataHandle InitialValue(
-      xla::ComputationBuilder* builder) = 0;
+  virtual xla::XlaOp InitialValue(xla::XlaBuilder* builder) = 0;
 
   // Implement the (scalar,scalar)->scalar lambda that should be
   // applied to each pair of elements to be reduced. The desired
   // computation should be added to 'builder' and
   // '(scalar_lhs,scalar_rhs)' are the function's inputs.
-  virtual void BuildReducer(xla::ComputationBuilder* builder,
-                            const xla::ComputationDataHandle& scalar_lhs,
-                            const xla::ComputationDataHandle& scalar_rhs) = 0;
+  virtual void BuildReducer(xla::XlaBuilder* builder,
+                            const xla::XlaOp& scalar_lhs,
+                            const xla::XlaOp& scalar_rhs) = 0;
 
   // Applies a transformation to the output of the reduction. The desired
   // computation should be added to 'builder'. Argument 'reduce_output' is the
   // output of the reduction. 'num_elements_reduced' is the number of elements
   // that contributed to the reduction. Returns the transformed reduction
   // output, Defaults to returning 'reduce_output' unchanged.
-  virtual xla::ComputationDataHandle BuildFinalizer(
-      xla::ComputationBuilder* builder,
-      const xla::ComputationDataHandle& reduce_output,
-      int64 num_elements_reduced);
+  virtual xla::XlaOp BuildFinalizer(xla::XlaBuilder* builder,
+                                    const xla::XlaOp& reduce_output,
+                                    int64 num_elements_reduced);
 
   void Compile(XlaOpKernelContext* ctx) override;
 
diff --git a/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc b/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc
index 64fe765ae9..4fd5bfd039 100644
--- a/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc
+++ b/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc
@@ -35,10 +35,9 @@ XlaReductionOp::XlaReductionOp(OpKernelConstruction* ctx,
 
 // Unless BuildFinalizer is overridden the reduction has no
 // finalizer.
-xla::ComputationDataHandle XlaReductionOp::BuildFinalizer(
-    xla::ComputationBuilder* builder,
-    const xla::ComputationDataHandle& reduce_output,
-    int64 num_elements_reduced) {
+xla::XlaOp XlaReductionOp::BuildFinalizer(xla::XlaBuilder* builder,
+                                          const xla::XlaOp& reduce_output,
+                                          int64 num_elements_reduced) {
   return reduce_output;
 }
 
@@ -96,9 +95,9 @@ void XlaReductionOp::Compile(XlaOpKernelContext* ctx) {
 
   string desc = ctx->op_kernel().name();
 
-  xla::ComputationBuilder* const b = ctx->builder();
+  xla::XlaBuilder* const b = ctx->builder();
   // Construct the builder for the reduction lambda.
-  xla::ComputationBuilder r(b->client(), strings::StrCat(desc, "-reduction"));
+  xla::XlaBuilder r(strings::StrCat(desc, "-reduction"));
   xla::PrimitiveType type;
   TF_CHECK_OK(DataTypeToPrimitiveType(reduction_type_, &type));
 
@@ -110,7 +109,7 @@ void XlaReductionOp::Compile(XlaOpKernelContext* ctx) {
   auto ry = r.Parameter(1, xla::ShapeUtil::MakeShape(type, {}), "y");
   // Call virtual method to build the reduction lambda.
   BuildReducer(&r, rx, ry);
-  xla::Computation reduction_computation = r.Build().ConsumeValueOrDie();
+  xla::XlaComputation reduction_computation = r.Build().ConsumeValueOrDie();
 
   auto reduce = b->Reduce(data, initial, reduction_computation, xla_axes);
   auto deconverted = XlaHelpers::ConvertElementType(b, reduce, input_type(0));
diff --git a/tensorflow/compiler/tf2xla/kernels/relu_op.cc b/tensorflow/compiler/tf2xla/kernels/relu_op.cc
index 12a3552999..ba7d484d53 100644
--- a/tensorflow/compiler/tf2xla/kernels/relu_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/relu_op.cc
@@ -18,7 +18,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/kernels/cwise_ops.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/core/framework/kernel_def_builder.h"
 #include "tensorflow/core/framework/types.h"
@@ -32,7 +32,7 @@ class ReluOp : public XlaOpKernel {
   explicit ReluOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
   // Computes the max of the scalar input x and 0.
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
     auto zero = XlaHelpers::Zero(builder, input_type(0));
     ctx->SetOutput(0, builder->Max(zero, ctx->Input(0)));
   }
@@ -43,7 +43,7 @@ class Relu6Op : public XlaOpKernel {
   explicit Relu6Op(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
   // Clamp the scalar input between 0 and 6.
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
     auto zero = XlaHelpers::Zero(builder, input_type(0));
     auto six = XlaHelpers::IntegerLiteral(builder, input_type(0), 6);
     ctx->SetOutput(0, builder->Clamp(zero, ctx->Input(0), six));
@@ -56,7 +56,7 @@ class ReluGradOp : public XlaOpKernel {
   // Return the lhs (incoming gradient) if the rhs (input feature) > 0,
   // otherwise return 0.
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
     const TensorShape shape = ctx->InputShape(0);
     const auto zero =
         b->Broadcast(XlaHelpers::Zero(b, input_type(0)), shape.dim_sizes());
@@ -71,7 +71,7 @@ class Relu6GradOp : public XlaOpKernel {
   // Return the lhs (incoming gradient) if the rhs (input feature) > 0,
   // otherwise return 0.
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
     const TensorShape shape = ctx->InputShape(0);
     const auto zero =
         b->Broadcast(XlaHelpers::Zero(b, input_type(0)), shape.dim_sizes());
diff --git a/tensorflow/compiler/tf2xla/kernels/retval_op.cc b/tensorflow/compiler/tf2xla/kernels/retval_op.cc
index c283e3b02c..70547290ea 100644
--- a/tensorflow/compiler/tf2xla/kernels/retval_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/retval_op.cc
@@ -16,7 +16,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/xla_context.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/core/framework/kernel_def_builder.h"
 #include "tensorflow/core/framework/op_kernel.h"
 
@@ -45,7 +45,7 @@ class RetvalOp : public XlaOpKernel {
       // compilation.
       OP_REQUIRES_OK(ctx, frame->SetRetval(index_, input));
     } else {
-      xla::ComputationDataHandle input = ctx->Input(0);
+      xla::XlaOp input = ctx->Input(0);
       const TensorShape input_shape = ctx->InputShape(0);
 
       auto is_constant = ctx->builder()->IsConstant(input);
diff --git a/tensorflow/compiler/tf2xla/kernels/reverse_op.cc b/tensorflow/compiler/tf2xla/kernels/reverse_op.cc
index e51d386926..2872a3c4d4 100644
--- a/tensorflow/compiler/tf2xla/kernels/reverse_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/reverse_op.cc
@@ -48,7 +48,7 @@ class ReverseOp : public XlaOpKernel {
       ctx->SetOutput(0, ctx->Input(0));
       return;
     }
-    // ComputationBuilder::Rev() requires concrete values for dimensions arg.
+    // XlaBuilder::Rev() requires concrete values for dimensions arg.
     xla::Literal lax;
     OP_REQUIRES_OK(ctx, ctx->ConstantInputReshaped(1, {x_shape.dims()}, &lax));
     std::vector<bool> revdims(x_shape.dims());
@@ -90,7 +90,7 @@ class ReverseV2Op : public XlaOpKernel {
       ctx->SetOutput(0, ctx->Input(0));
       return;
     }
-    // ComputationBuilder::Rev() requires concrete values for dimensions arg.
+    // XlaBuilder::Rev() requires concrete values for dimensions arg.
     std::vector<int64> axes;
     OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &axes));
 
diff --git a/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc b/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc
index 6bc5d3adb0..0ed4c4707d 100644
--- a/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc
@@ -54,7 +54,7 @@ class ReverseSequenceOp : public XlaOpKernel {
                                 "), ", "(", seq_lens_shape.num_elements(),
                                 " vs. ", input_shape.dim_size(batch_dim_)));
 
-    xla::ComputationBuilder* builder = context->builder();
+    xla::XlaBuilder* builder = context->builder();
     const auto input = context->Input(0);
     const auto seq_lens = context->Input(1);
 
@@ -155,7 +155,7 @@ class ReverseSequenceOp : public XlaOpKernel {
     auto output = builder->GetTupleElement(loop_output, 2);
 
     // Mask out elements after the sequence length.
-    xla::ComputationDataHandle iota;
+    xla::XlaOp iota;
     OP_REQUIRES_OK(
         context, XlaHelpers::Iota(builder, seq_lens_type, max_seq_len, &iota));
     std::vector<int64> dims(input_shape.dims(), 1);
diff --git a/tensorflow/compiler/tf2xla/kernels/scan_ops.cc b/tensorflow/compiler/tf2xla/kernels/scan_ops.cc
index 4cfa28a0ce..1819fb5433 100644
--- a/tensorflow/compiler/tf2xla/kernels/scan_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/scan_ops.cc
@@ -74,7 +74,7 @@ class ScanOp : public XlaOpKernel {
       return;
     }
 
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
 
     std::vector<int64> window_strides(input_shape.dims(), 1);
     std::vector<int64> window_dims(input_shape.dims(), 1);
@@ -91,8 +91,8 @@ class ScanOp : public XlaOpKernel {
       std::swap(padding[axis].first, padding[axis].second);
     }
 
-    xla::ComputationDataHandle init;
-    const xla::Computation* reducer;
+    xla::XlaOp init;
+    const xla::XlaComputation* reducer;
     if (sum_) {
       init = XlaHelpers::Zero(builder, dtype);
       reducer = ctx->GetOrCreateAdd(dtype);
diff --git a/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc b/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc
index 8433a29c4e..f2c63b4f90 100644
--- a/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc
@@ -102,7 +102,7 @@ class ScatterNdOp : public XlaOpKernel {
     OP_REQUIRES_OK(context, ValidateUpdateShape(buffer_shape, indices_shape,
                                                 updates_shape));
 
-    xla::ComputationBuilder* builder = context->builder();
+    xla::XlaBuilder* builder = context->builder();
     auto buffer = builder->Broadcast(XlaHelpers::Zero(builder, dtype),
                                      buffer_shape.dim_sizes());
     auto indices = context->Input(0);
diff --git a/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc b/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc
index 498342a988..664078ca16 100644
--- a/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc
@@ -17,7 +17,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 
 namespace tensorflow {
 namespace {
@@ -62,16 +62,16 @@ class UnsortedSegmentSum : public XlaOpKernel {
                       d, " differs ", data_shape.dim_size(d), " vs. ",
                       indices_shape.dim_size(d)));
     }
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
     TensorShape buffer_shape = data_shape;
     buffer_shape.RemoveDimRange(0, indices_shape.dims());
     buffer_shape.InsertDim(0, num_segments);
     auto buffer = builder->Broadcast(XlaHelpers::Zero(builder, dtype_),
                                      buffer_shape.dim_sizes());
 
-    auto combiner =
-        [](xla::ComputationDataHandle a, xla::ComputationDataHandle b,
-           xla::ComputationBuilder* builder) { return builder->Add(a, b); };
+    auto combiner = [](xla::XlaOp a, xla::XlaOp b, xla::XlaBuilder* builder) {
+      return builder->Add(a, b);
+    };
 
     auto result = XlaScatter(buffer, /*updates=*/data, indices,
                              /*indices_are_vectors=*/false, combiner, builder);
diff --git a/tensorflow/compiler/tf2xla/kernels/select_op.cc b/tensorflow/compiler/tf2xla/kernels/select_op.cc
index 8081d3c41c..f9f48164d6 100644
--- a/tensorflow/compiler/tf2xla/kernels/select_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/select_op.cc
@@ -40,7 +40,7 @@ class SelectOp : public XlaOpKernel {
             "'then' and 'else' must have the same size.  but received: ",
             then_shape.DebugString(), " vs. ", else_shape.DebugString()));
 
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
 
     auto cond_handle = ctx->Input(0);
     auto then_handle = ctx->Input(1);
diff --git a/tensorflow/compiler/tf2xla/kernels/sendrecv_ops.cc b/tensorflow/compiler/tf2xla/kernels/sendrecv_ops.cc
index d079b89861..9ce01d0d44 100644
--- a/tensorflow/compiler/tf2xla/kernels/sendrecv_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/sendrecv_ops.cc
@@ -18,7 +18,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
 #include "tensorflow/core/framework/kernel_def_builder.h"
 #include "tensorflow/core/framework/types.h"
diff --git a/tensorflow/compiler/tf2xla/kernels/softmax_op.cc b/tensorflow/compiler/tf2xla/kernels/softmax_op.cc
index 463788b8b4..bbf5ee8b12 100644
--- a/tensorflow/compiler/tf2xla/kernels/softmax_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/softmax_op.cc
@@ -43,8 +43,8 @@ class SoftmaxOp : public XlaOpKernel {
     const DataType type = input_type(0);
     auto logits = ctx->Input(0);
 
-    xla::ComputationBuilder* const b = ctx->builder();
-    const xla::Computation& max_func = *ctx->GetOrCreateMax(type);
+    xla::XlaBuilder* const b = ctx->builder();
+    const xla::XlaComputation& max_func = *ctx->GetOrCreateMax(type);
 
     // Find the max in each batch, resulting in a tensor of shape [batch]
     auto logits_max =
@@ -76,16 +76,15 @@ class SoftmaxOp : public XlaOpKernel {
 REGISTER_XLA_OP(Name("Softmax"), SoftmaxOp);
 REGISTER_XLA_OP(Name("LogSoftmax"), SoftmaxOp);
 
-std::pair<xla::ComputationDataHandle, xla::ComputationDataHandle>
-CrossEntropyWithLogits(XlaOpKernelContext* ctx, DataType type,
-                       const xla::ComputationDataHandle& logits,
-                       const xla::ComputationDataHandle& labels) {
-  const xla::Computation& max_func = *ctx->GetOrCreateMax(type);
+std::pair<xla::XlaOp, xla::XlaOp> CrossEntropyWithLogits(
+    XlaOpKernelContext* ctx, DataType type, const xla::XlaOp& logits,
+    const xla::XlaOp& labels) {
+  const xla::XlaComputation& max_func = *ctx->GetOrCreateMax(type);
 
   const int kBatchDim = 0;
   const int kClassDim = 1;
 
-  xla::ComputationBuilder* b = ctx->builder();
+  xla::XlaBuilder* b = ctx->builder();
   // Find the max in each batch, resulting in a tensor of shape [batch]
   auto logits_max =
       b->Reduce(logits, XlaHelpers::MinValue(b, type), max_func, {kClassDim});
@@ -123,7 +122,7 @@ CrossEntropyWithLogits(XlaOpKernelContext* ctx, DataType type,
   // backprop: prob - labels, where
   //   prob = exp(logits - max_logits) / sum(exp(logits - max_logits))
   //     (where the division broadcasts along the batch dimension)
-  xla::ComputationDataHandle backprop =
+  xla::XlaOp backprop =
       b->Sub(b->Div(exp_shifted_logits, sum_exp, {kBatchDim}), labels);
   return {loss, backprop};
 }
@@ -150,7 +149,7 @@ class SoftmaxXentWithLogitsOp : public XlaOpKernel {
     auto logits = ctx->Input(0);
     auto labels = ctx->Input(1);
 
-    xla::ComputationDataHandle loss, backprop;
+    xla::XlaOp loss, backprop;
     std::tie(loss, backprop) =
         CrossEntropyWithLogits(ctx, type, logits, labels);
     ctx->SetOutput(0, loss);
@@ -191,10 +190,10 @@ class SparseSoftmaxXentWithLogitsOp : public XlaOpKernel {
     DataType logits_type = input_type(0);
     DataType indices_type = input_type(1);
 
-    xla::ComputationDataHandle indices = ctx->Input(1);
+    xla::XlaOp indices = ctx->Input(1);
 
-    xla::ComputationBuilder* builder = ctx->builder();
-    xla::ComputationDataHandle labels;
+    xla::XlaBuilder* builder = ctx->builder();
+    xla::XlaOp labels;
     OP_REQUIRES_OK(ctx,
                    XlaHelpers::OneHot(
                        builder, depth, /*axis=*/1, input_type(1), labels_shape,
@@ -207,7 +206,7 @@ class SparseSoftmaxXentWithLogitsOp : public XlaOpKernel {
     // Builds a vector of {batch_size} that is 0 if the index is in range, or
     // NaN otherwise; then add that vector to the labels to force out-of-range
     // values to NaNs.
-    xla::ComputationDataHandle nan_or_zero = builder->Select(
+    xla::XlaOp nan_or_zero = builder->Select(
         builder->And(
             builder->Le(XlaHelpers::Zero(builder, indices_type), indices),
             builder->Lt(indices, XlaHelpers::IntegerLiteral(
@@ -218,7 +217,7 @@ class SparseSoftmaxXentWithLogitsOp : public XlaOpKernel {
                            {batch_size}));
     labels = builder->Add(labels, nan_or_zero, {0});
 
-    xla::ComputationDataHandle loss, backprop;
+    xla::XlaOp loss, backprop;
     std::tie(loss, backprop) =
         CrossEntropyWithLogits(ctx, logits_type, ctx->Input(0), labels);
     ctx->SetOutput(0, loss);
diff --git a/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc b/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc
index 01b46e160d..ec077924b5 100644
--- a/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc
@@ -20,9 +20,8 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
-void SpaceToBatch(XlaOpKernelContext* ctx,
-                  const xla::ComputationDataHandle& input, DataType input_dtype,
-                  const TensorShape& input_tensor_shape,
+void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp& input,
+                  DataType input_dtype, const TensorShape& input_tensor_shape,
                   gtl::ArraySlice<int64> block_shape,
                   const xla::Literal& paddings) {
   const int input_rank = input_tensor_shape.dims();
@@ -46,7 +45,7 @@ void SpaceToBatch(XlaOpKernelContext* ctx,
                               ", 2] instead of ",
                               xla::ShapeUtil::HumanString(paddings.shape())));
 
-  xla::ComputationBuilder* b = ctx->builder();
+  xla::XlaBuilder* b = ctx->builder();
 
   // 1. Zero-pad the start and end of dimensions `[1, ..., M]` of the
   //  input according to `paddings` to produce `padded` of shape `padded_shape`.
@@ -73,7 +72,7 @@ void SpaceToBatch(XlaOpKernelContext* ctx,
               errors::InvalidArgument(
                   "The product of the block dimensions must be positive"));
 
-  xla::ComputationDataHandle padded =
+  xla::XlaOp padded =
       b->Pad(input, XlaHelpers::Zero(b, input_dtype), padding_config);
 
   // 2. Reshape `padded` to `reshaped_padded` of shape:
@@ -101,8 +100,7 @@ void SpaceToBatch(XlaOpKernelContext* ctx,
   std::copy(remainder_shape.begin(), remainder_shape.end(),
             reshaped_padded_shape.begin() + 1 + 2 * block_rank);
 
-  xla::ComputationDataHandle reshaped_padded =
-      b->Reshape(padded, reshaped_padded_shape);
+  xla::XlaOp reshaped_padded = b->Reshape(padded, reshaped_padded_shape);
 
   // 3. Permute dimensions of `reshaped_padded` to produce
   //    `permuted_reshaped_padded` of shape:
@@ -121,7 +119,7 @@ void SpaceToBatch(XlaOpKernelContext* ctx,
   permutation[block_rank] = 0;
   std::iota(permutation.begin() + 1 + block_rank * 2, permutation.end(),
             1 + block_rank * 2);
-  xla::ComputationDataHandle permuted_reshaped_padded =
+  xla::XlaOp permuted_reshaped_padded =
       b->Transpose(reshaped_padded, permutation);
 
   // 4. Reshape `permuted_reshaped_padded` to flatten `block_shape` into the
@@ -142,8 +140,7 @@ void SpaceToBatch(XlaOpKernelContext* ctx,
   std::copy(remainder_shape.begin(), remainder_shape.end(),
             output_shape.begin() + 1 + block_rank);
 
-  xla::ComputationDataHandle output =
-      b->Reshape(permuted_reshaped_padded, output_shape);
+  xla::XlaOp output = b->Reshape(permuted_reshaped_padded, output_shape);
   ctx->SetOutput(0, output);
 }
 
diff --git a/tensorflow/compiler/tf2xla/kernels/spacetodepth_op.cc b/tensorflow/compiler/tf2xla/kernels/spacetodepth_op.cc
index 806fda632c..4c5886ee2a 100644
--- a/tensorflow/compiler/tf2xla/kernels/spacetodepth_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/spacetodepth_op.cc
@@ -50,8 +50,8 @@ class SpaceToDepthOp : public XlaOpKernel {
     const gtl::InlinedVector<int64, 4> input_shape =
         input_tensor_shape.dim_sizes();
 
-    xla::ComputationBuilder* b = ctx->builder();
-    xla::ComputationDataHandle input = ctx->Input(0);
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp input = ctx->Input(0);
 
     int feature_dim = GetTensorFeatureDimIndex(input_rank, data_format_);
     int num_spatial_dims = GetTensorSpatialDims(input_rank, data_format_);
@@ -135,7 +135,7 @@ class SpaceToDepthOp : public XlaOpKernel {
     //       input_shape[1] / block_size_, block_size_,
     //       input_shape[2] / block_size_, block_size_,
     //       depth]
-    xla::ComputationDataHandle reshaped = b->Reshape(input, reshaped_shape);
+    xla::XlaOp reshaped = b->Reshape(input, reshaped_shape);
 
     // 2. Permute dimensions of `reshaped` to produce
     //    `permuted_reshaped` of shape:
@@ -145,8 +145,7 @@ class SpaceToDepthOp : public XlaOpKernel {
     //       input_shape[2] / block_size_,
     //       block_size_, block_size_,
     //       depth]
-    xla::ComputationDataHandle permuted_reshaped =
-        b->Transpose(reshaped, transpose_order);
+    xla::XlaOp permuted_reshaped = b->Transpose(reshaped, transpose_order);
 
     // 3. Reshape `permuted_reshaped` to flatten `block_shape` into the
     //    batch dimension, producing an output tensor of shape:
@@ -156,8 +155,7 @@ class SpaceToDepthOp : public XlaOpKernel {
     //       input_shape[2] / block_size_,
     //       block_size_ * block_size_ * depth]
     //
-    xla::ComputationDataHandle output =
-        b->Reshape(permuted_reshaped, output_shape);
+    xla::XlaOp output = b->Reshape(permuted_reshaped, output_shape);
 
     ctx->SetOutput(0, output);
   }
diff --git a/tensorflow/compiler/tf2xla/kernels/split_op.cc b/tensorflow/compiler/tf2xla/kernels/split_op.cc
index 43c15e7538..8958b2e770 100644
--- a/tensorflow/compiler/tf2xla/kernels/split_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/split_op.cc
@@ -124,7 +124,7 @@ class SplitVOp : public XlaOpKernel {
                                         input_shape.dims(), "), but got ",
                                         split_dim_orig));
 
-    xla::ComputationDataHandle input = ctx->Input(0);
+    xla::XlaOp input = ctx->Input(0);
 
     OP_REQUIRES(ctx, input_shape.dims() > 0,
                 errors::InvalidArgument("Can't split a 0 dimensional input"));
diff --git a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc
index 1a78c7ab9b..0fb05a2be7 100644
--- a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc
@@ -38,13 +38,13 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
-Status GetStackShape(xla::ComputationBuilder* builder, XlaResource* resource,
+Status GetStackShape(xla::XlaBuilder* builder, XlaResource* resource,
                      TensorShape* stack_shape) {
   auto shape_or_status = builder->GetShape(resource->value());
   if (!shape_or_status.ok()) {
     return shape_or_status.status();
   }
-  xla::Shape shape = *shape_or_status.ValueOrDie();
+  xla::Shape shape = shape_or_status.ValueOrDie();
   TF_RET_CHECK(xla::ShapeUtil::IsTuple(shape));
   return XLAShapeToTensorShape(xla::ShapeUtil::GetTupleElementShape(shape, 0),
                                stack_shape);
@@ -60,9 +60,8 @@ Status GetStackShape(xla::ComputationBuilder* builder, XlaResource* resource,
 //
 // TODO(phawkins): consider changing the API of the stack operators to
 // allow an optional element shape at stack construction time.
-Status MaybeInitializeStack(xla::ComputationBuilder* builder,
-                            XlaResource* resource, DataType dtype,
-                            const TensorShape& elem_shape) {
+Status MaybeInitializeStack(xla::XlaBuilder* builder, XlaResource* resource,
+                            DataType dtype, const TensorShape& elem_shape) {
   if (resource->type() != dtype) {
     return errors::InvalidArgument(
         "Stack dtype is ", DataTypeString(resource->type()),
@@ -75,8 +74,6 @@ Status MaybeInitializeStack(xla::ComputationBuilder* builder,
 
   if (!resource->initialized()) {
     // Stack has not been initialized.
-    xla::ComputationDataHandle zero =
-        XlaHelpers::Zero(builder, resource->type());
     TF_RETURN_IF_ERROR(resource->SetTypeAndShape(dtype, elem_shape));
     TF_RETURN_IF_ERROR(resource->SetZeroValue(builder));
   } else {
@@ -111,7 +108,7 @@ class StackOp : public XlaOpKernel {
 
     // We defer initializing the Stack resource until we see the first push.
     // Otherwise we do not know the shape of the stack elements.
-    xla::ComputationDataHandle value;
+    xla::XlaOp value;
     XlaContext& xc = XlaContext::Get(ctx);
     XlaResource* resource;
     string name = strings::StrCat("Stack: ", stack_name_);
@@ -138,7 +135,7 @@ class StackPushOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
     TensorShape elem_shape = ctx->InputShape(1);
 
     XlaResource* resource;
@@ -147,9 +144,9 @@ class StackPushOp : public XlaOpKernel {
     // Initializes the Stack, if the element shape was not already known.
     OP_REQUIRES_OK(ctx, MaybeInitializeStack(b, resource, dtype_, elem_shape));
 
-    xla::ComputationDataHandle ta = b->GetTupleElement(resource->value(), 0);
-    xla::ComputationDataHandle index = b->GetTupleElement(resource->value(), 1);
-    xla::ComputationDataHandle value = ctx->Input(1);
+    xla::XlaOp ta = b->GetTupleElement(resource->value(), 0);
+    xla::XlaOp index = b->GetTupleElement(resource->value(), 1);
+    xla::XlaOp value = ctx->Input(1);
 
     // start_indices of the DynamicUpdateSlice are [index, 0, 0, ..., 0].
     auto start_indices =
@@ -184,7 +181,7 @@ class StackPopOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     XlaResource* resource;
     OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource));
@@ -199,9 +196,9 @@ class StackPopOp : public XlaOpKernel {
     TensorShape stack_shape;
     OP_REQUIRES_OK(ctx, GetStackShape(b, resource, &stack_shape));
 
-    xla::ComputationDataHandle state = resource->value();
-    xla::ComputationDataHandle ta = b->GetTupleElement(state, 0);
-    xla::ComputationDataHandle index = b->GetTupleElement(state, 1);
+    xla::XlaOp state = resource->value();
+    xla::XlaOp ta = b->GetTupleElement(state, 0);
+    xla::XlaOp index = b->GetTupleElement(state, 1);
 
     index = b->Sub(index, b->ConstantR0<int32>(1));
     OP_REQUIRES_OK(ctx, resource->SetValue(b->Tuple({ta, index})));
@@ -216,8 +213,7 @@ class StackPopOp : public XlaOpKernel {
 
     // TODO(phawkins): We don't check the index is in bounds --- there is no
     // error mechanism in XLA.
-    xla::ComputationDataHandle read =
-        b->DynamicSlice(ta, start_indices, slice_shape);
+    xla::XlaOp read = b->DynamicSlice(ta, start_indices, slice_shape);
 
     // Remove the leading '1' dimension.
     std::vector<int64> value_shape(slice_shape.begin() + 1, slice_shape.end());
diff --git a/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc b/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc
index 5bb773d97f..6340c22518 100644
--- a/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc
@@ -30,9 +30,8 @@ namespace tensorflow {
 namespace {
 
 // Rotates a 32-bit integer 'v' left by 'distance' bits.
-xla::ComputationDataHandle RotateLeftS32(xla::ComputationBuilder* builder,
-                                         const xla::ComputationDataHandle& v,
-                                         int distance) {
+xla::XlaOp RotateLeftS32(xla::XlaBuilder* builder, const xla::XlaOp& v,
+                         int distance) {
   return builder->Or(
       builder->ShiftLeft(v, builder->ConstantR0<int>(distance)),
       builder->ShiftRightLogical(v, builder->ConstantR0<int>(32 - distance)));
@@ -40,25 +39,24 @@ xla::ComputationDataHandle RotateLeftS32(xla::ComputationBuilder* builder,
 
 // TODO(b/65209188): add a primitive XOR to XLA and call it here, rather than
 // building XOR out of other bitwise operators.
-xla::ComputationDataHandle BitwiseXor(xla::ComputationBuilder* builder,
-                                      const xla::ComputationDataHandle& x,
-                                      const xla::ComputationDataHandle& y) {
+xla::XlaOp BitwiseXor(xla::XlaBuilder* builder, const xla::XlaOp& x,
+                      const xla::XlaOp& y) {
   return builder->Or(builder->And(x, builder->Not(y)),
                      builder->And(builder->Not(x), y));
 }
 
-using ThreeFry2x32State = std::array<xla::ComputationDataHandle, 2>;
+using ThreeFry2x32State = std::array<xla::XlaOp, 2>;
 
 // Implements the ThreeFry counter-based PRNG algorithm.
 // Salmon et al. SC 2011. Parallel random numbers: as easy as 1, 2, 3.
 // http://www.thesalmons.org/john/random123/papers/random123sc11.pdf
-ThreeFry2x32State ThreeFry2x32(xla::ComputationBuilder* builder,
+ThreeFry2x32State ThreeFry2x32(xla::XlaBuilder* builder,
                                ThreeFry2x32State input, ThreeFry2x32State key) {
   // Rotation distances specified by the Threefry2x32 algorithm.
   constexpr std::array<int, 8> rotations = {13, 15, 26, 6, 17, 29, 16, 24};
   ThreeFry2x32State x;
 
-  std::array<xla::ComputationDataHandle, 3> ks;
+  std::array<xla::XlaOp, 3> ks;
   // 0x1BD11BDA is a parity constant specified by the ThreeFry2x32 algorithm.
   ks[2] = builder->ConstantR0<int32>(0x1BD11BDA);
   for (int i = 0; i < 2; ++i) {
@@ -121,10 +119,9 @@ ThreeFry2x32State ThreeFry2x32(xla::ComputationBuilder* builder,
 
 // Returns a tensor of 'shape' random values uniformly distributed in the range
 // [minval, maxval)
-xla::ComputationDataHandle RandomUniform(xla::ComputationBuilder* builder,
-                                         const xla::ComputationDataHandle& seed,
-                                         const TensorShape& shape,
-                                         double minval, double maxval) {
+xla::XlaOp RandomUniform(xla::XlaBuilder* builder, const xla::XlaOp& seed,
+                         const TensorShape& shape, double minval,
+                         double maxval) {
   // Split the seed into two 32-bit scalars to form a key.
   auto seed0 = builder->Reshape(builder->Slice(seed, {0}, {1}, {1}), {});
   auto seed1 = builder->Reshape(builder->Slice(seed, {1}, {2}, {1}), {});
@@ -178,9 +175,8 @@ xla::ComputationDataHandle RandomUniform(xla::ComputationBuilder* builder,
 //     p = sum_{i=1}^n gq[i]*w^i
 //   }
 //   return p*x
-xla::ComputationDataHandle ErfInvF32(xla::ComputationBuilder* b,
-                                     const xla::ComputationDataHandle& x,
-                                     const TensorShape& shape) {
+xla::XlaOp ErfInvF32(xla::XlaBuilder* b, const xla::XlaOp& x,
+                     const TensorShape& shape) {
   constexpr int kDegree = 9;
   constexpr std::array<float, 9> w_less_than_5_constants = {
       2.81022636e-08f,  3.43273939e-07f, -3.5233877e-06f,
@@ -220,7 +216,7 @@ class StatelessRandomUniformOp : public XlaOpKernel {
       : XlaOpKernel(ctx) {}
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
 
     TensorShape shape;
     OP_REQUIRES_OK(ctx, ctx->ConstantInputAsShape(0, &shape));
@@ -229,7 +225,7 @@ class StatelessRandomUniformOp : public XlaOpKernel {
     OP_REQUIRES(ctx, seed_shape.dims() == 1 && seed_shape.dim_size(0) == 2,
                 errors::InvalidArgument("seed must have shape [2], not ",
                                         seed_shape.DebugString()));
-    xla::ComputationDataHandle seed = ctx->Input(1);
+    xla::XlaOp seed = ctx->Input(1);
     ctx->SetOutput(0, RandomUniform(builder, seed, shape, 0.0, 1.0));
   }
 
@@ -257,8 +253,8 @@ class StatelessRandomNormalOp : public XlaOpKernel {
     OP_REQUIRES(ctx, seed_shape == TensorShape({2}),
                 errors::InvalidArgument("seed must have shape [2], not ",
                                         seed_shape.DebugString()));
-    xla::ComputationDataHandle seed = ctx->Input(1);
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaOp seed = ctx->Input(1);
+    xla::XlaBuilder* builder = ctx->builder();
     auto uniform = RandomUniform(builder, seed, shape, -1.0, 1.0);
     // Convert uniform distribution to normal distribution by computing
     // sqrt(2) * erfinv(x)
diff --git a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
index 6204aa4e27..55254c746e 100644
--- a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
@@ -90,7 +90,7 @@ class StridedSliceOp : public XlaOpKernel {
       }
     }
 
-    xla::ComputationDataHandle slice = ctx->Input(0);
+    xla::XlaOp slice = ctx->Input(0);
     if (!dimensions_to_reverse.empty()) {
       slice = ctx->builder()->Rev(slice, dimensions_to_reverse);
     }
@@ -168,7 +168,7 @@ class StridedSliceGradOp : public XlaOpKernel {
 
     auto zero = XlaHelpers::Zero(ctx->builder(), ctx->expected_output_dtype(0));
 
-    xla::ComputationDataHandle grad = ctx->Input(4);
+    xla::XlaOp grad = ctx->Input(4);
 
     // Undo any new/shrink axes.
     grad = ctx->builder()->Reshape(grad, processing_shape.dim_sizes());
@@ -255,7 +255,7 @@ class StridedSliceAssignOp : public XlaOpKernel {
                                             &strides_tensor));
 
     TensorShape lhs_shape;
-    xla::ComputationDataHandle lhs;
+    xla::XlaOp lhs;
     OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, dtype_, &lhs_shape, &lhs));
 
     const TensorShape rhs_shape = ctx->InputShape(4);
@@ -284,7 +284,7 @@ class StridedSliceAssignOp : public XlaOpKernel {
                     " does not match r-value shape ", rhs_shape.DebugString(),
                     ". Automatic broadcasting not yet implemented."));
 
-    xla::ComputationDataHandle rhs = ctx->Input(4);
+    xla::XlaOp rhs = ctx->Input(4);
 
     gtl::InlinedVector<int64, 4> dimensions_to_reverse;
     gtl::InlinedVector<int64, 4> slice_begin, slice_dims;
diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc
index 000b50af6b..9adee78a1f 100644
--- a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc
@@ -47,7 +47,7 @@ namespace {
 // the TensorArray with elements of `elem_shape`. For both initialized and
 // uninitialized TensorArrays, checks that the tensor has a type compatible with
 // 'dtype' and shape compatible with 'elem_shape'.
-Status MaybeInitializeTensorArray(xla::ComputationBuilder* builder,
+Status MaybeInitializeTensorArray(xla::XlaBuilder* builder,
                                   XlaResource* resource, DataType dtype,
                                   const TensorShape& elem_shape) {
   if (resource->kind() != XlaResource::kTensorArray) {
@@ -64,9 +64,6 @@ Status MaybeInitializeTensorArray(xla::ComputationBuilder* builder,
       << resource->name() << " size " << resource->tensor_array_size();
 
   if (!resource->initialized()) {
-    xla::ComputationDataHandle zero =
-        XlaHelpers::Zero(builder, resource->type());
-
     TF_RETURN_IF_ERROR(resource->SetTypeAndShape(dtype, elem_shape));
     TF_RETURN_IF_ERROR(resource->SetZeroValue(builder));
   } else {
@@ -77,7 +74,7 @@ Status MaybeInitializeTensorArray(xla::ComputationBuilder* builder,
     }
     TensorShape shape;
     TF_RETURN_IF_ERROR(
-        XLAShapeToTensorShape(*shape_or_status.ValueOrDie(), &shape));
+        XLAShapeToTensorShape(shape_or_status.ValueOrDie(), &shape));
 
     TensorShape ta_shape;
     ta_shape.AddDim(resource->tensor_array_size());
@@ -114,23 +111,21 @@ Status CheckTensorArrayIsInitialized(const string& op_name,
 }
 
 Status GetTensorArrayShape(const XlaResource* resource,
-                           xla::ComputationBuilder* builder,
-                           TensorShape* shape) {
+                           xla::XlaBuilder* builder, TensorShape* shape) {
   *shape = resource->shape();
   shape->InsertDim(0, resource->tensor_array_size());
   return Status::OK();
 }
 
-// Like ComputationBuilder::DynamicUpdateSlice, but adds 'update' to the
+// Like XlaBuilder::DynamicUpdateSlice, but adds 'update' to the
 // relevant slice of 'operand'.
-xla::ComputationDataHandle DynamicAddSlice(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& operand,
-    const xla::ComputationDataHandle& update,
-    const gtl::ArraySlice<int64>& update_dims,
-    const xla::ComputationDataHandle& start_indices) {
-  xla::ComputationDataHandle current =
+xla::XlaOp DynamicAddSlice(xla::XlaBuilder* builder, const xla::XlaOp& operand,
+                           const xla::XlaOp& update,
+                           const gtl::ArraySlice<int64>& update_dims,
+                           const xla::XlaOp& start_indices) {
+  xla::XlaOp current =
       builder->DynamicSlice(operand, start_indices, update_dims);
-  xla::ComputationDataHandle sum = builder->Add(current, update);
+  xla::XlaOp sum = builder->Add(current, update);
   return builder->DynamicUpdateSlice(operand, sum, start_indices);
 }
 
@@ -155,18 +150,18 @@ class TensorArrayOp : public XlaOpKernel {
     OP_REQUIRES(ctx, size >= 0,
                 errors::InvalidArgument("TensorArray size must be >= 0"));
 
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     // Initializes the TensorArray value if we know the element shape.
     // Otherwise, defer initialization to the first write.
-    xla::ComputationDataHandle value;
+    xla::XlaOp value;
     TensorShape shape;
     if (element_shape_.IsFullyDefined()) {
       CHECK(element_shape_.AsTensorShape(&shape));
       TensorShape ta_shape;
       ta_shape.AddDim(size);
       ta_shape.AppendShape(shape);
-      xla::ComputationDataHandle zero = XlaHelpers::Zero(b, dtype_);
+      xla::XlaOp zero = XlaHelpers::Zero(b, dtype_);
       value = b->Broadcast(zero, ta_shape.dim_sizes());
     }
 
@@ -202,7 +197,7 @@ class TensorArrayWriteOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     TensorShape elem_shape = ctx->InputShape(2);
 
@@ -213,10 +208,10 @@ class TensorArrayWriteOp : public XlaOpKernel {
     OP_REQUIRES_OK(ctx,
                    MaybeInitializeTensorArray(b, resource, dtype_, elem_shape));
 
-    xla::ComputationDataHandle ta = resource->value();
-    xla::ComputationDataHandle index = ctx->Input(1);
-    xla::ComputationDataHandle value = ctx->Input(2);
-    xla::ComputationDataHandle flow = ctx->Input(3);
+    xla::XlaOp ta = resource->value();
+    xla::XlaOp index = ctx->Input(1);
+    xla::XlaOp value = ctx->Input(2);
+    xla::XlaOp flow = ctx->Input(3);
 
     // start_indices of the DynamicUpdateSlice are [index, 0, 0, ..., 0].
     auto start_indices =
@@ -227,7 +222,7 @@ class TensorArrayWriteOp : public XlaOpKernel {
     slice_shape.InsertDim(0, 1LL);
     auto update = b->Reshape(value, slice_shape.dim_sizes());
 
-    xla::ComputationDataHandle written =
+    xla::XlaOp written =
         DynamicAddSlice(b, ta, update, slice_shape.dim_sizes(), start_indices);
 
     OP_REQUIRES_OK(ctx, resource->SetValue(written));
@@ -249,7 +244,7 @@ class TensorArrayReadOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     XlaResource* resource;
     OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource));
@@ -259,8 +254,8 @@ class TensorArrayReadOp : public XlaOpKernel {
     TensorShape ta_shape;
     OP_REQUIRES_OK(ctx, GetTensorArrayShape(resource, b, &ta_shape));
 
-    xla::ComputationDataHandle ta = resource->value();
-    xla::ComputationDataHandle index = ctx->Input(1);
+    xla::XlaOp ta = resource->value();
+    xla::XlaOp index = ctx->Input(1);
 
     // start_indices of the DynamicSlice are [index, 0, 0, ..., 0].
     auto start_indices =
@@ -270,8 +265,7 @@ class TensorArrayReadOp : public XlaOpKernel {
     auto slice_shape = ta_shape.dim_sizes();
     slice_shape[0] = 1LL;
 
-    xla::ComputationDataHandle read =
-        b->DynamicSlice(ta, start_indices, slice_shape);
+    xla::XlaOp read = b->DynamicSlice(ta, start_indices, slice_shape);
 
     // Remove the leading '1' dimension.
     std::vector<int64> value_shape(slice_shape.begin() + 1, slice_shape.end());
@@ -293,7 +287,7 @@ class TensorArrayGatherOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     XlaResource* resource;
     OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource));
@@ -309,7 +303,7 @@ class TensorArrayGatherOp : public XlaOpKernel {
     auto indices = ctx->Input(1);
     DataType index_type = ctx->input_type(1);
 
-    xla::ComputationDataHandle ta = resource->value();
+    xla::XlaOp ta = resource->value();
 
     // Look for the case where the gather takes a simple slice from the
     // tensor array (0, 1, 2, 3, 4, ..., N)
@@ -337,7 +331,7 @@ class TensorArrayGatherOp : public XlaOpKernel {
       }
     }
 
-    xla::ComputationDataHandle gather;
+    xla::XlaOp gather;
     OP_REQUIRES_OK(
         ctx,
         XlaGather(ta, ta_shape, indices, indices_shape, /*axis=*/0,
@@ -360,7 +354,7 @@ class TensorArrayScatterOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     const TensorShape value_shape = ctx->InputShape(2);
 
@@ -375,11 +369,11 @@ class TensorArrayScatterOp : public XlaOpKernel {
     OP_REQUIRES(ctx, indices_shape.dims() >= 1,
                 errors::InvalidArgument("indices must be rank 1"));
     const int num_indices = indices_shape.dim_size(0);
-    const xla::ComputationDataHandle indices = ctx->Input(1);
+    const xla::XlaOp indices = ctx->Input(1);
 
-    xla::ComputationDataHandle ta = resource->value();
-    const xla::ComputationDataHandle value = ctx->Input(2);
-    const xla::ComputationDataHandle flow = ctx->Input(3);
+    xla::XlaOp ta = resource->value();
+    const xla::XlaOp value = ctx->Input(2);
+    const xla::XlaOp flow = ctx->Input(3);
 
     // Look for the case where the scatter is for each sub-tensor in order. The
     // tensor array implementation allows for this to be a straight addition.
@@ -443,7 +437,7 @@ class TensorArrayConcatOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     XlaResource* resource;
     OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource));
@@ -453,7 +447,7 @@ class TensorArrayConcatOp : public XlaOpKernel {
     TensorShape ta_shape;
     OP_REQUIRES_OK(ctx, GetTensorArrayShape(resource, b, &ta_shape));
 
-    xla::ComputationDataHandle ta = resource->value();
+    xla::XlaOp ta = resource->value();
 
     auto ta_dims = ta_shape.dim_sizes();
     std::vector<int64> shape(ta_dims.begin() + 1, ta_dims.end());
@@ -503,12 +497,12 @@ class TensorArraySplitOp : public XlaOpKernel {
     TensorShape elem_shape = value_shape;
     elem_shape.set_dim(0, length);
 
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
     XlaResource* resource;
     OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource));
     OP_REQUIRES_OK(ctx,
                    MaybeInitializeTensorArray(b, resource, dtype_, elem_shape));
-    xla::ComputationDataHandle ta = resource->value();
+    xla::XlaOp ta = resource->value();
 
     TensorShape ta_shape;
     ta_shape.AddDim(resource->tensor_array_size());
@@ -520,8 +514,8 @@ class TensorArraySplitOp : public XlaOpKernel {
             "TensorArray's size is not equal to the size of lengths (",
             lengths.size(), " vs. ", resource->tensor_array_size(), ")"));
 
-    const xla::ComputationDataHandle value = ctx->Input(1);
-    const xla::ComputationDataHandle flow = ctx->Input(3);
+    const xla::XlaOp value = ctx->Input(1);
+    const xla::XlaOp flow = ctx->Input(3);
 
     OP_REQUIRES(ctx, value_shape.num_elements() == ta_shape.num_elements(),
                 errors::InvalidArgument("mismatched element count ",
@@ -569,7 +563,7 @@ class TensorArrayGradOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     XlaResource* resource;
     OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource));
diff --git a/tensorflow/compiler/tf2xla/kernels/tile_ops.cc b/tensorflow/compiler/tf2xla/kernels/tile_ops.cc
index 9aefcd4fc7..e91075196b 100644
--- a/tensorflow/compiler/tf2xla/kernels/tile_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/tile_ops.cc
@@ -112,7 +112,7 @@ class TileOp : public XlaOpKernel {
       flattened.push_back(i);
       flattened.push_back(i + output_shape.size());
     }
-    xla::ComputationDataHandle output =
+    xla::XlaOp output =
         ctx->builder()->Reshape(broadcasted, flattened, output_shape);
 
     ctx->SetOutput(0, output);
diff --git a/tensorflow/compiler/tf2xla/kernels/training_ops.cc b/tensorflow/compiler/tf2xla/kernels/training_ops.cc
index f750f7003b..34caefa050 100644
--- a/tensorflow/compiler/tf2xla/kernels/training_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/training_ops.cc
@@ -16,7 +16,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/kernels/cwise_ops.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/core/framework/kernel_def_builder.h"
 #include "tensorflow/core/framework/types.h"
@@ -30,8 +30,8 @@ class ResourceApplyGradientDescent : public XlaOpKernel {
   explicit ResourceApplyGradientDescent(OpKernelConstruction* ctx)
       : XlaOpKernel(ctx) {}
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationDataHandle handle;
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaOp handle;
+    xla::XlaBuilder* b = ctx->builder();
     DataType type = ctx->input_type(1);
     TensorShape var_shape;
     OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, type, &var_shape, &handle));
@@ -63,12 +63,12 @@ class ResourceApplyMomentum : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     DataType type = ctx->input_type(2);
 
     TensorShape var_shape, accum_shape;
-    xla::ComputationDataHandle var, accum;
+    xla::XlaOp var, accum;
     OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, type, &var_shape, &var));
     OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, type, &accum_shape, &accum));
 
@@ -93,9 +93,9 @@ class ResourceApplyMomentum : public XlaOpKernel {
                 errors::InvalidArgument("momentum is not a scalar: ",
                                         momentum_shape.DebugString()));
 
-    xla::ComputationDataHandle lr = ctx->Input(2);
-    xla::ComputationDataHandle grad = ctx->Input(3);
-    xla::ComputationDataHandle momentum = ctx->Input(4);
+    xla::XlaOp lr = ctx->Input(2);
+    xla::XlaOp grad = ctx->Input(3);
+    xla::XlaOp momentum = ctx->Input(4);
 
     accum = b->Add(b->Mul(accum, momentum), grad);
     if (use_nesterov_) {
@@ -121,12 +121,12 @@ class ResourceApplyAdagrad : public XlaOpKernel {
   explicit ResourceApplyAdagrad(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     DataType type = ctx->input_type(2);
 
     TensorShape var_shape, accum_shape;
-    xla::ComputationDataHandle var, accum;
+    xla::XlaOp var, accum;
     OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, type, &var_shape, &var));
     OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, type, &accum_shape, &accum));
 
@@ -146,8 +146,8 @@ class ResourceApplyAdagrad : public XlaOpKernel {
                     "var and grad do not have the same shape",
                     var_shape.DebugString(), " ", grad_shape.DebugString()));
 
-    xla::ComputationDataHandle lr = ctx->Input(2);
-    xla::ComputationDataHandle grad = ctx->Input(3);
+    xla::XlaOp lr = ctx->Input(2);
+    xla::XlaOp grad = ctx->Input(3);
 
     accum = b->Add(accum, b->Pow(grad, XlaHelpers::FloatLiteral(b, type, 2.0)));
     var = b->Sub(
@@ -168,7 +168,7 @@ class ResourceApplyAdam : public XlaOpKernel {
 
   void Compile(XlaOpKernelContext* ctx) override {
     TensorShape var_shape, m_shape, v_shape;
-    xla::ComputationDataHandle var, m, v;
+    xla::XlaOp var, m, v;
     OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, dtype_, &var_shape, &var));
     OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, dtype_, &m_shape, &m));
     OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(2, dtype_, &v_shape, &v));
@@ -213,25 +213,25 @@ class ResourceApplyAdam : public XlaOpKernel {
                     "var and grad do not have the same shape",
                     var_shape.DebugString(), " ", grad_shape.DebugString()));
 
-    xla::ComputationDataHandle beta1_power = ctx->Input(3);
-    xla::ComputationDataHandle beta2_power = ctx->Input(4);
-    xla::ComputationDataHandle lr = ctx->Input(5);
-    xla::ComputationDataHandle beta1 = ctx->Input(6);
-    xla::ComputationDataHandle beta2 = ctx->Input(7);
-    xla::ComputationDataHandle epsilon = ctx->Input(8);
-    xla::ComputationDataHandle grad = ctx->Input(9);
+    xla::XlaOp beta1_power = ctx->Input(3);
+    xla::XlaOp beta2_power = ctx->Input(4);
+    xla::XlaOp lr = ctx->Input(5);
+    xla::XlaOp beta1 = ctx->Input(6);
+    xla::XlaOp beta2 = ctx->Input(7);
+    xla::XlaOp epsilon = ctx->Input(8);
+    xla::XlaOp grad = ctx->Input(9);
 
     // alpha <- learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
     // m_t <- beta1 * m_{t-1} + (1 - beta1) * g_t
     // v_t <- beta2 * v_{t-1} + (1 - beta2) * g_t * g_t
     // variable <- variable - alpha * m_t / (sqrt(v_t) + epsilon)
 
-    xla::ComputationBuilder* b = ctx->builder();
-    xla::ComputationDataHandle half = XlaHelpers::FloatLiteral(b, dtype_, 0.5);
-    xla::ComputationDataHandle one = XlaHelpers::FloatLiteral(b, dtype_, 1.0);
-    xla::ComputationDataHandle two = XlaHelpers::FloatLiteral(b, dtype_, 2.0);
+    xla::XlaBuilder* b = ctx->builder();
+    xla::XlaOp half = XlaHelpers::FloatLiteral(b, dtype_, 0.5);
+    xla::XlaOp one = XlaHelpers::FloatLiteral(b, dtype_, 1.0);
+    xla::XlaOp two = XlaHelpers::FloatLiteral(b, dtype_, 2.0);
 
-    xla::ComputationDataHandle alpha =
+    xla::XlaOp alpha =
         b->Div(b->Mul(lr, b->Pow(b->Sub(one, beta2_power), half)),
                b->Sub(one, beta1_power));
     m = b->Add(m, b->Mul(b->Sub(grad, m), b->Sub(one, beta1)));
@@ -255,12 +255,12 @@ class ResourceApplyRMSProp : public XlaOpKernel {
   explicit ResourceApplyRMSProp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* b = ctx->builder();
+    xla::XlaBuilder* b = ctx->builder();
 
     DataType type = ctx->input_type(3);
 
     TensorShape var_shape, ms_shape, mom_shape;
-    xla::ComputationDataHandle var, ms, mom;
+    xla::XlaOp var, ms, mom;
     OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, type, &var_shape, &var));
     OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, type, &ms_shape, &ms));
     OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(2, type, &mom_shape, &mom));
@@ -297,11 +297,11 @@ class ResourceApplyRMSProp : public XlaOpKernel {
                     "var and grad do not have the same shape",
                     var_shape.DebugString(), " ", grad_shape.DebugString()));
 
-    xla::ComputationDataHandle lr = ctx->Input(3);
-    xla::ComputationDataHandle rho = ctx->Input(4);
-    xla::ComputationDataHandle momentum = ctx->Input(5);
-    xla::ComputationDataHandle epsilon = ctx->Input(6);
-    xla::ComputationDataHandle grad = ctx->Input(7);
+    xla::XlaOp lr = ctx->Input(3);
+    xla::XlaOp rho = ctx->Input(4);
+    xla::XlaOp momentum = ctx->Input(5);
+    xla::XlaOp epsilon = ctx->Input(6);
+    xla::XlaOp grad = ctx->Input(7);
 
     // ms <- rho * ms_{t-1} + (1-rho) * grad * grad
     // mom <- momentum * mom_{t-1} + lr * grad / sqrt(ms + epsilon)
@@ -320,16 +320,16 @@ class ResourceApplyRMSProp : public XlaOpKernel {
     //    ms <- grad**2 (1 - rho) + ms * rho
     //
     // Which is the equation listed above.
-    xla::ComputationDataHandle new_ms = b->Add(
+    xla::XlaOp new_ms = b->Add(
         ms,
         b->Mul(b->Sub(b->Pow(grad, XlaHelpers::FloatLiteral(b, type, 2.0)), ms),
                b->Sub(XlaHelpers::FloatLiteral(b, type, 1.0), rho)));
-    xla::ComputationDataHandle new_mom =
+    xla::XlaOp new_mom =
         b->Add(b->Mul(mom, momentum),
                b->Mul(b->Mul(grad, lr),
                       b->Pow(b->Add(new_ms, epsilon),
                              XlaHelpers::FloatLiteral(b, type, -0.5))));
-    xla::ComputationDataHandle new_var = b->Sub(var, new_mom);
+    xla::XlaOp new_var = b->Sub(var, new_mom);
 
     OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, type, new_var));
     OP_REQUIRES_OK(ctx, ctx->AssignVariable(1, type, new_ms));
@@ -341,10 +341,10 @@ REGISTER_XLA_OP(Name("ResourceApplyRMSProp").TypeConstraint("T", kFloatTypes),
 
 void CompileFtrl(XlaOpKernelContext* ctx, DataType dtype,
                  bool has_l2_shrinkage) {
-  xla::ComputationBuilder* b = ctx->builder();
+  xla::XlaBuilder* b = ctx->builder();
 
   TensorShape var_shape, accum_shape, linear_shape;
-  xla::ComputationDataHandle var, accum, linear;
+  xla::XlaOp var, accum, linear;
   OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, dtype, &var_shape, &var));
   OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, dtype, &accum_shape, &accum));
   OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(2, dtype, &linear_shape, &linear));
@@ -399,12 +399,12 @@ void CompileFtrl(XlaOpKernelContext* ctx, DataType dtype,
               errors::InvalidArgument("lr_power is not a scalar: ",
                                       lr_power_shape.DebugString()));
 
-  xla::ComputationDataHandle grad = ctx->Input(3);
-  xla::ComputationDataHandle lr = ctx->Input(4);
-  xla::ComputationDataHandle l1 = ctx->Input(5);
-  xla::ComputationDataHandle l2 = ctx->Input(6);
-  xla::ComputationDataHandle l2_shrinkage;
-  xla::ComputationDataHandle lr_power;
+  xla::XlaOp grad = ctx->Input(3);
+  xla::XlaOp lr = ctx->Input(4);
+  xla::XlaOp l1 = ctx->Input(5);
+  xla::XlaOp l2 = ctx->Input(6);
+  xla::XlaOp l2_shrinkage;
+  xla::XlaOp lr_power;
   if (has_l2_shrinkage) {
     l2_shrinkage = ctx->Input(7);
     lr_power = ctx->Input(8);
@@ -421,26 +421,23 @@ void CompileFtrl(XlaOpKernelContext* ctx, DataType dtype,
   // var = (linear_clipped - linear) / quadratic
   // accum = new_accum
 
-  xla::ComputationDataHandle two = XlaHelpers::FloatLiteral(b, dtype, 2.0);
-  xla::ComputationDataHandle grad_to_use;
+  xla::XlaOp two = XlaHelpers::FloatLiteral(b, dtype, 2.0);
+  xla::XlaOp grad_to_use;
   if (has_l2_shrinkage) {
     grad_to_use = b->Add(grad, b->Mul(two, b->Mul(l2_shrinkage, var)));
   } else {
     grad_to_use = grad;
   }
 
-  xla::ComputationDataHandle new_accum =
-      b->Add(accum, b->Pow(grad_to_use, two));
-  xla::ComputationDataHandle new_accum_lr_pow =
-      b->Pow(new_accum, b->Neg(lr_power));
-  xla::ComputationDataHandle accum_lr_pow = b->Pow(accum, b->Neg(lr_power));
+  xla::XlaOp new_accum = b->Add(accum, b->Pow(grad_to_use, two));
+  xla::XlaOp new_accum_lr_pow = b->Pow(new_accum, b->Neg(lr_power));
+  xla::XlaOp accum_lr_pow = b->Pow(accum, b->Neg(lr_power));
   linear = b->Add(
       linear,
       b->Sub(grad_to_use,
              b->Mul(b->Div(b->Sub(new_accum_lr_pow, accum_lr_pow), lr), var)));
-  xla::ComputationDataHandle linear_clipped = b->Clamp(b->Neg(l1), linear, l1);
-  xla::ComputationDataHandle quadratic =
-      b->Add(b->Div(new_accum_lr_pow, lr), b->Mul(two, l2));
+  xla::XlaOp linear_clipped = b->Clamp(b->Neg(l1), linear, l1);
+  xla::XlaOp quadratic = b->Add(b->Div(new_accum_lr_pow, lr), b->Mul(two, l2));
   var = b->Div(b->Sub(linear_clipped, linear), quadratic);
   accum = new_accum;
 
diff --git a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc
index 7cb47f908d..a4f50f52eb 100644
--- a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc
@@ -19,7 +19,7 @@ limitations under the License.
 #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/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/core/framework/kernel_def_builder.h"
 
 namespace tensorflow {
@@ -33,9 +33,9 @@ namespace {
    public:                                                             \
     explicit NAME##Op(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} \
     void Compile(XlaOpKernelContext* ctx) {                            \
-      xla::ComputationBuilder* b = ctx->builder();                     \
-      xla::ComputationDataHandle x = ctx->Input(0);                    \
-      xla::ComputationDataHandle y = COMPUTATION;                      \
+      xla::XlaBuilder* b = ctx->builder();                             \
+      xla::XlaOp x = ctx->Input(0);                                    \
+      xla::XlaOp y = COMPUTATION;                                      \
       ctx->SetOutput(0, y);                                            \
     }                                                                  \
   };                                                                   \
@@ -124,9 +124,8 @@ XLAJIT_MAKE_UNARY(Neg, b->Neg(x));
 
 // Implements Banker's rounding: numbers that are equidistant between two
 // integers are rounded towards even.
-static xla::ComputationDataHandle Round(xla::ComputationBuilder* b,
-                                        DataType dtype,
-                                        const xla::ComputationDataHandle& x) {
+static xla::XlaOp Round(xla::XlaBuilder* b, DataType dtype,
+                        const xla::XlaOp& x) {
   auto half = XlaHelpers::FloatLiteral(b, dtype, 0.5);
   auto one = XlaHelpers::FloatLiteral(b, dtype, 1.0);
   auto two = XlaHelpers::FloatLiteral(b, dtype, 2.0);
@@ -148,9 +147,8 @@ XLAJIT_MAKE_UNARY(Rsqrt,
                   b->Pow(x, XlaHelpers::FloatLiteral(b, input_type(0), -0.5)));
 
 // Expresses sigmoid as a rescaled tanh: sigmoid(x) == (tanh(x/2) + 1) / 2.
-static xla::ComputationDataHandle Sigmoid(xla::ComputationBuilder* b,
-                                          DataType dtype,
-                                          const xla::ComputationDataHandle& x) {
+static xla::XlaOp Sigmoid(xla::XlaBuilder* b, DataType dtype,
+                          const xla::XlaOp& x) {
   auto half = XlaHelpers::FloatLiteral(b, dtype, 0.5);
   return b->Add(half, b->Mul(half, b->Tanh(b->Mul(half, x))));
 }
@@ -162,20 +160,18 @@ XLAJIT_MAKE_UNARY(Sinh,
                   b->Mul(b->Sub(b->Exp(x), b->Exp(b->Neg(x))),
                          XlaHelpers::FloatLiteral(b, input_type(0), 0.5)));
 
-static xla::ComputationDataHandle Softplus(
-    xla::ComputationBuilder* b, DataType dtype,
-    const xla::ComputationDataHandle& features) {
-  xla::ComputationDataHandle threshold =
-      b->Add(b->Log(XlaHelpers::Epsilon(b, dtype)),
-             XlaHelpers::FloatLiteral(b, dtype, 2.0));
+static xla::XlaOp Softplus(xla::XlaBuilder* b, DataType dtype,
+                           const xla::XlaOp& features) {
+  xla::XlaOp threshold = b->Add(b->Log(XlaHelpers::Epsilon(b, dtype)),
+                                XlaHelpers::FloatLiteral(b, dtype, 2.0));
   // Value above which exp(x) may overflow, but softplus(x) == x
   // is within machine epsilon.
-  xla::ComputationDataHandle too_large = b->Gt(features, b->Neg(threshold));
+  xla::XlaOp too_large = b->Gt(features, b->Neg(threshold));
   // Value below which exp(x) may underflow, but softplus(x) == exp(x)
   // is within machine epsilon.
-  xla::ComputationDataHandle too_small = b->Lt(features, threshold);
-  xla::ComputationDataHandle features_exp = b->Exp(features);
-  xla::ComputationDataHandle output = b->Select(
+  xla::XlaOp too_small = b->Lt(features, threshold);
+  xla::XlaOp features_exp = b->Exp(features);
+  xla::XlaOp output = b->Select(
       too_large, features,
       b->Select(too_small, features_exp,
                 b->Log(b->Add(features_exp, XlaHelpers::One(b, dtype)))));
diff --git a/tensorflow/compiler/tf2xla/kernels/variable_ops.cc b/tensorflow/compiler/tf2xla/kernels/variable_ops.cc
index 71173f5aea..6109db8e89 100644
--- a/tensorflow/compiler/tf2xla/kernels/variable_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/variable_ops.cc
@@ -19,7 +19,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/shape_util.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/core/framework/kernel_def_builder.h"
 #include "tensorflow/core/framework/types.h"
@@ -48,7 +48,7 @@ class ReadVariableOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationDataHandle handle;
+    xla::XlaOp handle;
     OP_REQUIRES_OK(
         ctx, ctx->ReadVariableInput(0, dtype_, /*shape=*/nullptr, &handle));
     ctx->SetOutput(0, handle);
@@ -74,7 +74,7 @@ class AssignAddVariableOp : public XlaOpKernel {
   explicit AssignAddVariableOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
   void Compile(XlaOpKernelContext* ctx) override {
     DataType type = ctx->input_type(1);
-    xla::ComputationDataHandle handle;
+    xla::XlaOp handle;
     OP_REQUIRES_OK(ctx,
                    ctx->ReadVariableInput(0, type, /*shape=*/nullptr, &handle));
     handle = ctx->builder()->Add(handle, ctx->Input(1));
@@ -90,7 +90,7 @@ class AssignSubVariableOp : public XlaOpKernel {
   explicit AssignSubVariableOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
   void Compile(XlaOpKernelContext* ctx) override {
     DataType type = ctx->input_type(1);
-    xla::ComputationDataHandle handle;
+    xla::XlaOp handle;
     OP_REQUIRES_OK(ctx,
                    ctx->ReadVariableInput(0, type, /*shape=*/nullptr, &handle));
     handle = ctx->builder()->Sub(handle, ctx->Input(1));
@@ -105,19 +105,19 @@ class ResourceGatherOp : public XlaOpKernel {
  public:
   explicit ResourceGatherOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
   void Compile(XlaOpKernelContext* ctx) override {
-    xla::ComputationBuilder* builder = ctx->builder();
+    xla::XlaBuilder* builder = ctx->builder();
 
     DataType type = ctx->expected_output_dtype(0);
 
     TensorShape resource_shape;
-    xla::ComputationDataHandle resource_handle;
+    xla::XlaOp resource_handle;
     OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, type, &resource_shape,
                                                &resource_handle));
 
     auto indices = ctx->Input(1);
     auto indices_shape = ctx->InputShape(1);
     DataType index_type = ctx->input_type(1);
-    xla::ComputationDataHandle gather;
+    xla::XlaOp gather;
     OP_REQUIRES_OK(
         ctx, XlaGather(resource_handle, resource_shape, indices, indices_shape,
                        /*axis=*/0, /*indices_are_nd=*/false, type, index_type,
diff --git a/tensorflow/compiler/tf2xla/kernels/while_op.cc b/tensorflow/compiler/tf2xla/kernels/while_op.cc
index 0ff1b65ae9..5467c5d994 100644
--- a/tensorflow/compiler/tf2xla/kernels/while_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/while_op.cc
@@ -21,7 +21,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/op_kernel.h"
@@ -101,7 +101,7 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) {
       ctx, MakeXlaCompilerArgumentsFromInputs(
                ctx, &arguments, &has_uninitialized_vars, &has_tensor_arrays));
 
-  xla::ComputationBuilder* builder = ctx->builder();
+  xla::XlaBuilder* builder = ctx->builder();
   XlaCompiler* compiler = ctx->compiler();
 
   VLOG(1) << "Compiling body";
@@ -234,7 +234,7 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) {
                   xla::ShapeUtil::HumanString(cond.xla_output_shape)));
 
   int num_inputs = body.input_mapping.size();
-  std::vector<xla::ComputationDataHandle> inputs(num_inputs);
+  std::vector<xla::XlaOp> inputs(num_inputs);
   for (int i = 0; i < num_inputs; ++i) {
     int input_num = body.input_mapping[i];
     if (ctx->input_type(input_num) == DT_RESOURCE) {
@@ -246,24 +246,24 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) {
     }
   }
 
-  xla::ComputationDataHandle init = builder->Tuple(inputs);
+  xla::XlaOp init = builder->Tuple(inputs);
 
   VLOG(1) << "Building while loop";
 
   // Wraps the condition in a computation that unpacks the output tuple.
-  xla::Computation cond_wrapper;
+  xla::XlaComputation cond_wrapper;
   {
-    std::unique_ptr<xla::ComputationBuilder> cb =
+    std::unique_ptr<xla::XlaBuilder> cb =
         builder->CreateSubBuilder("cond_wrapper");
     auto inputs = cb->Parameter(0, cond_input_shape, "inputs");
     auto outputs = cb->Call(*cond.computation, {inputs});
     cb->GetTupleElement(outputs, 0);
-    xla::StatusOr<xla::Computation> result = cb->Build();
+    xla::StatusOr<xla::XlaComputation> result = cb->Build();
     OP_REQUIRES_OK(ctx, result.status());
     cond_wrapper = std::move(result.ValueOrDie());
   }
 
-  xla::ComputationDataHandle while_result =
+  xla::XlaOp while_result =
       builder->While(cond_wrapper, *body.computation, init);
 
   // Sets non-variable outputs.
diff --git a/tensorflow/compiler/tf2xla/lib/BUILD b/tensorflow/compiler/tf2xla/lib/BUILD
index 12fdfb605d..04ad3694a0 100644
--- a/tensorflow/compiler/tf2xla/lib/BUILD
+++ b/tensorflow/compiler/tf2xla/lib/BUILD
@@ -25,8 +25,8 @@ cc_library(
         "//tensorflow/compiler/xla:shape_util",
         "//tensorflow/compiler/xla:status_macros",
         "//tensorflow/compiler/xla:statusor",
-        "//tensorflow/compiler/xla/client:computation",
-        "//tensorflow/compiler/xla/client:computation_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/core:lib",
     ],
 )
@@ -44,8 +44,8 @@ cc_library(
         "//tensorflow/compiler/xla:shape_util",
         "//tensorflow/compiler/xla:status_macros",
         "//tensorflow/compiler/xla:statusor",
-        "//tensorflow/compiler/xla/client:computation",
-        "//tensorflow/compiler/xla/client:computation_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/core:lib",
     ],
 )
@@ -62,9 +62,9 @@ cc_library(
         "//tensorflow/compiler/xla:status_macros",
         "//tensorflow/compiler/xla:statusor",
         "//tensorflow/compiler/xla:util",
-        "//tensorflow/compiler/xla/client:computation",
-        "//tensorflow/compiler/xla/client:computation_builder",
         "//tensorflow/compiler/xla/client/lib:arithmetic",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/core:lib",
     ],
 )
@@ -82,8 +82,8 @@ cc_library(
         "//tensorflow/compiler/xla:statusor",
         "//tensorflow/compiler/xla:types",
         "//tensorflow/compiler/xla:util",
-        "//tensorflow/compiler/xla/client:computation",
-        "//tensorflow/compiler/xla/client:computation_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/core:lib",
     ],
 )
@@ -101,9 +101,9 @@ xla_test(
         "//tensorflow/compiler/xla:test",
         "//tensorflow/compiler/xla:types",
         "//tensorflow/compiler/xla:xla_data_proto",
-        "//tensorflow/compiler/xla/client:computation_builder",
         "//tensorflow/compiler/xla/client:global_data",
         "//tensorflow/compiler/xla/client:local_client",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
         "//tensorflow/compiler/xla/tests:client_library_test_base",
         "//tensorflow/compiler/xla/tests:literal_test_util",
         "//tensorflow/compiler/xla/tests:xla_internal_test_main",
@@ -122,8 +122,8 @@ cc_library(
         "//tensorflow/compiler/xla:status_macros",
         "//tensorflow/compiler/xla:statusor",
         "//tensorflow/compiler/xla:util",
-        "//tensorflow/compiler/xla/client:computation",
-        "//tensorflow/compiler/xla/client:computation_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/core:lib",
     ],
 )
@@ -161,8 +161,8 @@ cc_library(
         "//tensorflow/compiler/xla:shape_util",
         "//tensorflow/compiler/xla:status_macros",
         "//tensorflow/compiler/xla:statusor",
-        "//tensorflow/compiler/xla/client:computation",
-        "//tensorflow/compiler/xla/client:computation_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/core:lib",
     ],
 )
diff --git a/tensorflow/compiler/tf2xla/lib/batch_dot.cc b/tensorflow/compiler/tf2xla/lib/batch_dot.cc
index 798f0fa780..526694d5a0 100644
--- a/tensorflow/compiler/tf2xla/lib/batch_dot.cc
+++ b/tensorflow/compiler/tf2xla/lib/batch_dot.cc
@@ -25,24 +25,22 @@ limitations under the License.
 
 namespace tensorflow {
 
-xla::StatusOr<xla::ComputationDataHandle> BatchDot(
-    xla::ComputationBuilder* builder, xla::ComputationDataHandle x,
-    xla::ComputationDataHandle y, bool transpose_x, bool transpose_y,
-    bool conjugate_x, bool conjugate_y) {
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> x_shape,
-                      builder->GetShape(x));
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> y_shape,
-                      builder->GetShape(y));
+xla::StatusOr<xla::XlaOp> BatchDot(xla::XlaBuilder* builder, xla::XlaOp x,
+                                   xla::XlaOp y, bool transpose_x,
+                                   bool transpose_y, bool conjugate_x,
+                                   bool conjugate_y) {
+  TF_ASSIGN_OR_RETURN(xla::Shape x_shape, builder->GetShape(x));
+  TF_ASSIGN_OR_RETURN(xla::Shape y_shape, builder->GetShape(y));
 
   // Check that both tensors have the same number of dimensions. There must be
   // at least two (the batch dimensions can be empty).
-  if (xla::ShapeUtil::Rank(*x_shape) != xla::ShapeUtil::Rank(*y_shape)) {
+  if (xla::ShapeUtil::Rank(x_shape) != xla::ShapeUtil::Rank(y_shape)) {
     return errors::InvalidArgument(
         "Arguments to BatchedDot have different ranks: ",
-        xla::ShapeUtil::HumanString(*x_shape), " vs. ",
-        xla::ShapeUtil::HumanString(*y_shape));
+        xla::ShapeUtil::HumanString(x_shape), " vs. ",
+        xla::ShapeUtil::HumanString(y_shape));
   }
-  const int ndims = xla::ShapeUtil::Rank(*x_shape);
+  const int ndims = xla::ShapeUtil::Rank(x_shape);
   if (ndims < 2) {
     return errors::InvalidArgument(
         "Arguments to BatchedDot must have rank >= 2: ", ndims);
@@ -52,46 +50,46 @@ xla::StatusOr<xla::ComputationDataHandle> BatchDot(
   // valid.
   std::vector<int64> batch_dimension_numbers;
   for (int i = 0; i < ndims - 2; ++i) {
-    if (x_shape->dimensions(i) != y_shape->dimensions(i)) {
+    if (x_shape.dimensions(i) != y_shape.dimensions(i)) {
       return errors::InvalidArgument(
           "Dimension ", i, " of inputs to BatchedDot must be equal: ",
-          xla::ShapeUtil::HumanString(*x_shape), " vs ",
-          xla::ShapeUtil::HumanString(*y_shape));
+          xla::ShapeUtil::HumanString(x_shape), " vs ",
+          xla::ShapeUtil::HumanString(y_shape));
     }
     batch_dimension_numbers.push_back(i);
   }
 
   int x_inner_dim = transpose_x ? (ndims - 2) : (ndims - 1);
   int y_inner_dim = transpose_y ? (ndims - 1) : (ndims - 2);
-  if (x_shape->dimensions(x_inner_dim) != y_shape->dimensions(y_inner_dim)) {
+  if (x_shape.dimensions(x_inner_dim) != y_shape.dimensions(y_inner_dim)) {
     return errors::InvalidArgument(
         "Dimensions ", x_inner_dim, " and ", y_inner_dim,
         " of arguments to BatchedDot must be equal: ",
-        xla::ShapeUtil::HumanString(*x_shape), " transpose: ", transpose_x,
-        " vs. ", xla::ShapeUtil::HumanString(*y_shape),
+        xla::ShapeUtil::HumanString(x_shape), " transpose: ", transpose_x,
+        " vs. ", xla::ShapeUtil::HumanString(y_shape),
         " transpose: ", transpose_y);
   }
 
   // Check for zero lhs/rhs dim size.
-  if (xla::ShapeUtil::HasZeroElements(*x_shape) ||
-      xla::ShapeUtil::HasZeroElements(*y_shape)) {
+  if (xla::ShapeUtil::HasZeroElements(x_shape) ||
+      xla::ShapeUtil::HasZeroElements(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]);
+      dimensions[i] = x_shape.dimensions(batch_dimension_numbers[i]);
     }
     int x_outer_dim = transpose_x ? (ndims - 1) : (ndims - 2);
     int y_outer_dim = transpose_y ? (ndims - 2) : (ndims - 1);
-    dimensions.push_back(x_shape->dimensions(x_outer_dim));
-    dimensions.push_back(y_shape->dimensions(y_outer_dim));
+    dimensions.push_back(x_shape.dimensions(x_outer_dim));
+    dimensions.push_back(y_shape.dimensions(y_outer_dim));
     return builder->Broadcast(
-        builder->ConstantLiteral(xla::Literal::Zero(x_shape->element_type())),
+        builder->ConstantLiteral(xla::Literal::Zero(x_shape.element_type())),
         dimensions);
   }
 
-  if (x_shape->element_type() == xla::C64 && conjugate_x) {
+  if (x_shape.element_type() == xla::C64 && conjugate_x) {
     x = builder->Conj(x);
   }
-  if (y_shape->element_type() == xla::C64 && conjugate_y) {
+  if (y_shape.element_type() == xla::C64 && conjugate_y) {
     y = builder->Conj(y);
   }
 
diff --git a/tensorflow/compiler/tf2xla/lib/batch_dot.h b/tensorflow/compiler/tf2xla/lib/batch_dot.h
index b230e885f1..1acc72033b 100644
--- a/tensorflow/compiler/tf2xla/lib/batch_dot.h
+++ b/tensorflow/compiler/tf2xla/lib/batch_dot.h
@@ -16,8 +16,8 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_TF2XLA_LIB_BATCH_DOT_H_
 #define TENSORFLOW_COMPILER_TF2XLA_LIB_BATCH_DOT_H_
 
-#include "tensorflow/compiler/xla/client/computation.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 
 namespace tensorflow {
 
@@ -43,10 +43,10 @@ namespace tensorflow {
 // It is computed as:
 //
 //     output[..., :, :] = matrix(x[..., :, :]) * matrix(y[..., :, :])
-xla::StatusOr<xla::ComputationDataHandle> BatchDot(
-    xla::ComputationBuilder* builder, xla::ComputationDataHandle x,
-    xla::ComputationDataHandle y, bool transpose_x, bool transpose_y,
-    bool conjugate_x = false, bool conjugate_y = false);
+xla::StatusOr<xla::XlaOp> BatchDot(xla::XlaBuilder* builder, xla::XlaOp x,
+                                   xla::XlaOp y, bool transpose_x,
+                                   bool transpose_y, bool conjugate_x = false,
+                                   bool conjugate_y = false);
 
 }  // namespace tensorflow
 
diff --git a/tensorflow/compiler/tf2xla/lib/cholesky.cc b/tensorflow/compiler/tf2xla/lib/cholesky.cc
index 203365e2ab..83e7382786 100644
--- a/tensorflow/compiler/tf2xla/lib/cholesky.cc
+++ b/tensorflow/compiler/tf2xla/lib/cholesky.cc
@@ -47,23 +47,21 @@ namespace {
 //     l[..., j+1:, j] = (a[..., j+1:, j] - np.dot(l[..., j+1:, :j], row_t)) /
 //                       l[..., j, j]
 //   return l
-xla::StatusOr<xla::ComputationDataHandle> CholeskyUnblocked(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& a) {
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> a_shape,
-                      builder->GetShape(a));
-  const int n_dims = xla::ShapeUtil::Rank(*a_shape);
-  const int64 n = xla::ShapeUtil::GetDimension(*a_shape, -1);
-  gtl::ArraySlice<int64> major_dims(xla::AsInt64Slice(a_shape->dimensions()),
+xla::StatusOr<xla::XlaOp> CholeskyUnblocked(xla::XlaBuilder* builder,
+                                            const xla::XlaOp& a) {
+  TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a));
+  const int n_dims = xla::ShapeUtil::Rank(a_shape);
+  const int64 n = xla::ShapeUtil::GetDimension(a_shape, -1);
+  gtl::ArraySlice<int64> major_dims(xla::AsInt64Slice(a_shape.dimensions()),
                                     /*pos=*/0,
                                     /*len=*/n_dims - 2);
 
-  xla::ComputationDataHandle l = Zeros(builder, *a_shape);
+  xla::XlaOp l = Zeros(builder, a_shape);
 
   // Construct the for loop body to iterate over rows.
-  auto body_fn = [&](xla::ComputationDataHandle i,
-                     gtl::ArraySlice<xla::ComputationDataHandle> loop_vars,
-                     xla::ComputationBuilder* body_builder)
-      -> xla::StatusOr<std::vector<xla::ComputationDataHandle>> {
+  auto body_fn = [&](xla::XlaOp i, gtl::ArraySlice<xla::XlaOp> loop_vars,
+                     xla::XlaBuilder* body_builder)
+      -> xla::StatusOr<std::vector<xla::XlaOp>> {
     xla::Shape col_shape;
     xla::Shape row_shape;
     for (int64 d : major_dims) {
@@ -72,12 +70,12 @@ xla::StatusOr<xla::ComputationDataHandle> CholeskyUnblocked(
     }
     row_shape.add_dimensions(1);
     row_shape.add_dimensions(n);
-    row_shape.set_element_type(a_shape->element_type());
+    row_shape.set_element_type(a_shape.element_type());
     auto mask_zeros_row = Zeros(body_builder, row_shape);
 
     col_shape.add_dimensions(n);
     col_shape.add_dimensions(1);
-    col_shape.set_element_type(a_shape->element_type());
+    col_shape.set_element_type(a_shape.element_type());
     auto mask_zeros_col = Zeros(body_builder, col_shape);
 
     std::vector<int32> mask_vector(n);
@@ -101,7 +99,7 @@ xla::StatusOr<xla::ComputationDataHandle> CholeskyUnblocked(
     TF_ASSIGN_OR_RETURN(auto a_ii, DynamicSliceInMinorDims(body_builder, body_a,
                                                            {i, i}, {1, 1}));
     // np.dot(row, np.swapaxes(row, -1, -2))
-    xla::ComputationDataHandle diag_dot;
+    xla::XlaOp diag_dot;
     TF_ASSIGN_OR_RETURN(diag_dot, BatchDot(body_builder, row, row,
                                            /*transpose_x=*/false,
                                            /*transpose_y=*/true));
@@ -109,7 +107,7 @@ xla::StatusOr<xla::ComputationDataHandle> CholeskyUnblocked(
     //                                              np.swapaxes(row, -1, -2)))
     auto l_ii = body_builder->Pow(
         body_builder->Sub(a_ii, diag_dot),
-        FloatLiteral(body_builder, a_shape->element_type(), 0.5));
+        FloatLiteral(body_builder, a_shape.element_type(), 0.5));
 
     // a[..., i+1:, i]
     auto ip1 = body_builder->Add(i, body_builder->ConstantR0<int32>(1));
@@ -140,7 +138,7 @@ xla::StatusOr<xla::ComputationDataHandle> CholeskyUnblocked(
     TF_ASSIGN_OR_RETURN(body_l, DynamicUpdateSliceInMinorDims(
                                     body_builder, body_l, l_ii, {i, i}));
 
-    return std::vector<xla::ComputationDataHandle>{body_a, body_l};
+    return std::vector<xla::XlaOp>{body_a, body_l};
   };
 
   TF_ASSIGN_OR_RETURN(
@@ -152,22 +150,20 @@ xla::StatusOr<xla::ComputationDataHandle> CholeskyUnblocked(
 
 }  // namespace
 
-xla::StatusOr<xla::ComputationDataHandle> Cholesky(
-    xla::ComputationBuilder* builder, xla::ComputationDataHandle a,
-    int64 block_size) {
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> a_shape,
-                      builder->GetShape(a));
-  const int ndims = xla::ShapeUtil::Rank(*a_shape);
+xla::StatusOr<xla::XlaOp> Cholesky(xla::XlaBuilder* builder, xla::XlaOp a,
+                                   int64 block_size) {
+  TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a));
+  const int ndims = xla::ShapeUtil::Rank(a_shape);
   if (ndims < 2) {
     return errors::InvalidArgument(
         "Arguments to Cholesky must have rank >= 2: ", ndims);
   }
 
-  const int64 n = xla::ShapeUtil::GetDimension(*a_shape, -1);
-  if (n != xla::ShapeUtil::GetDimension(*a_shape, -2)) {
+  const int64 n = xla::ShapeUtil::GetDimension(a_shape, -1);
+  if (n != xla::ShapeUtil::GetDimension(a_shape, -2)) {
     return errors::InvalidArgument(
         "Arguments to Cholesky must be square matrices: ",
-        xla::ShapeUtil::HumanString(*a_shape));
+        xla::ShapeUtil::HumanString(a_shape));
   }
 
   if (block_size < 1) {
@@ -179,7 +175,7 @@ xla::StatusOr<xla::ComputationDataHandle> Cholesky(
   // Algorithm 1 from
   // Haidar, Azzam, et al. "High-performance Cholesky factorization for GPU-only
   // execution." Proceedings of General Purpose GPUs. ACM, 2017.
-  xla::ComputationDataHandle l = Zeros(builder, *a_shape);
+  xla::XlaOp l = Zeros(builder, a_shape);
   for (int64 i = 0; i < n; i += block_size) {
     int64 k = std::min(block_size, n - i);
     if (i > 0) {
diff --git a/tensorflow/compiler/tf2xla/lib/cholesky.h b/tensorflow/compiler/tf2xla/lib/cholesky.h
index 17da8d8b22..20fca7969e 100644
--- a/tensorflow/compiler/tf2xla/lib/cholesky.h
+++ b/tensorflow/compiler/tf2xla/lib/cholesky.h
@@ -16,8 +16,8 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_TF2XLA_LIB_CHOLESKY_H_
 #define TENSORFLOW_COMPILER_TF2XLA_LIB_CHOLESKY_H_
 
-#include "tensorflow/compiler/xla/client/computation.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 
 namespace tensorflow {
 
@@ -30,9 +30,8 @@ namespace tensorflow {
 // TODO(phawkins): check for negative values on the diagonal and return an
 // error, instead of silently yielding NaNs.
 // TODO(znado): handle the complex Hermitian case
-xla::StatusOr<xla::ComputationDataHandle> Cholesky(
-    xla::ComputationBuilder* builder, xla::ComputationDataHandle a,
-    int64 block_size = 256);
+xla::StatusOr<xla::XlaOp> Cholesky(xla::XlaBuilder* builder, xla::XlaOp a,
+                                   int64 block_size = 256);
 
 }  // namespace tensorflow
 
diff --git a/tensorflow/compiler/tf2xla/lib/scatter.cc b/tensorflow/compiler/tf2xla/lib/scatter.cc
index 45699233ea..d5a27abb25 100644
--- a/tensorflow/compiler/tf2xla/lib/scatter.cc
+++ b/tensorflow/compiler/tf2xla/lib/scatter.cc
@@ -30,24 +30,19 @@ limitations under the License.
 
 namespace tensorflow {
 
-xla::StatusOr<xla::ComputationDataHandle> XlaScatter(
-    const xla::ComputationDataHandle& buffer,
-    const xla::ComputationDataHandle& updates,
-    const xla::ComputationDataHandle& indices, bool indices_are_vectors,
-    const std::function<xla::ComputationDataHandle(
-        xla::ComputationDataHandle, xla::ComputationDataHandle,
-        xla::ComputationBuilder*)>& combiner,
-    xla::ComputationBuilder* builder) {
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> buffer_shape,
-                      builder->GetShape(buffer));
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> updates_shape,
-                      builder->GetShape(updates));
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> indices_shape,
-                      builder->GetShape(indices));
+xla::StatusOr<xla::XlaOp> XlaScatter(
+    const xla::XlaOp& buffer, const xla::XlaOp& updates,
+    const xla::XlaOp& indices, bool indices_are_vectors,
+    const std::function<xla::XlaOp(xla::XlaOp, xla::XlaOp, xla::XlaBuilder*)>&
+        combiner,
+    xla::XlaBuilder* builder) {
+  TF_ASSIGN_OR_RETURN(xla::Shape buffer_shape, builder->GetShape(buffer));
+  TF_RETURN_IF_ERROR(builder->GetShape(updates).status());
+  TF_ASSIGN_OR_RETURN(xla::Shape indices_shape, builder->GetShape(indices));
   gtl::ArraySlice<int64> indices_dims =
-      xla::AsInt64Slice(indices_shape->dimensions());
+      xla::AsInt64Slice(indices_shape.dimensions());
   gtl::ArraySlice<int64> buffer_dims =
-      xla::AsInt64Slice(buffer_shape->dimensions());
+      xla::AsInt64Slice(buffer_shape.dimensions());
 
   // If the indices are N-dimensional, the minor dimension of indices contains
   // the indices to update. Otherwise the indices are all scalars.
@@ -55,12 +50,12 @@ xla::StatusOr<xla::ComputationDataHandle> XlaScatter(
   if (indices_are_vectors) {
     TF_RET_CHECK(!indices_dims.empty());
     num_index_dims = indices_dims.back();
-    if (num_index_dims > xla::ShapeUtil::Rank(*buffer_shape)) {
+    if (num_index_dims > xla::ShapeUtil::Rank(buffer_shape)) {
       return errors::InvalidArgument(
           "The size of the minor dimension of the indices (shape: ",
-          xla::ShapeUtil::HumanString(*indices_shape),
+          xla::ShapeUtil::HumanString(indices_shape),
           ") must be <= the rank of the buffer (shape: ",
-          xla::ShapeUtil::HumanString(*buffer_shape), ")");
+          xla::ShapeUtil::HumanString(buffer_shape), ")");
     }
     indices_dims.pop_back();
   }
@@ -78,10 +73,10 @@ xla::StatusOr<xla::ComputationDataHandle> XlaScatter(
   // If any of the indexed dimensions are zero in the buffer, the update cannot
   // succeed since it updates a slice of size 1.
   for (int64 i = 0; i < num_index_dims; ++i) {
-    if (xla::ShapeUtil::GetDimension(*buffer_shape, i) == 0) {
-      return errors::InvalidArgument(
-          "Scatter dimension ", i, " is of size zero in tensor with shape ",
-          xla::ShapeUtil::HumanString(*buffer_shape));
+    if (xla::ShapeUtil::GetDimension(buffer_shape, i) == 0) {
+      return errors::InvalidArgument("Scatter dimension ", i,
+                                     " is of size zero in tensor with shape ",
+                                     xla::ShapeUtil::HumanString(buffer_shape));
     }
   }
 
@@ -111,18 +106,17 @@ xla::StatusOr<xla::ComputationDataHandle> XlaScatter(
   //   index = dynamic-slice(indices, i)
   //   update = dynamic-slice(updates, i)
   //   buffer = dynamic-update-slice(buffer, update, index)
-  auto body_fn = [&](xla::ComputationDataHandle i,
-                     gtl::ArraySlice<xla::ComputationDataHandle> loop_vars,
-                     xla::ComputationBuilder* body_builder) {
+  auto body_fn = [&](xla::XlaOp i, gtl::ArraySlice<xla::XlaOp> loop_vars,
+                     xla::XlaBuilder* body_builder) {
     auto indices = loop_vars[0];
     auto updates = loop_vars[1];
     auto buffer = loop_vars[2];
 
     auto zero_index = body_builder->ConstantLiteral(
-        xla::Literal::Zero(indices_shape->element_type()));
+        xla::Literal::Zero(indices_shape.element_type()));
 
     // Slice the i-th index from the indices array.
-    xla::ComputationDataHandle index;
+    xla::XlaOp index;
     auto indices_offset = body_builder->Reshape(i, {1});
     if (indices_are_vectors) {
       indices_offset = body_builder->Pad(indices_offset, zero_index,
@@ -180,12 +174,12 @@ xla::StatusOr<xla::ComputationDataHandle> XlaScatter(
     // Apply the update.
     buffer = body_builder->DynamicUpdateSlice(buffer, update, index);
 
-    return std::vector<xla::ComputationDataHandle>{indices, updates, buffer};
+    return std::vector<xla::XlaOp>{indices, updates, buffer};
   };
 
-  TF_ASSIGN_OR_RETURN(
-      auto outputs, XlaForEachIndex(num_indices, indices_shape->element_type(),
-                                    body_fn, init, "scatter", builder));
+  TF_ASSIGN_OR_RETURN(auto outputs,
+                      XlaForEachIndex(num_indices, indices_shape.element_type(),
+                                      body_fn, init, "scatter", builder));
   return outputs[2];
 }
 
diff --git a/tensorflow/compiler/tf2xla/lib/scatter.h b/tensorflow/compiler/tf2xla/lib/scatter.h
index 41e6d3b195..87309e10ed 100644
--- a/tensorflow/compiler/tf2xla/lib/scatter.h
+++ b/tensorflow/compiler/tf2xla/lib/scatter.h
@@ -18,8 +18,8 @@ limitations under the License.
 
 #include <functional>
 
-#include "tensorflow/compiler/xla/client/computation.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 #include "tensorflow/compiler/xla/statusor.h"
 
 namespace tensorflow {
@@ -39,14 +39,12 @@ namespace tensorflow {
 // If a `combiner` is provided, updates are combined with the existing values in
 // the buffer using the combiner function. Otherwise, the updates replace the
 // existing values. The order of updates is implementation-defined.
-xla::StatusOr<xla::ComputationDataHandle> XlaScatter(
-    const xla::ComputationDataHandle& buffer,
-    const xla::ComputationDataHandle& updates,
-    const xla::ComputationDataHandle& indices, bool indices_are_vectors,
-    const std::function<xla::ComputationDataHandle(
-        xla::ComputationDataHandle, xla::ComputationDataHandle,
-        xla::ComputationBuilder*)>& combiner,
-    xla::ComputationBuilder* builder);
+xla::StatusOr<xla::XlaOp> XlaScatter(
+    const xla::XlaOp& buffer, const xla::XlaOp& updates,
+    const xla::XlaOp& indices, bool indices_are_vectors,
+    const std::function<xla::XlaOp(xla::XlaOp, xla::XlaOp, xla::XlaBuilder*)>&
+        combiner,
+    xla::XlaBuilder* builder);
 
 }  // namespace tensorflow
 
diff --git a/tensorflow/compiler/tf2xla/lib/triangular_solve.cc b/tensorflow/compiler/tf2xla/lib/triangular_solve.cc
index 9bf5821b54..d0279d4412 100644
--- a/tensorflow/compiler/tf2xla/lib/triangular_solve.cc
+++ b/tensorflow/compiler/tf2xla/lib/triangular_solve.cc
@@ -29,21 +29,20 @@ limitations under the License.
 
 namespace tensorflow {
 
-xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& a,
-    xla::ComputationDataHandle b, bool left_side, bool lower, bool transpose_a,
-    bool conjugate_a, int64 block_size) {
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> a_shape,
-                      builder->GetShape(a));
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> b_shape,
-                      builder->GetShape(b));
-  if (xla::ShapeUtil::Rank(*a_shape) != xla::ShapeUtil::Rank(*b_shape)) {
+xla::StatusOr<xla::XlaOp> TriangularSolve(xla::XlaBuilder* builder,
+                                          const xla::XlaOp& a, xla::XlaOp b,
+                                          bool left_side, bool lower,
+                                          bool transpose_a, bool conjugate_a,
+                                          int64 block_size) {
+  TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a));
+  TF_ASSIGN_OR_RETURN(xla::Shape b_shape, builder->GetShape(b));
+  if (xla::ShapeUtil::Rank(a_shape) != xla::ShapeUtil::Rank(b_shape)) {
     return errors::InvalidArgument(
         "Arguments to TriangularSolve have different ranks: ",
-        xla::ShapeUtil::HumanString(*a_shape), " vs. ",
-        xla::ShapeUtil::HumanString(*b_shape));
+        xla::ShapeUtil::HumanString(a_shape), " vs. ",
+        xla::ShapeUtil::HumanString(b_shape));
   }
-  const int ndims = xla::ShapeUtil::Rank(*a_shape);
+  const int ndims = xla::ShapeUtil::Rank(a_shape);
   if (ndims < 2) {
     return errors::InvalidArgument(
         "Arguments to TriangularSolve must have rank >= 2: ", ndims);
@@ -51,30 +50,30 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
   // The batch dimensions must be equal.
   std::vector<int64> batch_dimensions;
   for (int i = 0; i < ndims - 2; ++i) {
-    int64 a_size = a_shape->dimensions(i);
-    int64 b_size = b_shape->dimensions(i);
+    int64 a_size = a_shape.dimensions(i);
+    int64 b_size = b_shape.dimensions(i);
     if (a_size != b_size) {
       return errors::InvalidArgument(
           "Batch dimensions of arguments to TriangularSolve must be equal: ",
-          xla::ShapeUtil::HumanString(*a_shape), " vs ",
-          xla::ShapeUtil::HumanString(*b_shape));
+          xla::ShapeUtil::HumanString(a_shape), " vs ",
+          xla::ShapeUtil::HumanString(b_shape));
     }
     batch_dimensions.push_back(a_size);
   }
 
-  if (xla::ShapeUtil::GetDimension(*a_shape, -1) !=
-      xla::ShapeUtil::GetDimension(*a_shape, -2)) {
+  if (xla::ShapeUtil::GetDimension(a_shape, -1) !=
+      xla::ShapeUtil::GetDimension(a_shape, -2)) {
     return errors::InvalidArgument(
         "The 'a' arguments to TriangularSolve must be square matrices: ",
-        xla::ShapeUtil::HumanString(*a_shape));
+        xla::ShapeUtil::HumanString(a_shape));
   }
-  const int64 m = xla::ShapeUtil::GetDimension(*b_shape, -2);
-  const int64 n = xla::ShapeUtil::GetDimension(*b_shape, -1);
-  if ((left_side ? m : n) != xla::ShapeUtil::GetDimension(*a_shape, -1)) {
+  const int64 m = xla::ShapeUtil::GetDimension(b_shape, -2);
+  const int64 n = xla::ShapeUtil::GetDimension(b_shape, -1);
+  if ((left_side ? m : n) != xla::ShapeUtil::GetDimension(a_shape, -1)) {
     return errors::InvalidArgument(
         "Arguments to TriangularSolve have incompatible matrix shapes: ",
-        xla::ShapeUtil::HumanString(*a_shape), " vs ",
-        xla::ShapeUtil::HumanString(*b_shape));
+        xla::ShapeUtil::HumanString(a_shape), " vs ",
+        xla::ShapeUtil::HumanString(b_shape));
   }
 
   if (block_size < 1) {
@@ -85,24 +84,23 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
 
   // Applies a complex conjugation operation if `a` is complex and `conjugate_a`
   // is true, otherwise returns its argument.
-  auto maybe_conj = [&](xla::ComputationBuilder* builder,
-                        xla::ComputationDataHandle x) {
-    auto perform_conj = a_shape->element_type() == xla::C64 && conjugate_a;
+  auto maybe_conj = [&](xla::XlaBuilder* builder, xla::XlaOp x) {
+    auto perform_conj = a_shape.element_type() == xla::C64 && conjugate_a;
     return perform_conj ? builder->Conj(x) : x;
   };
 
-  std::map<int, xla::Computation> base_computations;
+  std::map<int, xla::XlaComputation> base_computations;
   auto get_base_triangular_solve =
-      [&](int k) -> xla::StatusOr<xla::Computation*> {
-    xla::Computation& computation = base_computations[k];
+      [&](int k) -> xla::StatusOr<xla::XlaComputation*> {
+    xla::XlaComputation& computation = base_computations[k];
     if (computation.IsNull()) {
-      std::unique_ptr<xla::ComputationBuilder> sub = builder->CreateSubBuilder(
+      std::unique_ptr<xla::XlaBuilder> sub = builder->CreateSubBuilder(
           tensorflow::strings::StrCat("trsm_base_", k));
 
       auto a_param = sub->Parameter(
           0,
           xla::ShapeUtil::MakeShape(
-              b_shape->element_type(),
+              b_shape.element_type(),
               PrependMajorDims(sub.get(), batch_dimensions, {k, k})),
           "a");
 
@@ -115,7 +113,7 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
       auto b_param = sub->Parameter(
           1,
           xla::ShapeUtil::MakeShape(
-              b_shape->element_type(),
+              b_shape.element_type(),
               PrependMajorDims(sub.get(), batch_dimensions, b_lastd)),
           "b");
 
@@ -142,7 +140,7 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
     return &computation;
   };
 
-  xla::ComputationDataHandle output = Zeros(builder, *b_shape);
+  xla::XlaOp output = Zeros(builder, b_shape);
 
   // Right-looking blocked triangular solve.
   // For an explanation of the algorithm, see the TRSM discussion in:
@@ -165,9 +163,9 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
                           SliceInMinorDims(builder, a, {i, i}, {i + k, i + k}));
       TF_ASSIGN_OR_RETURN(auto b_slice,
                           SliceInMinorDims(builder, b, {0, i}, {m, i + k}));
-      xla::ComputationDataHandle update;
+      xla::XlaOp update;
       if (k > 1) {
-        TF_ASSIGN_OR_RETURN(xla::Computation * solve,
+        TF_ASSIGN_OR_RETURN(xla::XlaComputation * solve,
                             get_base_triangular_solve(k));
         update = builder->Call(*solve, {a_slice, b_slice});
       } else {
@@ -181,7 +179,7 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
       //   a_slice_2 = T(a_slice_2) if transpose_a else a_slice_2
       //   b[..., :, i+k:] -= np.matmul(output[..., :, i:i+k], a_slice_2)
       if (i + k < n) {
-        xla::ComputationDataHandle a_slice_2;
+        xla::XlaOp a_slice_2;
         if (lower) {
           TF_ASSIGN_OR_RETURN(
               a_slice_2, SliceInMinorDims(builder, a, {i + k, i}, {n, i + k}));
@@ -215,9 +213,9 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
                           SliceInMinorDims(builder, a, {i, i}, {i + k, i + k}));
       TF_ASSIGN_OR_RETURN(auto b_slice,
                           SliceInMinorDims(builder, b, {i, 0}, {i + k, n}));
-      xla::ComputationDataHandle update;
+      xla::XlaOp update;
       if (k > 1) {
-        TF_ASSIGN_OR_RETURN(xla::Computation * solve,
+        TF_ASSIGN_OR_RETURN(xla::XlaComputation * solve,
                             get_base_triangular_solve(k));
         update = builder->Call(*solve, {a_slice, b_slice});
       } else {
@@ -231,7 +229,7 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
       //   a_slice_2 = T(a_slice_2) if transpose_a else a_slice_2
       //   b[..., i+k:, :] -= np.matmul(a_slice_2, output[..., i:i+k, :])
       if (i + k < m) {
-        xla::ComputationDataHandle a_slice_2;
+        xla::XlaOp a_slice_2;
         if (lower) {
           TF_ASSIGN_OR_RETURN(
               a_slice_2, SliceInMinorDims(builder, a, {i + k, i}, {m, i + k}));
@@ -264,9 +262,9 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
                           SliceInMinorDims(builder, a, {i, i}, {i + k, i + k}));
       TF_ASSIGN_OR_RETURN(auto b_slice,
                           SliceInMinorDims(builder, b, {0, i}, {m, i + k}));
-      xla::ComputationDataHandle update;
+      xla::XlaOp update;
       if (k > 1) {
-        TF_ASSIGN_OR_RETURN(xla::Computation * solve,
+        TF_ASSIGN_OR_RETURN(xla::XlaComputation * solve,
                             get_base_triangular_solve(k));
         update = builder->Call(*solve, {a_slice, b_slice});
       } else {
@@ -280,7 +278,7 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
       //   a_slice_2 = T(a_slice_2) if transpose_a else a_slice_2
       //   b[..., :, :i] -= np.matmul(out[..., :, i:i+k], a_slice_2)
       if (i - k >= 0) {
-        xla::ComputationDataHandle a_slice_2;
+        xla::XlaOp a_slice_2;
         if (lower) {
           TF_ASSIGN_OR_RETURN(a_slice_2,
                               SliceInMinorDims(builder, a, {i, 0}, {i + k, i}));
@@ -314,9 +312,9 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
                           SliceInMinorDims(builder, a, {i, i}, {i + k, i + k}));
       TF_ASSIGN_OR_RETURN(auto b_slice,
                           SliceInMinorDims(builder, b, {i, 0}, {i + k, n}));
-      xla::ComputationDataHandle update;
+      xla::XlaOp update;
       if (k > 1) {
-        TF_ASSIGN_OR_RETURN(xla::Computation * solve,
+        TF_ASSIGN_OR_RETURN(xla::XlaComputation * solve,
                             get_base_triangular_solve(k));
         update = builder->Call(*solve, {a_slice, b_slice});
       } else {
@@ -330,7 +328,7 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
       //   a_slice_2 = T(a_slice_2) if transpose_a else a_slice_2
       //   b[..., :i, :] -= np.matmul(a_slice_2, out[..., i:i+k, :])
       if (i - k >= 0) {
-        xla::ComputationDataHandle a_slice_2;
+        xla::XlaOp a_slice_2;
         if (lower) {
           TF_ASSIGN_OR_RETURN(a_slice_2,
                               SliceInMinorDims(builder, a, {i, 0}, {i + k, i}));
@@ -356,26 +354,25 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
   return output;
 }
 
-xla::StatusOr<xla::ComputationDataHandle> TriangularSolveLeftLooking(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& a,
-    const xla::ComputationDataHandle& b, bool transpose_a, bool conjugate_a) {
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> a_shape,
-                      builder->GetShape(a));
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> b_shape,
-                      builder->GetShape(b));
-  const int64 m = xla::ShapeUtil::GetDimension(*b_shape, -2);
-  const int64 n = xla::ShapeUtil::GetDimension(*b_shape, -1);
-  const int64 ndims = xla::ShapeUtil::Rank(*a_shape);
+xla::StatusOr<xla::XlaOp> TriangularSolveLeftLooking(xla::XlaBuilder* builder,
+                                                     const xla::XlaOp& a,
+                                                     const xla::XlaOp& b,
+                                                     bool transpose_a,
+                                                     bool conjugate_a) {
+  TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a));
+  TF_ASSIGN_OR_RETURN(xla::Shape b_shape, builder->GetShape(b));
+  const int64 m = xla::ShapeUtil::GetDimension(b_shape, -2);
+  const int64 n = xla::ShapeUtil::GetDimension(b_shape, -1);
+  const int64 ndims = xla::ShapeUtil::Rank(a_shape);
 
   std::vector<int64> batch_dimensions;
   for (int i = 0; i < ndims - 2; ++i) {
-    int64 a_size = a_shape->dimensions(i);
+    int64 a_size = a_shape.dimensions(i);
     batch_dimensions.push_back(a_size);
   }
 
-  auto maybe_conj = [&](xla::ComputationBuilder* builder,
-                        xla::ComputationDataHandle x) {
-    auto perform_conj = a_shape->element_type() == xla::C64 && conjugate_a;
+  auto maybe_conj = [&](xla::XlaBuilder* builder, xla::XlaOp x) {
+    auto perform_conj = a_shape.element_type() == xla::C64 && conjugate_a;
     return perform_conj ? builder->Conj(x) : x;
   };
 
@@ -387,7 +384,7 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolveLeftLooking(
   //   output[..., m-1:, :] = b[..., m-1:, :] / a[..., m-1:, m-1:]
   // else:
   //   output[..., :1, :] = b[..., :1, :] / a[..., :1, :1]
-  xla::ComputationDataHandle output = Zeros(builder, *b_shape);
+  xla::XlaOp output = Zeros(builder, b_shape);
   {
     auto i = transpose_a ? m - 1 : 0;
     TF_ASSIGN_OR_RETURN(auto a_slice,
@@ -408,11 +405,11 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolveLeftLooking(
       // The loop iteration counter is a scalar, incremented each iteration.
       xla::ShapeUtil::MakeShape(xla::S32, {}),
       // The output has the shape of b, with one row updated each iteration.
-      *b_shape,
+      b_shape,
       // The coefficient matrix a is a loop invariant.
-      *a_shape,
+      a_shape,
       // The right-hand-side matrix b is a loop invariant.
-      *b_shape};
+      b_shape};
   xla::Shape tuple_shape = xla::ShapeUtil::MakeTupleShape(tuple_shapes);
   auto init_i = builder->ConstantR0<int32>(transpose_a ? m - 2 : 1);
   auto init = builder->Tuple({init_i, output, a, b});
@@ -421,7 +418,7 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolveLeftLooking(
   // def cond_fun(loop_carry):
   //   i, output, a, b = loop_carry
   //   return i >= 0 if transpose_a else i < m
-  std::unique_ptr<xla::ComputationBuilder> condb =
+  std::unique_ptr<xla::XlaBuilder> condb =
       builder->CreateSubBuilder("TriangularSolveLeftLookingWhileCond");
   {
     auto i = condb->GetTupleElement(
@@ -451,7 +448,7 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolveLeftLooking(
   //     return (i + 1, output, a, b)
   // We have to do some extra FLOPs propagating zeros in the matrix multiply
   // because we can't have the size of its arguments depend on the loop counter.
-  std::unique_ptr<xla::ComputationBuilder> bodyb =
+  std::unique_ptr<xla::XlaBuilder> bodyb =
       builder->CreateSubBuilder("TriangularSolveLeftLookingWhileBody");
   {
     auto input_tuple = bodyb->Parameter(0, tuple_shape,
@@ -475,7 +472,7 @@ xla::StatusOr<xla::ComputationDataHandle> TriangularSolveLeftLooking(
     // But since we can't have intermediate array sizes depend on the loop
     // counter, we instead exploit the fact that we initialized the output to
     // all zeros and use that as zero-padding (doing unnecessary FLOPs).
-    xla::ComputationDataHandle a_row;
+    xla::XlaOp a_row;
     if (transpose_a) {
       TF_ASSIGN_OR_RETURN(a_row, DynamicSliceInMinorDims(bodyb.get(), body_a,
                                                          {zero, i}, {m, 1}));
diff --git a/tensorflow/compiler/tf2xla/lib/triangular_solve.h b/tensorflow/compiler/tf2xla/lib/triangular_solve.h
index e32223bfdd..fd8f2489d1 100644
--- a/tensorflow/compiler/tf2xla/lib/triangular_solve.h
+++ b/tensorflow/compiler/tf2xla/lib/triangular_solve.h
@@ -16,8 +16,8 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_TF2XLA_LIB_TRIANGULAR_SOLVE_H_
 #define TENSORFLOW_COMPILER_TF2XLA_LIB_TRIANGULAR_SOLVE_H_
 
-#include "tensorflow/compiler/xla/client/computation.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 
 namespace tensorflow {
 
@@ -57,14 +57,17 @@ namespace tensorflow {
 //
 // Uses a blocked algorithm if `block_size` is > 1; if block_size == 1 then no
 // blocking is used.
-xla::StatusOr<xla::ComputationDataHandle> TriangularSolve(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& a,
-    xla::ComputationDataHandle b, bool left_side, bool lower, bool transpose_a,
-    bool conjugate_a, int64 block_size = 256);
+xla::StatusOr<xla::XlaOp> TriangularSolve(xla::XlaBuilder* builder,
+                                          const xla::XlaOp& a, xla::XlaOp b,
+                                          bool left_side, bool lower,
+                                          bool transpose_a, bool conjugate_a,
+                                          int64 block_size = 256);
 
-xla::StatusOr<xla::ComputationDataHandle> TriangularSolveLeftLooking(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& a,
-    const xla::ComputationDataHandle& b, bool transpose_a, bool conjugate_a);
+xla::StatusOr<xla::XlaOp> TriangularSolveLeftLooking(xla::XlaBuilder* builder,
+                                                     const xla::XlaOp& a,
+                                                     const xla::XlaOp& b,
+                                                     bool transpose_a,
+                                                     bool conjugate_a);
 
 }  // namespace tensorflow
 
diff --git a/tensorflow/compiler/tf2xla/lib/triangular_solve_test.cc b/tensorflow/compiler/tf2xla/lib/triangular_solve_test.cc
index 6617070629..87ea4763f7 100644
--- a/tensorflow/compiler/tf2xla/lib/triangular_solve_test.cc
+++ b/tensorflow/compiler/tf2xla/lib/triangular_solve_test.cc
@@ -20,7 +20,7 @@ limitations under the License.
 #include <vector>
 
 #include "tensorflow/compiler/xla/array2d.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/compiler/xla/statusor.h"
 #include "tensorflow/compiler/xla/test.h"
@@ -80,9 +80,9 @@ xla::Array2D<float> AValsFull() {
 }
 
 XLA_TEST_F(TriangularSolveTest, SimpleRightLowerTranspose) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, b;
+  xla::XlaOp a, b;
   auto a_data = CreateR2Parameter<float>(AValsLower(), 0, "a", &builder, &a);
   auto b_data = CreateR2Parameter<float>(BValsRight(), 1, "b", &builder, &b);
   auto result = TriangularSolve(&builder, a, b,
@@ -102,9 +102,9 @@ XLA_TEST_F(TriangularSolveTest, SimpleRightLowerTranspose) {
 }
 
 XLA_TEST_F(TriangularSolveTest, SimpleRightLowerNotranspose) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, b;
+  xla::XlaOp a, b;
   auto a_data = CreateR2Parameter<float>(AValsLower(), 0, "a", &builder, &a);
   auto b_data = CreateR2Parameter<float>(BValsRight(), 1, "b", &builder, &b);
   auto result = TriangularSolve(&builder, a, b,
@@ -124,9 +124,9 @@ XLA_TEST_F(TriangularSolveTest, SimpleRightLowerNotranspose) {
 }
 
 XLA_TEST_F(TriangularSolveTest, SimpleRightUpperTranspose) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, b;
+  xla::XlaOp a, b;
   auto a_data = CreateR2Parameter<float>(AValsUpper(), 0, "a", &builder, &a);
   auto b_data = CreateR2Parameter<float>(BValsRight(), 1, "b", &builder, &b);
   auto result = TriangularSolve(&builder, a, b,
@@ -146,9 +146,9 @@ XLA_TEST_F(TriangularSolveTest, SimpleRightUpperTranspose) {
 }
 
 XLA_TEST_F(TriangularSolveTest, SimpleRightUpperNotranspose) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, b;
+  xla::XlaOp a, b;
   auto a_data = CreateR2Parameter<float>(AValsUpper(), 0, "a", &builder, &a);
   auto b_data = CreateR2Parameter<float>(BValsRight(), 1, "b", &builder, &b);
   auto result = TriangularSolve(&builder, a, b,
@@ -168,9 +168,9 @@ XLA_TEST_F(TriangularSolveTest, SimpleRightUpperNotranspose) {
 }
 
 XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerTranspose) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, b;
+  xla::XlaOp a, b;
   auto a_data = CreateR2Parameter<float>(AValsLower(), 0, "a", &builder, &a);
   auto b_data = CreateR2Parameter<float>(BValsLeft(), 1, "b", &builder, &b);
   auto result = TriangularSolve(&builder, a, b,
@@ -191,9 +191,9 @@ XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerTranspose) {
 }
 
 XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerNotranspose) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, b;
+  xla::XlaOp a, b;
   auto a_data = CreateR2Parameter<float>(AValsLower(), 0, "a", &builder, &a);
   auto b_data = CreateR2Parameter<float>(BValsLeft(), 1, "b", &builder, &b);
   auto result = TriangularSolve(&builder, a, b,
@@ -214,9 +214,9 @@ XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerNotranspose) {
 }
 
 XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperTranspose) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, b;
+  xla::XlaOp a, b;
   auto a_data = CreateR2Parameter<float>(AValsUpper(), 0, "a", &builder, &a);
   auto b_data = CreateR2Parameter<float>(BValsLeft(), 1, "b", &builder, &b);
   auto result = TriangularSolve(&builder, a, b,
@@ -237,9 +237,9 @@ XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperTranspose) {
 }
 
 XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperNotranspose) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, b;
+  xla::XlaOp a, b;
   auto a_data = CreateR2Parameter<float>(AValsUpper(), 0, "a", &builder, &a);
   auto b_data = CreateR2Parameter<float>(BValsLeft(), 1, "b", &builder, &b);
   auto result = TriangularSolve(&builder, a, b,
@@ -260,9 +260,9 @@ XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperNotranspose) {
 }
 
 XLA_TEST_F(TriangularSolveTest, SimpleRightLowerTransposeConjugate) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, b;
+  xla::XlaOp a, b;
   auto a_data =
       CreateR2Parameter<complex64>(AValsLowerComplex(), 0, "a", &builder, &a);
   auto b_data =
@@ -288,9 +288,9 @@ XLA_TEST_F(TriangularSolveTest, SimpleRightLowerTransposeConjugate) {
 }
 
 XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperTransposeNoconjugate) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, b;
+  xla::XlaOp a, b;
   auto a_data =
       CreateR2Parameter<complex64>(AValsUpperComplex(), 0, "a", &builder, &a);
   auto b_data =
@@ -318,9 +318,9 @@ XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperTransposeNoconjugate) {
 }
 
 XLA_TEST_F(TriangularSolveLeftLookingTest, Simple) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, b;
+  xla::XlaOp a, b;
   auto a_data = CreateR2Parameter<float>(AValsLower(), 0, "a", &builder, &a);
   auto b_data = CreateR2Parameter<float>(BValsLeft(), 1, "b", &builder, &b);
   auto result = TriangularSolveLeftLooking(&builder, a, b,
@@ -340,9 +340,9 @@ XLA_TEST_F(TriangularSolveLeftLookingTest, Simple) {
 }
 
 XLA_TEST_F(TriangularSolveLeftLookingTest, NonzeroUpperTriangle) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, b;
+  xla::XlaOp a, b;
   auto a_data = CreateR2Parameter<float>(AValsFull(), 0, "a", &builder, &a);
   auto b_data = CreateR2Parameter<float>(BValsLeft(), 1, "b", &builder, &b);
   auto result = TriangularSolveLeftLooking(&builder, a, b,
diff --git a/tensorflow/compiler/tf2xla/lib/util.cc b/tensorflow/compiler/tf2xla/lib/util.cc
index 31d823ca33..cc7b13571c 100644
--- a/tensorflow/compiler/tf2xla/lib/util.cc
+++ b/tensorflow/compiler/tf2xla/lib/util.cc
@@ -27,15 +27,14 @@ limitations under the License.
 
 namespace tensorflow {
 
-xla::ComputationDataHandle Zeros(xla::ComputationBuilder* builder,
-                                 const xla::Shape& shape) {
+xla::XlaOp Zeros(xla::XlaBuilder* builder, const xla::Shape& shape) {
   return builder->Broadcast(
       builder->ConstantLiteral(xla::Literal::Zero(shape.element_type())),
       xla::AsInt64Slice(shape.dimensions()));
 }
 
-xla::ComputationDataHandle FloatLiteral(xla::ComputationBuilder* builder,
-                                        xla::PrimitiveType type, double value) {
+xla::XlaOp FloatLiteral(xla::XlaBuilder* builder, xla::PrimitiveType type,
+                        double value) {
   switch (type) {
     case xla::F16:
       return builder->ConstantR0<xla::half>(static_cast<xla::half>(value));
@@ -57,9 +56,8 @@ xla::ComputationDataHandle FloatLiteral(xla::ComputationBuilder* builder,
   }
 }
 
-xla::ComputationDataHandle IntegerLiteral(xla::ComputationBuilder* builder,
-                                          xla::PrimitiveType type,
-                                          int64 value) {
+xla::XlaOp IntegerLiteral(xla::XlaBuilder* builder, xla::PrimitiveType type,
+                          int64 value) {
   xla::Literal literal;
   switch (type) {
     case xla::U8:
@@ -112,17 +110,18 @@ xla::ComputationDataHandle IntegerLiteral(xla::ComputationBuilder* builder,
   return builder->ConstantLiteral(literal);
 }
 
-xla::StatusOr<xla::ComputationDataHandle> SliceInMinorDims(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x,
-    gtl::ArraySlice<int64> start, gtl::ArraySlice<int64> end) {
+xla::StatusOr<xla::XlaOp> SliceInMinorDims(xla::XlaBuilder* builder,
+                                           const xla::XlaOp& x,
+                                           gtl::ArraySlice<int64> start,
+                                           gtl::ArraySlice<int64> end) {
   TF_RET_CHECK(start.size() == end.size());
   int64 n_minor_dims = start.size();
 
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> shape, builder->GetShape(x));
+  TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(x));
 
-  const int64 n_dims = xla::ShapeUtil::Rank(*shape);
+  const int64 n_dims = xla::ShapeUtil::Rank(shape);
   TF_RET_CHECK(n_minor_dims <= n_dims);
-  gtl::ArraySlice<int64> major_dims(xla::AsInt64Slice(shape->dimensions()),
+  gtl::ArraySlice<int64> major_dims(xla::AsInt64Slice(shape.dimensions()),
                                     /*pos=*/0,
                                     /*len=*/n_dims - n_minor_dims);
 
@@ -140,7 +139,7 @@ xla::StatusOr<xla::ComputationDataHandle> SliceInMinorDims(
   return builder->Slice(x, padded_start, padded_end, strides);
 }
 
-std::vector<int64> PrependMajorDims(xla::ComputationBuilder* builder,
+std::vector<int64> PrependMajorDims(xla::XlaBuilder* builder,
                                     const gtl::ArraySlice<int64>& major_dims,
                                     const gtl::ArraySlice<int64>& indices) {
   std::vector<int64> output(indices.size() + major_dims.size());
@@ -149,16 +148,16 @@ std::vector<int64> PrependMajorDims(xla::ComputationBuilder* builder,
   return output;
 }
 
-xla::StatusOr<xla::ComputationDataHandle> DynamicSliceInMinorDims(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x,
-    const std::vector<xla::ComputationDataHandle>& starts,
+xla::StatusOr<xla::XlaOp> DynamicSliceInMinorDims(
+    xla::XlaBuilder* builder, const xla::XlaOp& x,
+    const std::vector<xla::XlaOp>& starts,
     const gtl::ArraySlice<int64>& sizes) {
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> shape, builder->GetShape(x));
-  const int64 n_dims = xla::ShapeUtil::Rank(*shape);
+  TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(x));
+  const int64 n_dims = xla::ShapeUtil::Rank(shape);
   int64 n_minor_dims = starts.size();
   TF_RET_CHECK(n_minor_dims == sizes.size());
   TF_RET_CHECK(n_minor_dims <= n_dims);
-  gtl::ArraySlice<int64> major_dims(xla::AsInt64Slice(shape->dimensions()),
+  gtl::ArraySlice<int64> major_dims(xla::AsInt64Slice(shape.dimensions()),
                                     /*pos=*/0,
                                     /*len=*/n_dims - sizes.size());
   TF_ASSIGN_OR_RETURN(auto padded_starts,
@@ -167,27 +166,29 @@ xla::StatusOr<xla::ComputationDataHandle> DynamicSliceInMinorDims(
   return builder->DynamicSlice(x, padded_starts, padded_sizes);
 }
 
-xla::StatusOr<xla::ComputationDataHandle> UpdateSlice(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x,
-    const xla::ComputationDataHandle& update, gtl::ArraySlice<int64> start) {
+xla::StatusOr<xla::XlaOp> UpdateSlice(xla::XlaBuilder* builder,
+                                      const xla::XlaOp& x,
+                                      const xla::XlaOp& update,
+                                      gtl::ArraySlice<int64> start) {
   // TODO(phawkins): make int64 work on all backends, remove the int32 cast.
   std::vector<int32> start_as_int32(start.begin(), start.end());
   auto start_constant = builder->ConstantR1<int32>(start_as_int32);
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> shape, builder->GetShape(x));
-  const int64 n_dims = xla::ShapeUtil::Rank(*shape);
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> start_constant_shape,
+  TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(x));
+  const int64 n_dims = xla::ShapeUtil::Rank(shape);
+  TF_ASSIGN_OR_RETURN(xla::Shape start_constant_shape,
                       builder->GetShape(start_constant));
   const int64 start_length =
-      xla::ShapeUtil::GetDimension(*start_constant_shape, -1);
+      xla::ShapeUtil::GetDimension(start_constant_shape, -1);
   TF_RET_CHECK(start_length == n_dims);
   return builder->DynamicUpdateSlice(x, update, start_constant);
 }
 
-xla::StatusOr<xla::ComputationDataHandle> UpdateSliceInMinorDims(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x,
-    const xla::ComputationDataHandle& update, gtl::ArraySlice<int64> start) {
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> shape, builder->GetShape(x));
-  const int64 n_dims = xla::ShapeUtil::Rank(*shape);
+xla::StatusOr<xla::XlaOp> UpdateSliceInMinorDims(xla::XlaBuilder* builder,
+                                                 const xla::XlaOp& x,
+                                                 const xla::XlaOp& update,
+                                                 gtl::ArraySlice<int64> start) {
+  TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(x));
+  const int64 n_dims = xla::ShapeUtil::Rank(shape);
   const int64 n_minor_dims = start.size();
   TF_RET_CHECK(n_minor_dims <= n_dims);
   std::vector<int64> padded_start(n_dims, 0);
@@ -196,22 +197,21 @@ xla::StatusOr<xla::ComputationDataHandle> UpdateSliceInMinorDims(
   return UpdateSlice(builder, x, update, padded_start);
 }
 
-xla::StatusOr<xla::ComputationDataHandle> DynamicUpdateSliceInMinorDims(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x,
-    const xla::ComputationDataHandle& update,
-    const std::vector<xla::ComputationDataHandle>& starts) {
+xla::StatusOr<xla::XlaOp> DynamicUpdateSliceInMinorDims(
+    xla::XlaBuilder* builder, const xla::XlaOp& x, const xla::XlaOp& update,
+    const std::vector<xla::XlaOp>& starts) {
   TF_ASSIGN_OR_RETURN(auto padded_starts,
                       PrependZerosInMajorDims(builder, x, starts));
   return builder->DynamicUpdateSlice(x, update, padded_starts);
 }
 
-xla::StatusOr<xla::ComputationDataHandle> PrependZerosInMajorDims(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x,
-    const std::vector<xla::ComputationDataHandle>& starts) {
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> shape, builder->GetShape(x));
-  const int64 n_dims = xla::ShapeUtil::Rank(*shape);
+xla::StatusOr<xla::XlaOp> PrependZerosInMajorDims(
+    xla::XlaBuilder* builder, const xla::XlaOp& x,
+    const std::vector<xla::XlaOp>& starts) {
+  TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(x));
+  const int64 n_dims = xla::ShapeUtil::Rank(shape);
   auto zero = builder->Reshape(builder->ConstantR0<int32>(0), {1});
-  std::vector<xla::ComputationDataHandle> padded_starts(n_dims, zero);
+  std::vector<xla::XlaOp> padded_starts(n_dims, zero);
   for (int i = 0; i < starts.size(); ++i) {
     padded_starts[n_dims - starts.size() + i] =
         builder->Reshape(starts[i], {1});
@@ -219,10 +219,10 @@ xla::StatusOr<xla::ComputationDataHandle> PrependZerosInMajorDims(
   return builder->ConcatInDim(padded_starts, 0);
 }
 
-xla::StatusOr<xla::ComputationDataHandle> TransposeInMinorDims(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x) {
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::Shape> shape, builder->GetShape(x));
-  const int64 n_dims = xla::ShapeUtil::Rank(*shape);
+xla::StatusOr<xla::XlaOp> TransposeInMinorDims(xla::XlaBuilder* builder,
+                                               const xla::XlaOp& x) {
+  TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(x));
+  const int64 n_dims = xla::ShapeUtil::Rank(shape);
   TF_RET_CHECK(n_dims >= 2);
   std::vector<int64> permutation(n_dims);
   std::iota(permutation.begin(), permutation.end(), 0);
diff --git a/tensorflow/compiler/tf2xla/lib/util.h b/tensorflow/compiler/tf2xla/lib/util.h
index b684123f13..3df44ef035 100644
--- a/tensorflow/compiler/tf2xla/lib/util.h
+++ b/tensorflow/compiler/tf2xla/lib/util.h
@@ -16,75 +16,74 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_TF2XLA_LIB_UTIL_H_
 #define TENSORFLOW_COMPILER_TF2XLA_LIB_UTIL_H_
 
-#include "tensorflow/compiler/xla/client/computation.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 #include "tensorflow/compiler/xla/statusor.h"
 #include "tensorflow/core/lib/gtl/array_slice.h"
 
 namespace tensorflow {
 
 // Returns a zero-filled tensor with shape `shape`.
-xla::ComputationDataHandle Zeros(xla::ComputationBuilder* builder,
-                                 const xla::Shape& shape);
+xla::XlaOp Zeros(xla::XlaBuilder* builder, const xla::Shape& shape);
 
 // Returns a floating point scalar constant of 'type' with 'value'.
 // If 'type' is complex, returns a real value with zero imaginary component.
-xla::ComputationDataHandle FloatLiteral(xla::ComputationBuilder* builder,
-                                        xla::PrimitiveType type, double value);
+xla::XlaOp FloatLiteral(xla::XlaBuilder* builder, xla::PrimitiveType type,
+                        double value);
 
 // Makes a 1D tensor [0, ..., x, y] from two tensors x and y with zeros
 // prepended until the array is length n_dims.
-xla::ComputationDataHandle PrependZerosInMajorDims(
-    xla::ComputationBuilder* builder,
-    gtl::ArraySlice<xla::ComputationDataHandle> starts);
+xla::XlaOp PrependZerosInMajorDims(xla::XlaBuilder* builder,
+                                   gtl::ArraySlice<xla::XlaOp> starts);
 
 // Returns a integer scalar constant of 'type' with 'value'.
 // If 'type' is complex, returns a real value with zero imaginary component.
-xla::ComputationDataHandle IntegerLiteral(xla::ComputationBuilder* builder,
-                                          xla::PrimitiveType type, int64 value);
+xla::XlaOp IntegerLiteral(xla::XlaBuilder* builder, xla::PrimitiveType type,
+                          int64 value);
 
 // Builds a vector of zeros of length rank(x) with the last two values being
 // those in `starts`.
-xla::StatusOr<xla::ComputationDataHandle> PrependZerosInMajorDims(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x,
-    const std::vector<xla::ComputationDataHandle>& starts);
+xla::StatusOr<xla::XlaOp> PrependZerosInMajorDims(
+    xla::XlaBuilder* builder, const xla::XlaOp& x,
+    const std::vector<xla::XlaOp>& starts);
 
 // Performs a slice in the minor dimensions of a Tensor.
-xla::StatusOr<xla::ComputationDataHandle> SliceInMinorDims(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x,
-    gtl::ArraySlice<int64> start, gtl::ArraySlice<int64> end);
+xla::StatusOr<xla::XlaOp> SliceInMinorDims(xla::XlaBuilder* builder,
+                                           const xla::XlaOp& x,
+                                           gtl::ArraySlice<int64> start,
+                                           gtl::ArraySlice<int64> end);
 
 // Builds a 1-d vector out of a concatenation of `major_dims` and `starts`.
-std::vector<int64> PrependMajorDims(xla::ComputationBuilder* builder,
+std::vector<int64> PrependMajorDims(xla::XlaBuilder* builder,
                                     const gtl::ArraySlice<int64>& major_dims,
                                     const gtl::ArraySlice<int64>& indices);
 
 // Performs a dynamic slice in the minor dimensions of a Tensor.
-xla::StatusOr<xla::ComputationDataHandle> DynamicSliceInMinorDims(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x,
-    const std::vector<xla::ComputationDataHandle>& starts,
-    const gtl::ArraySlice<int64>& sizes);
+xla::StatusOr<xla::XlaOp> DynamicSliceInMinorDims(
+    xla::XlaBuilder* builder, const xla::XlaOp& x,
+    const std::vector<xla::XlaOp>& starts, const gtl::ArraySlice<int64>& sizes);
 
 // Updates a slice of 'x', i.e.,
 // x[start[0], ..., start[n]] = update
-xla::StatusOr<xla::ComputationDataHandle> UpdateSlice(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x,
-    const xla::ComputationDataHandle& update, gtl::ArraySlice<int64> start);
+xla::StatusOr<xla::XlaOp> UpdateSlice(xla::XlaBuilder* builder,
+                                      const xla::XlaOp& x,
+                                      const xla::XlaOp& update,
+                                      gtl::ArraySlice<int64> start);
 
 // Updates a slice of 'x', where 'start' contains a list of minor dimensions:
 // x[..., start[0], ..., start[n]] = update
-xla::StatusOr<xla::ComputationDataHandle> UpdateSliceInMinorDims(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x,
-    const xla::ComputationDataHandle& update, gtl::ArraySlice<int64> start);
+xla::StatusOr<xla::XlaOp> UpdateSliceInMinorDims(xla::XlaBuilder* builder,
+                                                 const xla::XlaOp& x,
+                                                 const xla::XlaOp& update,
+                                                 gtl::ArraySlice<int64> start);
 
-xla::StatusOr<xla::ComputationDataHandle> DynamicUpdateSliceInMinorDims(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x,
-    const xla::ComputationDataHandle& update,
-    const std::vector<xla::ComputationDataHandle>& starts);
+xla::StatusOr<xla::XlaOp> DynamicUpdateSliceInMinorDims(
+    xla::XlaBuilder* builder, const xla::XlaOp& x, const xla::XlaOp& update,
+    const std::vector<xla::XlaOp>& starts);
 
 // Transposes a stack of matrices `x` by swapping the last two dimensions.
-xla::StatusOr<xla::ComputationDataHandle> TransposeInMinorDims(
-    xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x);
+xla::StatusOr<xla::XlaOp> TransposeInMinorDims(xla::XlaBuilder* builder,
+                                               const xla::XlaOp& x);
 
 }  // namespace tensorflow
 
diff --git a/tensorflow/compiler/tf2xla/lib/util_test.cc b/tensorflow/compiler/tf2xla/lib/util_test.cc
index b6bd33af2e..265b39402c 100644
--- a/tensorflow/compiler/tf2xla/lib/util_test.cc
+++ b/tensorflow/compiler/tf2xla/lib/util_test.cc
@@ -21,7 +21,6 @@ limitations under the License.
 
 #include "tensorflow/compiler/tf2xla/lib/batch_dot.h"
 #include "tensorflow/compiler/xla/array2d.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/compiler/xla/statusor.h"
 #include "tensorflow/compiler/xla/test.h"
@@ -65,9 +64,9 @@ xla::Array3D<float> BatchedAValsFull() {
 }
 
 XLA_TEST_F(UtilTest, Simple2dLookup) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, x, y;
+  xla::XlaOp a, x, y;
   auto a_data = CreateR2Parameter<float>(BValsRight(), 0, "a", &builder, &a);
   auto x_data = CreateR0Parameter<int>(2, 1, "x", &builder, &x);
   auto y_data = CreateR0Parameter<int>(1, 2, "y", &builder, &y);
@@ -80,9 +79,9 @@ XLA_TEST_F(UtilTest, Simple2dLookup) {
 }
 
 XLA_TEST_F(UtilTest, Simple3dLookup) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, index;
+  xla::XlaOp a, index;
   auto a_data =
       CreateR3Parameter<float>(BatchedAValsFull(), 0, "a", &builder, &a);
   auto index_data = CreateR0Parameter<int>(1, 1, "index", &builder, &index);
@@ -97,9 +96,9 @@ XLA_TEST_F(UtilTest, Simple3dLookup) {
 }
 
 XLA_TEST_F(UtilTest, SimpleSliceUpdate) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
-  xla::ComputationDataHandle a, b, x, y;
+  xla::XlaOp a, b, x, y;
   auto a_data = CreateR2Parameter<float>(AValsFull(), 0, "a", &builder, &a);
   auto b_data = CreateR2Parameter<float>({{9, 1, -10}}, 1, "b", &builder, &b);
   auto x_data = CreateR0Parameter<int>(2, 2, "x", &builder, &x);
@@ -117,11 +116,11 @@ XLA_TEST_F(UtilTest, SimpleSliceUpdate) {
 }
 
 XLA_TEST_F(UtilTest, RowBatchDot) {
-  xla::ComputationBuilder builder(client_, TestName());
+  xla::XlaBuilder builder(TestName());
 
   int n = 4;
 
-  xla::ComputationDataHandle a, row, index;
+  xla::XlaOp a, row, index;
   auto a_data =
       CreateR3Parameter<float>(BatchedAValsFull(), 0, "a", &builder, &a);
   auto row_data = CreateR3Parameter<float>({{{9, 1, 0, 0}}, {{2, 4, 0, 0}}}, 1,
diff --git a/tensorflow/compiler/tf2xla/lib/while_loop.cc b/tensorflow/compiler/tf2xla/lib/while_loop.cc
index 495d9c6078..09ce594930 100644
--- a/tensorflow/compiler/tf2xla/lib/while_loop.cc
+++ b/tensorflow/compiler/tf2xla/lib/while_loop.cc
@@ -20,24 +20,24 @@ limitations under the License.
 
 namespace tensorflow {
 
-xla::StatusOr<std::vector<xla::ComputationDataHandle>> XlaWhileLoop(
+xla::StatusOr<std::vector<xla::XlaOp>> XlaWhileLoop(
     const LoopConditionFunction& condition_function,
     const LoopBodyFunction& body_function,
-    gtl::ArraySlice<xla::ComputationDataHandle> initial_values,
-    StringPiece name, xla::ComputationBuilder* builder) {
+    gtl::ArraySlice<xla::XlaOp> initial_values, StringPiece name,
+    xla::XlaBuilder* builder) {
   int arity = initial_values.size();
   std::vector<xla::Shape> var_shapes;
   var_shapes.reserve(arity);
-  for (const xla::ComputationDataHandle& input : initial_values) {
+  for (const xla::XlaOp& input : initial_values) {
     TF_ASSIGN_OR_RETURN(auto shape, builder->GetShape(input));
-    var_shapes.push_back(std::move(*shape));
+    var_shapes.push_back(std::move(shape));
   }
   xla::Shape tuple_shape = xla::ShapeUtil::MakeTupleShape(var_shapes);
 
   // Unpacks a tuple into its component parts.
-  auto unpack_tuple = [](xla::ComputationDataHandle tuple, int arity,
-                         xla::ComputationBuilder* builder) {
-    std::vector<xla::ComputationDataHandle> elements(arity);
+  auto unpack_tuple = [](xla::XlaOp tuple, int arity,
+                         xla::XlaBuilder* builder) {
+    std::vector<xla::XlaOp> elements(arity);
     for (int i = 0; i < arity; ++i) {
       elements[i] = builder->GetTupleElement(tuple, i);
     }
@@ -45,20 +45,20 @@ xla::StatusOr<std::vector<xla::ComputationDataHandle>> XlaWhileLoop(
   };
 
   // Build the condition.
-  std::unique_ptr<xla::ComputationBuilder> cond_builder =
+  std::unique_ptr<xla::XlaBuilder> cond_builder =
       builder->CreateSubBuilder(strings::StrCat(name, "_condition"));
   {
     auto parameter = cond_builder->Parameter(0, tuple_shape, "parameter");
 
-    TF_ASSIGN_OR_RETURN(
-        auto result,
+    TF_RETURN_IF_ERROR(
         condition_function(unpack_tuple(parameter, arity, cond_builder.get()),
-                           cond_builder.get()));
+                           cond_builder.get())
+            .status());
   }
   TF_ASSIGN_OR_RETURN(auto cond, cond_builder->Build());
 
   // Build the body.
-  std::unique_ptr<xla::ComputationBuilder> body_builder =
+  std::unique_ptr<xla::XlaBuilder> body_builder =
       builder->CreateSubBuilder(strings::StrCat(name, "_body"));
   {
     auto parameter = body_builder->Parameter(0, tuple_shape, "parameter");
@@ -78,38 +78,38 @@ xla::StatusOr<std::vector<xla::ComputationDataHandle>> XlaWhileLoop(
   return unpack_tuple(outputs, arity, builder);
 }
 
-xla::StatusOr<std::vector<xla::ComputationDataHandle>> XlaForEachIndex(
+xla::StatusOr<std::vector<xla::XlaOp>> XlaForEachIndex(
     int64 num_iterations, xla::PrimitiveType num_iterations_type,
     const ForEachIndexBodyFunction& body_function,
-    gtl::ArraySlice<xla::ComputationDataHandle> initial_values,
-    StringPiece name, xla::ComputationBuilder* builder) {
-  auto while_cond_fn = [&](gtl::ArraySlice<xla::ComputationDataHandle> values,
-                           xla::ComputationBuilder* cond_builder)
-      -> xla::StatusOr<xla::ComputationDataHandle> {
+    gtl::ArraySlice<xla::XlaOp> initial_values, StringPiece name,
+    xla::XlaBuilder* builder) {
+  auto while_cond_fn =
+      [&](gtl::ArraySlice<xla::XlaOp> values,
+          xla::XlaBuilder* cond_builder) -> xla::StatusOr<xla::XlaOp> {
     return cond_builder->Lt(
         values[0],
         IntegerLiteral(cond_builder, num_iterations_type, num_iterations));
   };
-  auto while_body_fn = [&](gtl::ArraySlice<xla::ComputationDataHandle> values,
-                           xla::ComputationBuilder* body_builder)
-      -> xla::StatusOr<std::vector<xla::ComputationDataHandle>> {
-    xla::ComputationDataHandle iteration = values[0];
+  auto while_body_fn = [&](gtl::ArraySlice<xla::XlaOp> values,
+                           xla::XlaBuilder* body_builder)
+      -> xla::StatusOr<std::vector<xla::XlaOp>> {
+    xla::XlaOp iteration = values[0];
 
-    std::vector<xla::ComputationDataHandle> updated_values;
+    std::vector<xla::XlaOp> updated_values;
     updated_values.reserve(values.size());
     updated_values.push_back(body_builder->Add(
         iteration,
         body_builder->ConstantLiteral(xla::Literal::One(num_iterations_type))));
 
     values.remove_prefix(1);
-    TF_ASSIGN_OR_RETURN(std::vector<xla::ComputationDataHandle> body_outputs,
+    TF_ASSIGN_OR_RETURN(std::vector<xla::XlaOp> body_outputs,
                         body_function(iteration, values, body_builder));
     updated_values.insert(updated_values.end(), body_outputs.begin(),
                           body_outputs.end());
     return updated_values;
   };
 
-  std::vector<xla::ComputationDataHandle> values;
+  std::vector<xla::XlaOp> values;
   values.reserve(initial_values.size() + 1);
   values.push_back(
       builder->ConstantLiteral(xla::Literal::Zero(num_iterations_type)));
diff --git a/tensorflow/compiler/tf2xla/lib/while_loop.h b/tensorflow/compiler/tf2xla/lib/while_loop.h
index 2e67a0c99b..5b6684c995 100644
--- a/tensorflow/compiler/tf2xla/lib/while_loop.h
+++ b/tensorflow/compiler/tf2xla/lib/while_loop.h
@@ -19,8 +19,8 @@ limitations under the License.
 #include <functional>
 #include <vector>
 
-#include "tensorflow/compiler/xla/client/computation.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 #include "tensorflow/compiler/xla/statusor.h"
 #include "tensorflow/core/lib/core/stringpiece.h"
 #include "tensorflow/core/lib/gtl/array_slice.h"
@@ -29,14 +29,14 @@ namespace tensorflow {
 
 // Function that builds a loop condition. Takes as input a sequence of input
 // values, and returns a boolean value representing if the condition succeeds.
-typedef std::function<xla::StatusOr<xla::ComputationDataHandle>(
-    gtl::ArraySlice<xla::ComputationDataHandle>, xla::ComputationBuilder*)>
+typedef std::function<xla::StatusOr<xla::XlaOp>(gtl::ArraySlice<xla::XlaOp>,
+                                                xla::XlaBuilder*)>
     LoopConditionFunction;
 
 // Function that builds a loop body. Takes as input a sequence of input values
 // and returns a sequence of output values.
-typedef std::function<xla::StatusOr<std::vector<xla::ComputationDataHandle>>(
-    gtl::ArraySlice<xla::ComputationDataHandle>, xla::ComputationBuilder*)>
+typedef std::function<xla::StatusOr<std::vector<xla::XlaOp>>(
+    gtl::ArraySlice<xla::XlaOp>, xla::XlaBuilder*)>
     LoopBodyFunction;
 
 // Helper function for building an XLA while loop, where the values carried by
@@ -47,27 +47,26 @@ typedef std::function<xla::StatusOr<std::vector<xla::ComputationDataHandle>>(
 //   init: (a, b, c)
 // )
 // 'name' is a descriptive name for the loop.
-xla::StatusOr<std::vector<xla::ComputationDataHandle>> XlaWhileLoop(
+xla::StatusOr<std::vector<xla::XlaOp>> XlaWhileLoop(
     const LoopConditionFunction& condition_function,
     const LoopBodyFunction& body_function,
-    gtl::ArraySlice<xla::ComputationDataHandle> initial_values,
-    StringPiece name, xla::ComputationBuilder* builder);
+    gtl::ArraySlice<xla::XlaOp> initial_values, StringPiece name,
+    xla::XlaBuilder* builder);
 
 // Builds an XLA loop that repeats a computation `num_iterations` times.
 //
 // The body function (ForEachIndexBodyFunction) takes as input a pair of
 // (current iteration number, loop-carried values), and returns an updated
 // vector of the loop-carried values.
-typedef std::function<xla::StatusOr<std::vector<xla::ComputationDataHandle>>(
-    xla::ComputationDataHandle, gtl::ArraySlice<xla::ComputationDataHandle>,
-    xla::ComputationBuilder*)>
+typedef std::function<xla::StatusOr<std::vector<xla::XlaOp>>(
+    xla::XlaOp, gtl::ArraySlice<xla::XlaOp>, xla::XlaBuilder*)>
     ForEachIndexBodyFunction;
 
-xla::StatusOr<std::vector<xla::ComputationDataHandle>> XlaForEachIndex(
+xla::StatusOr<std::vector<xla::XlaOp>> XlaForEachIndex(
     int64 num_iterations, xla::PrimitiveType num_iterations_type,
     const ForEachIndexBodyFunction& body_function,
-    gtl::ArraySlice<xla::ComputationDataHandle> initial_values,
-    StringPiece name, xla::ComputationBuilder* builder);
+    gtl::ArraySlice<xla::XlaOp> initial_values, StringPiece name,
+    xla::XlaBuilder* builder);
 
 }  // namespace tensorflow
 
diff --git a/tensorflow/compiler/tf2xla/tf2xla.cc b/tensorflow/compiler/tf2xla/tf2xla.cc
index 6051d7dffd..3a08aa8cf4 100644
--- a/tensorflow/compiler/tf2xla/tf2xla.cc
+++ b/tensorflow/compiler/tf2xla/tf2xla.cc
@@ -251,7 +251,7 @@ Status CreateXlaArgs(const Graph& graph,
 // Converts the TensorFlow graph into an XLA computation, by executing the
 // graph symbolically, with each op building up the XLA HLO.
 Status ConvertGraphToXla(std::unique_ptr<Graph> graph, xla::Client* client,
-                         xla::Computation* computation) {
+                         xla::XlaComputation* computation) {
   XlaOpRegistry::RegisterCompilationKernels();
   for (Node* node : graph->nodes()) {
     node->set_assigned_device_name(
@@ -303,7 +303,7 @@ Status ConvertGraphToXla(std::unique_ptr<Graph> graph, xla::Client* client,
 }
 
 // InitGraph creates a graph based on the graph_def, that may then be converted
-// to an xla::Computation via ConvertGraphToXla.
+// to an xla::XlaComputation via ConvertGraphToXla.
 //
 // The graph is rewritten with _Arg and _Retval nodes, representing the inputs
 // and outputs of the function that will be compiled.  Each feed id causes a new
@@ -348,7 +348,7 @@ Status InitGraph(const GraphDef& graph_def, const tf2xla::Config& config,
 
 Status ConvertGraphDefToXla(const GraphDef& graph_def,
                             const tf2xla::Config& config, xla::Client* client,
-                            xla::Computation* computation) {
+                            xla::XlaComputation* computation) {
   std::unique_ptr<Graph> graph;
   TF_RETURN_IF_ERROR(InitGraph(graph_def, config, &graph));
   TF_RETURN_IF_ERROR(ConvertGraphToXla(std::move(graph), client, computation));
diff --git a/tensorflow/compiler/tf2xla/tf2xla.h b/tensorflow/compiler/tf2xla/tf2xla.h
index 473c431b12..d02fc56c5b 100644
--- a/tensorflow/compiler/tf2xla/tf2xla.h
+++ b/tensorflow/compiler/tf2xla/tf2xla.h
@@ -18,21 +18,21 @@ limitations under the License.
 
 #include "tensorflow/compiler/tf2xla/tf2xla.pb.h"
 #include "tensorflow/compiler/xla/client/client.h"
-#include "tensorflow/compiler/xla/client/computation.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 #include "tensorflow/core/framework/graph.pb.h"
 
 namespace tensorflow {
 
-// Converts a tensorflow::GraphDef into an xla::Computation.  The given `config`
-// specifies the portion of the graph to convert, via feeds and fetches. Each
-// feed is a positional input argument for the generated computation, while each
-// fetch is a positional output argument.
+// Converts a tensorflow::GraphDef into an xla::XlaComputation. The given
+// `config` specifies the portion of the graph to convert, via feeds and
+// fetches. Each feed is a positional input argument for the generated
+// computation, while each fetch is a positional output argument.
 //
 // The computation is built in the context of the given `client`, which may
 // subsequently be used to compile or execute the computation.
 Status ConvertGraphDefToXla(const GraphDef& graph_def,
                             const tf2xla::Config& config, xla::Client* client,
-                            xla::Computation* computation);
+                            xla::XlaComputation* computation);
 
 }  // namespace tensorflow
 
diff --git a/tensorflow/compiler/tf2xla/tf2xla_test.cc b/tensorflow/compiler/tf2xla/tf2xla_test.cc
index b813668a9e..84c133ffab 100644
--- a/tensorflow/compiler/tf2xla/tf2xla_test.cc
+++ b/tensorflow/compiler/tf2xla/tf2xla_test.cc
@@ -69,7 +69,7 @@ TEST(ConvertGraphDefToXla, Sum) {
   tf2xla::Config config = SumConfig();
 
   xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie();
-  xla::Computation computation;
+  xla::XlaComputation computation;
   TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation));
 
   // Set up arguments.
diff --git a/tensorflow/compiler/tf2xla/xla_compilation_device.cc b/tensorflow/compiler/tf2xla/xla_compilation_device.cc
index fcb0a4e638..fe7ec633ec 100644
--- a/tensorflow/compiler/tf2xla/xla_compilation_device.cc
+++ b/tensorflow/compiler/tf2xla/xla_compilation_device.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/sharding_util.h"
 #include "tensorflow/compiler/tf2xla/xla_context.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/core/common_runtime/local_device.h"
 #include "tensorflow/core/framework/device_base.h"
 #include "tensorflow/core/platform/mem.h"
@@ -108,7 +109,7 @@ void XlaCompilationDevice::Compute(OpKernel* op_kernel,
   // If no sharding metadata is found, XLA is free to use whatever device it
   // wants. In practice this usually has the effect of placing things on device
   // 0.
-  xla::ScopedShardingAssignment assign_sharding(b, op_sharding);
+  xla::XlaScopedShardingAssignment assign_sharding(b, op_sharding);
   op_kernel->Compute(context);
 
   b->ClearOpMetadata();
@@ -126,9 +127,7 @@ Status XlaCompilationDevice::MakeTensorFromProto(
 
 XlaExpression::XlaExpression() = default;
 
-void XlaExpression::set_handle(const xla::ComputationDataHandle& h) {
-  handle_ = h;
-}
+void XlaExpression::set_handle(const xla::XlaOp& h) { handle_ = h; }
 
 void XlaExpression::set_constant_value(Tensor value) {
   has_constant_value_ = true;
diff --git a/tensorflow/compiler/tf2xla/xla_compilation_device.h b/tensorflow/compiler/tf2xla/xla_compilation_device.h
index 0243ee332f..d0b9e34e16 100644
--- a/tensorflow/compiler/tf2xla/xla_compilation_device.h
+++ b/tensorflow/compiler/tf2xla/xla_compilation_device.h
@@ -19,7 +19,7 @@ limitations under the License.
 #include <memory>
 
 #include "tensorflow/compiler/tf2xla/xla_resource.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
 #include "tensorflow/core/common_runtime/local_device.h"
 #include "tensorflow/core/framework/device_base.h"
@@ -69,7 +69,7 @@ class XlaCompilationDevice : public LocalDevice {
 
 // A XlaExpression wraps an XLA computation. Each Tensor on an
 // XlaCompilationDevice contains an XlaExpression, and the shape of the Tensor
-// matches the shape of the subcomputation in the ComputationDataHandle. Each
+// matches the shape of the subcomputation in the XlaOp. Each
 // expression is either a constant, or a function of previously-compiled
 // expressions.
 class XlaExpression {
@@ -78,8 +78,8 @@ class XlaExpression {
 
   // handle() stores the XLA handle of the computation that the
   // expression represents.
-  void set_handle(const xla::ComputationDataHandle& h);
-  const xla::ComputationDataHandle& handle() const { return handle_; }
+  void set_handle(const xla::XlaOp& h);
+  const xla::XlaOp& handle() const { return handle_; }
 
   void set_constant_value(Tensor value);
   bool has_constant_value() const { return has_constant_value_; }
@@ -90,7 +90,7 @@ class XlaExpression {
 
  private:
   // The XLA handle of the expression's computation.
-  xla::ComputationDataHandle handle_;
+  xla::XlaOp handle_;
 
   // If this expression is a constant with a known value, 'constant_value' is a
   // host-memory Tensor containing the value. Used to avoid invoking XLA for
diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc
index c0e9967684..3d1946c332 100644
--- a/tensorflow/compiler/tf2xla/xla_compiler.cc
+++ b/tensorflow/compiler/tf2xla/xla_compiler.cc
@@ -339,11 +339,11 @@ Status BuildComputation(
     const std::vector<int>& arg_cores,
     const std::vector<XlaExpression>& retvals,
     const std::vector<std::unique_ptr<XlaResource>>& resources,
-    bool return_updated_values_for_all_resources,
-    xla::ComputationBuilder* builder, xla::Computation* computation,
-    int* num_computation_outputs, int* num_nonconst_outputs,
+    bool return_updated_values_for_all_resources, xla::XlaBuilder* builder,
+    xla::XlaComputation* computation, int* num_computation_outputs,
+    int* num_nonconst_outputs,
     std::vector<XlaCompiler::ResourceUpdate>* resource_updates) {
-  std::vector<xla::ComputationDataHandle> elems;
+  std::vector<xla::XlaOp> elems;
   elems.reserve(retvals.size());
   for (const XlaExpression& retval : retvals) {
     if (!retval.has_constant_value()) {
@@ -376,14 +376,12 @@ Status BuildComputation(
     const XlaCompiler::Argument& arg = args[resource->arg_num()];
     const int core = arg_cores[resource->arg_num()];
     DCHECK_LT(resource->arg_num(), arg_cores.size());
-    bool modified =
-        resource->value().handle() != resource->initial_value().handle();
+    bool modified = resource->value() != resource->initial_value();
     // TensorArray gradients were modified if their values changed or there are
     // any newly created gradients.
     for (const auto& grad : resource->tensor_array_gradients()) {
       modified = modified ||
-                 grad.second->value().handle() !=
-                     grad.second->initial_value().handle() ||
+                 grad.second->value() != grad.second->initial_value() ||
                  arg.tensor_array_gradients.count(grad.first) == 0;
     }
     if (return_updated_values_for_all_resources || modified) {
@@ -398,11 +396,11 @@ Status BuildComputation(
       }
 
       // Request that the value be returned on a specific core.
-      xla::ScopedShardingAssignment assign_sharding(
+      xla::XlaScopedShardingAssignment assign_sharding(
           builder, core == -1 ? tensorflow::gtl::optional<xla::OpSharding>()
                               : xla::sharding_builder::AssignDevice(core));
 
-      xla::ComputationDataHandle handle;
+      xla::XlaOp handle;
       TF_RETURN_IF_ERROR(resource->Pack(&handle, builder));
 
       // Since we can't change the sharding metadata of <value> as this point,
@@ -421,7 +419,7 @@ Status BuildComputation(
   builder->Tuple(elems);
   builder->ClearOpMetadata();
 
-  xla::StatusOr<xla::Computation> computation_status = builder->Build();
+  xla::StatusOr<xla::XlaComputation> computation_status = builder->Build();
   if (!computation_status.ok()) {
     return computation_status.status();
   }
@@ -435,7 +433,7 @@ Status BuildComputation(
 // `args` are the arguments to the computation.
 Status XlaCompiler::BuildArguments(
     const Graph& graph, const std::vector<XlaCompiler::Argument>& args,
-    bool use_tuple_arg, xla::ComputationBuilder* builder, XlaContext* context,
+    bool use_tuple_arg, xla::XlaBuilder* builder, XlaContext* context,
     std::vector<int>* arg_cores, std::vector<XlaExpression>* arg_expressions,
     std::vector<int>* input_mapping, std::vector<xla::Shape>* input_shapes,
     bool is_entry_computation) {
@@ -461,8 +459,7 @@ Status XlaCompiler::BuildArguments(
         // alias.
         XlaResource* resource;
         TF_RETURN_IF_ERROR(context->CreateResource(
-            arg.resource_kind, i, arg.name, arg.type, arg.shape,
-            xla::ComputationDataHandle(),
+            arg.resource_kind, i, arg.name, arg.type, arg.shape, xla::XlaOp(),
             /*tensor_array_size=*/arg.tensor_array_size,
             /*tensor_array_gradients=*/arg.tensor_array_gradients, &resource));
         arg_expression.set_resource(resource);
@@ -531,9 +528,9 @@ Status XlaCompiler::BuildArguments(
   builder->SetOpMetadata(arg_metadata);
 
   // Build parameter handles for non-constant arguments.
-  std::vector<xla::ComputationDataHandle> arg_handles(input_mapping->size());
+  std::vector<xla::XlaOp> arg_handles(input_mapping->size());
   if (use_tuple_arg) {
-    xla::ComputationDataHandle tuple;
+    xla::XlaOp tuple;
     if (is_entry_computation) {
       xla::OpSharding tuple_sharding;
       tuple_sharding.set_type(xla::OpSharding::Type::OpSharding_Type_TUPLE);
@@ -544,15 +541,15 @@ Status XlaCompiler::BuildArguments(
             core == -1 ? xla::sharding_builder::AssignDevice(root_device)
                        : xla::sharding_builder::AssignDevice(core);
       }
-      xla::ScopedShardingAssignment assign_tuple_sharding(builder,
-                                                          tuple_sharding);
+      xla::XlaScopedShardingAssignment assign_tuple_sharding(builder,
+                                                             tuple_sharding);
       tuple = builder->Parameter(0, (*input_shapes)[0], "arg_tuple");
     } else {
       tuple = builder->Parameter(0, (*input_shapes)[0], "arg_tuple");
     }
     for (std::vector<int>::size_type i = 0; i < input_mapping->size(); ++i) {
       const int core = (*arg_cores)[input_mapping->at(i)];
-      xla::ScopedShardingAssignment assign_sharding(
+      xla::XlaScopedShardingAssignment assign_sharding(
           builder, core == -1 ? tensorflow::gtl::optional<xla::OpSharding>()
                               : xla::sharding_builder::AssignDevice(core));
       arg_handles[i] = builder->GetTupleElement(tuple, i);
@@ -560,7 +557,7 @@ Status XlaCompiler::BuildArguments(
   } else {
     for (std::vector<int>::size_type i = 0; i < input_mapping->size(); ++i) {
       const int core = (*arg_cores)[input_mapping->at(i)];
-      xla::ScopedShardingAssignment assign_sharding(
+      xla::XlaScopedShardingAssignment assign_sharding(
           builder, core == -1 ? tensorflow::gtl::optional<xla::OpSharding>()
                               : xla::sharding_builder::AssignDevice(core));
       arg_handles[i] =
@@ -647,7 +644,7 @@ Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options,
                                  std::unique_ptr<Graph> graph,
                                  const std::vector<XlaCompiler::Argument>& args,
                                  CompilationResult* result) {
-  VLOG(1) << "Executing graph symbolically to populate ComputationBuilder.";
+  VLOG(1) << "Executing graph symbolically to populate XlaBuilder.";
 
   if (VLOG_IS_ON(2)) {
     VLOG(2) << "XlaCompiler::CompileGraph: "
@@ -663,7 +660,7 @@ Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options,
   TF_RETURN_IF_ERROR(
       FunctionalizeControlFlow(graph.get(), local_flib_def_.get()));
 
-  xla::ComputationBuilder builder(client(), name);
+  xla::XlaBuilder builder(name);
   XlaContext* context =
       new XlaContext(this, &builder, options_.allow_cpu_custom_calls,
                      options.resolve_compile_time_constants,
@@ -683,7 +680,7 @@ Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options,
 
   int num_nonconst_outputs;
   int num_computation_outputs;
-  result->computation = std::make_shared<xla::Computation>();
+  result->computation = std::make_shared<xla::XlaComputation>();
   TF_RETURN_IF_ERROR(BuildComputation(
       args, arg_cores, context->retvals(), context->resources(),
       options.return_updated_values_for_all_resources, &builder,
@@ -814,7 +811,7 @@ Status XlaCompiler::SetHostToDeviceMetadata(
 }
 
 Status XlaCompiler::GetHostComputeControlDependency(
-    const string& host_compute_name, xla::ComputationDataHandle* handle) {
+    const string& host_compute_name, xla::XlaOp* handle) {
   const auto iter = host_compute_control_output_.find(host_compute_name);
   if (iter == host_compute_control_output_.end()) {
     return errors::InvalidArgument(
@@ -827,7 +824,7 @@ Status XlaCompiler::GetHostComputeControlDependency(
 }
 
 Status XlaCompiler::SetHostComputeControlDependency(
-    const string& host_compute_name, const xla::ComputationDataHandle& handle) {
+    const string& host_compute_name, const xla::XlaOp& handle) {
   if (host_compute_control_output_.find(host_compute_name) !=
       host_compute_control_output_.end()) {
     return errors::InvalidArgument(
diff --git a/tensorflow/compiler/tf2xla/xla_compiler.h b/tensorflow/compiler/tf2xla/xla_compiler.h
index 8f564f35ec..ca6cd822ef 100644
--- a/tensorflow/compiler/tf2xla/xla_compiler.h
+++ b/tensorflow/compiler/tf2xla/xla_compiler.h
@@ -227,7 +227,7 @@ class XlaCompiler {
     std::vector<ResourceUpdate> resource_updates;
 
     // The XLA computation built from the tensorflow subgraph.
-    std::shared_ptr<xla::Computation> computation;
+    std::shared_ptr<xla::XlaComputation> computation;
   };
 
   struct Options {
@@ -281,7 +281,7 @@ class XlaCompiler {
                          const NameAttrList& fn_name_attrs,
                          std::vector<Argument> args, CompilationResult* result);
 
-  // Compiles a tensorflow::Graph into an xla::Computation.
+  // Compiles a tensorflow::Graph into an xla::XlaComputation.
   // Similar to CompileFunction, but takes a Graph as input rather than a
   // function.
   Status CompileGraph(const CompileOptions& options, string const& name,
@@ -290,7 +290,7 @@ class XlaCompiler {
                       CompilationResult* result);
 
   // Compiles a single Op, given by an OpKernelContext, into an
-  // xla::Computation. Similar to CompileFunction but takes a single Op as
+  // xla::XlaComputation. Similar to CompileFunction but takes a single Op as
   // input.
   Status CompileSingleOp(const CompileOptions& options, string const& name,
                          OpKernelContext* ctx,
@@ -337,10 +337,9 @@ class XlaCompiler {
   // a given HostCompute Op as long as the names are unique within the
   // compilation.
   Status GetHostComputeControlDependency(const string& host_compute_name,
-                                         xla::ComputationDataHandle* handle);
-  Status SetHostComputeControlDependency(
-      const string& host_compute_name,
-      const xla::ComputationDataHandle& handle);
+                                         xla::XlaOp* handle);
+  Status SetHostComputeControlDependency(const string& host_compute_name,
+                                         const xla::XlaOp& handle);
 
   const Options& options() const { return options_; }
   xla::Client* client() const { return options_.client; }
@@ -358,7 +357,7 @@ class XlaCompiler {
   // `args` are the arguments to the computation.
   Status BuildArguments(const Graph& graph,
                         const std::vector<XlaCompiler::Argument>& args,
-                        bool use_tuple_arg, xla::ComputationBuilder* builder,
+                        bool use_tuple_arg, xla::XlaBuilder* builder,
                         XlaContext* context, std::vector<int>* arg_cores,
                         std::vector<XlaExpression>* arg_expressions,
                         std::vector<int>* input_mapping,
@@ -408,8 +407,7 @@ class XlaCompiler {
   std::unordered_map<string, tf2xla::HostTransferMetadata> host_compute_sends_;
   std::unordered_map<string, tf2xla::HostTransferMetadata> host_compute_recvs_;
 
-  std::unordered_map<string, xla::ComputationDataHandle>
-      host_compute_control_output_;
+  std::unordered_map<string, xla::XlaOp> host_compute_control_output_;
 
   TF_DISALLOW_COPY_AND_ASSIGN(XlaCompiler);
 };
diff --git a/tensorflow/compiler/tf2xla/xla_compiler_test.cc b/tensorflow/compiler/tf2xla/xla_compiler_test.cc
index 096dc7160b..6b8918b261 100644
--- a/tensorflow/compiler/tf2xla/xla_compiler_test.cc
+++ b/tensorflow/compiler/tf2xla/xla_compiler_test.cc
@@ -164,7 +164,6 @@ REGISTER_XLA_OP(Name("DummyDuplicateOp").Device(DEVICE_CPU_XLA_JIT),
 REGISTER_XLA_OP(Name("DummyDuplicateOp").Device(DEVICE_GPU_XLA_JIT),
                 DummyDuplicateOp);
 
-
 // Tests compilation and execution of an empty graph.
 TEST_F(XlaCompilerTest, EmptyReturnValues) {
   XlaCompiler compiler(DefaultOptions());
@@ -433,21 +432,26 @@ TEST_F(XlaCompilerTest, DeterministicCompilation) {
   }
 
   for (int64 i = 1; i < test_count; ++i) {
-    auto m1 =
-        results[i - 1].computation->Snapshot().ValueOrDie()->entry().requests();
-    auto m2 =
-        results[i].computation->Snapshot().ValueOrDie()->entry().requests();
-    // Check if every entry is the same.
-    for (auto& entry1 : m1) {
-      int64 key = entry1.first;
-      auto value1 = entry1.second;
-      auto entry2 = m2.find(key);
-      auto value2 = entry2->second;
-      EXPECT_TRUE(entry2 != m2.end());
-      string str1, str2;
-      value1.AppendToString(&str1);
-      value2.AppendToString(&str2);
-      EXPECT_EQ(str1, str2);
+    const auto& m1 = results[i - 1].computation->proto();
+    const auto& m2 = results[i].computation->proto();
+    ASSERT_EQ(m1.computations_size(), m2.computations_size());
+    // Check if every hlo computation is the same.
+    for (int k = 0; k < m1.computations_size(); k++) {
+      const auto& c1 = m1.computations(k);
+      const auto& c2 = m2.computations(k);
+      ASSERT_EQ(c1.instructions_size(), c2.instructions_size());
+      for (int j = 0; j < c1.instructions_size(); j++) {
+        auto instr1 = c1.instructions(j);
+        auto instr2 = c2.instructions(j);
+        instr1.clear_name();
+        instr2.clear_name();
+        // The names of instructions were uniquified by the XlaBuilder, the rest
+        // of the fields should be identical.
+        string str1, str2;
+        instr1.AppendPartialToString(&str1);
+        instr2.AppendPartialToString(&str2);
+        EXPECT_EQ(str1, str2);
+      }
     }
   }
 }
diff --git a/tensorflow/compiler/tf2xla/xla_context.cc b/tensorflow/compiler/tf2xla/xla_context.cc
index 8423921086..3dd2d183f3 100644
--- a/tensorflow/compiler/tf2xla/xla_context.cc
+++ b/tensorflow/compiler/tf2xla/xla_context.cc
@@ -25,7 +25,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/xla/client/client_library.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/layout_util.h"
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/compiler/xla/statusor.h"
@@ -63,7 +63,7 @@ void XlaContext::set_args(std::vector<XlaExpression> args) {
 }
 
 XlaContext::XlaContext(
-    XlaCompiler* compiler, xla::ComputationBuilder* builder,
+    XlaCompiler* compiler, xla::XlaBuilder* builder,
     bool allow_cpu_custom_calls, bool resolve_compile_time_constants,
     const std::function<TensorShape(const TensorShape&, DataType)>*
         variable_representation_shape_fn)
@@ -78,7 +78,7 @@ string XlaContext::DebugString() { return "TLA JIT context"; }
 // This is called by the Retval Op to associate a computed value
 // with a specific return value of the subgraph.
 void XlaContext::AddRetval(int retval_index, DataType type,
-                           const xla::ComputationDataHandle& handle) {
+                           const xla::XlaOp& handle) {
   VLOG(1) << "Added retval index " << retval_index << " to XLA computation";
   // Add the return value to the list being built up.
   if (retvals_.size() <= retval_index) {
@@ -104,13 +104,12 @@ Status XlaContext::AddConstRetval(int retval_index, DataType dtype,
   return Status::OK();
 }
 
-xla::ComputationBuilder* XlaContext::builder() { return builder_; }
+xla::XlaBuilder* XlaContext::builder() { return builder_; }
 
 Status XlaContext::CreateResource(
     XlaResource::Kind kind, int arg_num, string name, DataType type,
-    TensorShape shape, const xla::ComputationDataHandle& handle,
-    int64 tensor_array_size, const std::set<string>& tensor_array_gradients,
-    XlaResource** resource) {
+    TensorShape shape, const xla::XlaOp& handle, int64 tensor_array_size,
+    const std::set<string>& tensor_array_gradients, XlaResource** resource) {
   resources_.emplace_back(
       new XlaResource(kind, arg_num, std::move(name), type, std::move(shape),
                       handle, tensor_array_size, tensor_array_gradients));
@@ -123,11 +122,11 @@ TensorShape XlaContext::VariableRepresentationShape(const TensorShape& shape,
   return (*variable_representation_shape_fn_)(shape, type);
 }
 
-const xla::Computation* XlaContext::GetOrCreateMax(const DataType type) {
+const xla::XlaComputation* XlaContext::GetOrCreateMax(const DataType type) {
   return LookupOrCreate(type, &max_func_, [this, type] {
     const string type_string = DataTypeString(type);
     VLOG(1) << "Building Max() for " << type_string;
-    xla::ComputationBuilder b(builder()->client(), "max<" + type_string + ">");
+    xla::XlaBuilder b("max<" + type_string + ">");
     xla::PrimitiveType xla_type;
     TF_CHECK_OK(DataTypeToPrimitiveType(type, &xla_type));
     auto x = b.Parameter(0, xla::ShapeUtil::MakeShape(xla_type, {}), "x");
@@ -137,11 +136,11 @@ const xla::Computation* XlaContext::GetOrCreateMax(const DataType type) {
   });
 }
 
-const xla::Computation* XlaContext::GetOrCreateMin(const DataType type) {
+const xla::XlaComputation* XlaContext::GetOrCreateMin(const DataType type) {
   return LookupOrCreate(type, &min_func_, [this, type] {
     const string type_string = DataTypeString(type);
     VLOG(1) << "Building Min() for " << type_string;
-    xla::ComputationBuilder b(builder()->client(), "min<" + type_string + ">");
+    xla::XlaBuilder b("min<" + type_string + ">");
     xla::PrimitiveType xla_type;
     TF_CHECK_OK(DataTypeToPrimitiveType(type, &xla_type));
     auto x = b.Parameter(0, xla::ShapeUtil::MakeShape(xla_type, {}), "x");
@@ -151,11 +150,11 @@ const xla::Computation* XlaContext::GetOrCreateMin(const DataType type) {
   });
 }
 
-const xla::Computation* XlaContext::GetOrCreateAdd(const DataType type) {
+const xla::XlaComputation* XlaContext::GetOrCreateAdd(const DataType type) {
   return LookupOrCreate(type, &add_func_, [this, type] {
     const string type_string = DataTypeString(type);
     VLOG(1) << "Building Add() for " << type_string;
-    xla::ComputationBuilder b(builder()->client(), "add<" + type_string + ">");
+    xla::XlaBuilder b("add<" + type_string + ">");
     xla::PrimitiveType xla_type;
     TF_CHECK_OK(DataTypeToPrimitiveType(type, &xla_type));
     auto x = b.Parameter(0, xla::ShapeUtil::MakeShape(xla_type, {}), "x");
@@ -165,11 +164,11 @@ const xla::Computation* XlaContext::GetOrCreateAdd(const DataType type) {
   });
 }
 
-const xla::Computation* XlaContext::GetOrCreateMul(const DataType type) {
+const xla::XlaComputation* XlaContext::GetOrCreateMul(const DataType type) {
   return LookupOrCreate(type, &mul_func_, [this, type] {
     const string type_string = DataTypeString(type);
     VLOG(1) << "Building Mul() for " << type_string;
-    xla::ComputationBuilder b(builder()->client(), "mul<" + type_string + ">");
+    xla::XlaBuilder b("mul<" + type_string + ">");
     xla::PrimitiveType xla_type;
     TF_CHECK_OK(DataTypeToPrimitiveType(type, &xla_type));
     auto x = b.Parameter(0, xla::ShapeUtil::MakeShape(xla_type, {}), "x");
@@ -179,9 +178,9 @@ const xla::Computation* XlaContext::GetOrCreateMul(const DataType type) {
   });
 }
 
-const xla::Computation* XlaContext::LookupOrCreate(
+const xla::XlaComputation* XlaContext::LookupOrCreate(
     DataType type, ComputationMap* out,
-    const std::function<xla::Computation()>& create) {
+    const std::function<xla::XlaComputation()>& create) {
   {
     const auto& entry = (*out)[type];
     if (!entry.IsNull()) {
diff --git a/tensorflow/compiler/tf2xla/xla_context.h b/tensorflow/compiler/tf2xla/xla_context.h
index 00fbaba37c..1136ffe507 100644
--- a/tensorflow/compiler/tf2xla/xla_context.h
+++ b/tensorflow/compiler/tf2xla/xla_context.h
@@ -22,8 +22,8 @@ limitations under the License.
 
 #include "tensorflow/compiler/tf2xla/xla_compilation_device.h"
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
-#include "tensorflow/compiler/xla/client/computation.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/resource_mgr.h"
@@ -43,7 +43,7 @@ class XlaContext : public ResourceBase {
   static XlaContext& Get(const XlaOpKernelContext* ctx);
 
   // Creates a new XlaContext.
-  XlaContext(XlaCompiler* compiler, xla::ComputationBuilder* builder,
+  XlaContext(XlaCompiler* compiler, xla::XlaBuilder* builder,
              bool allow_cpu_custom_calls, bool resolve_compile_time_constants,
              const std::function<TensorShape(const TensorShape&, DataType)>*
                  variable_representation_shape_fn);
@@ -53,9 +53,8 @@ class XlaContext : public ResourceBase {
 
   XlaCompiler* compiler() const { return compiler_; }
 
-  // Returns the ComputationBuilder that Ops use for compiling new
-  // expressions.
-  xla::ComputationBuilder* builder();
+  // Returns the XlaBuilder that Ops use for compiling new expressions.
+  xla::XlaBuilder* builder();
 
   bool allow_cpu_custom_calls() const { return allow_cpu_custom_calls_; }
 
@@ -66,8 +65,7 @@ class XlaContext : public ResourceBase {
 
   // This is called by the Retval Op to associate a computed value
   // with a specific return value of the subgraph.
-  void AddRetval(int retval_index, DataType type,
-                 const xla::ComputationDataHandle& handle);
+  void AddRetval(int retval_index, DataType type, const xla::XlaOp& handle);
 
   // As for Retval, but for return values that are compile-time constants.
   Status AddConstRetval(int retval_index, DataType dtype,
@@ -79,8 +77,7 @@ class XlaContext : public ResourceBase {
   // Fails if the resource already exists.
   Status CreateResource(XlaResource::Kind kind, int arg_num, string name,
                         DataType type, TensorShape shape,
-                        const xla::ComputationDataHandle& handle,
-                        int64 tensor_array_size,
+                        const xla::XlaOp& handle, int64 tensor_array_size,
                         const std::set<string>& tensor_array_gradients,
                         XlaResource** resource);
 
@@ -96,22 +93,22 @@ class XlaContext : public ResourceBase {
   // Get an XLA lambda to compute Max. This is cached in the
   // XlaContext since it may be used by multiple Ops. There is a
   // separate specialization of the computation for each DataType.
-  const xla::Computation* GetOrCreateMax(const DataType type);
+  const xla::XlaComputation* GetOrCreateMax(const DataType type);
 
   // Get an XLA lambda to compute Min. This is cached in the
   // XlaContext since it may be used by multiple Ops. There is a
   // separate specialization of the computation for each DataType.
-  const xla::Computation* GetOrCreateMin(const DataType type);
+  const xla::XlaComputation* GetOrCreateMin(const DataType type);
 
   // Get an XLA lambda to compute Add. This is cached in the
   // XlaContext since it may be used by multiple Ops. There is a
   // separate specialization of the computation for each DataType.
-  const xla::Computation* GetOrCreateAdd(const DataType type);
+  const xla::XlaComputation* GetOrCreateAdd(const DataType type);
 
   // Get an XLA lambda to compute Mul. This is cached in the
   // XlaContext since it may be used by multiple Ops. There is a
   // separate specialization of the computation for each DataType.
-  const xla::Computation* GetOrCreateMul(const DataType type);
+  const xla::XlaComputation* GetOrCreateMul(const DataType type);
 
   // The name of the XlaContext resource during symbolic graph execution.
   static const char kXlaContextResourceName[];
@@ -119,9 +116,8 @@ class XlaContext : public ResourceBase {
  private:
   XlaCompiler* const compiler_;
 
-  // The ComputationBuilder used to construct the subgraph's compiled
-  // representation.
-  xla::ComputationBuilder* builder_;
+  // The XlaBuilder used to construct the subgraph's compiled representation.
+  xla::XlaBuilder* builder_;
 
   // Allow ops to emit CustomCall operations for CPU.
   const bool allow_cpu_custom_calls_;
@@ -146,14 +142,14 @@ class XlaContext : public ResourceBase {
       variable_representation_shape_fn_;
 
   // Cache of prebuilt computations indexed by their type.
-  using ComputationMap = std::map<DataType, xla::Computation>;
+  using ComputationMap = std::map<DataType, xla::XlaComputation>;
 
   // Finds the value for the given type in out map if it already
   // exists or makes a new value with create function and keeps it the
   // map. The returned value != nullptr and is owned by the map.
-  const xla::Computation* LookupOrCreate(
+  const xla::XlaComputation* LookupOrCreate(
       DataType type, ComputationMap* out,
-      const std::function<xla::Computation()>& create);
+      const std::function<xla::XlaComputation()>& create);
 
   // Cached computation to compute Max of two elements, specialized by type.
   ComputationMap max_func_;
diff --git a/tensorflow/compiler/tf2xla/xla_helpers.cc b/tensorflow/compiler/tf2xla/xla_helpers.cc
index a3deb02a1f..f1594193af 100644
--- a/tensorflow/compiler/tf2xla/xla_helpers.cc
+++ b/tensorflow/compiler/tf2xla/xla_helpers.cc
@@ -22,7 +22,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/type_util.h"
 #include "tensorflow/compiler/tf2xla/xla_context.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/types.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/lib/core/status.h"
@@ -32,13 +32,12 @@ namespace tensorflow {
 
 namespace {
 
-Status ArgMinMax(xla::ComputationBuilder* builder, XlaOpKernelContext* ctx,
-                 const xla::ComputationDataHandle& input,
-                 const TensorShape& input_shape, DataType input_type,
-                 DataType output_type, int axis, bool is_min,
-                 xla::ComputationDataHandle* argminmax) {
-  xla::ComputationDataHandle init_value;
-  const xla::Computation* reducer;
+Status ArgMinMax(xla::XlaBuilder* builder, XlaOpKernelContext* ctx,
+                 const xla::XlaOp& input, const TensorShape& input_shape,
+                 DataType input_type, DataType output_type, int axis,
+                 bool is_min, xla::XlaOp* argminmax) {
+  xla::XlaOp init_value;
+  const xla::XlaComputation* reducer;
   if (is_min) {
     init_value = XlaHelpers::MaxValue(builder, input_type);
     reducer = ctx->GetOrCreateMin(input_type);
@@ -50,13 +49,13 @@ Status ArgMinMax(xla::ComputationBuilder* builder, XlaOpKernelContext* ctx,
   xla::PrimitiveType xla_output_type;
   TF_RETURN_IF_ERROR(DataTypeToPrimitiveType(output_type, &xla_output_type));
 
-  xla::ComputationDataHandle input_max = builder->Reduce(
-      input, init_value, *reducer, /*dimensions_to_reduce=*/{axis});
+  xla::XlaOp input_max = builder->Reduce(input, init_value, *reducer,
+                                         /*dimensions_to_reduce=*/{axis});
   std::vector<int64> broadcast_dims(input_shape.dims() - 1);
   std::iota(broadcast_dims.begin(), broadcast_dims.begin() + axis, 0);
   std::iota(broadcast_dims.begin() + axis, broadcast_dims.end(), axis + 1);
   // Compute a mask that has 1s for elements equal to the maximum.
-  xla::ComputationDataHandle partial_mask = builder->ConvertElementType(
+  xla::XlaOp partial_mask = builder->ConvertElementType(
       builder->Eq(input, input_max, broadcast_dims), xla_output_type);
 
   // In order to make identity elements for a bitwise And, we:
@@ -65,23 +64,23 @@ Status ArgMinMax(xla::ComputationBuilder* builder, XlaOpKernelContext* ctx,
   //   0xFF...F
   int32 bits_in_type =
       xla::ShapeUtil::ByteSizeOfPrimitiveType(xla_output_type) * 8 - 1;
-  xla::ComputationDataHandle shift_amount =
+  xla::XlaOp shift_amount =
       XlaHelpers::IntegerLiteral(builder, output_type, bits_in_type);
-  xla::ComputationDataHandle full_mask = builder->ShiftRightArithmetic(
+  xla::XlaOp full_mask = builder->ShiftRightArithmetic(
       builder->ShiftLeft(partial_mask, shift_amount), shift_amount);
 
   // And with the vector [0, 1, 2, ...] to convert each 0xFF...F into its
   // index.
-  xla::ComputationDataHandle iota;
+  xla::XlaOp iota;
 
   const int64 axis_size = input_shape.dim_size(axis);
   TF_RETURN_IF_ERROR(XlaHelpers::Iota(builder, output_type, axis_size, &iota));
-  xla::ComputationDataHandle product =
+  xla::XlaOp product =
       builder->And(full_mask, iota, /*broadcast_dimensions=*/{axis});
 
   // If there are multiple maximum elements, choose the one with the highest
   // index.
-  xla::ComputationDataHandle output =
+  xla::XlaOp output =
       builder->Reduce(product, XlaHelpers::MinValue(builder, output_type),
                       *ctx->GetOrCreateMax(output_type),
                       /*dimensions_to_reduce=*/{axis});
@@ -91,36 +90,31 @@ Status ArgMinMax(xla::ComputationBuilder* builder, XlaOpKernelContext* ctx,
 
 }  // namespace
 
-xla::ComputationDataHandle XlaHelpers::MinValue(xla::ComputationBuilder* b,
-                                                DataType data_type) {
+xla::XlaOp XlaHelpers::MinValue(xla::XlaBuilder* b, DataType data_type) {
   xla::PrimitiveType type;
   TF_CHECK_OK(DataTypeToPrimitiveType(data_type, &type));
   return b->ConstantLiteral(xla::Literal::MinValue(type));
 }
 
-xla::ComputationDataHandle XlaHelpers::MaxValue(xla::ComputationBuilder* b,
-                                                DataType data_type) {
+xla::XlaOp XlaHelpers::MaxValue(xla::XlaBuilder* b, DataType data_type) {
   xla::PrimitiveType type;
   TF_CHECK_OK(DataTypeToPrimitiveType(data_type, &type));
   return b->ConstantLiteral(xla::Literal::MaxValue(type));
 }
 
-xla::ComputationDataHandle XlaHelpers::Zero(xla::ComputationBuilder* b,
-                                            DataType data_type) {
+xla::XlaOp XlaHelpers::Zero(xla::XlaBuilder* b, DataType data_type) {
   xla::PrimitiveType type;
   TF_CHECK_OK(DataTypeToPrimitiveType(data_type, &type));
   return b->ConstantLiteral(xla::Literal::Zero(type));
 }
 
-xla::ComputationDataHandle XlaHelpers::One(xla::ComputationBuilder* b,
-                                           DataType data_type) {
+xla::XlaOp XlaHelpers::One(xla::XlaBuilder* b, DataType data_type) {
   xla::PrimitiveType type;
   TF_CHECK_OK(DataTypeToPrimitiveType(data_type, &type));
   return b->ConstantLiteral(xla::Literal::One(type));
 }
 
-xla::ComputationDataHandle XlaHelpers::Epsilon(xla::ComputationBuilder* b,
-                                               DataType data_type) {
+xla::XlaOp XlaHelpers::Epsilon(xla::XlaBuilder* b, DataType data_type) {
   switch (data_type) {
     case DT_HALF:
       return b->ConstantR0<Eigen::half>(
@@ -137,16 +131,15 @@ xla::ComputationDataHandle XlaHelpers::Epsilon(xla::ComputationBuilder* b,
   }
 }
 
-xla::ComputationDataHandle XlaHelpers::IntegerLiteral(
-    xla::ComputationBuilder* b, DataType data_type, int64 value) {
+xla::XlaOp XlaHelpers::IntegerLiteral(xla::XlaBuilder* b, DataType data_type,
+                                      int64 value) {
   xla::PrimitiveType type;
   TF_CHECK_OK(DataTypeToPrimitiveType(data_type, &type));
   return ::tensorflow::IntegerLiteral(b, type, value);
 }
 
-xla::ComputationDataHandle XlaHelpers::FloatLiteral(xla::ComputationBuilder* b,
-                                                    DataType data_type,
-                                                    double value) {
+xla::XlaOp XlaHelpers::FloatLiteral(xla::XlaBuilder* b, DataType data_type,
+                                    double value) {
   xla::PrimitiveType type;
   TF_CHECK_OK(DataTypeToPrimitiveType(data_type, &type));
   return ::tensorflow::FloatLiteral(b, type, value);
@@ -183,28 +176,24 @@ static Tensor MakeLinspaceTensor(const TensorShape& shape, int64 depth) {
   return linspace;
 }
 
-Status XlaHelpers::ArgMax(xla::ComputationBuilder* builder,
-                          XlaOpKernelContext* ctx,
-                          const xla::ComputationDataHandle& input,
+Status XlaHelpers::ArgMax(xla::XlaBuilder* builder, XlaOpKernelContext* ctx,
+                          const xla::XlaOp& input,
                           const TensorShape& input_shape, DataType input_type,
-                          DataType output_type, int axis,
-                          xla::ComputationDataHandle* argmax) {
+                          DataType output_type, int axis, xla::XlaOp* argmax) {
   return ArgMinMax(builder, ctx, input, input_shape, input_type, output_type,
                    axis, /*is_min=*/false, argmax);
 }
 
-Status XlaHelpers::ArgMin(xla::ComputationBuilder* builder,
-                          XlaOpKernelContext* ctx,
-                          const xla::ComputationDataHandle& input,
+Status XlaHelpers::ArgMin(xla::XlaBuilder* builder, XlaOpKernelContext* ctx,
+                          const xla::XlaOp& input,
                           const TensorShape& input_shape, DataType input_type,
-                          DataType output_type, int axis,
-                          xla::ComputationDataHandle* argmin) {
+                          DataType output_type, int axis, xla::XlaOp* argmin) {
   return ArgMinMax(builder, ctx, input, input_shape, input_type, output_type,
                    axis, /*is_min=*/true, argmin);
 }
 
-Status XlaHelpers::Iota(xla::ComputationBuilder* builder, DataType dtype,
-                        int64 size, xla::ComputationDataHandle* iota) {
+Status XlaHelpers::Iota(xla::XlaBuilder* builder, DataType dtype, int64 size,
+                        xla::XlaOp* iota) {
   TensorShape linspace_shape({size});
   Tensor linspace;
   switch (dtype) {
@@ -227,13 +216,10 @@ Status XlaHelpers::Iota(xla::ComputationBuilder* builder, DataType dtype,
   return Status::OK();
 }
 
-Status XlaHelpers::OneHot(xla::ComputationBuilder* builder, int64 depth,
-                          int axis, DataType index_type,
-                          const TensorShape& indices_shape,
-                          const xla::ComputationDataHandle& indices,
-                          const xla::ComputationDataHandle& on_value,
-                          const xla::ComputationDataHandle& off_value,
-                          xla::ComputationDataHandle* one_hot) {
+Status XlaHelpers::OneHot(xla::XlaBuilder* builder, int64 depth, int axis,
+                          DataType index_type, const TensorShape& indices_shape,
+                          const xla::XlaOp& indices, const xla::XlaOp& on_value,
+                          const xla::XlaOp& off_value, xla::XlaOp* one_hot) {
   const int indices_dims = indices_shape.dims();
   const int output_dims = indices_dims + 1;
 
@@ -267,7 +253,7 @@ Status XlaHelpers::OneHot(xla::ComputationBuilder* builder, int64 depth,
   std::vector<int64> broadcast_dims(indices_shape.dims());
   std::iota(broadcast_dims.begin(), broadcast_dims.begin() + axis, 0);
   std::iota(broadcast_dims.begin() + axis, broadcast_dims.end(), axis + 1);
-  xla::ComputationDataHandle one_hot_bool = builder->Eq(
+  xla::XlaOp one_hot_bool = builder->Eq(
       indices, builder->ConstantLiteral(linspace_literal), broadcast_dims);
 
   // Selects the user-provided off_value and on_value values.
@@ -284,10 +270,9 @@ DataType XlaHelpers::SumAccumulationType(const DataType& dtype) {
   return dtype;
 }
 
-xla::ComputationDataHandle XlaHelpers::ConvertElementType(
-    xla::ComputationBuilder* const builder,
-    const xla::ComputationDataHandle& operand,
-    const DataType new_element_type) {
+xla::XlaOp XlaHelpers::ConvertElementType(xla::XlaBuilder* const builder,
+                                          const xla::XlaOp& operand,
+                                          const DataType new_element_type) {
   xla::PrimitiveType convert_to;
   TF_CHECK_OK(DataTypeToPrimitiveType(new_element_type, &convert_to));
   return builder->ConvertElementType(operand, convert_to);
diff --git a/tensorflow/compiler/tf2xla/xla_helpers.h b/tensorflow/compiler/tf2xla/xla_helpers.h
index 68ab93b64a..c3fdc5252e 100644
--- a/tensorflow/compiler/tf2xla/xla_helpers.h
+++ b/tensorflow/compiler/tf2xla/xla_helpers.h
@@ -19,7 +19,7 @@ limitations under the License.
 #define TENSORFLOW_COMPILER_TF2XLA_XLA_HELPERS_H_
 
 #include "tensorflow/compiler/tf2xla/xla_context.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/lib/gtl/array_slice.h"
 
@@ -30,41 +30,34 @@ class XlaHelpers {
  public:
   // Returns a handle representing the minimum value of a scalar
   // element of data_type.
-  static xla::ComputationDataHandle MinValue(xla::ComputationBuilder* b,
-                                             DataType data_type);
+  static xla::XlaOp MinValue(xla::XlaBuilder* b, DataType data_type);
 
   // Returns a handle representing the maximum value of a scalar
   // element of data_type.
-  static xla::ComputationDataHandle MaxValue(xla::ComputationBuilder* b,
-                                             DataType data_type);
+  static xla::XlaOp MaxValue(xla::XlaBuilder* b, DataType data_type);
 
   // Returns a handle representing the zero value of a scalar
   // element of data_type.
-  static xla::ComputationDataHandle Zero(xla::ComputationBuilder* b,
-                                         DataType data_type);
+  static xla::XlaOp Zero(xla::XlaBuilder* b, DataType data_type);
 
   // Returns a handle representing the one value of a scalar
   // element of data_type.
-  static xla::ComputationDataHandle One(xla::ComputationBuilder* b,
-                                        DataType data_type);
+  static xla::XlaOp One(xla::XlaBuilder* b, DataType data_type);
 
   // Returns the machine epsilon for floating-point type `data_type`, i.e.,
   // the difference between 1.0 and the next representable value.
-  static xla::ComputationDataHandle Epsilon(xla::ComputationBuilder* b,
-                                            DataType data_type);
+  static xla::XlaOp Epsilon(xla::XlaBuilder* b, DataType data_type);
 
   // Returns a handle representing the given value of an integer scalar
   // element of data_type.
   // Note that unlike One and Zero, does not work on boolean types.
-  static xla::ComputationDataHandle IntegerLiteral(xla::ComputationBuilder* b,
-                                                   DataType data_type,
-                                                   int64 value);
+  static xla::XlaOp IntegerLiteral(xla::XlaBuilder* b, DataType data_type,
+                                   int64 value);
 
   // Returns a handle representing the given value of a floating-point scalar
   // element of data_type.
-  static xla::ComputationDataHandle FloatLiteral(xla::ComputationBuilder* b,
-                                                 DataType data_type,
-                                                 double value);
+  static xla::XlaOp FloatLiteral(xla::XlaBuilder* b, DataType data_type,
+                                 double value);
 
   // Reshapes literal 'input' to have 'shape'. Both the original shape and
   // 'shape' must contain the same number of elements.
@@ -75,38 +68,32 @@ class XlaHelpers {
   // Sets `argmax` to the argmax of `input` along `axis`. `input_shape` and
   // `input_dtype` are the shape and dtype of `input` respectively, and
   // `output_type` is the dtype to use for `argmax`.
-  static Status ArgMax(xla::ComputationBuilder* builder,
-                       XlaOpKernelContext* ctx,
-                       const xla::ComputationDataHandle& input,
-                       const TensorShape& input_shape, DataType input_type,
-                       DataType output_type, int axis,
-                       xla::ComputationDataHandle* argmax);
+  static Status ArgMax(xla::XlaBuilder* builder, XlaOpKernelContext* ctx,
+                       const xla::XlaOp& input, const TensorShape& input_shape,
+                       DataType input_type, DataType output_type, int axis,
+                       xla::XlaOp* argmax);
 
   // Sets `argmin` to the argmin of `input` along `axis`. `input_shape` and
   // `input_dtype` are the shape and dtype of `input` respectively, and
   // `output_type` is the dtype to use for `argmin`.
-  static Status ArgMin(xla::ComputationBuilder* builder,
-                       XlaOpKernelContext* ctx,
-                       const xla::ComputationDataHandle& input,
-                       const TensorShape& input_shape, DataType input_type,
-                       DataType output_type, int axis,
-                       xla::ComputationDataHandle* argmin);
+  static Status ArgMin(xla::XlaBuilder* builder, XlaOpKernelContext* ctx,
+                       const xla::XlaOp& input, const TensorShape& input_shape,
+                       DataType input_type, DataType output_type, int axis,
+                       xla::XlaOp* argmin);
 
   // Sets *iota to a rank 1 tensor with values [0, 1, 2, ...] of `dtype`.
-  static Status Iota(xla::ComputationBuilder* builder, DataType dtype,
-                     int64 size, xla::ComputationDataHandle* iota);
+  static Status Iota(xla::XlaBuilder* builder, DataType dtype, int64 size,
+                     xla::XlaOp* iota);
 
   // Converts `indices` into a one-hot representation. `depth` is the size
   // of the new axis to add. `axis` is the position at which to add the new
   // axis. `indices_shape` is the shape of `indices`. `on_value` and
   // `off_value` represent the values to use for the on and off positions,
   // respectively.
-  static Status OneHot(xla::ComputationBuilder* builder, int64 depth, int axis,
+  static Status OneHot(xla::XlaBuilder* builder, int64 depth, int axis,
                        DataType index_type, const TensorShape& indices_shape,
-                       const xla::ComputationDataHandle& indices,
-                       const xla::ComputationDataHandle& on_value,
-                       const xla::ComputationDataHandle& off_value,
-                       xla::ComputationDataHandle* one_hot);
+                       const xla::XlaOp& indices, const xla::XlaOp& on_value,
+                       const xla::XlaOp& off_value, xla::XlaOp* one_hot);
 
   // Certain DataTypes should use increased precision DataTypes when performing
   // reductions.  This function remaps a given DataType to a higher precision
@@ -115,10 +102,9 @@ class XlaHelpers {
 
   // A helper for creating a ConvertElementType xla op given a DataType rather
   // than the xla::PrimitiveType.
-  static xla::ComputationDataHandle ConvertElementType(
-      xla::ComputationBuilder* const builder,
-      const xla::ComputationDataHandle& operand,
-      const DataType new_element_type);
+  static xla::XlaOp ConvertElementType(xla::XlaBuilder* const builder,
+                                       const xla::XlaOp& operand,
+                                       const DataType new_element_type);
 };
 
 }  // end namespace tensorflow
diff --git a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc
index 1fe6e69ff2..9e17756b27 100644
--- a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc
+++ b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc
@@ -112,10 +112,10 @@ void CollectNames(const T& entries, std::vector<string>* nonempty_names,
 XlaJitCompiledCpuFunction::Compile(
     const GraphDef& graph_def, const tf2xla::Config& config,
     const xla::ExecutableBuildOptions& build_options) {
-  // Convert the graph_def into an xla::Computation.
+  // Convert the graph_def into an xla::XlaComputation.
   TF_ASSIGN_OR_RETURN(xla::LocalClient * client,
                       xla::ClientLibrary::GetOrCreateLocalClient());
-  xla::Computation computation;
+  xla::XlaComputation computation;
   TF_RETURN_IF_ERROR(tensorflow::ConvertGraphDefToXla(graph_def, config, client,
                                                       &computation));
 
diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.cc b/tensorflow/compiler/tf2xla/xla_op_kernel.cc
index c4bb90d587..2b65f4d5d5 100644
--- a/tensorflow/compiler/tf2xla/xla_op_kernel.cc
+++ b/tensorflow/compiler/tf2xla/xla_op_kernel.cc
@@ -30,7 +30,7 @@ bool XlaOpKernelContext::ValidateInputsAreSameShape(OpKernel* op) {
   return context_->ValidateInputsAreSameShape(op);
 }
 
-xla::ComputationBuilder* XlaOpKernelContext::builder() const {
+xla::XlaBuilder* XlaOpKernelContext::builder() const {
   return XlaContext::Get(this).builder();
 }
 
@@ -38,9 +38,9 @@ xla::ComputationBuilder* XlaOpKernelContext::builder() const {
 static const XlaExpression* CastExpressionFromTensor(const Tensor& tensor) {
   const XlaExpression* expression =
       reinterpret_cast<const XlaExpression*>(tensor.tensor_data().data());
-  CHECK(expression->handle().handle() != 0 ||
+  CHECK(expression->handle().builder() != nullptr ||
         expression->resource() != nullptr);
-  VLOG(1) << "Fetched T" << expression->handle().handle();
+  VLOG(1) << "Fetched T" << expression->handle();
   return expression;
 }
 
@@ -48,20 +48,18 @@ static const XlaExpression* CastExpressionFromTensor(const Tensor& tensor) {
 static XlaExpression* CastExpressionFromUninitializedTensor(Tensor* tensor) {
   const XlaExpression* expression =
       reinterpret_cast<const XlaExpression*>(tensor->tensor_data().data());
-  CHECK_EQ(expression->handle().handle(), 0);
+  CHECK_EQ(expression->handle().builder(), nullptr);
   return const_cast<XlaExpression*>(expression);
 }
 
-// Retrieves the ComputationDataHandle from an input Tensor to an Op. This
-// computation was constructed by an Op that executed previously and
-// created the output Tensor using CreateOutputTensorFromComputation
-// or CreateConstantOutputTensor.
-static const xla::ComputationDataHandle& GetComputationFromTensor(
-    const Tensor& tensor) {
+// Retrieves the XlaOp from an input Tensor to an Op. This computation was
+// constructed by an Op that executed previously and created the output Tensor
+// using CreateOutputTensorFromComputation or CreateConstantOutputTensor.
+static const xla::XlaOp& GetComputationFromTensor(const Tensor& tensor) {
   return CastExpressionFromTensor(tensor)->handle();
 }
 
-const xla::ComputationDataHandle& XlaOpKernelContext::Input(int index) {
+const xla::XlaOp& XlaOpKernelContext::Input(int index) {
   return GetComputationFromTensor(context_->input(index));
 }
 
@@ -106,7 +104,7 @@ Status XlaOpKernelContext::ConstantInputReshaped(
     return HostTensorToLiteral(temp, constant_literal);
   }
 
-  xla::ComputationDataHandle handle = expression->handle();
+  xla::XlaOp handle = expression->handle();
   if (new_shape != tensor.shape()) {
     // Reshape the handle to the desired shape.
     handle = builder()->Reshape(handle, new_shape.dim_sizes());
@@ -141,8 +139,17 @@ Status XlaOpKernelContext::ConstantInputReshaped(
   }
 
   // Ask the XLA compiler to evaluate the data handle to a literal.
+  xla::StatusOr<xla::XlaComputation> constant_graph =
+      builder()->BuildConstantSubGraph(handle);
+  if (!constant_graph.ok()) {
+    return errors::Internal(
+        "Error getting a compile-time constant graph for ",
+        context_->op_kernel().name(), " input ", index,
+        ".\nError: ", constant_graph.status().error_message());
+  }
   xla::StatusOr<std::unique_ptr<xla::Literal>> computed =
-      builder()->ComputeConstant(handle, &layout);
+      compiler()->client()->ComputeConstant(constant_graph.ValueOrDie(),
+                                            &layout);
   if (!computed.ok()) {
     return errors::Internal("Error evaluating ", context_->op_kernel().name(),
                             " input ", index,
@@ -260,9 +267,9 @@ Status XlaOpKernelContext::ConstantInputAsShape(int index, TensorShape* shape) {
   return Status::OK();
 }
 
-Status XlaOpKernelContext::InputList(
-    StringPiece name, std::vector<xla::ComputationDataHandle>* handles,
-    std::vector<TensorShape>* shapes) {
+Status XlaOpKernelContext::InputList(StringPiece name,
+                                     std::vector<xla::XlaOp>* handles,
+                                     std::vector<TensorShape>* shapes) {
   OpInputList inputs;
   TF_RETURN_IF_ERROR(context_->input_list(name, &inputs));
   handles->clear();
@@ -285,9 +292,9 @@ Status XlaOpKernelContext::ConstantInputList(
   return Status::OK();
 }
 
-Status XlaOpKernelContext::ReadVariableInput(
-    int index, DataType type, TensorShape* shape,
-    xla::ComputationDataHandle* value) {
+Status XlaOpKernelContext::ReadVariableInput(int index, DataType type,
+                                             TensorShape* shape,
+                                             xla::XlaOp* value) {
   const Tensor& tensor = context_->input(index);
   const XlaExpression* expression = CastExpressionFromTensor(tensor);
   XlaResource* variable = expression->resource();
@@ -334,8 +341,7 @@ Status XlaOpKernelContext::GetVariableTypeAndShape(int index, DataType* type,
   return Status::OK();
 }
 
-void XlaOpKernelContext::SetOutput(int index,
-                                   const xla::ComputationDataHandle& handle) {
+void XlaOpKernelContext::SetOutput(int index, const xla::XlaOp& handle) {
   // Makes the host Tensor that will refer to the expression.
   Tensor* output = nullptr;
   auto shape = builder()->GetShape(handle);
@@ -349,7 +355,7 @@ void XlaOpKernelContext::SetOutput(int index,
   // corresponds.
   TensorShape tensor_shape;
   OP_REQUIRES_OK(context_,
-                 XLAShapeToTensorShape(*shape.ValueOrDie(), &tensor_shape));
+                 XLAShapeToTensorShape(shape.ValueOrDie(), &tensor_shape));
   OP_REQUIRES_OK(context_,
                  context_->allocate_output(index, tensor_shape, &output));
 
@@ -364,8 +370,8 @@ void XlaOpKernelContext::SetConstantOutput(int index, const Tensor& constant) {
 
   xla::Literal literal;
   OP_REQUIRES_OK(context_, HostTensorToLiteral(constant, &literal));
-  xla::ComputationDataHandle handle = builder()->ConstantLiteral(literal);
-  CHECK_NE(handle.handle(), 0);
+  xla::XlaOp handle = builder()->ConstantLiteral(literal);
+  CHECK_NE(handle.builder(), nullptr);
 
   // Make the Tensor that will refer to the expression.
   Tensor* output = nullptr;
@@ -386,8 +392,7 @@ void XlaOpKernelContext::SetInvalidOutput(int index) {
   OP_REQUIRES_OK(context_,
                  context_->allocate_output(index, TensorShape({}), &output));
   XlaExpression* expression = CastExpressionFromUninitializedTensor(output);
-  xla::ComputationDataHandle handle;
-  handle.set_handle(0);
+  xla::XlaOp handle;
   expression->set_handle(handle);
 }
 
@@ -410,8 +415,8 @@ Status XlaOpKernelContext::GetResourceInput(int index, XlaResource** resource) {
 }
 
 Status XlaOpKernelContext::AssignVariable(int input_index, DataType type,
-                                          xla::ComputationDataHandle handle) {
-  TF_RET_CHECK(handle.handle() != 0);
+                                          xla::XlaOp handle) {
+  TF_RET_CHECK(handle.builder() != nullptr);
 
   const XlaExpression* expression =
       CastExpressionFromTensor(context_->input(input_index));
@@ -425,7 +430,7 @@ Status XlaOpKernelContext::AssignVariable(int input_index, DataType type,
   }
   TensorShape shape;
   TF_RETURN_IF_ERROR(
-      XLAShapeToTensorShape(*shape_or_status.ValueOrDie(), &shape));
+      XLAShapeToTensorShape(shape_or_status.ValueOrDie(), &shape));
 
   TF_RETURN_IF_ERROR(variable->SetTypeAndShape(type, shape));
 
@@ -457,22 +462,22 @@ void XlaOpKernelContext::CtxFailureWithWarning(const char* file, int line,
   context_->CtxFailureWithWarning(file, line, s);
 }
 
-const xla::Computation* XlaOpKernelContext::GetOrCreateMax(
+const xla::XlaComputation* XlaOpKernelContext::GetOrCreateMax(
     const DataType type) {
   return XlaContext::Get(context_).GetOrCreateMax(type);
 }
 
-const xla::Computation* XlaOpKernelContext::GetOrCreateMin(
+const xla::XlaComputation* XlaOpKernelContext::GetOrCreateMin(
     const DataType type) {
   return XlaContext::Get(context_).GetOrCreateMin(type);
 }
 
-const xla::Computation* XlaOpKernelContext::GetOrCreateAdd(
+const xla::XlaComputation* XlaOpKernelContext::GetOrCreateAdd(
     const DataType type) {
   return XlaContext::Get(context_).GetOrCreateAdd(type);
 }
 
-const xla::Computation* XlaOpKernelContext::GetOrCreateMul(
+const xla::XlaComputation* XlaOpKernelContext::GetOrCreateMul(
     const DataType type) {
   return XlaContext::Get(context_).GetOrCreateMul(type);
 }
diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.h b/tensorflow/compiler/tf2xla/xla_op_kernel.h
index 4e4b97e0ce..667dc262ca 100644
--- a/tensorflow/compiler/tf2xla/xla_op_kernel.h
+++ b/tensorflow/compiler/tf2xla/xla_op_kernel.h
@@ -17,7 +17,7 @@ limitations under the License.
 #define TENSORFLOW_COMPILER_TF2XLA_XLA_OP_KERNEL_H_
 
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/platform/macros.h"
 
@@ -58,8 +58,8 @@ class XlaOpKernelContext {
  public:
   explicit XlaOpKernelContext(OpKernelContext* context);
 
-  // Returns the XLA ComputationBuilder containing the output of compilation.
-  xla::ComputationBuilder* builder() const;
+  // Returns the XLA XlaBuilder containing the output of compilation.
+  xla::XlaBuilder* builder() const;
 
   // Inputs
 
@@ -72,10 +72,10 @@ class XlaOpKernelContext {
   // Returns the shape of input 'index'.
   TensorShape InputShape(int index);
 
-  // Returns input 'index' as a ComputationDataHandle. Unlike
+  // Returns input 'index' as a XlaOp. Unlike
   // OpKernelContext::Input returns a symbolic value rather than a concrete
   // Tensor.
-  const xla::ComputationDataHandle& Input(int index);
+  const xla::XlaOp& Input(int index);
 
   // Returns true if all inputs are the same shape, otherwise sets the
   // status to a non-OK value and returns false.
@@ -85,8 +85,7 @@ class XlaOpKernelContext {
   // Returns the named list-valued immutable input in "list", as
   // defined in the OpDef.  If the named output is not list-valued,
   // returns a one-element list.
-  Status InputList(StringPiece name,
-                   std::vector<xla::ComputationDataHandle>* handles,
+  Status InputList(StringPiece name, std::vector<xla::XlaOp>* handles,
                    std::vector<TensorShape>* shapes);
 
   // Helper methods for constant inputs.
@@ -132,10 +131,10 @@ class XlaOpKernelContext {
     return context_->expected_output_dtype(index);
   }
 
-  // Sets output 'index' to the ComputationDataHandle 'handle'.
+  // Sets output 'index' to the XlaOp 'handle'.
   // All outputs should be set using SetOutput and SetConstantOutput, not
   // via the underlying OpKernelContext.
-  void SetOutput(int index, const xla::ComputationDataHandle& handle);
+  void SetOutput(int index, const xla::XlaOp& handle);
 
   // Sets output 'index' to compile-time constant 'host_tensor', where
   // 'host_tensor' is a tensor in host memory. It is preferable to use
@@ -168,14 +167,13 @@ class XlaOpKernelContext {
   // variable. Returns an error if the variable has not been initialized, or if
   // its type does not match `type`.
   Status ReadVariableInput(int index, DataType type, TensorShape* shape,
-                           xla::ComputationDataHandle* value);
+                           xla::XlaOp* value);
 
   // Assigns the value `handle` to the variable referenced by input
   // `input_index`. The variable must be of `type`. Returns an error if the
   // variable has been initialized with a different type or with a
   // different shape.
-  Status AssignVariable(int input_index, DataType type,
-                        xla::ComputationDataHandle handle);
+  Status AssignVariable(int input_index, DataType type, xla::XlaOp handle);
 
   // Helper routines for the OP_REQUIRES macros
   void CtxFailure(const Status& s);
@@ -205,22 +203,22 @@ class XlaOpKernelContext {
   // Gets an XLA lambda to compute Max. This is cached in the
   // XlaContext since it may be used by multiple Ops. There is a
   // separate specialization of the computation for each DataType.
-  const xla::Computation* GetOrCreateMax(const DataType type);
+  const xla::XlaComputation* GetOrCreateMax(const DataType type);
 
   // Gets an XLA lambda to compute Min. This is cached in the
   // XlaContext since it may be used by multiple Ops. There is a
   // separate specialization of the computation for each DataType.
-  const xla::Computation* GetOrCreateMin(const DataType type);
+  const xla::XlaComputation* GetOrCreateMin(const DataType type);
 
   // Gets an XLA lambda to compute Add. This is cached in the
   // XlaContext since it may be used by multiple Ops. There is a
   // separate specialization of the computation for each DataType.
-  const xla::Computation* GetOrCreateAdd(const DataType type);
+  const xla::XlaComputation* GetOrCreateAdd(const DataType type);
 
   // Gets an XLA lambda to compute Mul. This is cached in the
   // XlaContext since it may be used by multiple Ops. There is a
   // separate specialization of the computation for each DataType.
-  const xla::Computation* GetOrCreateMul(const DataType type);
+  const xla::XlaComputation* GetOrCreateMul(const DataType type);
 
  private:
   OpKernelContext* const context_;
diff --git a/tensorflow/compiler/tf2xla/xla_resource.cc b/tensorflow/compiler/tf2xla/xla_resource.cc
index c2075b44b8..540c65c597 100644
--- a/tensorflow/compiler/tf2xla/xla_resource.cc
+++ b/tensorflow/compiler/tf2xla/xla_resource.cc
@@ -26,8 +26,7 @@ limitations under the License.
 namespace tensorflow {
 
 XlaResource::XlaResource(Kind kind, int arg_num, string name, DataType type,
-                         TensorShape shape,
-                         const xla::ComputationDataHandle& initial_value,
+                         TensorShape shape, const xla::XlaOp& initial_value,
                          int64 tensor_array_size,
                          const std::set<string>& tensor_array_gradients)
     : kind_(kind),
@@ -41,11 +40,10 @@ XlaResource::XlaResource(Kind kind, int arg_num, string name, DataType type,
   CHECK(kind_ != kInvalid);
 
   for (const string& gradient : tensor_array_gradients) {
-    tensor_array_gradients_[gradient].reset(
-        new XlaResource(/*kind=*/kTensorArray, /*arg_num=*/-1,
-                        /*name=*/strings::StrCat("TensorArrayGrad: ", name_),
-                        type_, shape_, xla::ComputationDataHandle(),
-                        tensor_array_size_, /*tensor_array_gradients=*/{}));
+    tensor_array_gradients_[gradient].reset(new XlaResource(
+        /*kind=*/kTensorArray, /*arg_num=*/-1,
+        /*name=*/strings::StrCat("TensorArrayGrad: ", name_), type_, shape_,
+        xla::XlaOp(), tensor_array_size_, /*tensor_array_gradients=*/{}));
   }
 }
 
@@ -73,7 +71,7 @@ Status XlaResource::SetTypeAndShape(DataType type, const TensorShape& shape) {
   return Status::OK();
 }
 
-Status XlaResource::SetValue(const xla::ComputationDataHandle& value) {
+Status XlaResource::SetValue(const xla::XlaOp& value) {
   if (type_ == DT_INVALID) {
     return errors::InvalidArgument(
         "Resource '", name_,
@@ -83,7 +81,7 @@ Status XlaResource::SetValue(const xla::ComputationDataHandle& value) {
   return Status::OK();
 }
 
-Status XlaResource::SetZeroValue(xla::ComputationBuilder* builder) {
+Status XlaResource::SetZeroValue(xla::XlaBuilder* builder) {
   if (type_ == DT_INVALID) {
     return errors::InvalidArgument(
         "Resource '", name_,
@@ -121,9 +119,9 @@ Status XlaResource::SetZeroValue(xla::ComputationBuilder* builder) {
   return Status::OK();
 }
 
-Status XlaResource::GetOrCreateTensorArrayGradient(
-    const string& source, xla::ComputationBuilder* builder,
-    XlaResource** gradient_out) {
+Status XlaResource::GetOrCreateTensorArrayGradient(const string& source,
+                                                   xla::XlaBuilder* builder,
+                                                   XlaResource** gradient_out) {
   VLOG(2) << "Gradient lookup for resource: " << name_
           << " gradient: " << source;
   TF_RET_CHECK(kind_ == kTensorArray);
@@ -132,7 +130,7 @@ Status XlaResource::GetOrCreateTensorArrayGradient(
     TensorShape ta_shape;
     ta_shape.AddDim(tensor_array_size_);
     ta_shape.AppendShape(shape_);
-    xla::ComputationDataHandle gradient_value = builder->Broadcast(
+    xla::XlaOp gradient_value = builder->Broadcast(
         XlaHelpers::Zero(builder, type_), ta_shape.dim_sizes());
     gradient.reset(
         new XlaResource(/*kind=*/kTensorArray, /*arg_num=*/-1,
@@ -144,13 +142,12 @@ Status XlaResource::GetOrCreateTensorArrayGradient(
   return Status::OK();
 }
 
-Status XlaResource::Pack(xla::ComputationDataHandle* pack,
-                         xla::ComputationBuilder* builder) const {
+Status XlaResource::Pack(xla::XlaOp* pack, xla::XlaBuilder* builder) const {
   if (tensor_array_gradients_.empty()) {
     *pack = value_;
   } else {
     TF_RET_CHECK(kind_ == kTensorArray);
-    std::vector<xla::ComputationDataHandle> elems;
+    std::vector<xla::XlaOp> elems;
     elems.push_back(value_);
     for (const auto& gradient : tensor_array_gradients_) {
       elems.push_back(gradient.second->value_);
@@ -161,8 +158,8 @@ Status XlaResource::Pack(xla::ComputationDataHandle* pack,
 }
 
 Status XlaResource::SetFromPack(const std::set<string>& gradient_sources,
-                                const xla::ComputationDataHandle& pack,
-                                xla::ComputationBuilder* builder) {
+                                const xla::XlaOp& pack,
+                                xla::XlaBuilder* builder) {
   if (gradient_sources.empty()) {
     if (!initialized()) {
       initial_value_ = pack;
diff --git a/tensorflow/compiler/tf2xla/xla_resource.h b/tensorflow/compiler/tf2xla/xla_resource.h
index 1bb2c7274e..9ce36d1aa7 100644
--- a/tensorflow/compiler/tf2xla/xla_resource.h
+++ b/tensorflow/compiler/tf2xla/xla_resource.h
@@ -18,7 +18,7 @@ limitations under the License.
 
 #include <memory>
 
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.pb.h"
@@ -37,8 +37,7 @@ class XlaResource {
   };
 
   XlaResource(Kind kind, int arg_num, string name, DataType type,
-              TensorShape shape,
-              const xla::ComputationDataHandle& initial_value,
+              TensorShape shape, const xla::XlaOp& initial_value,
               int64 tensor_array_size,
               const std::set<string>& tensor_array_gradients);
 
@@ -69,16 +68,14 @@ class XlaResource {
   // this is the shape of each entry in the TensorArray/Stack.
   const TensorShape& shape() const { return shape_; }
 
-  const xla::ComputationDataHandle& value() const { return value_; }
+  const xla::XlaOp& value() const { return value_; }
 
   // Value of the resource at computation entry. Used to detect which
   // variables have new values that need to be written back.
-  const xla::ComputationDataHandle& initial_value() const {
-    return initial_value_;
-  }
+  const xla::XlaOp& initial_value() const { return initial_value_; }
 
   // A variable is initialized if it has a value.
-  bool initialized() const { return value_.handle() > 0; }
+  bool initialized() const { return value_.builder() != nullptr; }
 
   // Sets the type and shape of the resource. The type and shape of a resource
   // must not change once the variable has been initialized.
@@ -86,17 +83,17 @@ class XlaResource {
 
   // Sets the current value of the resource. Returns an error if the type is not
   // set to a valid value.
-  Status SetValue(const xla::ComputationDataHandle& value);
+  Status SetValue(const xla::XlaOp& value);
 
   // Sets the current value of the resource to an all-zero value.
-  Status SetZeroValue(xla::ComputationBuilder* builder);
+  Status SetZeroValue(xla::XlaBuilder* builder);
 
   // Looks up the gradient for `source`, or creates it if it does not already
   // exist. The call target must be an initialized TensorArray resource. A
   // TensorArray can have multiple named gradients; see the operator
   // documentation for TensorArrayGradV3 for details.
   Status GetOrCreateTensorArrayGradient(const string& source,
-                                        xla::ComputationBuilder* builder,
+                                        xla::XlaBuilder* builder,
                                         XlaResource** gradient_out);
 
   // Packs a resource into a single XLA value `pack`, suitable for use as
@@ -104,8 +101,7 @@ class XlaResource {
   // gradients, sets `*pack` to `value`.
   // For TensorArrays with gradients, packs the value and its gradient values in
   // a tuple; the gradients values are packed in order by source name.
-  Status Pack(xla::ComputationDataHandle* pack,
-              xla::ComputationBuilder* builder) const;
+  Status Pack(xla::XlaOp* pack, xla::XlaBuilder* builder) const;
 
   // Updates the resource with values from `pack`. If `gradient_sources` is
   // non-empty, treats `pack` as a tuple that represents a TensorArray and
@@ -114,8 +110,7 @@ class XlaResource {
   // values.
   // Opposite of Pack().
   Status SetFromPack(const std::set<string>& gradient_sources,
-                     const xla::ComputationDataHandle& pack,
-                     xla::ComputationBuilder* builder);
+                     const xla::XlaOp& pack, xla::XlaBuilder* builder);
 
   // TensorArray and Stack specific fields
 
@@ -144,8 +139,8 @@ class XlaResource {
 
   DataType type_;
   TensorShape shape_;
-  xla::ComputationDataHandle value_;
-  xla::ComputationDataHandle initial_value_;
+  xla::XlaOp value_;
+  xla::XlaOp initial_value_;
 
   int64 tensor_array_size_ = -1;
 
diff --git a/tensorflow/compiler/xla/client/BUILD b/tensorflow/compiler/xla/client/BUILD
index 286d06d12f..aac3273d5f 100644
--- a/tensorflow/compiler/xla/client/BUILD
+++ b/tensorflow/compiler/xla/client/BUILD
@@ -106,6 +106,7 @@ cc_library(
         "//tensorflow/compiler/xla:statusor",
         "//tensorflow/compiler/xla:util",
         "//tensorflow/compiler/xla:xla_data_proto",
+        "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/compiler/xla/service:backend",
         "//tensorflow/compiler/xla/service:compiler",
         "//tensorflow/compiler/xla/service:device_memory_allocator",
diff --git a/tensorflow/compiler/xla/client/local_client.h b/tensorflow/compiler/xla/client/local_client.h
index f306c520ed..4ce7059f7e 100644
--- a/tensorflow/compiler/xla/client/local_client.h
+++ b/tensorflow/compiler/xla/client/local_client.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/client/client.h"
 #include "tensorflow/compiler/xla/client/computation.h"
 #include "tensorflow/compiler/xla/client/executable_build_options.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 #include "tensorflow/compiler/xla/executable_run_options.h"
 #include "tensorflow/compiler/xla/service/compiler.h"
 #include "tensorflow/compiler/xla/service/device_memory_allocator.h"
-- 
GitLab


From c8a0f6e92b3197c76c5aac1d2a7612e2f4b3fc56 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 17:46:35 -0700
Subject: [PATCH 077/395] Go: Update generated wrapper functions for TensorFlow
 ops. PiperOrigin-RevId: 194874988

---
 tensorflow/go/op/wrappers.go | 190 +++++++++++++++++------------------
 1 file changed, 95 insertions(+), 95 deletions(-)

diff --git a/tensorflow/go/op/wrappers.go b/tensorflow/go/op/wrappers.go
index 83de1c5a92..c12ea51563 100644
--- a/tensorflow/go/op/wrappers.go
+++ b/tensorflow/go/op/wrappers.go
@@ -6655,6 +6655,101 @@ func FusedBatchNormV2(scope *Scope, x tf.Output, scale tf.Output, offset tf.Outp
 	return op.Output(0), op.Output(1), op.Output(2), op.Output(3), op.Output(4)
 }
 
+// Reverses specific dimensions of a tensor.
+//
+// NOTE `tf.reverse` has now changed behavior in preparation for 1.0.
+// `tf.reverse_v2` is currently an alias that will be deprecated before TF 1.0.
+//
+// Given a `tensor`, and a `int32` tensor `axis` representing the set of
+// dimensions of `tensor` to reverse. This operation reverses each dimension
+// `i` for which there exists `j` s.t. `axis[j] == i`.
+//
+// `tensor` can have up to 8 dimensions. The number of dimensions specified
+// in `axis` may be 0 or more entries. If an index is specified more than
+// once, a InvalidArgument error is raised.
+//
+// For example:
+//
+// ```
+// # tensor 't' is [[[[ 0,  1,  2,  3],
+// #                  [ 4,  5,  6,  7],
+// #                  [ 8,  9, 10, 11]],
+// #                 [[12, 13, 14, 15],
+// #                  [16, 17, 18, 19],
+// #                  [20, 21, 22, 23]]]]
+// # tensor 't' shape is [1, 2, 3, 4]
+//
+// # 'dims' is [3] or 'dims' is [-1]
+// reverse(t, dims) ==> [[[[ 3,  2,  1,  0],
+//                         [ 7,  6,  5,  4],
+//                         [ 11, 10, 9, 8]],
+//                        [[15, 14, 13, 12],
+//                         [19, 18, 17, 16],
+//                         [23, 22, 21, 20]]]]
+//
+// # 'dims' is '[1]' (or 'dims' is '[-3]')
+// reverse(t, dims) ==> [[[[12, 13, 14, 15],
+//                         [16, 17, 18, 19],
+//                         [20, 21, 22, 23]
+//                        [[ 0,  1,  2,  3],
+//                         [ 4,  5,  6,  7],
+//                         [ 8,  9, 10, 11]]]]
+//
+// # 'dims' is '[2]' (or 'dims' is '[-2]')
+// reverse(t, dims) ==> [[[[8, 9, 10, 11],
+//                         [4, 5, 6, 7],
+//                         [0, 1, 2, 3]]
+//                        [[20, 21, 22, 23],
+//                         [16, 17, 18, 19],
+//                         [12, 13, 14, 15]]]]
+// ```
+//
+// Arguments:
+//	tensor: Up to 8-D.
+//	axis: 1-D. The indices of the dimensions to reverse. Must be in the range
+// `[-rank(tensor), rank(tensor))`.
+//
+// Returns The same shape as `tensor`.
+func ReverseV2(scope *Scope, tensor tf.Output, axis tf.Output) (output tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	opspec := tf.OpSpec{
+		Type: "ReverseV2",
+		Input: []tf.Input{
+			tensor, axis,
+		},
+	}
+	op := scope.AddOperation(opspec)
+	return op.Output(0)
+}
+
+// Adds `bias` to `value`.
+//
+// This is a deprecated version of BiasAdd and will be soon removed.
+//
+// This is a special case of `tf.add` where `bias` is restricted to be 1-D.
+// Broadcasting is supported, so `value` may have any number of dimensions.
+//
+// Arguments:
+//	value: Any number of dimensions.
+//	bias: 1-D with size the last dimension of `value`.
+//
+// Returns Broadcasted sum of `value` and `bias`.
+func BiasAddV1(scope *Scope, value tf.Output, bias tf.Output) (output tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	opspec := tf.OpSpec{
+		Type: "BiasAddV1",
+		Input: []tf.Input{
+			value, bias,
+		},
+	}
+	op := scope.AddOperation(opspec)
+	return op.Output(0)
+}
+
 // Transforms a Tensor into a serialized TensorProto proto.
 //
 // Arguments:
@@ -13816,101 +13911,6 @@ func ResourceApplyCenteredRMSProp(scope *Scope, var_ tf.Output, mg tf.Output, ms
 	return scope.AddOperation(opspec)
 }
 
-// Adds `bias` to `value`.
-//
-// This is a deprecated version of BiasAdd and will be soon removed.
-//
-// This is a special case of `tf.add` where `bias` is restricted to be 1-D.
-// Broadcasting is supported, so `value` may have any number of dimensions.
-//
-// Arguments:
-//	value: Any number of dimensions.
-//	bias: 1-D with size the last dimension of `value`.
-//
-// Returns Broadcasted sum of `value` and `bias`.
-func BiasAddV1(scope *Scope, value tf.Output, bias tf.Output) (output tf.Output) {
-	if scope.Err() != nil {
-		return
-	}
-	opspec := tf.OpSpec{
-		Type: "BiasAddV1",
-		Input: []tf.Input{
-			value, bias,
-		},
-	}
-	op := scope.AddOperation(opspec)
-	return op.Output(0)
-}
-
-// Reverses specific dimensions of a tensor.
-//
-// NOTE `tf.reverse` has now changed behavior in preparation for 1.0.
-// `tf.reverse_v2` is currently an alias that will be deprecated before TF 1.0.
-//
-// Given a `tensor`, and a `int32` tensor `axis` representing the set of
-// dimensions of `tensor` to reverse. This operation reverses each dimension
-// `i` for which there exists `j` s.t. `axis[j] == i`.
-//
-// `tensor` can have up to 8 dimensions. The number of dimensions specified
-// in `axis` may be 0 or more entries. If an index is specified more than
-// once, a InvalidArgument error is raised.
-//
-// For example:
-//
-// ```
-// # tensor 't' is [[[[ 0,  1,  2,  3],
-// #                  [ 4,  5,  6,  7],
-// #                  [ 8,  9, 10, 11]],
-// #                 [[12, 13, 14, 15],
-// #                  [16, 17, 18, 19],
-// #                  [20, 21, 22, 23]]]]
-// # tensor 't' shape is [1, 2, 3, 4]
-//
-// # 'dims' is [3] or 'dims' is [-1]
-// reverse(t, dims) ==> [[[[ 3,  2,  1,  0],
-//                         [ 7,  6,  5,  4],
-//                         [ 11, 10, 9, 8]],
-//                        [[15, 14, 13, 12],
-//                         [19, 18, 17, 16],
-//                         [23, 22, 21, 20]]]]
-//
-// # 'dims' is '[1]' (or 'dims' is '[-3]')
-// reverse(t, dims) ==> [[[[12, 13, 14, 15],
-//                         [16, 17, 18, 19],
-//                         [20, 21, 22, 23]
-//                        [[ 0,  1,  2,  3],
-//                         [ 4,  5,  6,  7],
-//                         [ 8,  9, 10, 11]]]]
-//
-// # 'dims' is '[2]' (or 'dims' is '[-2]')
-// reverse(t, dims) ==> [[[[8, 9, 10, 11],
-//                         [4, 5, 6, 7],
-//                         [0, 1, 2, 3]]
-//                        [[20, 21, 22, 23],
-//                         [16, 17, 18, 19],
-//                         [12, 13, 14, 15]]]]
-// ```
-//
-// Arguments:
-//	tensor: Up to 8-D.
-//	axis: 1-D. The indices of the dimensions to reverse. Must be in the range
-// `[-rank(tensor), rank(tensor))`.
-//
-// Returns The same shape as `tensor`.
-func ReverseV2(scope *Scope, tensor tf.Output, axis tf.Output) (output tf.Output) {
-	if scope.Err() != nil {
-		return
-	}
-	opspec := tf.OpSpec{
-		Type: "ReverseV2",
-		Input: []tf.Input{
-			tensor, axis,
-		},
-	}
-	op := scope.AddOperation(opspec)
-	return op.Output(0)
-}
-
 // RealAttr is an optional argument to Real.
 type RealAttr func(optionalAttr)
 
-- 
GitLab


From 40721422bfc9cec546537799e16dd75f443d2db2 Mon Sep 17 00:00:00 2001
From: Yifei Feng <yifeif@google.com>
Date: Mon, 30 Apr 2018 18:01:23 -0700
Subject: [PATCH 078/395] Remove proto header import from
 core/framework/tracking_allocator.h

The goal is to make kernels mostly independent of proto headers, which will let us lock down our .so import.

PiperOrigin-RevId: 194876569
---
 tensorflow/core/common_runtime/eager/kernel_and_device.cc | 1 +
 tensorflow/core/common_runtime/eager/kernel_and_device.h  | 4 ++++
 tensorflow/core/framework/tracking_allocator.h            | 1 -
 3 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/common_runtime/eager/kernel_and_device.cc b/tensorflow/core/common_runtime/eager/kernel_and_device.cc
index 0a4895a938..a63b2b9711 100644
--- a/tensorflow/core/common_runtime/eager/kernel_and_device.cc
+++ b/tensorflow/core/common_runtime/eager/kernel_and_device.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/rendezvous_mgr.h"
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/step_stats.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/gtl/map_util.h"
 #include "tensorflow/core/lib/gtl/stl_util.h"
diff --git a/tensorflow/core/common_runtime/eager/kernel_and_device.h b/tensorflow/core/common_runtime/eager/kernel_and_device.h
index 46ec550c78..f78d197fd5 100644
--- a/tensorflow/core/common_runtime/eager/kernel_and_device.h
+++ b/tensorflow/core/common_runtime/eager/kernel_and_device.h
@@ -32,6 +32,10 @@ limitations under the License.
 
 namespace tensorflow {
 
+// Forward declaration for proto class NodeExecStats so we do not need to
+// include the proto header
+class NodeExecStats;
+
 // KernelAndDevice encapsulates an instantiated kernel and the device it is on.
 //
 // Also see:
diff --git a/tensorflow/core/framework/tracking_allocator.h b/tensorflow/core/framework/tracking_allocator.h
index f6c3c0b71b..661c28969e 100644
--- a/tensorflow/core/framework/tracking_allocator.h
+++ b/tensorflow/core/framework/tracking_allocator.h
@@ -18,7 +18,6 @@ limitations under the License.
 
 #include <unordered_map>
 #include "tensorflow/core/framework/allocator.h"
-#include "tensorflow/core/framework/step_stats.pb.h"
 #include "tensorflow/core/lib/core/refcount.h"
 #include "tensorflow/core/lib/gtl/inlined_vector.h"
 #include "tensorflow/core/platform/mutex.h"
-- 
GitLab


From 85d30bfcf412bd1ca06fa33548344bf40eedb4ac Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 18:05:37 -0700
Subject: [PATCH 079/395] Internal change.

PiperOrigin-RevId: 194877173
---
 .../kernels/bidirectional_sequence_lstm.cc    | 70 ++++++++++++-------
 .../bidirectional_sequence_lstm_test.cc       |  8 ---
 tensorflow/contrib/lite/kernels/lstm.cc       | 49 ++++++++-----
 tensorflow/contrib/lite/kernels/lstm_test.cc  |  4 --
 .../lite/kernels/optional_tensor_test.cc      |  4 --
 .../kernels/unidirectional_sequence_lstm.cc   | 47 +++++++++----
 .../unidirectional_sequence_lstm_test.cc      |  4 --
 tensorflow/contrib/lite/models/speech_test.cc | 16 ++---
 .../testdata/speech_asr_lm_model.test_spec    | 20 +++---
 .../identify_lstm_merge_inputs.cc             |  9 ++-
 .../identify_lstm_split_inputs.cc             |  7 +-
 .../toco/graph_transformations/lstm_utils.h   |  8 +--
 12 files changed, 139 insertions(+), 107 deletions(-)

diff --git a/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm.cc b/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm.cc
index a64ac42bc4..3ac0210f36 100644
--- a/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm.cc
+++ b/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm.cc
@@ -96,15 +96,23 @@ constexpr int kBwProjectionWeightsTensor = 33;  // Optional
 constexpr int kBwProjectionBiasTensor = 34;  // Optional
 
 // Output tensors.
-constexpr int kFwScratchBufferTensor = 0;
-constexpr int kFwOutputStateTensor = 1;
-constexpr int kFwCellStateTensor = 2;
-constexpr int kFwOutputTensor = 3;
+constexpr int kFwOutputStateTensor = 0;
+constexpr int kFwCellStateTensor = 1;
+constexpr int kFwOutputTensor = 2;
+
+constexpr int kBwOutputStateTensor = 3;
+constexpr int kBwCellStateTensor = 4;
+constexpr int kBwOutputTensor = 5;
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* scratch_tensor_index = new int;
+  context->AddTensors(context, 2, scratch_tensor_index);
+  return scratch_tensor_index;
+}
 
-constexpr int kBwScratchBufferTensor = 4;
-constexpr int kBwOutputStateTensor = 5;
-constexpr int kBwCellStateTensor = 6;
-constexpr int kBwOutputTensor = 7;
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<int*>(buffer);
+}
 
 // Check that input tensor dimensions matches with each other.
 TfLiteStatus CheckLstmTensorDimensions(
@@ -296,9 +304,11 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
 // Resize the output, state and scratch tensors based on the sizes of the input
 // tensors. Also check that the size of the input tensors match each other.
 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  int* scratch_tensor_index = reinterpret_cast<int*>(node->user_data);
+
   // Check we have all the inputs and outputs we need.
   TF_LITE_ENSURE_EQ(context, node->inputs->size, 35);
-  TF_LITE_ENSURE_EQ(context, node->outputs->size, 8);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 6);
 
   // Inferring batch size, number of outputs and sequence length and
   // number of cells from the input tensors.
@@ -330,12 +340,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TfLiteTensor* fw_output_state =
       GetOutput(context, node, kFwOutputStateTensor);
   TfLiteTensor* fw_cell_state = GetOutput(context, node, kFwCellStateTensor);
-  // TODO(ghodrat): Modify this as soon as we have a finalized method for
-  // scratch buffers.
-  TfLiteTensor* fw_scratch_buffer =
-      GetOutput(context, node, kFwScratchBufferTensor);
 
-  // Resize the output and output_state tensors.
+  // Resize the output, output_state and cell_state tensors.
   TfLiteIntArray* fw_output_size = TfLiteIntArrayCreate(3);
   fw_output_size->data[0] = max_time;
   fw_output_size->data[1] = n_batch;
@@ -349,13 +355,21 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, fw_output_state,
                                                    fw_output_state_size));
 
-  // Resize the scratch buffer tensor.
   TfLiteIntArray* fw_cell_size = TfLiteIntArrayCreate(2);
   fw_cell_size->data[0] = n_batch;
   fw_cell_size->data[1] = n_fw_cell;
   TF_LITE_ENSURE_OK(
       context, context->ResizeTensor(context, fw_cell_state, fw_cell_size));
 
+  // Create a scratch buffer tensor.
+  TfLiteIntArrayFree(node->temporaries);
+  node->temporaries = TfLiteIntArrayCreate(2);
+  node->temporaries->data[0] = *scratch_tensor_index;
+  TfLiteTensor* fw_scratch_buffer =
+      &context->tensors[node->temporaries->data[0]];
+  fw_scratch_buffer->type = input->type;
+  fw_scratch_buffer->allocation_type = kTfLiteArenaRw;
+
   // Mark state tensors as persistent tensors.
   fw_output_state->allocation_type = kTfLiteArenaRwPersistent;
   fw_cell_state->allocation_type = kTfLiteArenaRwPersistent;
@@ -392,17 +406,13 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   // Check that input tensor dimensions matches with each other.
   CheckInputTensorDimensions(context, node, n_input, n_bw_output, n_bw_cell);
 
-  // Get the pointer to output, state and scratch buffer tensors.
+  // Get the pointer to output, output_state and cell_state buffer tensors.
   TfLiteTensor* bw_output = GetOutput(context, node, kBwOutputTensor);
   TfLiteTensor* bw_output_state =
       GetOutput(context, node, kBwOutputStateTensor);
   TfLiteTensor* bw_cell_state = GetOutput(context, node, kBwCellStateTensor);
-  // TODO(ghodrat): Modify this as soon as we have a finalized method for
-  // scratch buffers.
-  TfLiteTensor* bw_scratch_buffer =
-      GetOutput(context, node, kBwScratchBufferTensor);
 
-  // Resize the output and output_state tensors.
+  // Resize the output, output_state and cell_state tensors.
   TfLiteIntArray* bw_output_size = TfLiteIntArrayCreate(3);
   bw_output_size->data[0] = max_time;
   bw_output_size->data[1] = n_batch;
@@ -416,13 +426,19 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, bw_output_state,
                                                    bw_output_state_size));
 
-  // Resize the scratch buffer tensor.
   TfLiteIntArray* bw_cell_size = TfLiteIntArrayCreate(2);
   bw_cell_size->data[0] = n_batch;
   bw_cell_size->data[1] = n_bw_cell;
   TF_LITE_ENSURE_OK(
       context, context->ResizeTensor(context, bw_cell_state, bw_cell_size));
 
+  // Create a scratch buffer tensor.
+  node->temporaries->data[1] = *(scratch_tensor_index) + 1;
+  TfLiteTensor* bw_scratch_buffer =
+      &context->tensors[node->temporaries->data[1]];
+  bw_scratch_buffer->type = input->type;
+  bw_scratch_buffer->allocation_type = kTfLiteArenaRw;
+
   // Mark state tensors as persistent tensors.
   bw_output_state->allocation_type = kTfLiteArenaRwPersistent;
   bw_cell_state->allocation_type = kTfLiteArenaRwPersistent;
@@ -553,7 +569,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
 
   // Index the scratch buffers pointers to the global scratch buffer.
   TfLiteTensor* fw_scratch_buffer =
-      GetOutput(context, node, kFwScratchBufferTensor);
+      &context->tensors[node->temporaries->data[0]];
   float* fw_input_gate_scratch = nullptr;
   float* fw_cell_scratch = nullptr;
   float* fw_forget_gate_scratch = nullptr;
@@ -624,7 +640,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
 
   // Index the scratch buffers pointers to the global scratch buffer.
   TfLiteTensor* bw_scratch_buffer =
-      GetOutput(context, node, kBwScratchBufferTensor);
+      &context->tensors[node->temporaries->data[1]];
   float* bw_input_gate_scratch = nullptr;
   float* bw_cell_scratch = nullptr;
   float* bw_forget_gate_scratch = nullptr;
@@ -691,9 +707,9 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
 }  // namespace bidirectional_sequence_lstm
 
 TfLiteRegistration* Register_BIDIRECTIONAL_SEQUENCE_LSTM() {
-  static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr,
-                                 bidirectional_sequence_lstm::Prepare,
-                                 bidirectional_sequence_lstm::Eval};
+  static TfLiteRegistration r = {
+      bidirectional_sequence_lstm::Init, bidirectional_sequence_lstm::Free,
+      bidirectional_sequence_lstm::Prepare, bidirectional_sequence_lstm::Eval};
   return &r;
 }
 
diff --git a/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm_test.cc b/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm_test.cc
index cca857bac0..a18e1bce34 100644
--- a/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm_test.cc
+++ b/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm_test.cc
@@ -102,9 +102,6 @@ class BidirectionalLSTMOpModel : public SingleOpModel {
       fw_projection_bias_ = AddNullInput();
     }
 
-    fw_scratch_buffer_ = AddOutput(TensorType_FLOAT32);
-    // TODO(ghodrat): Modify these states when we have a permanent solution for
-    // persistent buffer.
     fw_output_state_ = AddOutput(TensorType_FLOAT32);
     fw_cell_state_ = AddOutput(TensorType_FLOAT32);
     fw_output_ = AddOutput(TensorType_FLOAT32);
@@ -164,9 +161,6 @@ class BidirectionalLSTMOpModel : public SingleOpModel {
       bw_projection_bias_ = AddNullInput();
     }
 
-    bw_scratch_buffer_ = AddOutput(TensorType_FLOAT32);
-    // TODO(ghodrat): Modify these states when we have a permanent solution for
-    // persistent buffer.
     bw_output_state_ = AddOutput(TensorType_FLOAT32);
     bw_cell_state_ = AddOutput(TensorType_FLOAT32);
     bw_output_ = AddOutput(TensorType_FLOAT32);
@@ -349,12 +343,10 @@ class BidirectionalLSTMOpModel : public SingleOpModel {
   int fw_output_;
   int fw_output_state_;
   int fw_cell_state_;
-  int fw_scratch_buffer_;
 
   int bw_output_;
   int bw_output_state_;
   int bw_cell_state_;
-  int bw_scratch_buffer_;
 
   int n_batch_;
   int n_input_;
diff --git a/tensorflow/contrib/lite/kernels/lstm.cc b/tensorflow/contrib/lite/kernels/lstm.cc
index 8cf1165135..668226e674 100644
--- a/tensorflow/contrib/lite/kernels/lstm.cc
+++ b/tensorflow/contrib/lite/kernels/lstm.cc
@@ -66,10 +66,19 @@ constexpr int kProjectionWeightsTensor = 16;  // Optional
 constexpr int kProjectionBiasTensor = 17;  // Optional
 
 // Output tensors.
-constexpr int kScratchBufferTensor = 0;
-constexpr int kOutputStateTensor = 1;
-constexpr int kCellStateTensor = 2;
-constexpr int kOutputTensor = 3;
+constexpr int kOutputStateTensor = 0;
+constexpr int kCellStateTensor = 1;
+constexpr int kOutputTensor = 2;
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* scratch_tensor_index = new int;
+  context->AddTensors(context, 1, scratch_tensor_index);
+  return scratch_tensor_index;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<int*>(buffer);
+}
 
 // Check that input tensor dimensions matches with each other.
 TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
@@ -220,12 +229,15 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   return kTfLiteOk;
 }
 
-// Resize the output, state and scratch tensors based on the sizes of the input
-// tensors. Also check that the size of the input tensors match each other.
+// Resize the output, state tensors based on the sizes of the input tensors.
+// Allocate a temporary scratch tensor. Also check that the sizes of the input
+// tensors match each other.
 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  int* scratch_tensor_index = reinterpret_cast<int*>(node->user_data);
+
   // Check we have all the inputs and outputs we need.
   TF_LITE_ENSURE_EQ(context, node->inputs->size, 18);
-  TF_LITE_ENSURE_EQ(context, node->outputs->size, 4);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 3);
 
   // Inferring batch size, number of outputs and number of cells from the
   // input tensors.
@@ -250,15 +262,12 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   // Check that input tensor dimensions matches with each other.
   CheckInputTensorDimensions(context, node, n_input, n_output, n_cell);
 
-  // Get the pointer to output, state and scratch buffer tensors.
+  // Get the pointer to output, output_state and cell_state tensors.
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   TfLiteTensor* output_state = GetOutput(context, node, kOutputStateTensor);
   TfLiteTensor* cell_state = GetOutput(context, node, kCellStateTensor);
-  // TODO(ghodrat): Modify this as soon as we have a finalized method for
-  // scratch buffers.
-  TfLiteTensor* scratch_buffer = GetOutput(context, node, kScratchBufferTensor);
 
-  // Resize the output and output_state tensors.
+  // Resize the output, output_state and cell_state tensors.
   TfLiteIntArray* output_size = TfLiteIntArrayCreate(2);
   output_size->data[0] = n_batch;
   output_size->data[1] = n_output;
@@ -271,13 +280,20 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_OK(
       context, context->ResizeTensor(context, output_state, output_state_size));
 
-  // Resize the output, state and scratch buffer tensors.
   TfLiteIntArray* cell_size = TfLiteIntArrayCreate(2);
   cell_size->data[0] = n_batch;
   cell_size->data[1] = n_cell;
   TF_LITE_ENSURE_OK(context,
                     context->ResizeTensor(context, cell_state, cell_size));
 
+  // Create a scratch buffer tensor.
+  TfLiteIntArrayFree(node->temporaries);
+  node->temporaries = TfLiteIntArrayCreate(1);
+  node->temporaries->data[0] = *scratch_tensor_index;
+  TfLiteTensor* scratch_buffer = &context->tensors[node->temporaries->data[0]];
+  scratch_buffer->type = input->type;
+  scratch_buffer->allocation_type = kTfLiteArenaRw;
+
   // Mark state tensors as persistent tensors.
   output_state->allocation_type = kTfLiteArenaRwPersistent;
   cell_state->allocation_type = kTfLiteArenaRwPersistent;
@@ -362,7 +378,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   const bool use_peephole = (cell_to_output_weights != nullptr);
 
   // Index the scratch buffers pointers to the global scratch buffer.
-  TfLiteTensor* scratch_buffer = GetOutput(context, node, kScratchBufferTensor);
+  TfLiteTensor* scratch_buffer = &context->tensors[node->temporaries->data[0]];
+
   float* input_gate_scratch = nullptr;
   float* cell_scratch = nullptr;
   float* forget_gate_scratch = nullptr;
@@ -433,8 +450,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
 }  // namespace lstm
 
 TfLiteRegistration* Register_LSTM() {
-  static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr,
-                                 lstm::Prepare, lstm::Eval};
+  static TfLiteRegistration r = {lstm::Init, lstm::Free, lstm::Prepare,
+                                 lstm::Eval};
   return &r;
 }
 
diff --git a/tensorflow/contrib/lite/kernels/lstm_test.cc b/tensorflow/contrib/lite/kernels/lstm_test.cc
index c068286b0d..d81220d8d3 100644
--- a/tensorflow/contrib/lite/kernels/lstm_test.cc
+++ b/tensorflow/contrib/lite/kernels/lstm_test.cc
@@ -97,9 +97,6 @@ class LSTMOpModel : public SingleOpModel {
       projection_bias_ = AddNullInput();
     }
 
-    scratch_buffer_ = AddOutput(TensorType_FLOAT32);
-    // TODO(ghodrat): Modify these states when we have a permanent solution for
-    // persistent buffer.
     output_state_ = AddOutput(TensorType_FLOAT32);
     cell_state_ = AddOutput(TensorType_FLOAT32);
     output_ = AddOutput(TensorType_FLOAT32);
@@ -233,7 +230,6 @@ class LSTMOpModel : public SingleOpModel {
   int output_;
   int output_state_;
   int cell_state_;
-  int scratch_buffer_;
 
   int n_batch_;
   int n_input_;
diff --git a/tensorflow/contrib/lite/kernels/optional_tensor_test.cc b/tensorflow/contrib/lite/kernels/optional_tensor_test.cc
index cee3ec6197..bcad58406a 100644
--- a/tensorflow/contrib/lite/kernels/optional_tensor_test.cc
+++ b/tensorflow/contrib/lite/kernels/optional_tensor_test.cc
@@ -95,9 +95,6 @@ class LSTMOpModel : public SingleOpModel {
       projection_bias_ = AddNullInput();
     }
 
-    scratch_buffer_ = AddOutput(TensorType_FLOAT32);
-    // TODO(ghodrat): Modify these states when we have a permanent solution for
-    // persistent buffer.
     output_state_ = AddOutput(TensorType_FLOAT32);
     cell_state_ = AddOutput(TensorType_FLOAT32);
     output_ = AddOutput(TensorType_FLOAT32);
@@ -235,7 +232,6 @@ class LSTMOpModel : public SingleOpModel {
   int output_;
   int output_state_;
   int cell_state_;
-  int scratch_buffer_;
 
   int n_batch_;
   int n_input_;
diff --git a/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.cc b/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.cc
index 42941a97db..3c1256d3a6 100644
--- a/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.cc
+++ b/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.cc
@@ -66,10 +66,19 @@ constexpr int kProjectionWeightsTensor = 16;  // Optional
 constexpr int kProjectionBiasTensor = 17;  // Optional
 
 // Output tensors.
-constexpr int kScratchBufferTensor = 0;
-constexpr int kOutputStateTensor = 1;
-constexpr int kCellStateTensor = 2;
-constexpr int kOutputTensor = 3;
+constexpr int kOutputStateTensor = 0;
+constexpr int kCellStateTensor = 1;
+constexpr int kOutputTensor = 2;
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* scratch_tensor_index = new int;
+  context->AddTensors(context, 1, scratch_tensor_index);
+  return scratch_tensor_index;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<int*>(buffer);
+}
 
 // Check that input tensor dimensions matches with each other.
 TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
@@ -220,12 +229,15 @@ TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
   return kTfLiteOk;
 }
 
-// Resize the output, state and scratch tensors based on the sizes of the input
-// tensors. Also check that the size of the input tensors match each other.
+// Resize the output and  state tensors based on the sizes of the input tensors.
+// Allocate a temprory scratch tensor. Also check that the sizes of the input
+// tensors match each other.
 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  int* scratch_tensor_index = reinterpret_cast<int*>(node->user_data);
+
   // Check we have all the inputs and outputs we need.
   TF_LITE_ENSURE_EQ(context, node->inputs->size, 18);
-  TF_LITE_ENSURE_EQ(context, node->outputs->size, 4);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 3);
 
   // Inferring batch size, number of outputs and sequence length and
   // number of cells from the input tensors.
@@ -251,15 +263,12 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   // Check that input tensor dimensions matches with each other.
   CheckInputTensorDimensions(context, node, n_input, n_output, n_cell);
 
-  // Get the pointer to output, state and scratch buffer tensors.
+  // Get the pointer to output, output_state and cell_state buffer tensors.
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   TfLiteTensor* output_state = GetOutput(context, node, kOutputStateTensor);
   TfLiteTensor* cell_state = GetOutput(context, node, kCellStateTensor);
-  // TODO(ghodrat): Modify this as soon as we have a finalized method for
-  // scratch buffers.
-  TfLiteTensor* scratch_buffer = GetOutput(context, node, kScratchBufferTensor);
 
-  // Resize the output and output_state tensors.
+  // Resize the output, output_state and cell_state tensors.
   TfLiteIntArray* output_size = TfLiteIntArrayCreate(3);
   output_size->data[0] = max_time;
   output_size->data[1] = n_batch;
@@ -273,13 +282,20 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_OK(
       context, context->ResizeTensor(context, output_state, output_state_size));
 
-  // Resize the scratch buffer tensor.
   TfLiteIntArray* cell_size = TfLiteIntArrayCreate(2);
   cell_size->data[0] = n_batch;
   cell_size->data[1] = n_cell;
   TF_LITE_ENSURE_OK(context,
                     context->ResizeTensor(context, cell_state, cell_size));
 
+  // Create a scratch buffer tensor.
+  TfLiteIntArrayFree(node->temporaries);
+  node->temporaries = TfLiteIntArrayCreate(1);
+  node->temporaries->data[0] = *scratch_tensor_index;
+  TfLiteTensor* scratch_buffer = &context->tensors[node->temporaries->data[0]];
+  scratch_buffer->type = input->type;
+  scratch_buffer->allocation_type = kTfLiteArenaRw;
+
   // Mark state tensors as persistent tensors.
   output_state->allocation_type = kTfLiteArenaRwPersistent;
   cell_state->allocation_type = kTfLiteArenaRwPersistent;
@@ -365,7 +381,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   const bool use_peephole = (cell_to_output_weights != nullptr);
 
   // Index the scratch buffers pointers to the global scratch buffer.
-  TfLiteTensor* scratch_buffer = GetOutput(context, node, kScratchBufferTensor);
+  TfLiteTensor* scratch_buffer = &context->tensors[node->temporaries->data[0]];
   float* input_gate_scratch = nullptr;
   float* cell_scratch = nullptr;
   float* forget_gate_scratch = nullptr;
@@ -439,7 +455,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
 }  // namespace unidirectional_sequence_lstm
 
 TfLiteRegistration* Register_UNIDIRECTIONAL_SEQUENCE_LSTM() {
-  static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr,
+  static TfLiteRegistration r = {unidirectional_sequence_lstm::Init,
+                                 unidirectional_sequence_lstm::Free,
                                  unidirectional_sequence_lstm::Prepare,
                                  unidirectional_sequence_lstm::Eval};
   return &r;
diff --git a/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm_test.cc b/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm_test.cc
index 93b635ae57..5881ced7c7 100644
--- a/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm_test.cc
+++ b/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm_test.cc
@@ -100,9 +100,6 @@ class UnidirectionalLSTMOpModel : public SingleOpModel {
       projection_bias_ = AddNullInput();
     }
 
-    scratch_buffer_ = AddOutput(TensorType_FLOAT32);
-    // TODO(ghodrat): Modify these states when we have a permanent solution for
-    // persistent buffer.
     output_state_ = AddOutput(TensorType_FLOAT32);
     cell_state_ = AddOutput(TensorType_FLOAT32);
     output_ = AddOutput(TensorType_FLOAT32);
@@ -238,7 +235,6 @@ class UnidirectionalLSTMOpModel : public SingleOpModel {
   int output_;
   int output_state_;
   int cell_state_;
-  int scratch_buffer_;
 
   int n_batch_;
   int n_input_;
diff --git a/tensorflow/contrib/lite/models/speech_test.cc b/tensorflow/contrib/lite/models/speech_test.cc
index a354179a94..206de1962d 100644
--- a/tensorflow/contrib/lite/models/speech_test.cc
+++ b/tensorflow/contrib/lite/models/speech_test.cc
@@ -131,8 +131,8 @@ TEST_P(SpeechTest, SpeakerIdOkGoogleTest) {
   ASSERT_TRUE(ConvertCsvData(
       "speech_speakerid_model.tflite", "speech_speakerid_model_in.csv",
       "speech_speakerid_model_out.csv", /*input_tensor=*/"0",
-      /*output_tensor=*/"66",
-      /*persistent_tensors=*/"19,20,40,41,61,62",
+      /*output_tensor=*/"63",
+      /*persistent_tensors=*/"18,19,38,39,58,59",
       /*sequence_size=*/80, &os));
   testing::TfLiteDriver test_driver(/*use_nnapi=*/false);
   ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver, GetMaxInvocations()))
@@ -144,8 +144,8 @@ TEST_P(SpeechTest, AsrAmTest) {
   ASSERT_TRUE(
       ConvertCsvData("speech_asr_am_model.tflite", "speech_asr_am_model_in.csv",
                      "speech_asr_am_model_out.csv", /*input_tensor=*/"0",
-                     /*output_tensor=*/"109",
-                     /*persistent_tensors=*/"19,20,40,41,61,62,82,83,103,104",
+                     /*output_tensor=*/"104",
+                     /*persistent_tensors=*/"18,19,38,39,58,59,78,79,98,99",
                      /*sequence_size=*/320, &os));
   testing::TfLiteDriver test_driver(/*use_nnapi=*/false);
   ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver, GetMaxInvocations()))
@@ -170,8 +170,8 @@ TEST_P(SpeechTest, EndpointerTest) {
   ASSERT_TRUE(ConvertCsvData(
       "speech_endpointer_model.tflite", "speech_endpointer_model_in.csv",
       "speech_endpointer_model_out.csv", /*input_tensor=*/"0",
-      /*output_tensor=*/"58",
-      /*persistent_tensors=*/"28,29,49,50",
+      /*output_tensor=*/"56",
+      /*persistent_tensors=*/"27,28,47,48",
       /*sequence_size=*/320, &os));
   testing::TfLiteDriver test_driver(/*use_nnapi=*/false);
   ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver, GetMaxInvocations()))
@@ -183,8 +183,8 @@ TEST_P(SpeechTest, TtsTest) {
   ASSERT_TRUE(ConvertCsvData("speech_tts_model.tflite",
                              "speech_tts_model_in.csv",
                              "speech_tts_model_out.csv", /*input_tensor=*/"0",
-                             /*output_tensor=*/"74",
-                             /*persistent_tensors=*/"25,26,46,47,67,68,73",
+                             /*output_tensor=*/"71",
+                             /*persistent_tensors=*/"24,25,44,45,64,65,70",
                              /*sequence_size=*/334, &os));
   testing::TfLiteDriver test_driver(/*use_nnapi=*/false);
   ASSERT_TRUE(testing::ParseAndRunTests(&os, &test_driver, GetMaxInvocations()))
diff --git a/tensorflow/contrib/lite/models/testdata/speech_asr_lm_model.test_spec b/tensorflow/contrib/lite/models/testdata/speech_asr_lm_model.test_spec
index 5812de4b30..f7f518b75f 100644
--- a/tensorflow/contrib/lite/models/testdata/speech_asr_lm_model.test_spec
+++ b/tensorflow/contrib/lite/models/testdata/speech_asr_lm_model.test_spec
@@ -1,5 +1,5 @@
 load_model: "speech_asr_lm_model.tflite"
-init_state: "21,22,42,43,63,64"
+init_state: "20,21,40,41,60,61"
 invoke {
   id: 3
   input: "63982"
@@ -18,7 +18,7 @@ invoke {
   input: "63981"
   output: "-0.314846"
 }
-init_state: "21,22,42,43,63,64"
+init_state: "20,21,40,41,60,61"
 invoke {
   id: 6
   input: "63982"
@@ -31,7 +31,7 @@ invoke {
   input: "3082"
   output: "-3.63721"
 }
-init_state: "21,22,42,43,63,64"
+init_state: "20,21,40,41,60,61"
 invoke {
   id: 8
   input: "63982"
@@ -44,7 +44,7 @@ invoke {
   input: "18965"
   output: "-6.93985"
 }
-init_state: "21,22,42,43,63,64"
+init_state: "20,21,40,41,60,61"
 invoke {
   id: 13
   input: "63982"
@@ -63,7 +63,7 @@ invoke {
   input: "63981"
   output: "-3.82091"
 }
-init_state: "21,22,42,43,63,64"
+init_state: "20,21,40,41,60,61"
 invoke {
   id: 19
   input: "63982"
@@ -88,7 +88,7 @@ invoke {
   input: "63981"
   output: "-0.677399"
 }
-init_state: "21,22,42,43,63,64"
+init_state: "20,21,40,41,60,61"
 invoke {
   id: 26
   input: "63982"
@@ -113,7 +113,7 @@ invoke {
   input: "63981"
   output: "0.415889"
 }
-init_state: "21,22,42,43,63,64"
+init_state: "20,21,40,41,60,61"
 invoke {
   id: 30
   input: "63982"
@@ -131,7 +131,7 @@ invoke {
   input: "51923"
   output: "-14.1147"
 }
-init_state: "21,22,42,43,63,64"
+init_state: "20,21,40,41,60,61"
 invoke {
   id: 34
   input: "63982"
@@ -144,7 +144,7 @@ invoke {
   input: "16318"
   output: "-1.54815"
 }
-init_state: "21,22,42,43,63,64"
+init_state: "20,21,40,41,60,61"
 invoke {
   id: 36
   input: "63982"
@@ -157,7 +157,7 @@ invoke {
   input: "28303"
   output: "-14.0947"
 }
-init_state: "21,22,42,43,63,64"
+init_state: "20,21,40,41,60,61"
 invoke {
   id: 38
   input: "63982"
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm_merge_inputs.cc b/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm_merge_inputs.cc
index 45335fd78c..3f768bfee1 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm_merge_inputs.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm_merge_inputs.cc
@@ -146,16 +146,19 @@ bool MergeLstmCellInputs::Run(Model* model, std::size_t op_index) {
   lstm_cell_op->inputs[LstmCellOperator::PREV_ACTIV_INPUT] = prev_activ_input;
   lstm_cell_op->inputs[LstmCellOperator::PREV_STATE_INPUT] = prev_state_input;
 
-  // Reorder LstmCell's 4 outputs.
+  // Reorder LstmCell's 3 outputs.
   lstm_cell_op->outputs.resize(LstmCellOperator::NUM_OUTPUTS);
   lstm_cell_op->outputs[LstmCellOperator::ACTIV_OUTPUT] =
       src_op->outputs[kOutputTensor];
   lstm_cell_op->outputs[LstmCellOperator::STATE_OUTPUT] =
       src_op->outputs[kCellStateTensor];
-  lstm_cell_op->outputs[LstmCellOperator::CONCAT_TEMP] =
-      src_op->outputs[kScratchBufferTensor];
   lstm_cell_op->outputs[LstmCellOperator::ACTIV_TEMP] =
       src_op->outputs[kOutputStateTensor];
+  // Create a new temp array for the fourth output.
+  const string& concat_temp_array_name =
+      AvailableArrayName(*model, base_name + "concat_temp");
+  model->GetOrCreateArray(concat_temp_array_name);
+  lstm_cell_op->outputs[LstmCellOperator::CONCAT_TEMP] = concat_temp_array_name;
 
   // Add the op into model.
   model->operators.emplace(op_it, std::move(lstm_cell_op));
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm_split_inputs.cc b/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm_split_inputs.cc
index eca717680a..8e66323bd7 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm_split_inputs.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/identify_lstm_split_inputs.cc
@@ -138,10 +138,9 @@ bool SplitLstmCellInputs::Run(Model* model, std::size_t op_index) {
   CreateOptionalArray(model, &(lstm_cell_op->inputs[kProjectionBiasTensor]),
                       base_name + "proj_bias");
 
-  // Reorder LstmCell's outputs.
-  lstm_cell_op->outputs.resize(LstmCellOperator::NUM_OUTPUTS);
-  lstm_cell_op->outputs[kScratchBufferTensor] =
-      curr_op->outputs[LstmCellOperator::CONCAT_TEMP];
+  // Reorder and resize LstmCell's outputs.
+  lstm_cell_op->outputs.resize(
+      ExtendedLstmCellOutputs::kExtendedLstmOutputCount);
   lstm_cell_op->outputs[kOutputStateTensor] =
       curr_op->outputs[LstmCellOperator::ACTIV_TEMP];
   lstm_cell_op->outputs[kCellStateTensor] =
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/lstm_utils.h b/tensorflow/contrib/lite/toco/graph_transformations/lstm_utils.h
index 4a9974ed4e..1c32a78169 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/lstm_utils.h
+++ b/tensorflow/contrib/lite/toco/graph_transformations/lstm_utils.h
@@ -51,10 +51,10 @@ enum ExtendedLstmCellInputs {
 };
 
 enum ExtendedLstmCellOutputs {
-  kScratchBufferTensor = 0,
-  kOutputStateTensor = 1,
-  kCellStateTensor = 2,
-  kOutputTensor = 3
+  kOutputStateTensor = 0,
+  kCellStateTensor = 1,
+  kOutputTensor = 2,
+  kExtendedLstmOutputCount = 3
 };
 
 // Create optional array used for optional tensor in ExtendedLstmCell inputs.
-- 
GitLab


From 88821b0e41f59ae60d02a6880706aef8a1aba024 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 18:41:36 -0700
Subject: [PATCH 080/395] [XLA] Redesign: dump HloSnapshot at the point where
 it used to dump the SessionModule.

PiperOrigin-RevId: 194880385
---
 tensorflow/compiler/xla/service/executable.cc | 13 +++
 tensorflow/compiler/xla/service/executable.h  | 13 ++-
 tensorflow/compiler/xla/service/service.cc    | 91 +++++++++++++++++++
 3 files changed, 116 insertions(+), 1 deletion(-)

diff --git a/tensorflow/compiler/xla/service/executable.cc b/tensorflow/compiler/xla/service/executable.cc
index 021f09d310..8119478ce9 100644
--- a/tensorflow/compiler/xla/service/executable.cc
+++ b/tensorflow/compiler/xla/service/executable.cc
@@ -143,6 +143,19 @@ Status Executable::DumpSessionModule() {
                                      *session_module_);
 }
 
+Status Executable::DumpHloSnapshot() {
+  TF_RET_CHECK(dumping_snapshot());
+  TF_RET_CHECK(hlo_snapshot_->has_hlo() &&
+               hlo_snapshot_->hlo().has_hlo_module());
+  const string& directory_path =
+      module_config().debug_options().xla_dump_executions_to();
+  const auto& module = hlo_snapshot_->hlo().hlo_module();
+  string filename = tensorflow::strings::Printf(
+      "computation_%lld__%s__execution_%lld", module.id(),
+      module.entry_computation_name().c_str(), ++execution_count_);
+  return Executable::DumpToDirectory(directory_path, filename, *hlo_snapshot_);
+}
+
 /* static */ Status Executable::DumpToDirectory(
     const string& directory_path, string filename,
     const SessionModule& session_module) {
diff --git a/tensorflow/compiler/xla/service/executable.h b/tensorflow/compiler/xla/service/executable.h
index f7af1ca574..99762f4586 100644
--- a/tensorflow/compiler/xla/service/executable.h
+++ b/tensorflow/compiler/xla/service/executable.h
@@ -144,7 +144,7 @@ class Executable {
     return hlo_module_->config().entry_computation_layout().result_shape();
   }
 
-  // Dumping helpers.
+  // TODO(b/74197823): Delete the session module dumping helpers.
   void set_session_module(std::unique_ptr<xla::SessionModule> session_module) {
     session_module_ = std::move(session_module);
   }
@@ -152,6 +152,14 @@ class Executable {
   SessionModule* session_module() const { return session_module_.get(); }
   Status DumpSessionModule();
 
+  // Dumping helpers.
+  void set_hlo_snapshot(std::unique_ptr<xla::HloSnapshot> hlo_snapshot) {
+    hlo_snapshot_ = std::move(hlo_snapshot);
+  }
+  bool dumping_snapshot() const { return hlo_snapshot_ != nullptr; }
+  HloSnapshot* hlo_snapshot() const { return hlo_snapshot_.get(); }
+  Status DumpHloSnapshot();
+
   // Dump session_module to directory_path/filename.
   static Status DumpToDirectory(const string& directory_path, string filename,
                                 const SessionModule& session_module);
@@ -174,6 +182,9 @@ class Executable {
   // SessionModule this was compiled from. Null if not dumping executions.
   std::unique_ptr<SessionModule> session_module_;
 
+  // HloSnapshot this was compiled from. Null if not dumping executions.
+  std::unique_ptr<HloSnapshot> hlo_snapshot_;
+
   // Execution count, used to generate a unique filename for each dumped
   // execution.
   int64 execution_count_ = 0;
diff --git a/tensorflow/compiler/xla/service/service.cc b/tensorflow/compiler/xla/service/service.cc
index 849488f4f9..175ee96bbc 100644
--- a/tensorflow/compiler/xla/service/service.cc
+++ b/tensorflow/compiler/xla/service/service.cc
@@ -91,6 +91,34 @@ tensorflow::Status RecordResult(const ShapedBuffer& result,
   return tensorflow::Status::OK();
 }
 
+// Records the arguments used to invoke a computation in an HloSnapshot proto.
+tensorflow::Status RecordArguments(
+    const tensorflow::gtl::ArraySlice<const ShapedBuffer*> arguments,
+    se::StreamExecutor* executor, TransferManager* transfer_manager,
+    HloSnapshot* module) {
+  module->clear_arguments();
+  for (const ShapedBuffer* argument : arguments) {
+    TF_ASSIGN_OR_RETURN(
+        std::unique_ptr<Literal> literal,
+        transfer_manager->TransferLiteralFromDevice(executor, *argument));
+    *module->add_arguments() = literal->ToProto();
+  }
+  return tensorflow::Status::OK();
+}
+
+// Records the result of a computation in a HloSnapshot proto.
+tensorflow::Status RecordResult(const ShapedBuffer& result,
+                                se::StreamExecutor* executor,
+                                TransferManager* transfer_manager,
+                                HloSnapshot* module) {
+  module->clear_result();
+  TF_ASSIGN_OR_RETURN(
+      std::unique_ptr<Literal> literal,
+      transfer_manager->TransferLiteralFromDevice(executor, result));
+  *module->mutable_result() = literal->ToProto();
+  return tensorflow::Status::OK();
+}
+
 }  // namespace
 
 ServiceOptions& ServiceOptions::set_platform(se::Platform* platform) {
@@ -409,6 +437,28 @@ StatusOr<std::vector<std::unique_ptr<Executable>>> Service::BuildExecutables(
     DeviceMemoryAllocator* device_allocator) {
   VLOG(1) << Printf("BuildExecutable on service %p", this);
 
+  // Dump computation proto state if flag is set.
+  std::vector<std::unique_ptr<HloSnapshot>> hlo_snapshots;
+  for (int64 i = 0; i < module_protos.size(); ++i) {
+    const string& directory_path =
+        module_configs[i]->debug_options().xla_dump_computations_to();
+    const string& execution_directory_path =
+        module_configs[i]->debug_options().xla_dump_executions_to();
+    if (directory_path.empty() && execution_directory_path.empty()) {
+      continue;
+    }
+    auto hlo_snapshot = MakeUnique<HloSnapshot>();
+    *hlo_snapshot->mutable_hlo()->mutable_hlo_module() = *module_protos[i];
+    if (!directory_path.empty()) {
+      string filename =
+          Printf("computation_%lld__%s", module_protos[i]->id(),
+                 module_protos[i]->entry_computation_name().c_str());
+      TF_RETURN_IF_ERROR(
+          Executable::DumpToDirectory(directory_path, filename, *hlo_snapshot));
+      hlo_snapshots.push_back(std::move(hlo_snapshot));
+    }
+  }
+
   VLOG(1) << "Computations:";
   for (const HloModuleProto* proto : module_protos) {
     VLOG(1) << proto->name();
@@ -429,6 +479,12 @@ StatusOr<std::vector<std::unique_ptr<Executable>>> Service::BuildExecutables(
       backend->compiler()->Compile(std::move(modules), std::move(executors),
                                    device_allocator));
 
+  for (size_t i = 0; i < module_protos.size(); ++i) {
+    if (!module_configs[i]->debug_options().xla_dump_executions_to().empty()) {
+      executables[i]->set_hlo_snapshot(std::move(hlo_snapshots[i]));
+    }
+  }
+
   return std::move(executables);
 }
 
@@ -1132,6 +1188,22 @@ StatusOr<std::unique_ptr<Executable>> Service::BuildExecutable(
       "BuildExecutable on service %p with serialized module proto: %s", this,
       module_proto.name().c_str());
 
+  // Dump computation proto state if flag is set.
+  auto hlo_snapshot = MakeUnique<HloSnapshot>();
+  const string& directory_path =
+      module_config->debug_options().xla_dump_computations_to();
+  const string& execution_directory_path =
+      module_config->debug_options().xla_dump_executions_to();
+  if (!directory_path.empty() || !execution_directory_path.empty()) {
+    *hlo_snapshot->mutable_hlo()->mutable_hlo_module() = module_proto;
+    if (!directory_path.empty()) {
+      string filename = Printf("computation_%lld__%s", module_proto.id(),
+                               module_proto.entry_computation_name().c_str());
+      TF_RETURN_IF_ERROR(
+          Executable::DumpToDirectory(directory_path, filename, *hlo_snapshot));
+    }
+  }
+
   TF_ASSIGN_OR_RETURN(std::unique_ptr<HloModule> module,
                       HloModule::CreateFromProto(module_proto, *module_config));
 
@@ -1182,12 +1254,31 @@ tensorflow::Status Service::ExecuteGraph(const ExecuteGraphRequest* arg,
                       execute_backend_->default_stream_executor(),
                       /*device_allocator=*/nullptr));
 
+  if (executable->dumping_snapshot()) {
+    executable->hlo_snapshot()->set_execution_platform(
+        execute_backend_->platform()->Name());
+    TF_RETURN_IF_ERROR(RecordArguments(
+        replicated_arguments.front(),
+        execute_backend_->default_stream_executor(),
+        execute_backend_->transfer_manager(), executable->hlo_snapshot()));
+  }
+
   TF_ASSIGN_OR_RETURN(
       *result->mutable_output(),
       ExecuteAndRegisterResult(
           executable.get(), replicated_arguments, execute_backend_.get(),
           "result of " + arg->computation().name(), result->mutable_profile()));
 
+  if (executable->dumping_snapshot()) {
+    TF_ASSIGN_OR_RETURN(
+        const ShapedBuffer* result_buffer,
+        allocation_tracker_.ResolveForReplica(result->output(), 0));
+    TF_RETURN_IF_ERROR(RecordResult(
+        *result_buffer, execute_backend_->default_stream_executor(),
+        execute_backend_->transfer_manager(), executable->hlo_snapshot()));
+    TF_RETURN_IF_ERROR(executable->DumpHloSnapshot());
+  }
+
   VLOG(1) << "successfully completed 'execute-graph' request";
   return tensorflow::Status::OK();
 }
-- 
GitLab


From 79ccb99e9396a7b480615c9ee4b924e851f67163 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 18:51:48 -0700
Subject: [PATCH 081/395] Move LinearOperatorKronecker and
 LinearOperatorBlockDiag to core.

PiperOrigin-RevId: 194881237
---
 tensorflow/contrib/linalg/BUILD               |  44 ------
 tensorflow/contrib/linalg/__init__.py         |   4 +-
 tensorflow/python/kernel_tests/linalg/BUILD   |  44 ++++++
 .../linear_operator_block_diag_test.py        |   2 +-
 .../linalg}/linear_operator_kronecker_test.py |   2 +-
 tensorflow/python/ops/linalg/linalg.py        |   2 +
 .../ops/linalg}/linear_operator_block_diag.py |   6 +
 .../ops/linalg}/linear_operator_kronecker.py  |   6 +
 ...ar-operator-block-diag.__metaclass__.pbtxt |  14 ++
 ...w.linalg.-linear-operator-block-diag.pbtxt | 134 ++++++++++++++++++
 ...ear-operator-kronecker.__metaclass__.pbtxt |  14 ++
 ...ow.linalg.-linear-operator-kronecker.pbtxt | 134 ++++++++++++++++++
 .../tools/api/golden/tensorflow.linalg.pbtxt  |   8 ++
 13 files changed, 366 insertions(+), 48 deletions(-)
 rename tensorflow/{contrib/linalg/python/kernel_tests => python/kernel_tests/linalg}/linear_operator_block_diag_test.py (98%)
 rename tensorflow/{contrib/linalg/python/kernel_tests => python/kernel_tests/linalg}/linear_operator_kronecker_test.py (98%)
 rename tensorflow/{contrib/linalg/python/ops => python/ops/linalg}/linear_operator_block_diag.py (98%)
 rename tensorflow/{contrib/linalg/python/ops => python/ops/linalg}/linear_operator_kronecker.py (99%)
 create mode 100644 tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-block-diag.__metaclass__.pbtxt
 create mode 100644 tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-block-diag.pbtxt
 create mode 100644 tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-kronecker.__metaclass__.pbtxt
 create mode 100644 tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-kronecker.pbtxt

diff --git a/tensorflow/contrib/linalg/BUILD b/tensorflow/contrib/linalg/BUILD
index 2e92ad6eb3..78b7970069 100644
--- a/tensorflow/contrib/linalg/BUILD
+++ b/tensorflow/contrib/linalg/BUILD
@@ -42,47 +42,3 @@ cuda_py_test(
         "//tensorflow/python:platform_test",
     ],
 )
-
-cuda_py_test(
-    name = "linear_operator_block_diag_test",
-    size = "medium",
-    srcs = ["python/kernel_tests/linear_operator_block_diag_test.py"],
-    additional_deps = [
-        ":linalg_py",
-        "//third_party/py/numpy",
-        "//tensorflow/python:array_ops",
-        "//tensorflow/python:client_testlib",
-        "//tensorflow/python:framework",
-        "//tensorflow/python:framework_for_generated_wrappers",
-        "//tensorflow/python:framework_test_lib",
-        "//tensorflow/python:math_ops",
-        "//tensorflow/python:platform_test",
-    ],
-    shard_count = 5,
-    tags = [
-        "noasan",
-        "optonly",
-    ],
-)
-
-cuda_py_test(
-    name = "linear_operator_kronecker_test",
-    size = "medium",
-    srcs = ["python/kernel_tests/linear_operator_kronecker_test.py"],
-    additional_deps = [
-        ":linalg_py",
-        "//third_party/py/numpy",
-        "//tensorflow/python:array_ops",
-        "//tensorflow/python:client_testlib",
-        "//tensorflow/python:framework",
-        "//tensorflow/python:framework_for_generated_wrappers",
-        "//tensorflow/python:framework_test_lib",
-        "//tensorflow/python:math_ops",
-        "//tensorflow/python:platform_test",
-    ],
-    shard_count = 8,
-    tags = [
-        "noasan",
-        "optonly",
-    ],
-)
diff --git a/tensorflow/contrib/linalg/__init__.py b/tensorflow/contrib/linalg/__init__.py
index 554854da84..a262a099cf 100644
--- a/tensorflow/contrib/linalg/__init__.py
+++ b/tensorflow/contrib/linalg/__init__.py
@@ -39,14 +39,14 @@ from __future__ import print_function
 # pylint: disable=unused-import,wildcard-import,line-too-long,g-importing-member
 
 from tensorflow.contrib.linalg.python.ops.linear_operator_addition import *
-from tensorflow.contrib.linalg.python.ops.linear_operator_block_diag import *
-from tensorflow.contrib.linalg.python.ops.linear_operator_kronecker import *
 from tensorflow.python.ops.linalg.linear_operator import *
+from tensorflow.python.ops.linalg.linear_operator_block_diag import *
 from tensorflow.python.ops.linalg.linear_operator_circulant import *
 from tensorflow.python.ops.linalg.linear_operator_composition import *
 from tensorflow.python.ops.linalg.linear_operator_diag import *
 from tensorflow.python.ops.linalg.linear_operator_full_matrix import *
 from tensorflow.python.ops.linalg.linear_operator_identity import *
+from tensorflow.python.ops.linalg.linear_operator_kronecker import *
 from tensorflow.python.ops.linalg.linear_operator_low_rank_update import *
 from tensorflow.python.ops.linalg.linear_operator_lower_triangular import *
 
diff --git a/tensorflow/python/kernel_tests/linalg/BUILD b/tensorflow/python/kernel_tests/linalg/BUILD
index faeccc8fba..6573cb9a1a 100644
--- a/tensorflow/python/kernel_tests/linalg/BUILD
+++ b/tensorflow/python/kernel_tests/linalg/BUILD
@@ -24,6 +24,28 @@ cuda_py_test(
     ],
 )
 
+cuda_py_test(
+    name = "linear_operator_block_diag_test",
+    size = "medium",
+    srcs = ["linear_operator_block_diag_test.py"],
+    additional_deps = [
+        "//tensorflow/python/ops/linalg",
+        "//third_party/py/numpy",
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:framework",
+        "//tensorflow/python:framework_for_generated_wrappers",
+        "//tensorflow/python:framework_test_lib",
+        "//tensorflow/python:math_ops",
+        "//tensorflow/python:platform_test",
+    ],
+    shard_count = 6,
+    tags = [
+        "noasan",
+        "optonly",
+    ],
+)
+
 cuda_py_test(
     name = "linear_operator_composition_test",
     size = "medium",
@@ -114,6 +136,28 @@ cuda_py_test(
     shard_count = 5,
 )
 
+cuda_py_test(
+    name = "linear_operator_kronecker_test",
+    size = "medium",
+    srcs = ["linear_operator_kronecker_test.py"],
+    additional_deps = [
+        "//tensorflow/python/ops/linalg",
+        "//third_party/py/numpy",
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:framework",
+        "//tensorflow/python:framework_for_generated_wrappers",
+        "//tensorflow/python:framework_test_lib",
+        "//tensorflow/python:math_ops",
+        "//tensorflow/python:platform_test",
+    ],
+    shard_count = 8,
+    tags = [
+        "noasan",
+        "optonly",
+    ],
+)
+
 cuda_py_test(
     name = "linear_operator_lower_triangular_test",
     size = "medium",
diff --git a/tensorflow/contrib/linalg/python/kernel_tests/linear_operator_block_diag_test.py b/tensorflow/python/kernel_tests/linalg/linear_operator_block_diag_test.py
similarity index 98%
rename from tensorflow/contrib/linalg/python/kernel_tests/linear_operator_block_diag_test.py
rename to tensorflow/python/kernel_tests/linalg/linear_operator_block_diag_test.py
index e7407ede11..2b80f01b73 100644
--- a/tensorflow/contrib/linalg/python/kernel_tests/linear_operator_block_diag_test.py
+++ b/tensorflow/python/kernel_tests/linalg/linear_operator_block_diag_test.py
@@ -19,11 +19,11 @@ from __future__ import print_function
 
 import numpy as np
 
-from tensorflow.contrib.linalg.python.ops import linear_operator_block_diag as block_diag
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import random_seed
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops.linalg import linalg as linalg_lib
+from tensorflow.python.ops.linalg import linear_operator_block_diag as block_diag
 from tensorflow.python.ops.linalg import linear_operator_test_util
 from tensorflow.python.ops.linalg import linear_operator_util
 from tensorflow.python.platform import test
diff --git a/tensorflow/contrib/linalg/python/kernel_tests/linear_operator_kronecker_test.py b/tensorflow/python/kernel_tests/linalg/linear_operator_kronecker_test.py
similarity index 98%
rename from tensorflow/contrib/linalg/python/kernel_tests/linear_operator_kronecker_test.py
rename to tensorflow/python/kernel_tests/linalg/linear_operator_kronecker_test.py
index 6574da22a1..cce1ecd45e 100644
--- a/tensorflow/contrib/linalg/python/kernel_tests/linear_operator_kronecker_test.py
+++ b/tensorflow/python/kernel_tests/linalg/linear_operator_kronecker_test.py
@@ -19,12 +19,12 @@ from __future__ import print_function
 
 import numpy as np
 
-from tensorflow.contrib.linalg.python.ops import linear_operator_kronecker as kronecker
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import random_seed
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops.linalg import linalg as linalg_lib
+from tensorflow.python.ops.linalg import linear_operator_kronecker as kronecker
 from tensorflow.python.ops.linalg import linear_operator_test_util
 from tensorflow.python.ops.linalg import linear_operator_util
 from tensorflow.python.platform import test
diff --git a/tensorflow/python/ops/linalg/linalg.py b/tensorflow/python/ops/linalg/linalg.py
index d73c21cdc0..a7ba0bbe9c 100644
--- a/tensorflow/python/ops/linalg/linalg.py
+++ b/tensorflow/python/ops/linalg/linalg.py
@@ -22,11 +22,13 @@ from __future__ import print_function
 # pylint: disable=wildcard-import,unused-import
 from tensorflow.python.ops.linalg.linalg_impl import *
 from tensorflow.python.ops.linalg.linear_operator import *
+from tensorflow.python.ops.linalg.linear_operator_block_diag import *
 from tensorflow.python.ops.linalg.linear_operator_circulant import *
 from tensorflow.python.ops.linalg.linear_operator_composition import *
 from tensorflow.python.ops.linalg.linear_operator_diag import *
 from tensorflow.python.ops.linalg.linear_operator_full_matrix import *
 from tensorflow.python.ops.linalg.linear_operator_identity import *
+from tensorflow.python.ops.linalg.linear_operator_kronecker import *
 from tensorflow.python.ops.linalg.linear_operator_low_rank_update import *
 from tensorflow.python.ops.linalg.linear_operator_lower_triangular import *
 # pylint: enable=wildcard-import
diff --git a/tensorflow/contrib/linalg/python/ops/linear_operator_block_diag.py b/tensorflow/python/ops/linalg/linear_operator_block_diag.py
similarity index 98%
rename from tensorflow/contrib/linalg/python/ops/linear_operator_block_diag.py
rename to tensorflow/python/ops/linalg/linear_operator_block_diag.py
index 9d3af66c92..438c3496bd 100644
--- a/tensorflow/contrib/linalg/python/ops/linear_operator_block_diag.py
+++ b/tensorflow/python/ops/linalg/linear_operator_block_diag.py
@@ -27,8 +27,14 @@ from tensorflow.python.ops import check_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops.linalg import linear_operator
 from tensorflow.python.ops.linalg import linear_operator_util
+from tensorflow.python.util.tf_export import tf_export
 
+__all__ = [
+    "LinearOperatorBlockDiag",
+]
 
+
+@tf_export("linalg.LinearOperatorBlockDiag")
 class LinearOperatorBlockDiag(linear_operator.LinearOperator):
   """Combines one or more `LinearOperators` in to a Block Diagonal matrix.
 
diff --git a/tensorflow/contrib/linalg/python/ops/linear_operator_kronecker.py b/tensorflow/python/ops/linalg/linear_operator_kronecker.py
similarity index 99%
rename from tensorflow/contrib/linalg/python/ops/linear_operator_kronecker.py
rename to tensorflow/python/ops/linalg/linear_operator_kronecker.py
index 79080d194f..da959f9a1c 100644
--- a/tensorflow/contrib/linalg/python/ops/linear_operator_kronecker.py
+++ b/tensorflow/python/ops/linalg/linear_operator_kronecker.py
@@ -28,6 +28,11 @@ from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops.linalg import linalg_impl as linalg
 from tensorflow.python.ops.linalg import linear_operator
+from tensorflow.python.util.tf_export import tf_export
+
+__all__ = [
+    "LinearOperatorKronecker",
+]
 
 
 def _vec(x):
@@ -59,6 +64,7 @@ def _rotate_last_dim(x, rotate_right=False):
   return array_ops.transpose(x, transpose_perm)
 
 
+@tf_export("linalg.LinearOperatorKronecker")
 class LinearOperatorKronecker(linear_operator.LinearOperator):
   """Kronecker product between two `LinearOperators`.
 
diff --git a/tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-block-diag.__metaclass__.pbtxt b/tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-block-diag.__metaclass__.pbtxt
new file mode 100644
index 0000000000..b6dee63176
--- /dev/null
+++ b/tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-block-diag.__metaclass__.pbtxt
@@ -0,0 +1,14 @@
+path: "tensorflow.linalg.LinearOperatorBlockDiag.__metaclass__"
+tf_class {
+  is_instance: "<class \'abc.ABCMeta\'>"
+  member_method {
+    name: "__init__"
+  }
+  member_method {
+    name: "mro"
+  }
+  member_method {
+    name: "register"
+    argspec: "args=[\'cls\', \'subclass\'], varargs=None, keywords=None, defaults=None"
+  }
+}
diff --git a/tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-block-diag.pbtxt b/tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-block-diag.pbtxt
new file mode 100644
index 0000000000..973705dae2
--- /dev/null
+++ b/tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-block-diag.pbtxt
@@ -0,0 +1,134 @@
+path: "tensorflow.linalg.LinearOperatorBlockDiag"
+tf_class {
+  is_instance: "<class \'tensorflow.python.ops.linalg.linear_operator_block_diag.LinearOperatorBlockDiag\'>"
+  is_instance: "<class \'tensorflow.python.ops.linalg.linear_operator.LinearOperator\'>"
+  is_instance: "<type \'object\'>"
+  member {
+    name: "batch_shape"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "domain_dimension"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "dtype"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "graph_parents"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "is_non_singular"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "is_positive_definite"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "is_self_adjoint"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "is_square"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "name"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "operators"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "range_dimension"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "shape"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "tensor_rank"
+    mtype: "<type \'property\'>"
+  }
+  member_method {
+    name: "__init__"
+    argspec: "args=[\'self\', \'operators\', \'is_non_singular\', \'is_self_adjoint\', \'is_positive_definite\', \'is_square\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\', \'None\'], "
+  }
+  member_method {
+    name: "add_to_tensor"
+    argspec: "args=[\'self\', \'x\', \'name\'], varargs=None, keywords=None, defaults=[\'add_to_tensor\'], "
+  }
+  member_method {
+    name: "assert_non_singular"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'assert_non_singular\'], "
+  }
+  member_method {
+    name: "assert_positive_definite"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'assert_positive_definite\'], "
+  }
+  member_method {
+    name: "assert_self_adjoint"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'assert_self_adjoint\'], "
+  }
+  member_method {
+    name: "batch_shape_tensor"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'batch_shape_tensor\'], "
+  }
+  member_method {
+    name: "determinant"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'det\'], "
+  }
+  member_method {
+    name: "diag_part"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'diag_part\'], "
+  }
+  member_method {
+    name: "domain_dimension_tensor"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'domain_dimension_tensor\'], "
+  }
+  member_method {
+    name: "log_abs_determinant"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'log_abs_det\'], "
+  }
+  member_method {
+    name: "matmul"
+    argspec: "args=[\'self\', \'x\', \'adjoint\', \'adjoint_arg\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'matmul\'], "
+  }
+  member_method {
+    name: "matvec"
+    argspec: "args=[\'self\', \'x\', \'adjoint\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'matvec\'], "
+  }
+  member_method {
+    name: "range_dimension_tensor"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'range_dimension_tensor\'], "
+  }
+  member_method {
+    name: "shape_tensor"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'shape_tensor\'], "
+  }
+  member_method {
+    name: "solve"
+    argspec: "args=[\'self\', \'rhs\', \'adjoint\', \'adjoint_arg\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'solve\'], "
+  }
+  member_method {
+    name: "solvevec"
+    argspec: "args=[\'self\', \'rhs\', \'adjoint\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'solve\'], "
+  }
+  member_method {
+    name: "tensor_rank_tensor"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'tensor_rank_tensor\'], "
+  }
+  member_method {
+    name: "to_dense"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'to_dense\'], "
+  }
+  member_method {
+    name: "trace"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'trace\'], "
+  }
+}
diff --git a/tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-kronecker.__metaclass__.pbtxt b/tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-kronecker.__metaclass__.pbtxt
new file mode 100644
index 0000000000..5c6784dd02
--- /dev/null
+++ b/tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-kronecker.__metaclass__.pbtxt
@@ -0,0 +1,14 @@
+path: "tensorflow.linalg.LinearOperatorKronecker.__metaclass__"
+tf_class {
+  is_instance: "<class \'abc.ABCMeta\'>"
+  member_method {
+    name: "__init__"
+  }
+  member_method {
+    name: "mro"
+  }
+  member_method {
+    name: "register"
+    argspec: "args=[\'cls\', \'subclass\'], varargs=None, keywords=None, defaults=None"
+  }
+}
diff --git a/tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-kronecker.pbtxt b/tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-kronecker.pbtxt
new file mode 100644
index 0000000000..c11d390829
--- /dev/null
+++ b/tensorflow/tools/api/golden/tensorflow.linalg.-linear-operator-kronecker.pbtxt
@@ -0,0 +1,134 @@
+path: "tensorflow.linalg.LinearOperatorKronecker"
+tf_class {
+  is_instance: "<class \'tensorflow.python.ops.linalg.linear_operator_kronecker.LinearOperatorKronecker\'>"
+  is_instance: "<class \'tensorflow.python.ops.linalg.linear_operator.LinearOperator\'>"
+  is_instance: "<type \'object\'>"
+  member {
+    name: "batch_shape"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "domain_dimension"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "dtype"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "graph_parents"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "is_non_singular"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "is_positive_definite"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "is_self_adjoint"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "is_square"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "name"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "operators"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "range_dimension"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "shape"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "tensor_rank"
+    mtype: "<type \'property\'>"
+  }
+  member_method {
+    name: "__init__"
+    argspec: "args=[\'self\', \'operators\', \'is_non_singular\', \'is_self_adjoint\', \'is_positive_definite\', \'is_square\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\'], "
+  }
+  member_method {
+    name: "add_to_tensor"
+    argspec: "args=[\'self\', \'x\', \'name\'], varargs=None, keywords=None, defaults=[\'add_to_tensor\'], "
+  }
+  member_method {
+    name: "assert_non_singular"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'assert_non_singular\'], "
+  }
+  member_method {
+    name: "assert_positive_definite"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'assert_positive_definite\'], "
+  }
+  member_method {
+    name: "assert_self_adjoint"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'assert_self_adjoint\'], "
+  }
+  member_method {
+    name: "batch_shape_tensor"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'batch_shape_tensor\'], "
+  }
+  member_method {
+    name: "determinant"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'det\'], "
+  }
+  member_method {
+    name: "diag_part"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'diag_part\'], "
+  }
+  member_method {
+    name: "domain_dimension_tensor"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'domain_dimension_tensor\'], "
+  }
+  member_method {
+    name: "log_abs_determinant"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'log_abs_det\'], "
+  }
+  member_method {
+    name: "matmul"
+    argspec: "args=[\'self\', \'x\', \'adjoint\', \'adjoint_arg\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'matmul\'], "
+  }
+  member_method {
+    name: "matvec"
+    argspec: "args=[\'self\', \'x\', \'adjoint\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'matvec\'], "
+  }
+  member_method {
+    name: "range_dimension_tensor"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'range_dimension_tensor\'], "
+  }
+  member_method {
+    name: "shape_tensor"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'shape_tensor\'], "
+  }
+  member_method {
+    name: "solve"
+    argspec: "args=[\'self\', \'rhs\', \'adjoint\', \'adjoint_arg\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'False\', \'solve\'], "
+  }
+  member_method {
+    name: "solvevec"
+    argspec: "args=[\'self\', \'rhs\', \'adjoint\', \'name\'], varargs=None, keywords=None, defaults=[\'False\', \'solve\'], "
+  }
+  member_method {
+    name: "tensor_rank_tensor"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'tensor_rank_tensor\'], "
+  }
+  member_method {
+    name: "to_dense"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'to_dense\'], "
+  }
+  member_method {
+    name: "trace"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'trace\'], "
+  }
+}
diff --git a/tensorflow/tools/api/golden/tensorflow.linalg.pbtxt b/tensorflow/tools/api/golden/tensorflow.linalg.pbtxt
index 7a5c533872..00b9238543 100644
--- a/tensorflow/tools/api/golden/tensorflow.linalg.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.linalg.pbtxt
@@ -4,6 +4,10 @@ tf_module {
     name: "LinearOperator"
     mtype: "<class \'abc.ABCMeta\'>"
   }
+  member {
+    name: "LinearOperatorBlockDiag"
+    mtype: "<class \'abc.ABCMeta\'>"
+  }
   member {
     name: "LinearOperatorCirculant"
     mtype: "<class \'abc.ABCMeta\'>"
@@ -32,6 +36,10 @@ tf_module {
     name: "LinearOperatorIdentity"
     mtype: "<class \'abc.ABCMeta\'>"
   }
+  member {
+    name: "LinearOperatorKronecker"
+    mtype: "<class \'abc.ABCMeta\'>"
+  }
   member {
     name: "LinearOperatorLowRankUpdate"
     mtype: "<class \'abc.ABCMeta\'>"
-- 
GitLab


From 37191e98117c959fe5599df8b6f0d49b005b5782 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 30 Apr 2018 19:18:40 -0700
Subject: [PATCH 082/395] Update ops-related pbtxt files.

PiperOrigin-RevId: 194883351
---
 .../core/ops/compat/ops_history.v1.pbtxt      | 76 +++++++++++++++++++
 tensorflow/core/ops/ops.pbtxt                 | 76 +++++++++++++++++++
 2 files changed, 152 insertions(+)

diff --git a/tensorflow/core/ops/compat/ops_history.v1.pbtxt b/tensorflow/core/ops/compat/ops_history.v1.pbtxt
index 71ba5f016a..cb466ef817 100644
--- a/tensorflow/core/ops/compat/ops_history.v1.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history.v1.pbtxt
@@ -24166,6 +24166,82 @@ op {
     }
   }
 }
+op {
+  name: "GroupByReducerDataset"
+  input_arg {
+    name: "input_dataset"
+    type: DT_VARIANT
+  }
+  input_arg {
+    name: "key_func_other_arguments"
+    type_list_attr: "Tkey_func_other_arguments"
+  }
+  input_arg {
+    name: "init_func_other_arguments"
+    type_list_attr: "Tinit_func_other_arguments"
+  }
+  input_arg {
+    name: "reduce_func_other_arguments"
+    type_list_attr: "Treduce_func_other_arguments"
+  }
+  input_arg {
+    name: "finalize_func_other_arguments"
+    type_list_attr: "Tfinalize_func_other_arguments"
+  }
+  output_arg {
+    name: "handle"
+    type: DT_VARIANT
+  }
+  attr {
+    name: "key_func"
+    type: "func"
+  }
+  attr {
+    name: "init_func"
+    type: "func"
+  }
+  attr {
+    name: "reduce_func"
+    type: "func"
+  }
+  attr {
+    name: "finalize_func"
+    type: "func"
+  }
+  attr {
+    name: "Tkey_func_other_arguments"
+    type: "list(type)"
+    has_minimum: true
+  }
+  attr {
+    name: "Tinit_func_other_arguments"
+    type: "list(type)"
+    has_minimum: true
+  }
+  attr {
+    name: "Treduce_func_other_arguments"
+    type: "list(type)"
+    has_minimum: true
+  }
+  attr {
+    name: "Tfinalize_func_other_arguments"
+    type: "list(type)"
+    has_minimum: true
+  }
+  attr {
+    name: "output_types"
+    type: "list(type)"
+    has_minimum: true
+    minimum: 1
+  }
+  attr {
+    name: "output_shapes"
+    type: "list(shape)"
+    has_minimum: true
+    minimum: 1
+  }
+  is_stateful: true
+}
 op {
   name: "GroupByWindowDataset"
   input_arg {
diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 90368fe614..207dd1c3d7 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -11536,6 +11536,82 @@ op {
     }
   }
 }
+op {
+  name: "GroupByReducerDataset"
+  input_arg {
+    name: "input_dataset"
+    type: DT_VARIANT
+  }
+  input_arg {
+    name: "key_func_other_arguments"
+    type_list_attr: "Tkey_func_other_arguments"
+  }
+  input_arg {
+    name: "init_func_other_arguments"
+    type_list_attr: "Tinit_func_other_arguments"
+  }
+  input_arg {
+    name: "reduce_func_other_arguments"
+    type_list_attr: "Treduce_func_other_arguments"
+  }
+  input_arg {
+    name: "finalize_func_other_arguments"
+    type_list_attr: "Tfinalize_func_other_arguments"
+  }
+  output_arg {
+    name: "handle"
+    type: DT_VARIANT
+  }
+  attr {
+    name: "key_func"
+    type: "func"
+  }
+  attr {
+    name: "init_func"
+    type: "func"
+  }
+  attr {
+    name: "reduce_func"
+    type: "func"
+  }
+  attr {
+    name: "finalize_func"
+    type: "func"
+  }
+  attr {
+    name: "Tkey_func_other_arguments"
+    type: "list(type)"
+    has_minimum: true
+  }
+  attr {
+    name: "Tinit_func_other_arguments"
+    type: "list(type)"
+    has_minimum: true
+  }
+  attr {
+    name: "Treduce_func_other_arguments"
+    type: "list(type)"
+    has_minimum: true
+  }
+  attr {
+    name: "Tfinalize_func_other_arguments"
+    type: "list(type)"
+    has_minimum: true
+  }
+  attr {
+    name: "output_types"
+    type: "list(type)"
+    has_minimum: true
+    minimum: 1
+  }
+  attr {
+    name: "output_shapes"
+    type: "list(shape)"
+    has_minimum: true
+    minimum: 1
+  }
+  is_stateful: true
+}
 op {
   name: "GroupByWindowDataset"
   input_arg {
-- 
GitLab


From 7525a48ebf6f8175cd2845f0fa7ae8ae2a10e1c1 Mon Sep 17 00:00:00 2001
From: Sami Kama <skama@nvidia.com>
Date: Mon, 30 Apr 2018 20:22:51 -0700
Subject: [PATCH 083/395] Fixes for review comments

---
 .../contrib/tensorrt/convert/convert_graph.cc |   4 +-
 .../contrib/tensorrt/convert/convert_graph.h  |   2 -
 .../contrib/tensorrt/convert/convert_nodes.h  |   6 +-
 .../tensorrt/convert/trt_optimization_pass.cc | 113 ++++----
 .../tensorrt/convert/trt_optimization_pass.h  |  14 +-
 .../contrib/tensorrt/kernels/trt_engine_op.cc |  36 +--
 .../contrib/tensorrt/kernels/trt_engine_op.h  |   1 +
 .../tensorrt/resources/trt_allocator.h        |   7 +-
 .../contrib/tensorrt/segment/segment.cc       | 252 ++++++++++++------
 tensorflow/contrib/tensorrt/segment/segment.h |  81 +-----
 10 files changed, 265 insertions(+), 251 deletions(-)

diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.cc b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
index 632908f078..c1979afcf8 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_graph.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
@@ -349,7 +349,6 @@ tensorflow::Status ConvertGraphDefToTensorRT(
 
   // Layout optimization
   item.graph = graph_def;
-  tensorflow::grappler::LayoutOptimizer optimizer;
   tensorflow::grappler::Cluster* cluster;
 
   // virtual cluster
@@ -417,6 +416,7 @@ tensorflow::Status ConvertAfterShapes(
   for (auto s : segments) {
     total_num_nodes_in_segments += s.first.size();
   }
+  // Cluster may not be available
   std::map<string, tensorflow::Device*> name_to_device_map;
   if (cluster) {
     for (const auto dm : cluster->GetDeviceSet()->devices()) {
@@ -454,6 +454,8 @@ tensorflow::Status ConvertAfterShapes(
         cuda_device_id = cuda_gpu_id.value();
       }
       tensorflow::GPUOptions gpuoptions;
+      // we need to us PM here since in python path there is no way to get to
+      // allocators
       auto pm = tensorflow::ProcessState::singleton();
       // this should be instantiated by now
       auto dev_allocator = pm->GetGPUAllocator(gpuoptions, tf_gpu_id, 1);
diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.h b/tensorflow/contrib/tensorrt/convert/convert_graph.h
index 23a83b5094..65a67d7e73 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_graph.h
+++ b/tensorflow/contrib/tensorrt/convert/convert_graph.h
@@ -17,9 +17,7 @@ limitations under the License.
 
 #include <vector>
 
-#include "tensorflow/contrib/tensorrt/segment/segment.h"
 #include "tensorflow/core/framework/graph.pb.h"
-#include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/grappler/clusters/cluster.h"
 #include "tensorflow/core/grappler/costs/graph_properties.h"
 #include "tensorflow/core/lib/core/status.h"
diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.h b/tensorflow/contrib/tensorrt/convert/convert_nodes.h
index 50b0c37094..3f6592cd25 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_nodes.h
+++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.h
@@ -22,14 +22,14 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h"
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/grappler/costs/graph_properties.h"
 #include "tensorflow/core/lib/core/status.h"
-
 #if GOOGLE_CUDA
 #if GOOGLE_TENSORRT
-#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h"
+
 namespace tensorflow {
 namespace tensorrt {
 namespace convert {
@@ -49,7 +49,7 @@ struct SubGraphParams {
       std::unordered_map<string, std::pair<int, string>>* output_edges,
       tensorflow::NodeDef* constructed_trt_node,
       int engine_precision_mode = FP32MODE, const string& device_name = "",
-      std::shared_ptr<nvinfer1::IGpuAllocator> allocator = 0,
+      std::shared_ptr<nvinfer1::IGpuAllocator> allocator = nullptr,
       int cuda_gpu_id = 0)
       : graph(inp_graph),
         subgraph_node_ids(subgraph_node_id_numbers),
diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
index 743750998c..21013fbf9e 100644
--- a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
+++ b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
@@ -22,18 +22,19 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/public/session_options.h"
 
-using tensorflow::str_util::Uppercase;
-using tensorflow::strings::StrAppend;
-using tensorflow::strings::StrCat;
 #if GOOGLE_CUDA
 #if GOOGLE_TENSORRT
 namespace tensorflow {
 namespace tensorrt {
 namespace convert {
 // TODO(sami): Remove VLOG messages once the code matures
+using tensorflow::str_util::Uppercase;
+using tensorflow::strings::StrAppend;
+using tensorflow::strings::StrCat;
+
 tensorflow::Status TRTOptimizationPass::Init(
     const tensorflow::RewriterConfig_CustomGraphOptimizer* config) {
-  VLOG(1) << "Called INIT for " << m_name_ << " with config = " << config;
+  VLOG(1) << "Called INIT for " << name_ << " with config = " << config;
   if (config == nullptr) {
     maximum_workspace_size_ = 2 << 30;
     return tensorflow::Status::OK();
@@ -65,10 +66,9 @@ tensorflow::Status TRTOptimizationPass::Init(
   return tensorflow::Status::OK();
 };
 
-tensorflow::Status TRTOptimizationPass::Optimize(
+void TRTOptimizationPass::PrintDebugInfo(
     tensorflow::grappler::Cluster* cluster,
-    const tensorflow::grappler::GrapplerItem& item, GraphDef* optimized_graph) {
-  VLOG(1) << "Called TRTOptimization Pass " << m_name_;
+    const tensorflow::grappler::GrapplerItem& item) {
   VLOG(1) << "Cluster = " << cluster;
   string offset("  ");
   string offset2 = StrCat(offset, offset);
@@ -77,10 +77,10 @@ tensorflow::Status TRTOptimizationPass::Optimize(
   if (cluster) {
     VLOG(1) << offset << "type             = " << cluster->type();
     VLOG(1) << offset << "num warmup steps = " << cluster->NumWarmupSteps();
-    const auto devNames = cluster->GetDeviceNames();
-    if (devNames.size()) {
+    const auto dev_names = cluster->GetDeviceNames();
+    if (dev_names.size()) {
       VLOG(1) << offset << " Device names:";
-      for (const auto s : devNames) {
+      for (const auto s : dev_names) {
         VLOG(1) << offset2 << s;
       }
     }
@@ -122,38 +122,15 @@ tensorflow::Status TRTOptimizationPass::Optimize(
     }
   }
   VLOG(1) << "item: " << item.id;
-  int max_dim = -1;
   if (item.feed.size()) {
     VLOG(1) << offset << "Feeds  :";
     for (const auto& f : item.feed) {
       const auto& shape = f.second.shape();
-      if (shape.dims() > 0) {
-        if (shape.dim_size(0) > max_dim) max_dim = shape.dim_size(0);
-      }
-      VLOG(1) << offset2 << f.first << " = shaped "
-              << f.second.shape().DebugString();
+      VLOG(1) << offset2 << f.first << " = shaped " << shape.DebugString();
     }
   } else {
     VLOG(1) << offset << "No Feeds";
   }
-  if (maximum_batch_size_ < 0) {  // automatic batch size from input
-    if (max_dim > 0) {
-      maximum_batch_size_ = max_dim;
-      VLOG(1) << "Setting maximum batch size to " << max_dim;
-    } else {
-      maximum_batch_size_ = 128;
-      LOG(WARNING) << "Maximum batch size is not set"
-                      " and can't be deduced from inputs setting it to"
-                   << maximum_batch_size_
-                   << ". Suggest configuring it from configuration parameters";
-    }
-  } else {
-    if (max_dim > maximum_batch_size_) {
-      LOG(WARNING) << "Configured batch size " << maximum_batch_size_
-                   << " is less than input batch size " << max_dim
-                   << " adjusting maximum batch size to match input batch size";
-    }
-  }
   if (item.fetch.size()) {
     VLOG(1) << offset << "Fetches  :";
     for (const auto& f : item.fetch) {
@@ -182,9 +159,7 @@ tensorflow::Status TRTOptimizationPass::Optimize(
   } else {
     VLOG(1) << offset << "No keep ops";
   }
-  VLOG(1) << item.graph.DebugString();
-  tensorflow::grappler::GraphProperties static_graph_properties(item);
-  TF_RETURN_IF_ERROR(static_graph_properties.InferStatically(true));
+  VLOG(3) << item.graph.DebugString();
   for (const auto dev : cluster->GetDeviceSet()->devices()) {
     const auto& pname = dev->parsed_name();
     VLOG(1) << "Device name= " << dev->name()
@@ -192,6 +167,44 @@ tensorflow::Status TRTOptimizationPass::Optimize(
             << " has_id: " << pname.has_id << " has_job: " << pname.has_job
             << "has_type: " << pname.has_type << " type =" << pname.type;
   }
+}
+
+tensorflow::Status TRTOptimizationPass::Optimize(
+    tensorflow::grappler::Cluster* cluster,
+    const tensorflow::grappler::GrapplerItem& item, GraphDef* optimized_graph) {
+  VLOG(1) << "Called TRTOptimization Pass " << name_;
+  if (VLOG_IS_ON(1)) {
+    PrintDebugInfo(cluster, item);
+  }
+  int max_dim = -1;
+  if (item.feed.size()) {
+    for (const auto& f : item.feed) {
+      const auto& shape = f.second.shape();
+      if (shape.dims() > 0) {
+        if (shape.dim_size(0) > max_dim) max_dim = shape.dim_size(0);
+      }
+    }
+  }
+  if (maximum_batch_size_ < 0) {  // automatic batch size from input
+    if (max_dim > 0) {
+      maximum_batch_size_ = max_dim;
+      VLOG(1) << "Setting maximum batch size to " << max_dim;
+    } else {
+      maximum_batch_size_ = 128;
+      LOG(WARNING) << "Maximum batch size is not set"
+                      " and can't be deduced from inputs setting it to"
+                   << maximum_batch_size_
+                   << ". Suggest configuring it from configuration parameters";
+    }
+  } else {
+    if (max_dim > maximum_batch_size_) {
+      LOG(WARNING) << "Configured batch size " << maximum_batch_size_
+                   << " is less than input batch size " << max_dim
+                   << " adjusting maximum batch size to match input batch size";
+    }
+  }
+  tensorflow::grappler::GraphProperties static_graph_properties(item);
+  TF_RETURN_IF_ERROR(static_graph_properties.InferStatically(true));
   auto status = tensorflow::tensorrt::convert::ConvertAfterShapes(
       item.graph, item.fetch, maximum_batch_size_, maximum_workspace_size_,
       optimized_graph, precision_mode_, minimum_segment_size_,
@@ -205,20 +218,25 @@ void TRTOptimizationPass::Feedback(
     const tensorflow::grappler::GrapplerItem& item,
     const GraphDef& optimized_graph, double result) {}
 
-using tensorflow::grappler::CustomGraphOptimizerRegistrar;
-namespace {
 
-class samiReg : public CustomGraphOptimizerRegistrar {
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+
+class VerboseCustomGraphOptimizerRegistrar
+    : public tensorflow::grappler::CustomGraphOptimizerRegistrar {
  public:
-  samiReg(const tensorflow::grappler::CustomGraphOptimizerRegistry::Creator& cr,
-          const string& name)
-      : CustomGraphOptimizerRegistrar(cr, name) {
+  VerboseCustomGraphOptimizerRegistrar(
+      const tensorflow::grappler::CustomGraphOptimizerRegistry::Creator& cr,
+      const tensorflow::string& name)
+      : tensorflow::grappler::CustomGraphOptimizerRegistrar(cr, name) {
     VLOG(1) << "Constructing a CustomOptimizationPass registration object for "
             << name;
   }
 };
-// static CustomGraphOptimizerRegistrar TRTOptimizationPass_Registrar([]() {
-static samiReg TRTOptimizationPass_Registrar(
+
+static VerboseCustomGraphOptimizerRegistrar TRTOptimizationPass_Registrar(
     []() {
       VLOG(1)
           << "Instantiating CustomOptimizationPass object TensorRTOptimizer";
@@ -226,11 +244,6 @@ static samiReg TRTOptimizationPass_Registrar(
           "TensorRTOptimizer");
     },
     ("TensorRTOptimizer"));
-}  // namespace
-
-}  // namespace convert
-}  // namespace tensorrt
-}  // namespace tensorflow
 
 #endif
 #endif
diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
index aa9f289550..c554a5d784 100644
--- a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
+++ b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
@@ -16,11 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_CONTRIB_TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_
 #define TENSORFLOW_CONTRIB_TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_
 
-#include <set>
 #include <string>
-#include <unordered_map>
-#include <utility>
-#include <vector>
 
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h"
@@ -35,14 +31,14 @@ namespace convert {
 class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer {
  public:
   TRTOptimizationPass(const string& name = "TRTOptimizationPass")
-      : m_name_(name),
+      : name_(name),
         minimum_segment_size_(3),
         precision_mode_(0),
         maximum_batch_size_(-1),
         maximum_workspace_size_(-1) {
-    VLOG(1) << "Constructing " << m_name_;
+    VLOG(1) << "Constructing " << name_;
   };
-  string name() const override { return m_name_; };
+  string name() const override { return name_; };
   tensorflow::Status Init(const tensorflow::RewriterConfig_CustomGraphOptimizer*
                               config = nullptr) override;
 
@@ -52,9 +48,11 @@ class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer {
   void Feedback(tensorflow::grappler::Cluster* cluster,
                 const tensorflow::grappler::GrapplerItem& item,
                 const GraphDef& optimized_graph, double result) override;
+  void PrintDebugInfo(tensorflow::grappler::Cluster* cluster,
+                      const tensorflow::grappler::GrapplerItem& item);
 
  private:
-  string m_name_;
+  string name_;
   int minimum_segment_size_;
   int precision_mode_;
   int maximum_batch_size_;
diff --git a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
index 15a3bbd0d2..f10b10edec 100644
--- a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
+++ b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
@@ -15,9 +15,6 @@ limitations under the License.
 #include "tensorflow/contrib/tensorrt/kernels/trt_engine_op.h"
 
 #include "tensorflow/contrib/tensorrt/log/trt_logger.h"
-#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
-#include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h"
-#include "tensorflow/core/common_runtime/gpu/process_state.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/stream_executor.h"
 #include "tensorflow/core/platform/types.h"
@@ -42,38 +39,23 @@ TRTEngineOp::TRTEngineOp(OpKernelConstruction* context) : OpKernel(context) {
   OP_REQUIRES_OK(context, context->GetAttr("input_nodes", &input_nodes_));
   OP_REQUIRES_OK(context, context->GetAttr("output_nodes", &output_nodes_));
 
-  // TODO(samikama) runtime should be taken from a resourcemanager as well.
-  // Only engine should be in the op and context and runtime should be taken
-  // from resourcemanager
-  // TODO(jie): cudaSetDevice make sure trt engine is allocated on the same
-  // gpu where the input/output is also located.
-  // int gpu_id = context->device()->tensorflow_gpu_device_info()->gpu_id;
-  // cudaSetDevice(gpu_id);
-  // int device;
-  // cudaGetDevice(&device);
-  // if (gpu_id != device) LOG(FATAL) << "set device failed!";
+}
 
+void TRTEngineOp::Compute(OpKernelContext* context) {
   // TODO(samikama) runtime should be taken from a resourcemanager as well.
   // Only engine should be in the op and context and runtime should be taken
   // from resourcemanager
 
-  // IRuntime* infer = nvinfer1::createInferRuntime(logger);
-  // trt_engine_ptr_.reset(infer->deserializeCudaEngine(
-  //     serialized_engine.c_str(), serialized_engine.size(), nullptr));
-  // trt_execution_context_ptr_.reset(trt_engine_ptr_->createExecutionContext());
-  // Runtime is safe to delete after engine creation
-  // infer->destroy();
-}
-
-void TRTEngineOp::Compute(OpKernelContext* context) {
   if (!trt_execution_context_ptr_) {
     IRuntime* infer = nvinfer1::createInferRuntime(logger);
 #if NV_TENSORRT_MAJOR > 3
-    tensorflow::TfGpuId tf_gpu_id(
-        context->device()->tensorflow_gpu_device_info()->gpu_id);
-    tensorflow::GPUOptions gpuoptions;
-    auto pm = tensorflow::ProcessState::singleton();
-    auto dev_allocator = pm->GetGPUAllocator(gpuoptions, tf_gpu_id, 1);
+    auto device=context->device();
+    auto dev_allocator=device->getAllocator(tensorflow::AllocatorAttributes())
+    // tensorflow::TfGpuId tf_gpu_id(
+    //     context->device()->tensorflow_gpu_device_info()->gpu_id);
+    // tensorflow::GPUOptions gpuoptions;
+    // auto pm = tensorflow::ProcessState::singleton();
+    // auto dev_allocator = pm->GetGPUAllocator(gpuoptions, tf_gpu_id, 1);
     if (!dev_allocator) {
       LOG(FATAL) << "Can't find device allocator for gpu device" << tf_gpu_id;
     }
diff --git a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h
index 38ceec4704..fec4bd728b 100644
--- a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h
+++ b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h
@@ -32,6 +32,7 @@ namespace tensorflow {
 namespace tensorrt {
 class Logger;
 
+//  TODO(Sami): Remove this file?
 class TRTEngineOp : public OpKernel {
  public:
   explicit TRTEngineOp(OpKernelConstruction* context);
diff --git a/tensorflow/contrib/tensorrt/resources/trt_allocator.h b/tensorflow/contrib/tensorrt/resources/trt_allocator.h
index 05dcb7cde6..dd4f8c7943 100644
--- a/tensorflow/contrib/tensorrt/resources/trt_allocator.h
+++ b/tensorflow/contrib/tensorrt/resources/trt_allocator.h
@@ -40,17 +40,20 @@ class IGpuAllocator {
 namespace tensorflow {
 namespace tensorrt {
 class TRTCudaAllocator : public nvinfer1::IGpuAllocator {
+  // Allocator implementation that is using cuda allocator instead of device
+  // allocator in case we can't get device allocator from TF.
  public:
   TRTCudaAllocator() {}
-  virtual ~TRTCudaAllocator() {};
+  virtual ~TRTCudaAllocator(){};
   void* allocate(uint64_t size, uint64_t alignment, uint32_t flags) override;
   void free(void* memory) override;
 };
 
 class TRTDeviceAllocator : public nvinfer1::IGpuAllocator {
+  // Allocator implementation wrapping TF device allocators.
  public:
   TRTDeviceAllocator(tensorflow::Allocator* allocator);
-  virtual ~TRTDeviceAllocator() {};
+  virtual ~TRTDeviceAllocator(){};
   void* allocate(uint64_t size, uint64_t alignment, uint32_t flags) override;
   void free(void* memory) override;
 
diff --git a/tensorflow/contrib/tensorrt/segment/segment.cc b/tensorflow/contrib/tensorrt/segment/segment.cc
index a76d170236..7e094f552d 100644
--- a/tensorflow/contrib/tensorrt/segment/segment.cc
+++ b/tensorflow/contrib/tensorrt/segment/segment.cc
@@ -32,74 +32,92 @@ namespace tensorflow {
 namespace tensorrt {
 namespace segment {
 using ::tensorflow::strings::StrAppend;
-namespace {
-
-bool CheckCycles(const SimpleGraph* g, const SimpleNode* src,
-                 const std::vector<SimpleNode*>& start) {
-  //  copied from TF ReverseDFS
-  struct Work {
-    SimpleNode* node;
-    bool leave;  // Are we entering or leaving n?
-  };
-
-  std::vector<Work> stack(start.size());
-  for (int i = 0; i < start.size(); ++i) {
-    stack[i] = Work{start[i], false};
-  }
-
-  std::vector<bool> visited(g->num_node_ids(), false);
-  while (!stack.empty()) {
-    Work w = stack.back();
-    stack.pop_back();
-
-    auto n = w.node;
-    if (w.leave) {
-      if (n == src) {
-        return true;
-      }
-      continue;
-    }
-
-    if (visited[n->id()]) continue;
-    visited[n->id()] = true;
-    // Arrange to call leave(n) when all done with descendants.
-    stack.push_back(Work{n, true});
-
-    auto nodes = n->in_nodes();
-    for (const auto node : nodes) {
-      if (!visited[node->id()]) {
-        stack.push_back(Work{node, false});
-      }
+// A simple graph representation to mirror tensorflow::Graph. This structure
+// helps saving memory since segmenter modifies the graph in place, preventing
+// the need to create a copy of the graph. It is composed of edges and nodes.
+// Nodes keep pointers to original TF nodes.
+class SimpleNode;
+class SimpleGraph;
+class SimpleEdge {
+ public:
+  SimpleEdge(int id, SimpleNode* src, int src_port, SimpleNode* dst,
+             int dst_port, bool is_control = false)
+      : id_(id),
+        src_(src),
+        src_port_(src_port),
+        dst_(dst),
+        dst_port_(dst_port),
+        control_(is_control){};
+  SimpleNode* src() const { return src_; }
+  SimpleNode* dst() const { return dst_; }
+  int src_output() const { return src_port_; }
+  int dst_input() const { return dst_port_; }
+  int id() const { return id_; }
+  bool IsControlEdge() const { return control_; }
+  ~SimpleEdge() {}
+
+ private:
+  int id_;
+  SimpleNode* src_;
+  int src_port_;
+  SimpleNode* dst_;
+  int dst_port_;
+  bool control_;
+};
+class SimpleNode {
+ public:
+  SimpleNode(const tensorflow::Node* node, const int id);
+  const std::vector<SimpleEdge*>& in_edges() const { return in_edges_; };
+  const std::vector<SimpleEdge*>& out_edges() const { return out_edges_; };
+  std::vector<SimpleNode*> in_nodes() const {
+    std::vector<SimpleNode*> res;
+    res.reserve(in_edges_.size());
+    for (const auto e : in_edges_) {
+      if (e) res.push_back(e->src());
     }
+    return res;
   }
-  return false;
-}
-
-bool CanContractEdge(const SimpleEdge* edge, const SimpleGraph* graph) {
-  const auto src = edge->src();
-  const auto dst = edge->dst();
-
-  // Can't contract edge if doing so would cause a cycle in the
-  // graph. So, if there is a directed path from 'src' to 'dst', other
-  // than 'edge' (or any other direct edge from 'src' to 'dst'), then
-  // combining 'src' and 'dst' will cause a cycle along that path.
-  //
-  // In practice, to avoid modifying the graph and to take advantage
-  // of existing graph functions, we perform an equivalent.
-  //   1. Get all nodes incoming to 'dst', excluding 'src'
-  //   2. Reverse DFS from those nodes
-  //   3. If reverse DFS reaches 'src' then we have a cycle
-  std::vector<SimpleNode*> dfs_start_nodes;
-  for (SimpleNode* node : dst->in_nodes()) {
-    if (node != src) {
-      dfs_start_nodes.push_back(node);
-    }
+  const string& name() const { return node_->name(); }
+  const tensorflow::Node* tf_node() const { return node_; }
+  int id() const { return id_; }
+
+ private:
+  const tensorflow::Node* node_;
+  std::vector<SimpleEdge*> in_edges_;
+  std::vector<SimpleEdge*> out_edges_;
+  int id_;
+
+  friend class SimpleGraph;
+};
+
+class SimpleGraph {
+ public:
+  SimpleGraph(const tensorflow::Graph* g);
+  void AddControlEdge(SimpleNode* src, SimpleNode* dst);
+  void AddEdge(SimpleNode* src, int out_port, SimpleNode* dst, int in_port);
+  void RemoveEdge(const SimpleEdge*);
+  SimpleNode* FindNodeId(int node_id) {
+    if (node_id < 0 || node_id > (int)nodes_.size()) return nullptr;
+    return nodes_[node_id];
+  }
+  ~SimpleGraph();
+  int num_node_ids() const { return nodes_.size(); }
+  const SimpleNode* source_node() const {
+    return nodes_[tensorflow::Graph::kSourceId];
+  }
+  const SimpleNode* sink_node() const {
+    return nodes_[tensorflow::Graph::kSinkId];
   }
 
-  bool is_cycle = CheckCycles(graph, src, dfs_start_nodes);
-  return !is_cycle;
-}
-}  // namespace
+ private:
+  const tensorflow::Graph* g_;
+  std::vector<SimpleNode*> nodes_;
+  std::vector<SimpleEdge*> edges_;
+  // edge_ids_ and node_ids_ contain freed indices.
+  std::set<int> free_edge_ids_;
+  std::set<int> free_node_ids_;
+};
+
 SimpleNode::SimpleNode(const tensorflow::Node* node, const int id)
     : node_(node), id_(id) {
   if (node_) {
@@ -120,7 +138,7 @@ SimpleGraph::SimpleGraph(const tensorflow::Graph* g) : g_(g) {
     if (n) {
       nodes_[i] = new SimpleNode(n, i);
     } else {
-      node_ids_.insert(i);
+      free_node_ids_.insert(i);
     }
   }
   for (int i = 0; i < n_edges; i++) {
@@ -137,7 +155,7 @@ SimpleGraph::SimpleGraph(const tensorflow::Graph* g) : g_(g) {
       src->out_edges_.push_back(edge);
       dst->in_edges_.push_back(edge);
     } else {
-      edge_ids_.insert(i);
+      free_edge_ids_.insert(i);
     }
   }
 }
@@ -145,12 +163,12 @@ SimpleGraph::SimpleGraph(const tensorflow::Graph* g) : g_(g) {
 void SimpleGraph::AddEdge(SimpleNode* src, int out_port, SimpleNode* dst,
                           int in_port) {
   int i = edges_.size();
-  if (edge_ids_.size()) {
-    auto it = edge_ids_.begin();
+  if (free_edge_ids_.size()) {
+    auto it = free_edge_ids_.begin();
     i = *it;
-    edge_ids_.erase(it);
+    free_edge_ids_.erase(it);
   } else {
-    edges_.push_back(0);
+    edges_.push_back(nullptr);
   }
   bool is_control = (out_port == tensorflow::Graph::kControlSlot);
   is_control |= (in_port == tensorflow::Graph::kControlSlot);
@@ -187,7 +205,77 @@ SimpleGraph::~SimpleGraph() {
   for (auto x : edges_) delete x;
 }
 
-void ContractEdge(SimpleEdge* edge, SimpleGraph* graph,
+namespace {
+
+bool CheckCycles(const std::unique_ptr<SimpleGraph>& g, const SimpleNode* src,
+                 const std::vector<SimpleNode*>& start) {
+  // copied from TF ReverseDFS.
+  struct Work {
+    SimpleNode* node;
+    bool leave;  // Are we entering or leaving n?
+  };
+
+  std::vector<Work> stack(start.size());
+  for (int i = 0; i < start.size(); ++i) {
+    stack[i] = Work{start[i], false};
+  }
+
+  std::vector<bool> visited(g->num_node_ids(), false);
+  while (!stack.empty()) {
+    Work w = stack.back();
+    stack.pop_back();
+
+    auto n = w.node;
+    if (w.leave) {
+      if (n == src) {
+        return true;
+      }
+      continue;
+    }
+
+    if (visited[n->id()]) continue;
+    visited[n->id()] = true;
+    // Arrange to call leave(n) when all done with descendants.
+    stack.push_back(Work{n, true});
+
+    auto nodes = n->in_nodes();
+    for (const auto node : nodes) {
+      if (!visited[node->id()]) {
+        stack.push_back(Work{node, false});
+      }
+    }
+  }
+  return false;
+}
+
+bool CanContractEdge(const SimpleEdge* edge,
+                     const std::unique_ptr<SimpleGraph>& graph) {
+  const auto src = edge->src();
+  const auto dst = edge->dst();
+
+  // Can't contract edge if doing so would cause a cycle in the
+  // graph. So, if there is a directed path from 'src' to 'dst', other
+  // than 'edge' (or any other direct edge from 'src' to 'dst'), then
+  // combining 'src' and 'dst' will cause a cycle along that path.
+  //
+  // In practice, to avoid modifying the graph and to take advantage
+  // of existing graph functions, we perform an equivalent.
+  //   1. Get all nodes incoming to 'dst', excluding 'src'
+  //   2. Reverse DFS from those nodes
+  //   3. If reverse DFS reaches 'src' then we have a cycle
+  std::vector<SimpleNode*> dfs_start_nodes;
+  for (SimpleNode* node : dst->in_nodes()) {
+    if (node != src) {
+      dfs_start_nodes.push_back(node);
+    }
+  }
+
+  bool is_cycle = CheckCycles(graph, src, dfs_start_nodes);
+  return !is_cycle;
+}
+}  // namespace
+
+void ContractEdge(SimpleEdge* edge, std::unique_ptr<SimpleGraph>& graph,
                   std::vector<const SimpleEdge*>* remove_edges) {
   // Transfer all inputs and outputs of 'dst' to 'src' except edges
   // connecting the two.
@@ -265,7 +353,7 @@ tensorflow::Status SegmentGraph(
     const std::function<bool(const tensorflow::Node*)>& candidate_fn,
     const SegmentOptions& options, SegmentNodesVector* segments) {
   // tensorflow::DumpGraph("Pre-Segment", &graph);
-  SimpleGraph* graph = new SimpleGraph(tf_graph);
+  auto graph = std::unique_ptr<SimpleGraph>(new SimpleGraph(tf_graph));
   // Use a union-find to collect the nodes that belong to the same
   // segment. A node value of nullptr indicates that the node is not a candidate
   // for TRT.
@@ -370,6 +458,11 @@ tensorflow::Status SegmentGraph(
     if ((u.Value() != nullptr) && (u.ParentValue() != nullptr)) {
       sg_map[u.ParentValue()->name()].insert(u.Value()->name());
       auto tf_node = u.Value()->tf_node();
+      // has_assigned_device_name() is expected to return true
+      // when called from optimization pass. However, since graph
+      // is converted back and forth between graph and graphdef,
+      // assigned devices demoted to requested devices. If the graph
+      // is passed directly to this module, assigned devices will be set.
       if (tf_node->has_assigned_device_name()) {
         device_maps[u.ParentValue()->name()].insert(
             tf_node->assigned_device_name());
@@ -421,15 +514,16 @@ tensorflow::Status SegmentGraph(
           std::make_pair(segment_node_names, *(dev_itr->second.begin())));
     }
   }
-  for (const auto& d : device_maps) {
-    string s("Segment ");
-    StrAppend(&s, ": '", d.first, "' ");
-    for (const auto& dd : d.second) {
-      StrAppend(&s, dd, ", ");
+  if (VLOG_IS_ON(1)) {
+    for (const auto& d : device_maps) {
+      string s("Segment ");
+      StrAppend(&s, ": '", d.first, "' ");
+      for (const auto& dd : d.second) {
+        StrAppend(&s, dd, ", ");
+      }
+      VLOG(1) << "Devices " << s;
     }
-    VLOG(1) << "Devices " << s;
   }
-  delete graph;
   return tensorflow::Status::OK();
 }
 
diff --git a/tensorflow/contrib/tensorrt/segment/segment.h b/tensorflow/contrib/tensorrt/segment/segment.h
index 44a84cbd38..c5aca4bf04 100644
--- a/tensorflow/contrib/tensorrt/segment/segment.h
+++ b/tensorflow/contrib/tensorrt/segment/segment.h
@@ -29,87 +29,10 @@ namespace tensorflow {
 namespace tensorrt {
 namespace segment {
 
+// vector of segments, each entry contains a device name and a set of nodes in
+// segment
 using SegmentNodesVector = std::vector<std::pair<std::set<string>, string>>;
-class SimpleNode;
-class SimpleGraph;
-class SimpleEdge {
- public:
-  SimpleEdge(int id, SimpleNode* src, int src_port, SimpleNode* dst,
-             int dst_port, bool is_control = false)
-      : id_(id),
-        src_(src),
-        src_port_(src_port),
-        dst_(dst),
-        dst_port_(dst_port),
-        control_(is_control){};
-  SimpleNode* src() const { return src_; }
-  SimpleNode* dst() const { return dst_; }
-  int src_output() const { return src_port_; }
-  int dst_input() const { return dst_port_; }
-  int id() const { return id_; }
-  bool IsControlEdge() const { return control_; }
-  ~SimpleEdge() {}
 
- private:
-  int id_;
-  SimpleNode* src_;
-  int src_port_;
-  SimpleNode* dst_;
-  int dst_port_;
-  bool control_;
-};
-class SimpleNode {
- public:
-  SimpleNode(const tensorflow::Node* node, const int id);
-  const std::vector<SimpleEdge*>& in_edges() const { return in_edges_; };
-  const std::vector<SimpleEdge*>& out_edges() const { return out_edges_; };
-  std::vector<SimpleNode*> in_nodes() const {
-    std::vector<SimpleNode*> res;
-    res.reserve(in_edges_.size());
-    for (const auto e : in_edges_) {
-      if (e) res.push_back(e->src());
-    }
-    return res;
-  }
-  const string& name() const { return node_->name(); }
-  const tensorflow::Node* tf_node() const { return node_; }
-  int id() const { return id_; }
-
- private:
-  const tensorflow::Node* node_;
-  std::vector<SimpleEdge*> in_edges_;
-  std::vector<SimpleEdge*> out_edges_;
-  int id_;
-
-  friend class SimpleGraph;
-};
-
-class SimpleGraph {
- public:
-  SimpleGraph(const tensorflow::Graph* g);
-  void AddControlEdge(SimpleNode* src, SimpleNode* dst);
-  void AddEdge(SimpleNode* src, int out_port, SimpleNode* dst, int in_port);
-  void RemoveEdge(const SimpleEdge*);
-  SimpleNode* FindNodeId(int node_id) {
-    if (node_id < 0 || node_id > (int)nodes_.size()) return nullptr;
-    return nodes_[node_id];
-  }
-  ~SimpleGraph();
-  int num_node_ids() const { return nodes_.size(); }
-  const SimpleNode* source_node() const {
-    return nodes_[tensorflow::Graph::kSourceId];
-  }
-  const SimpleNode* sink_node() const {
-    return nodes_[tensorflow::Graph::kSinkId];
-  }
-
- private:
-  const tensorflow::Graph* g_;
-  std::vector<SimpleNode*> nodes_;
-  std::vector<SimpleEdge*> edges_;
-  std::set<int> edge_ids_;
-  std::set<int> node_ids_;
-};
 struct SegmentOptions {
   // Segment must contain at least this many nodes.
   int minimum_segment_size = 2;
-- 
GitLab


From 9e197152c04ebb81f055067534bd93322d182f0e Mon Sep 17 00:00:00 2001
From: Ben <bstriner@users.noreply.github.com>
Date: Mon, 30 Apr 2018 23:30:22 -0400
Subject: [PATCH 084/395] Fix MSVC openmp flag (#18973)

* Fix MSVC openmp flag
---
 tensorflow/contrib/cmake/CMakeLists.txt | 27 +++++++++++++++----------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/tensorflow/contrib/cmake/CMakeLists.txt b/tensorflow/contrib/cmake/CMakeLists.txt
index d81f6a0ae8..0708d6b7b9 100644
--- a/tensorflow/contrib/cmake/CMakeLists.txt
+++ b/tensorflow/contrib/cmake/CMakeLists.txt
@@ -172,19 +172,20 @@ if (tensorflow_OPTIMIZE_FOR_NATIVE_ARCH)
   endif()
 endif()
 
+include(CheckCXXCompilerFlag)
+
+# OpenMP Support
+CHECK_CXX_COMPILER_FLAG("-fopenmp" GCC_OPENMP_SUPPORT)
+if (GCC_OPENMP_SUPPORT)
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fopenmp")
+endif()
+CHECK_CXX_COMPILER_FLAG("/openmp" MSVC_OPENMP_SUPPORT)
+if (MSVC_OPENMP_SUPPORT)
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /openmp")
+endif()
+
 # MSVC SIMD instructions
 if (tensorflow_WIN_CPU_SIMD_OPTIONS)
-  include(CheckCXXCompilerFlag)
-  if (tensorflow_ENABLE_MKL_SUPPORT)
-    add_definitions(-DINTEL_MKL -DEIGEN_USE_VML)
-    if (NOT tensorflow_ENABLE_MKLDNN_SUPPORT)
-      add_definitions(-DINTEL_MKL_ML)
-    endif()
-  endif()
-  CHECK_CXX_COMPILER_FLAG("-fopenmp" COMPILER_OPT_OPENMP_SUPPORT)
-  if (COMPILER_OPT_OPENMP_SUPPORT)
-    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fopenmp")
-  endif()
   if (WIN32)
     CHECK_CXX_COMPILER_FLAG(${tensorflow_WIN_CPU_SIMD_OPTIONS} COMPILER_OPT_WIN_CPU_SIMD_SUPPORTED)
     if(COMPILER_OPT_WIN_CPU_SIMD_SUPPORTED)
@@ -323,7 +324,9 @@ if(HAIKU)
   list(APPEND tensorflow_EXTERNAL_LIBRARIES network)
 endif()
 
+# MKL Support
 if (tensorflow_ENABLE_MKL_SUPPORT)
+  add_definitions(-DINTEL_MKL -DEIGEN_USE_VML)
   if (WIN32)
     find_path(MKL_HOME_PLATFORM mkl
       PATHS ${MKL_HOME} ${MKL_HOME}/../ ${MKL_HOME}/../../
@@ -359,6 +362,8 @@ if (tensorflow_ENABLE_MKL_SUPPORT)
     list(APPEND tensorflow_EXTERNAL_LIBRARIES ${mkldnn_STATIC_LIBRARIES})
     list(APPEND tensorflow_EXTERNAL_DEPENDENCIES mkldnn)
     include_directories(${mkldnn_INCLUDE_DIRS})
+  else (tensorflow_ENABLE_MKLDNN_SUPPORT)
+    add_definitions(-DINTEL_MKL_ML)
   endif()
 endif (tensorflow_ENABLE_MKL_SUPPORT)
 
-- 
GitLab


From d0f5bc17560fc97bcc7de9164aa3b237a8d5221d Mon Sep 17 00:00:00 2001
From: Maciej <mbajkowski@hotmail.com>
Date: Mon, 30 Apr 2018 22:30:58 -0500
Subject: [PATCH 085/395]  Remove whitespace characters from
 tf_cuda_compute_capabilities user string (#18986)

Remove all whitespace characters from the user specified tf_cuda_compute_capabilities string as this can results in errors during the split operation, and is easy for users to do as it is natural to insert a space after a comma
---
 configure.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/configure.py b/configure.py
index b745e374a2..fe15bfc1a4 100644
--- a/configure.py
+++ b/configure.py
@@ -1226,6 +1226,9 @@ def set_tf_cuda_compute_capabilities(environ_cp):
         ask_cuda_compute_capabilities, default_cuda_compute_capabilities)
     # Check whether all capabilities from the input is valid
     all_valid = True
+    # Remove all whitespace characters before splitting the string
+    # that users may insert by accident, as this will result in error 
+    tf_cuda_compute_capabilities = ''.join(tf_cuda_compute_capabilities.split())
     for compute_capability in tf_cuda_compute_capabilities.split(','):
       m = re.match('[0-9]+.[0-9]+', compute_capability)
       if not m:
-- 
GitLab


From 95b36432d2c04a8355d2de2aeb4817fb3042d639 Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Tue, 1 May 2018 00:42:56 -0700
Subject: [PATCH 086/395] [XLA:CPU] Open source some tests.

PiperOrigin-RevId: 194903752
---
 .../compiler/xla/service/cpu/tests/BUILD      | 126 +++++++
 .../service/cpu/tests/cpu_bytesizeof_test.cc  |  37 ++
 .../xla/service/cpu/tests/cpu_codegen_test.h  |  30 ++
 .../cpu/tests/cpu_eigen_dot_operation_test.cc | 113 ++++++
 .../cpu/tests/cpu_external_constants_test.cc  |  73 ++++
 .../xla/service/cpu/tests/cpu_fusion_test.cc  | 330 ++++++++++++++++++
 .../service/cpu/tests/cpu_intrinsic_test.cc   | 151 ++++++++
 .../xla/service/cpu/tests/cpu_noalias_test.cc | 136 ++++++++
 tensorflow/compiler/xla/tests/filecheck.cc    |   7 +-
 9 files changed, 1002 insertions(+), 1 deletion(-)
 create mode 100644 tensorflow/compiler/xla/service/cpu/tests/BUILD
 create mode 100644 tensorflow/compiler/xla/service/cpu/tests/cpu_bytesizeof_test.cc
 create mode 100644 tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h
 create mode 100644 tensorflow/compiler/xla/service/cpu/tests/cpu_eigen_dot_operation_test.cc
 create mode 100644 tensorflow/compiler/xla/service/cpu/tests/cpu_external_constants_test.cc
 create mode 100644 tensorflow/compiler/xla/service/cpu/tests/cpu_fusion_test.cc
 create mode 100644 tensorflow/compiler/xla/service/cpu/tests/cpu_intrinsic_test.cc
 create mode 100644 tensorflow/compiler/xla/service/cpu/tests/cpu_noalias_test.cc

diff --git a/tensorflow/compiler/xla/service/cpu/tests/BUILD b/tensorflow/compiler/xla/service/cpu/tests/BUILD
new file mode 100644
index 0000000000..9425b948c1
--- /dev/null
+++ b/tensorflow/compiler/xla/service/cpu/tests/BUILD
@@ -0,0 +1,126 @@
+# Description:
+#    Tests for LLVM-based CPU backend for XLA.
+
+licenses(["notice"])  # Apache 2.0
+
+package(
+    default_visibility = [":friends"],
+)
+
+package_group(
+    name = "friends",
+    includes = [
+        "//tensorflow/compiler/xla:friends",
+    ],
+)
+
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+
+# Filegroup used to collect source files for dependency checking.
+filegroup(
+    name = "c_srcs",
+    data = glob([
+        "**/*.cc",
+        "**/*.h",
+    ]),
+)
+
+cc_library(
+    name = "cpu_codegen_test",
+    testonly = True,
+    hdrs = ["cpu_codegen_test.h"],
+    deps = [
+        "//tensorflow/compiler/xla/service:cpu_plugin",
+        "//tensorflow/compiler/xla/tests:llvm_irgen_test_base",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
+
+tf_cc_test(
+    name = "cpu_fusion_test",
+    srcs = ["cpu_fusion_test.cc"],
+    deps = [
+        "//tensorflow/compiler/xla:literal_util",
+        "//tensorflow/compiler/xla:shape_util",
+        "//tensorflow/compiler/xla:util",
+        "//tensorflow/compiler/xla:xla_data_proto",
+        "//tensorflow/compiler/xla/service:cpu_plugin",
+        "//tensorflow/compiler/xla/service:hlo",
+        "//tensorflow/compiler/xla/service/cpu:cpu_instruction_fusion",
+        "//tensorflow/compiler/xla/tests:hlo_test_base",
+        "//tensorflow/compiler/xla/tests:literal_test_util",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
+
+tf_cc_test(
+    name = "cpu_bytesizeof_test",
+    srcs = ["cpu_bytesizeof_test.cc"],
+    deps = [
+        "//tensorflow/compiler/xla:shape_util",
+        "//tensorflow/compiler/xla/service/llvm_ir:llvm_util",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
+
+tf_cc_test(
+    name = "cpu_external_constants_test",
+    srcs = ["cpu_external_constants_test.cc"],
+    deps = [
+        "//tensorflow/compiler/xla:array2d",
+        "//tensorflow/compiler/xla:shape_util",
+        "//tensorflow/compiler/xla/service:hlo",
+        "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test",
+        "//tensorflow/compiler/xla/tests:filecheck",
+        "//tensorflow/core:test",
+    ],
+)
+
+tf_cc_test(
+    name = "cpu_noalias_test",
+    srcs = ["cpu_noalias_test.cc"],
+    deps = [
+        "//tensorflow/compiler/xla:literal_util",
+        "//tensorflow/compiler/xla:shape_util",
+        "//tensorflow/compiler/xla:util",
+        "//tensorflow/compiler/xla:xla_data_proto",
+        "//tensorflow/compiler/xla/service:buffer_assignment",
+        "//tensorflow/compiler/xla/service:hlo",
+        "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test",
+        "//tensorflow/compiler/xla/service/llvm_ir:alias_analysis",
+        "//tensorflow/compiler/xla/service/llvm_ir:llvm_util",
+        "//tensorflow/compiler/xla/tests:filecheck",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "@llvm//:core",
+    ],
+)
+
+tf_cc_test(
+    name = "cpu_intrinsic_test",
+    srcs = ["cpu_intrinsic_test.cc"],
+    deps = [
+        "//tensorflow/compiler/xla/service:hlo",
+        "//tensorflow/compiler/xla/service/cpu:cpu_compiler",
+        "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
+
+tf_cc_test(
+    name = "cpu_eigen_dot_operation_test",
+    srcs = ["cpu_eigen_dot_operation_test.cc"],
+    deps = [
+        "//tensorflow/compiler/xla/service:hlo",
+        "//tensorflow/compiler/xla/service/cpu:cpu_compiler",
+        "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_bytesizeof_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_bytesizeof_test.cc
new file mode 100644
index 0000000000..d5bbe7677a
--- /dev/null
+++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_bytesizeof_test.cc
@@ -0,0 +1,37 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h"
+#include "tensorflow/compiler/xla/shape_util.h"
+#include "tensorflow/core/platform/test.h"
+
+class CpuByteSizeOfTest : public ::testing::Test {};
+
+TEST_F(CpuByteSizeOfTest, ARM32) {
+  llvm::DataLayout data_layout(
+      "e-m:e-p:32:32-i64:64-v128:64:128-a:0:32-n32-S64");
+  auto tuple_shape =
+      xla::ShapeUtil::MakeTupleShape({xla::ShapeUtil::MakeShape(xla::F32, {})});
+  EXPECT_EQ(xla::llvm_ir::ByteSizeOf(tuple_shape, data_layout),
+            data_layout.getPointerSize(0 /* default address space */));
+}
+
+TEST_F(CpuByteSizeOfTest, ARM64) {
+  llvm::DataLayout data_layout("e-m:e-i64:64-i128:128-n32:64-S128");
+  auto tuple_shape =
+      xla::ShapeUtil::MakeTupleShape({xla::ShapeUtil::MakeShape(xla::F32, {})});
+  EXPECT_EQ(xla::llvm_ir::ByteSizeOf(tuple_shape, data_layout),
+            data_layout.getPointerSize(0 /* default address space */));
+}
diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h b/tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h
new file mode 100644
index 0000000000..7c8d07a10b
--- /dev/null
+++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h
@@ -0,0 +1,30 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_TESTS_CPU_CODEGEN_TEST_H_
+#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_TESTS_CPU_CODEGEN_TEST_H_
+
+#include "tensorflow/compiler/xla/tests/llvm_irgen_test_base.h"
+
+namespace xla {
+namespace cpu {
+
+// Tests that verify IR emitted by the CPU backend is as expected.
+class CpuCodegenTest : public LLVMIRGenTestBase {};
+
+}  // namespace cpu
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_TESTS_CPU_CODEGEN_TEST_H_
diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_eigen_dot_operation_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_eigen_dot_operation_test.cc
new file mode 100644
index 0000000000..6fcce42eaa
--- /dev/null
+++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_eigen_dot_operation_test.cc
@@ -0,0 +1,113 @@
+/* Copyright 2018 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 that we call into Eigen for dot operations as needed.
+
+#include <algorithm>
+#include <cctype>
+#include <string>
+
+#include "tensorflow/compiler/xla/service/cpu/cpu_compiler.h"
+#include "tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h"
+#include "tensorflow/compiler/xla/service/hlo_computation.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace xla {
+namespace cpu {
+namespace {
+
+struct DotTestSpec {
+  PrimitiveType primitive_type;
+  string filecheck_lines;
+};
+
+string DotTestSpecToString(const ::testing::TestParamInfo<DotTestSpec>& info) {
+  return PrimitiveType_Name(info.param.primitive_type);
+}
+
+class CpuEigenDotOperationTest
+    : public CpuCodegenTest,
+      public ::testing::WithParamInterface<DotTestSpec> {
+ protected:
+  void CompileAndCheck(std::unique_ptr<HloComputation> entry_computation,
+                       const string& filecheck_lines) {
+    CpuAotCompilationOptions options{
+        /*triple=*/"x86_64", /*cpu_name=*/"", /*features=*/"",
+        /*entry_point_name=*/"entry",
+        /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static};
+
+    auto hlo_module = CreateNewModule();
+    hlo_module->AddEntryComputation(std::move(entry_computation));
+
+    CompileAheadOfTimeAndVerifyIr(std::move(hlo_module), options,
+                                  filecheck_lines,
+                                  /*match_optimized_ir=*/true);
+  }
+};
+
+TEST_P(CpuEigenDotOperationTest, SimpleDotOp) {
+  HloComputation::Builder builder(TestName());
+  DotTestSpec spec = GetParam();
+
+  auto param_shape = ShapeUtil::MakeShape(spec.primitive_type, {128, 128});
+
+  HloInstruction* lhs = builder.AddInstruction(
+      HloInstruction::CreateParameter(0, param_shape, "input"));
+  HloInstruction* rhs = builder.AddInstruction(
+      HloInstruction::CreateParameter(1, param_shape, "input"));
+
+  builder.AddInstruction(
+      HloInstruction::CreateCanonicalDot(param_shape, lhs, rhs));
+  CompileAndCheck(builder.Build(), spec.filecheck_lines);
+}
+
+TEST_P(CpuEigenDotOperationTest, DotTransposeOp) {
+  HloComputation::Builder builder(TestName());
+  DotTestSpec spec = GetParam();
+
+  auto param_shape = ShapeUtil::MakeShape(spec.primitive_type, {128, 128});
+
+  HloInstruction* lhs = builder.AddInstruction(
+      HloInstruction::CreateParameter(0, param_shape, "input"));
+  HloInstruction* rhs = builder.AddInstruction(
+      HloInstruction::CreateParameter(1, param_shape, "input"));
+  HloInstruction* lhs_transposed = builder.AddInstruction(
+      HloInstruction::CreateTranspose(param_shape, lhs, {1, 0}));
+
+  builder.AddInstruction(
+      HloInstruction::CreateCanonicalDot(param_shape, lhs_transposed, rhs));
+  CompileAndCheck(builder.Build(), spec.filecheck_lines);
+}
+
+std::vector<DotTestSpec> GetDotTestCases() {
+  std::vector<DotTestSpec> result;
+  result.push_back(
+      {F16, R"(CHECK: call void @__xla_cpu_runtime_EigenMatMulF16)"});
+  result.push_back(
+      {F32, R"(CHECK: call void @__xla_cpu_runtime_EigenMatMulF32)"});
+  result.push_back(
+      {F64, R"(CHECK: call void @__xla_cpu_runtime_EigenMatMulF64)"});
+  return result;
+}
+
+INSTANTIATE_TEST_CASE_P(CpuEigenDotOperationTestInstantiation,
+                        CpuEigenDotOperationTest,
+                        ::testing::ValuesIn(GetDotTestCases()),
+                        DotTestSpecToString);
+
+}  // namespace
+}  // namespace cpu
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_external_constants_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_external_constants_test.cc
new file mode 100644
index 0000000000..ed8f375bd6
--- /dev/null
+++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_external_constants_test.cc
@@ -0,0 +1,73 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include <memory>
+#include <utility>
+
+#include "tensorflow/compiler/xla/array2d.h"
+#include "tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h"
+#include "tensorflow/compiler/xla/service/hlo_computation.h"
+#include "tensorflow/compiler/xla/service/hlo_instruction.h"
+#include "tensorflow/compiler/xla/service/hlo_module.h"
+#include "tensorflow/compiler/xla/shape_util.h"
+#include "tensorflow/compiler/xla/tests/filecheck.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace xla {
+namespace cpu {
+namespace {
+class CpuExternalConstantsTest : public CpuCodegenTest {
+ public:
+  void TestWithArray(int64 rows, int64 cols, const char* filecheck_pattern) {
+    HloComputation::Builder builder(TestName());
+
+    Array2D<float> backing_array(rows, cols);
+    backing_array.FillUnique();
+
+    auto shape = ShapeUtil::MakeShape(F32, {rows, cols});
+
+    HloInstruction* constant =
+        builder.AddInstruction(HloInstruction::CreateConstant(
+            Literal::CreateR2FromArray2D(backing_array)));
+    HloInstruction* param =
+        builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "x"));
+    builder.AddInstruction(
+        HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param, constant));
+
+    std::unique_ptr<HloModule> module = CreateNewModule();
+    module->AddEntryComputation(builder.Build());
+
+    CompileAndVerifyIr(std::move(module), filecheck_pattern,
+                       /*match_optimized_ir=*/false);
+  }
+};
+
+TEST_F(CpuExternalConstantsTest, Basic) {
+  TestWithArray(/*rows=*/1024, /*cols=*/1024, R"(
+CHECK: @constant_global_0 = external constant [1024 x [1024 x float]], align 16
+)");
+}
+
+TEST_F(CpuExternalConstantsTest, BasicNegative) {
+  // The constant array in this test case is small enough that there is no need
+  // to externalize it.
+  TestWithArray(/*rows=*/4, /*cols=*/4, R"(
+CHECK-NOT: @constant_global_0 = external constant [4 x [4 x float]], align 8
+CHECK: @0 = private constant [4 x [4 x float]] {{.*}}, align 8
+)");
+}
+}  // namespace
+}  // namespace cpu
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_fusion_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_fusion_test.cc
new file mode 100644
index 0000000000..23e7a3de4d
--- /dev/null
+++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_fusion_test.cc
@@ -0,0 +1,330 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/compiler/xla/literal_util.h"
+#include "tensorflow/compiler/xla/ptr_util.h"
+#include "tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.h"
+#include "tensorflow/compiler/xla/service/hlo_computation.h"
+#include "tensorflow/compiler/xla/service/hlo_instruction.h"
+#include "tensorflow/compiler/xla/service/hlo_module.h"
+#include "tensorflow/compiler/xla/service/hlo_opcode.h"
+#include "tensorflow/compiler/xla/shape_util.h"
+#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
+#include "tensorflow/compiler/xla/tests/literal_test_util.h"
+#include "tensorflow/compiler/xla/xla_data.pb.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace xla {
+namespace cpu {
+namespace {
+
+class CpuFusionTest : public HloTestBase {
+ protected:
+  CpuFusionTest() {}
+
+  ErrorSpec error_spec_{0.0001, 1e-5};
+};
+
+TEST_F(CpuFusionTest, FuseTwoElementwiseOps) {
+  auto builder = HloComputation::Builder(TestName());
+  auto input_literal1 = Literal::CreateR1<float>({1.0, 2.0, 3.0});
+  auto input_literal2 = Literal::CreateR1<float>({-2.0, -42.0, 2.0});
+  Shape vshape = input_literal1->shape();
+
+  auto input1 = builder.AddInstruction(
+      HloInstruction::CreateConstant(std::move(input_literal1)));
+  auto input2 = builder.AddInstruction(
+      HloInstruction::CreateConstant(std::move(input_literal2)));
+
+  auto add1 = builder.AddInstruction(
+      HloInstruction::CreateBinary(vshape, HloOpcode::kAdd, input1, input2));
+  builder.AddInstruction(
+      HloInstruction::CreateUnary(vshape, HloOpcode::kNegate, add1));
+
+  auto module = CreateNewModule();
+  module->AddEntryComputation(builder.Build());
+
+  CpuInstructionFusion fusion;
+  EXPECT_TRUE(fusion.Run(module.get()).ValueOrDie());
+
+  // The computation root instruction was fused. Verify the fusion instruction
+  // is now the root.
+  auto computation = module->entry_computation();
+  auto fusion_instruction = computation->root_instruction();
+  EXPECT_EQ(HloOpcode::kFusion, fusion_instruction->opcode());
+  EXPECT_EQ(HloOpcode::kNegate,
+            fusion_instruction->fused_expression_root()->opcode());
+  // There should be four fused instructions: 2 parameters, the add, and the
+  // negate.
+  EXPECT_EQ(4, fusion_instruction->fused_instruction_count());
+
+  // Compile and execute the computation.
+  auto result = ExecuteAndTransfer(std::move(module), {});
+
+  // Check the output correctness.
+  LiteralTestUtil::ExpectR1Near<float>({1.0, 40.0, -5.0}, *result, error_spec_);
+}
+
+TEST_F(CpuFusionTest, FuseElementwiseOpChain) {
+  auto builder = HloComputation::Builder(TestName());
+  auto input_literal = Literal::CreateR1<float>({-1.5, -2.5, -3.0});
+  Shape vshape = input_literal->shape();
+
+  auto input = builder.AddInstruction(
+      HloInstruction::CreateConstant(std::move(input_literal)));
+  auto negate = builder.AddInstruction(
+      HloInstruction::CreateUnary(vshape, HloOpcode::kNegate, input));
+  auto ceil = builder.AddInstruction(
+      HloInstruction::CreateUnary(vshape, HloOpcode::kCeil, negate));
+  auto exp = builder.AddInstruction(
+      HloInstruction::CreateUnary(vshape, HloOpcode::kExp, ceil));
+  auto floor = builder.AddInstruction(
+      HloInstruction::CreateUnary(vshape, HloOpcode::kFloor, exp));
+  auto two = builder.AddInstruction(
+      HloInstruction::CreateConstant(Literal::CreateR0<float>(2.0)));
+  builder.AddInstruction(
+      HloInstruction::CreateBinary(vshape, HloOpcode::kMultiply, two, floor));
+
+  auto module = CreateNewModule();
+  module->AddEntryComputation(builder.Build());
+
+  CpuInstructionFusion fusion;
+  EXPECT_TRUE(fusion.Run(module.get()).ValueOrDie());
+
+  // The computation root instruction was fused. Verify the fusion instruction
+  // is now the root.
+  auto computation = module->entry_computation();
+  auto fusion_instruction = computation->root_instruction();
+  EXPECT_EQ(HloOpcode::kFusion, fusion_instruction->opcode());
+  EXPECT_EQ(HloOpcode::kMultiply,
+            fusion_instruction->fused_expression_root()->opcode());
+  // There should be 7 fused instructions: 2 parameters and the fused
+  // operations.
+  EXPECT_EQ(7, fusion_instruction->fused_instruction_count());
+
+  // Compile and execute the computation.
+  auto result = ExecuteAndTransfer(std::move(module), {});
+
+  // Check the output correctness.
+  LiteralTestUtil::ExpectR1Near<float>({14.0, 40.0, 40.0}, *result,
+                                       error_spec_);
+}
+
+TEST_F(CpuFusionTest, ElementwiseOpChainWithNonfusableInstruction) {
+  // Test a chain of fusable ops with a non-fusable op (a reduce) thrown in the
+  // middle.
+  auto module = CreateNewModule();
+  auto builder = HloComputation::Builder(TestName());
+  auto input_literal = Literal::CreateR1<float>({-1.5, -2.5, -3.0});
+  Shape vshape = input_literal->shape();
+
+  auto input = builder.AddInstruction(
+      HloInstruction::CreateConstant(std::move(input_literal)));
+  auto negate = builder.AddInstruction(
+      HloInstruction::CreateUnary(vshape, HloOpcode::kNegate, input));
+  auto ceil = builder.AddInstruction(
+      HloInstruction::CreateUnary(vshape, HloOpcode::kCeil, negate));
+
+  auto cshape = ShapeUtil::MakeShape(F32, {6});
+  auto concatenate = builder.AddInstruction(
+      HloInstruction::CreateConcatenate(cshape, {ceil, ceil}, /*dimension=*/0));
+
+  // Build an x+y computation to use in a reduce.
+  Shape r0f32 = ShapeUtil::MakeShape(F32, {});
+  auto embedded_builder = HloComputation::Builder("f32+f32");
+  embedded_builder.AddInstruction(HloInstruction::CreateBinary(
+      r0f32, HloOpcode::kAdd,
+      embedded_builder.AddInstruction(
+          HloInstruction::CreateParameter(0, r0f32, "x")),
+      embedded_builder.AddInstruction(
+          HloInstruction::CreateParameter(1, r0f32, "y"))));
+  auto add_f32 = module->AddEmbeddedComputation(embedded_builder.Build());
+
+  // This is a nop reduction.
+  auto reduce = builder.AddInstruction(HloInstruction::CreateReduce(
+      cshape,
+      builder.AddInstruction(HloInstruction::CreateReshape(
+          ShapeUtil::MakeShape(F32, {6, 1}), concatenate)),
+      /*init_value=*/
+      builder.AddInstruction(
+          HloInstruction::CreateConstant(Literal::CreateR0<float>(0))),
+      /*dimensions_to_reduce=*/{1}, add_f32));
+
+  auto exp = builder.AddInstruction(
+      HloInstruction::CreateUnary(cshape, HloOpcode::kExp, reduce));
+  auto floor = builder.AddInstruction(
+      HloInstruction::CreateUnary(cshape, HloOpcode::kFloor, exp));
+  auto two = builder.AddInstruction(
+      HloInstruction::CreateConstant(Literal::CreateR0<float>(2.0)));
+  builder.AddInstruction(
+      HloInstruction::CreateBinary(cshape, HloOpcode::kMultiply, two, floor));
+
+  module->AddEntryComputation(builder.Build());
+
+  CpuInstructionFusion fusion;
+  EXPECT_TRUE(fusion.Run(module.get()).ValueOrDie());
+
+  // The computation root instruction was fused. Verify the fusion instruction
+  // is now the root.
+  auto computation = module->entry_computation();
+
+  auto fusion_instruction1 = computation->root_instruction();
+  EXPECT_EQ(HloOpcode::kFusion, fusion_instruction1->opcode());
+  EXPECT_EQ(HloOpcode::kMultiply,
+            fusion_instruction1->fused_expression_root()->opcode());
+  // There should be 5 fused instructions in the root fusion instruction: 2
+  // parameters, multiply, floor, and exp.
+  EXPECT_EQ(5, fusion_instruction1->fused_instruction_count())
+      << fusion_instruction1->fused_instructions_computation()->ToString();
+
+  auto fusion_instruction2 = reduce->operand(0);
+  EXPECT_EQ(HloOpcode::kFusion, fusion_instruction1->opcode());
+  EXPECT_EQ(HloOpcode::kReshape,
+            fusion_instruction2->fused_expression_root()->opcode());
+  // There should be 5 fused instructions in the second fusion instruction: 1
+  // parameter, negate, ceil, concat, and reshape.
+  EXPECT_EQ(5, fusion_instruction2->fused_instruction_count())
+      << fusion_instruction2->fused_instructions_computation()->ToString();
+
+  // Compile and execute the computation.
+  auto result = ExecuteAndTransfer(std::move(module), {});
+
+  // Check the output correctness.
+  LiteralTestUtil::ExpectR1Near<float>({14.0, 40.0, 40.0, 14.0, 40.0, 40.0},
+                                       *result, error_spec_);
+}
+
+TEST_F(CpuFusionTest, TestOperandOrderToAvoidDuplication) {
+  // Test that the operands of an instruction to be fused are considered in the
+  // proper order to avoid duplication. Test input:
+  //
+  //   constant = {...}
+  //   negate    = neg(constant)
+  //   ceil      = ceil(negate)
+  //   add1      = add(negate, ceil)
+  //   add2      = add(ceil, negate)
+  //
+  // In this example, the operands of both add1 and add2 should be fused in the
+  // order {ceil, negate} even though they have different orders in their
+  // operand vectors. Test for this problem by counting the number of nodes in
+  // each fusion instruction to ensure that negate is not duplicated.
+  auto builder = HloComputation::Builder(TestName());
+  auto input_literal = Literal::CreateR1<float>({1.0, 2.0, 3.0});
+  Shape vshape = input_literal->shape();
+
+  auto constant = builder.AddInstruction(
+      HloInstruction::CreateConstant(std::move(input_literal)));
+  auto negate = builder.AddInstruction(
+      HloInstruction::CreateUnary(vshape, HloOpcode::kNegate, constant));
+  auto ceil = builder.AddInstruction(
+      HloInstruction::CreateUnary(vshape, HloOpcode::kCeil, negate));
+
+  auto add1 = builder.AddInstruction(
+      HloInstruction::CreateBinary(vshape, HloOpcode::kMultiply, negate, ceil));
+  auto add2 = builder.AddInstruction(
+      HloInstruction::CreateBinary(vshape, HloOpcode::kMultiply, ceil, negate));
+
+  // Tie together the two adds with a tuple to create a single root.
+  auto result =
+      builder.AddInstruction(HloInstruction::CreateTuple({add1, add2}));
+
+  // Create computation and module.
+  auto module = CreateNewModule();
+  module->AddEntryComputation(builder.Build());
+
+  // Run fusion.
+  CpuInstructionFusion fusion;
+  EXPECT_TRUE(fusion.Run(module.get()).ValueOrDie());
+
+  auto fusion1 = result->operand(0);
+  auto fusion2 = result->operand(1);
+  EXPECT_EQ(HloOpcode::kFusion, fusion1->opcode());
+  EXPECT_EQ(HloOpcode::kFusion, fusion2->opcode());
+
+  // Each fusion instruction should have 4 fused instruction inside: add, ceil,
+  // negate, and the fused parameter.
+  EXPECT_EQ(4, fusion1->fused_instruction_count());
+  EXPECT_EQ(4, fusion2->fused_instruction_count());
+
+  // Each fusion instruction should have one parameter and the parameter should
+  // be the constant.
+  EXPECT_EQ(1, fusion1->operand_count());
+  EXPECT_EQ(constant, fusion1->operand(0));
+  EXPECT_EQ(1, fusion2->operand_count());
+  EXPECT_EQ(constant, fusion2->operand(0));
+}
+
+TEST_F(CpuFusionTest, DoNotDuplicateExpensiveOps) {
+  // Verify that expensive operations will not be fused if the fusion results in
+  // duplication. Test code:
+  //
+  //   constant = 42.0
+  //   exp1 = exp(constant)
+  //   negate1 = negate(exp1)
+  //   exp2 = exp(constant)
+  //   negate2 = negate(exp2)
+  //   tuple = tuple(negate1, negate2, exp2)
+  //
+  // exp1 should be fused down into negate1, but exp2 will not be fused into
+  // negate2 because this will result in duplication of the expensive exp
+  // computation. The duplication is caused by the other use of exp2 in the
+  // tuple.
+  auto builder = HloComputation::Builder(TestName());
+  auto input_literal1 = Literal::CreateR1<float>({1.0, 2.0, 3.0});
+  auto input_literal2 = Literal::CreateR1<float>({-2.0, -42.0, 2.0});
+  auto constant = builder.AddInstruction(
+      HloInstruction::CreateConstant(Literal::CreateR0<float>(42.0)));
+  Shape shape = constant->shape();
+
+  auto exp1 = builder.AddInstruction(
+      HloInstruction::CreateUnary(shape, HloOpcode::kExp, constant));
+  auto negate1 = builder.AddInstruction(
+      HloInstruction::CreateUnary(shape, HloOpcode::kNegate, exp1));
+
+  auto exp2 = builder.AddInstruction(
+      HloInstruction::CreateUnary(shape, HloOpcode::kExp, constant));
+  auto negate2 = builder.AddInstruction(
+      HloInstruction::CreateUnary(shape, HloOpcode::kNegate, exp2));
+
+  auto tuple = builder.AddInstruction(
+      HloInstruction::CreateTuple({negate1, negate2, exp2}));
+
+  auto module = CreateNewModule();
+  module->AddEntryComputation(builder.Build());
+
+  CpuInstructionFusion fusion;
+  EXPECT_TRUE(fusion.Run(module.get()).ValueOrDie());
+
+  // The only fusion instruction should be operand 0 of the tuple (formerly
+  // negate1).
+  EXPECT_EQ(HloOpcode::kFusion, tuple->operand(0)->opcode());
+  EXPECT_EQ(HloOpcode::kNegate, tuple->operand(1)->opcode());
+  EXPECT_EQ(HloOpcode::kExp, tuple->operand(2)->opcode());
+
+  auto fusion_inst = tuple->operand(0);
+  // There should be three fused instructions: negate2, exp2, and the fused
+  // parameter.
+  EXPECT_EQ(3, fusion_inst->fused_instruction_count());
+  EXPECT_EQ(1, fusion_inst->operand_count());
+  EXPECT_EQ(constant, fusion_inst->operand(0));
+}
+
+}  // namespace
+}  // namespace cpu
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_intrinsic_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_intrinsic_test.cc
new file mode 100644
index 0000000000..973aac8766
--- /dev/null
+++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_intrinsic_test.cc
@@ -0,0 +1,151 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include <algorithm>
+#include <cctype>
+#include <string>
+
+#include "tensorflow/compiler/xla/service/cpu/cpu_compiler.h"
+#include "tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h"
+#include "tensorflow/compiler/xla/service/hlo_computation.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace xla {
+namespace cpu {
+namespace {
+
+const char* const kTriple_x86_64 = "x86_64-pc-linux";
+const char* const kTriple_android_arm = "armv7-none-android";
+
+struct IntrinsicTestSpec {
+  HloOpcode opcode;
+  tensorflow::StringPiece triple;
+  tensorflow::StringPiece features;
+  tensorflow::StringPiece check_lines;
+};
+
+// Tests that unary functions get lowered using intrinsic calls.
+class CpuUnaryIntrinsicTest
+    : public CpuCodegenTest,
+      public ::testing::WithParamInterface<IntrinsicTestSpec> {
+ public:
+  static string Name(const ::testing::TestParamInfo<IntrinsicTestSpec>& info) {
+    auto spec = info.param;
+
+    string opcode = HloOpcodeString(spec.opcode);
+    opcode[0] = toupper(opcode[0]);
+
+    string triple{spec.triple.data(), spec.triple.size()};
+    if (triple == kTriple_x86_64) {
+      triple = "x86_64";
+    } else if (triple == kTriple_android_arm) {
+      triple = "android_arm";
+    } else {
+      triple = "Unknown";
+    }
+
+    string features{spec.features.data(), spec.features.size()};
+    if (!features.empty()) {
+      std::replace_if(features.begin(), features.end(),
+                      [](char c) { return c != '_' && !isalnum(c); }, '_');
+    } else {
+      features = "";
+    }
+
+    return tensorflow::strings::StrCat(opcode.c_str(), "_On_", triple.c_str(),
+                                       features.empty() ? "" : "_With",
+                                       features.c_str());
+  }
+};
+
+// Creates a module with a call to the unary op, and tests if the
+// compiler replaced it with a call to the intrinsic.
+TEST_P(CpuUnaryIntrinsicTest, DoIt) {
+  HloComputation::Builder builder(TestName());
+  IntrinsicTestSpec spec = GetParam();
+
+  auto param_shape = ShapeUtil::MakeShape(F32, {1024});
+  HloInstruction* param = builder.AddInstruction(
+      HloInstruction::CreateParameter(0, param_shape, "input"));
+  builder.AddInstruction(
+      HloInstruction::CreateUnary(param_shape, spec.opcode, param));
+  std::unique_ptr<HloComputation> computation = builder.Build();
+
+  string triple{spec.triple.data(), spec.triple.size()};
+  string features{spec.features.data(), spec.features.size()};
+
+  CpuAotCompilationOptions options{
+      /*triple=*/triple, /*cpu_name=*/"", /*features=*/features,
+      /*entry_point_name=*/"entry",
+      /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static};
+
+  auto hlo_module = CreateNewModule();
+  hlo_module->AddEntryComputation(std::move(computation));
+
+  string check_lines{spec.check_lines.data(), spec.check_lines.size()};
+
+  CompileAheadOfTimeAndVerifyIr(std::move(hlo_module), options, check_lines,
+                                /*match_optimized_ir=*/true);
+}
+
+IntrinsicTestSpec CpuUnaryIntrinsicTestCases[] = {
+    // The intrinsics are always inlined, so we match a line from it instead of
+    // a function call.
+
+    IntrinsicTestSpec{
+        HloOpcode::kExp, kTriple_x86_64, "",
+        R"(CHECK: fmul fast <4 x float> <float 0xBF2BD01060000000, float 0xBF2BD01060000000, float 0xBF2BD01060000000, float 0xBF2BD01060000000>)"},
+
+    IntrinsicTestSpec{
+        HloOpcode::kExp, kTriple_x86_64, "+avx",
+        R"(CHECK: fmul fast <8 x float> <float 0xBF2BD01060000000, float 0xBF2BD01060000000, float 0xBF2BD01060000000, float 0xBF2BD01060000000, float 0xBF2BD01060000000, float 0xBF2BD01060000000, float 0xBF2BD01060000000, float 0xBF2BD01060000000>)"},
+
+    IntrinsicTestSpec{
+        HloOpcode::kExp, kTriple_android_arm, "+neon",
+        R"(CHECK: fmul fast <4 x float> <float 0xBF2BD01060000000, float 0xBF2BD01060000000, float 0xBF2BD01060000000, float 0xBF2BD01060000000>)"},
+
+    IntrinsicTestSpec{
+        HloOpcode::kTanh, kTriple_x86_64, "",
+        R"(CHECK: fcmp fast uge <4 x float> %wide.load, <float -9.000000e+00, float -9.000000e+00, float -9.000000e+00, float -9.000000e+00>)"},
+
+    IntrinsicTestSpec{
+        HloOpcode::kTanh, kTriple_x86_64, "+avx",
+        R"(CHECK: fcmp fast uge <8 x float> %wide.load, <float -9.000000e+00, float -9.000000e+00, float -9.000000e+00, float -9.000000e+00, float -9.000000e+00, float -9.000000e+00, float -9.000000e+00, float -9.000000e+00>)"},
+
+    IntrinsicTestSpec{
+        HloOpcode::kTanh, kTriple_android_arm, "",
+        R"(CHECK: fcmp fast uge <4 x float> %wide.load, <float -9.000000e+00, float -9.000000e+00, float -9.000000e+00, float -9.000000e+00>)"},
+
+    IntrinsicTestSpec{
+        HloOpcode::kLog, kTriple_x86_64, "",
+        R"(CHECK: fadd fast <4 x float> <float 0x3FBDE4A340000000, float 0x3FBDE4A340000000, float 0x3FBDE4A340000000, float 0x3FBDE4A340000000>)"},
+
+    IntrinsicTestSpec{
+        HloOpcode::kLog, kTriple_x86_64, "+avx",
+        R"(CHECK: fadd fast <8 x float> <float 0x3FBDE4A340000000, float 0x3FBDE4A340000000, float 0x3FBDE4A340000000, float 0x3FBDE4A340000000, float 0x3FBDE4A340000000, float 0x3FBDE4A340000000, float 0x3FBDE4A340000000, float 0x3FBDE4A340000000>)"},
+
+    IntrinsicTestSpec{
+        HloOpcode::kLog, kTriple_android_arm, "",
+        R"(CHECK: fadd fast <4 x float> <float 0x3FBDE4A340000000, float 0x3FBDE4A340000000, float 0x3FBDE4A340000000, float 0x3FBDE4A340000000>)"}};
+
+INSTANTIATE_TEST_CASE_P(CpuUnaryIntrinsicTestInstantiation,
+                        CpuUnaryIntrinsicTest,
+                        ::testing::ValuesIn(CpuUnaryIntrinsicTestCases),
+                        CpuUnaryIntrinsicTest::Name);
+
+}  // namespace
+}  // namespace cpu
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_noalias_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_noalias_test.cc
new file mode 100644
index 0000000000..3b6b0ed740
--- /dev/null
+++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_noalias_test.cc
@@ -0,0 +1,136 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include <memory>
+#include <utility>
+
+#include "llvm/IR/Module.h"
+#include "tensorflow/compiler/xla/literal_util.h"
+#include "tensorflow/compiler/xla/ptr_util.h"
+#include "tensorflow/compiler/xla/service/buffer_assignment.h"
+#include "tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h"
+#include "tensorflow/compiler/xla/service/hlo_computation.h"
+#include "tensorflow/compiler/xla/service/hlo_instruction.h"
+#include "tensorflow/compiler/xla/service/hlo_module.h"
+#include "tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h"
+#include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h"
+#include "tensorflow/compiler/xla/shape_util.h"
+#include "tensorflow/compiler/xla/tests/filecheck.h"
+#include "tensorflow/compiler/xla/xla_data.pb.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace xla {
+namespace cpu {
+
+class CpuNoAliasTest : public CpuCodegenTest {};
+
+// Creates a simple HLO ir_module (runs concat(concat(x, y), x)), and then
+// inspects the aliasing information for loads to its buffers.
+TEST_F(CpuNoAliasTest, Concat) {
+  HloComputation::Builder builder(TestName());
+
+  std::unique_ptr<Literal> literal =
+      Literal::CreateR2<float>({{1.0, 2.0}, {3.0, 4.0}});
+  auto param_shape = ShapeUtil::MakeShape(F32, {2, 2});
+  HloInstruction* param_x = builder.AddInstruction(
+      HloInstruction::CreateParameter(0, param_shape, "x"));
+  HloInstruction* param_y = builder.AddInstruction(
+      HloInstruction::CreateParameter(1, param_shape, "y"));
+  HloInstruction* concat1 =
+      builder.AddInstruction(HloInstruction::CreateConcatenate(
+          ShapeUtil::MakeShape(F32, {2, 4}), {param_x, param_y}, 1));
+  HloInstruction* concat2 =
+      builder.AddInstruction(HloInstruction::CreateConcatenate(
+          ShapeUtil::MakeShape(F32, {2, 6}), {concat1, param_x}, 1));
+
+  std::unique_ptr<HloComputation> computation = builder.Build();
+
+  auto hlo_module = CreateNewModule();
+  hlo_module->AddEntryComputation(std::move(computation));
+
+  // Now that we have an HLO module, build an llvm_ir::AliasAnalysis for it.
+  auto status_or_buffer_assn = BufferAssigner::Run(
+      hlo_module.get(), MakeUnique<DependencyHloOrdering>(hlo_module.get()),
+      backend().compiler()->BufferSizeBytesFunction(),
+      [](LogicalBuffer::Color) { return /*alignment=*/1; });
+  ASSERT_EQ(status_or_buffer_assn.status(), Status::OK());
+
+  llvm::LLVMContext context;
+  llvm_ir::AliasAnalysis aa(*hlo_module, *status_or_buffer_assn.ValueOrDie(),
+                            &context);
+
+  // Construct an LLVM module containing loads that we annotate as being from
+  // the buffers in the HLO module.  We'll inspect these loads to ensure that
+  // they have the expected alias information.
+  llvm::Module ir_module("test", context);
+  llvm::Function* func = llvm::cast<llvm::Function>(
+      ir_module.getOrInsertFunction("test_fn", llvm::Type::getVoidTy(context)));
+  llvm::BasicBlock* bb = llvm::BasicBlock::Create(context, "body", func);
+  llvm::IRBuilder<> ir_builder(bb);
+  auto* zero = llvm::ConstantInt::get(llvm::Type::getInt32Ty(context), 0);
+  llvm_ir::IrArray::Index zero2D({zero, zero});
+
+  llvm::ArrayType* array2d_type = llvm::ArrayType::get(
+      llvm::ArrayType::get(llvm::Type::getFloatTy(context), 100), 100);
+
+  {
+    llvm::Value* param_x_val =
+        ir_module.getOrInsertGlobal("param_x", array2d_type);
+    llvm_ir::IrArray param_x_array(param_x_val, param_shape);
+    aa.AddAliasingInformationToIrArray(*param_x, &param_x_array);
+    param_x_array.EmitReadArrayElement(zero2D, &ir_builder)
+        ->setName("read_param_x_array");
+  }
+
+  {
+    llvm::Value* concat1_val =
+        ir_module.getOrInsertGlobal("concat1", array2d_type);
+    auto shape = ShapeUtil::MakeShape(F32, {2, 4});
+    llvm_ir::IrArray concat1_array(concat1_val, shape);
+    aa.AddAliasingInformationToIrArray(*concat1, &concat1_array);
+    concat1_array.EmitReadArrayElement(zero2D, &ir_builder)
+        ->setName("read_concat1_array");
+  }
+
+  {
+    llvm::Value* concat2_val =
+        ir_module.getOrInsertGlobal("concat2", array2d_type);
+    auto shape = ShapeUtil::MakeShape(F32, {2, 6});
+    llvm_ir::IrArray concat2_array(concat2_val, shape);
+    aa.AddAliasingInformationToIrArray(*concat2, &concat2_array);
+    concat2_array.EmitReadArrayElement(zero2D, &ir_builder)
+        ->setName("read_concat2_array");
+  }
+
+  // Check the AA info in the loads.
+  const char* filecheck_pattern = R"(
+    CHECK: %read_param_x_array = load {{.*}} !noalias [[param_x_noalias:![0-9]+]]
+    CHECK: %read_concat1_array = load {{.*}} !alias.scope [[concat1_scope:![0-9]+]], !noalias [[concat1_noalias:![0-9]+]]
+    CHECK: %read_concat2_array = load {{.*}} !alias.scope [[concat1_noalias]], !noalias [[concat1_scope]]
+    CHECK-DAG: [[buf_size32:![0-9]+]] = !{!"buffer:{{.*}} size:32
+    CHECK-DAG: [[buf_size48:![0-9]+]] = !{!"buffer:{{.*}} size:48
+    CHECK-DAG: [[param_x_noalias]] = !{[[buf_size32]], [[buf_size48]]}
+    CHECK-DAG: [[concat1_scope]] = !{[[buf_size32]]}
+    CHECK-DAG: [[concat1_noalias]] = !{[[buf_size48]]}
+  )";
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      bool filecheck_match,
+      RunFileCheck(llvm_ir::DumpModuleToString(ir_module), filecheck_pattern));
+  EXPECT_TRUE(filecheck_match);
+}
+
+}  // namespace cpu
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/tests/filecheck.cc b/tensorflow/compiler/xla/tests/filecheck.cc
index a5f6872c46..93d1c921c4 100644
--- a/tensorflow/compiler/xla/tests/filecheck.cc
+++ b/tensorflow/compiler/xla/tests/filecheck.cc
@@ -38,7 +38,7 @@ StatusOr<bool> RunFileCheck(const string& input, const string& pattern) {
   TF_RETURN_IF_ERROR(tensorflow::WriteStringToFile(env, pattern_path, pattern));
 
   // Invoke FileCheck to check whether input matches `pattern`.
-  const char* file_check_path_suffix = "external/llvm/FileCheck";
+  const char* file_check_path_suffix = "org_tensorflow/external/llvm/FileCheck";
   string file_check_path;
   if (const char* test_srcdir = getenv("TEST_SRCDIR")) {
     file_check_path = JoinPath(test_srcdir, file_check_path_suffix);
@@ -66,6 +66,11 @@ StatusOr<bool> RunFileCheck(const string& input, const string& pattern) {
   // the error message generated by FileCheck and the inputs.
   bool succeeded = (exit_status == 0);
   if (!succeeded) {
+    LOG(WARNING) << "Tried to execute FileCheck at " << file_check_path;
+    if (!env->FileExists(file_check_path).ok()) {
+      LOG(WARNING) << "NOTE: FileCheck binary does not exist!";
+    }
+
     LOG(WARNING) << "FileCheck error: " << standard_error;
     LOG(WARNING) << "FileCheck input was:";
     XLA_LOG_LINES(tensorflow::WARNING, input);
-- 
GitLab


From a4343eb6cd10fa6c0fdfaa18585706d78e8c9d26 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 03:25:26 -0700
Subject: [PATCH 087/395] Protocol buffer classes now list their fields in
 dir(cls)

PiperOrigin-RevId: 194917415
---
 .../tensorflow.-attr-value.-list-value.pbtxt  |  36 +++---
 .../api/golden/tensorflow.-attr-value.pbtxt   |  48 +++----
 ...ow.-config-proto.-device-count-entry.pbtxt |   8 +-
 .../api/golden/tensorflow.-config-proto.pbtxt |  72 +++++------
 .../tools/api/golden/tensorflow.-event.pbtxt  |  36 +++---
 .../golden/tensorflow.-g-p-u-options.pbtxt    |  48 +++----
 .../api/golden/tensorflow.-graph-def.pbtxt    |  16 +--
 .../golden/tensorflow.-graph-options.pbtxt    |  44 +++----
 .../golden/tensorflow.-histogram-proto.pbtxt  |  36 +++---
 .../api/golden/tensorflow.-log-message.pbtxt  |  16 +--
 ...meta-graph-def.-collection-def-entry.pbtxt |   8 +-
 ...rflow.-meta-graph-def.-meta-info-def.pbtxt |  32 ++---
 ...-meta-graph-def.-signature-def-entry.pbtxt |   8 +-
 .../golden/tensorflow.-meta-graph-def.pbtxt   |  40 +++---
 ...nsorflow.-name-attr-list.-attr-entry.pbtxt |   8 +-
 .../golden/tensorflow.-name-attr-list.pbtxt   |  12 +-
 .../tensorflow.-node-def.-attr-entry.pbtxt    |   8 +-
 .../api/golden/tensorflow.-node-def.pbtxt     |  28 ++--
 .../tensorflow.-optimizer-options.pbtxt       |  44 +++----
 .../api/golden/tensorflow.-run-metadata.pbtxt |  16 +--
 .../api/golden/tensorflow.-run-options.pbtxt  |  36 +++---
 .../api/golden/tensorflow.-session-log.pbtxt  |  24 ++--
 ...rflow.-summary-metadata.-plugin-data.pbtxt |  12 +-
 .../golden/tensorflow.-summary-metadata.pbtxt |  20 +--
 .../golden/tensorflow.-summary.-audio.pbtxt   |  28 ++--
 .../golden/tensorflow.-summary.-image.pbtxt   |  24 ++--
 .../golden/tensorflow.-summary.-value.pbtxt   |  40 +++---
 .../api/golden/tensorflow.-summary.pbtxt      |   8 +-
 .../tensorflow.-tensor-info.-coo-sparse.pbtxt |  16 +--
 .../api/golden/tensorflow.-tensor-info.pbtxt  |  24 ++--
 ...flow.profiler.-advice-proto.-checker.pbtxt |   4 +-
 ...ofiler.-advice-proto.-checkers-entry.pbtxt |   8 +-
 .../tensorflow.profiler.-advice-proto.pbtxt   |   8 +-
 ...graph-node-proto.-input-shapes-entry.pbtxt |   8 +-
 ...ensorflow.profiler.-graph-node-proto.pbtxt | 120 +++++++++---------
 ...low.profiler.-multi-graph-node-proto.pbtxt |  92 +++++++-------
 ...er.-op-log-proto.-id-to-string-entry.pbtxt |   8 +-
 .../tensorflow.profiler.-op-log-proto.pbtxt   |  12 +-
 .../golden/tensorflow.summary.-event.pbtxt    |  36 +++---
 .../tensorflow.summary.-session-log.pbtxt     |  24 ++--
 ...sorflow.summary.-summary-description.pbtxt |   4 +-
 .../tensorflow.summary.-summary.-audio.pbtxt  |  28 ++--
 .../tensorflow.summary.-summary.-image.pbtxt  |  24 ++--
 .../tensorflow.summary.-summary.-value.pbtxt  |  40 +++---
 .../golden/tensorflow.summary.-summary.pbtxt  |   8 +-
 ...sorflow.summary.-tagged-run-metadata.pbtxt |   8 +-
 .../golden/tensorflow.train.-bytes-list.pbtxt |   4 +-
 .../tensorflow.train.-cluster-def.pbtxt       |   4 +-
 .../golden/tensorflow.train.-example.pbtxt    |   4 +-
 .../tensorflow.train.-feature-list.pbtxt      |   4 +-
 ...n.-feature-lists.-feature-list-entry.pbtxt |   8 +-
 .../tensorflow.train.-feature-lists.pbtxt     |   8 +-
 .../golden/tensorflow.train.-feature.pbtxt    |  16 +--
 ...rflow.train.-features.-feature-entry.pbtxt |   8 +-
 .../golden/tensorflow.train.-features.pbtxt   |   8 +-
 .../golden/tensorflow.train.-float-list.pbtxt |   4 +-
 .../golden/tensorflow.train.-int64-list.pbtxt |   4 +-
 ...nsorflow.train.-job-def.-tasks-entry.pbtxt |   8 +-
 .../golden/tensorflow.train.-job-def.pbtxt    |  12 +-
 .../golden/tensorflow.train.-saver-def.pbtxt  |  34 ++---
 .../tensorflow.train.-sequence-example.pbtxt  |  12 +-
 .../golden/tensorflow.train.-server-def.pbtxt |  28 ++--
 62 files changed, 697 insertions(+), 697 deletions(-)

diff --git a/tensorflow/tools/api/golden/tensorflow.-attr-value.-list-value.pbtxt b/tensorflow/tools/api/golden/tensorflow.-attr-value.-list-value.pbtxt
index 0fb1aaba28..004d716954 100644
--- a/tensorflow/tools/api/golden/tensorflow.-attr-value.-list-value.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-attr-value.-list-value.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.AttrValue.ListValue"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.attr_value_pb2.ListValue\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "B_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -15,32 +11,36 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "FUNC_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "b"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "f"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "F_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "func"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "I_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "i"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SHAPE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "s"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "S_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "shape"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TENSOR_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "tensor"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TYPE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "type"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-attr-value.pbtxt b/tensorflow/tools/api/golden/tensorflow.-attr-value.pbtxt
index e7a3a1f02f..2996e02483 100644
--- a/tensorflow/tools/api/golden/tensorflow.-attr-value.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-attr-value.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.AttrValue"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.attr_value_pb2.AttrValue\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "B_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -15,44 +11,48 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "FUNC_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "ListValue"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+  }
+  member {
+    name: "b"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "F_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "f"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "I_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "func"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "LIST_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "i"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "ListValue"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "list"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "PLACEHOLDER_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "placeholder"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SHAPE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "s"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "S_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "shape"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TENSOR_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "tensor"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TYPE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "type"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-config-proto.-device-count-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.-config-proto.-device-count-entry.pbtxt
index 29bb3be35c..c7022e7593 100644
--- a/tensorflow/tools/api/golden/tensorflow.-config-proto.-device-count-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-config-proto.-device-count-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "KEY_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "key"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-config-proto.pbtxt b/tensorflow/tools/api/golden/tensorflow.-config-proto.pbtxt
index 009d64aed0..ca9530de85 100644
--- a/tensorflow/tools/api/golden/tensorflow.-config-proto.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-config-proto.pbtxt
@@ -3,76 +3,76 @@ tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.config_pb2.ConfigProto\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "ALLOW_SOFT_PLACEMENT_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "CLUSTER_DEF_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "DeviceCountEntry"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "DEVICE_COUNT_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "allow_soft_placement"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "DEVICE_FILTERS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "cluster_def"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "DeviceCountEntry"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "device_count"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "device_filters"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "GPU_OPTIONS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "gpu_options"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "GRAPH_OPTIONS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "graph_options"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "INTER_OP_PARALLELISM_THREADS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "inter_op_parallelism_threads"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "INTRA_OP_PARALLELISM_THREADS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "intra_op_parallelism_threads"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "ISOLATE_SESSION_STATE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "isolate_session_state"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "LOG_DEVICE_PLACEMENT_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "log_device_placement"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "OPERATION_TIMEOUT_IN_MS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "operation_timeout_in_ms"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "PLACEMENT_PERIOD_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "placement_period"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "RPC_OPTIONS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "rpc_options"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SESSION_INTER_OP_THREAD_POOL_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "session_inter_op_thread_pool"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "USE_PER_SESSION_THREADS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "use_per_session_threads"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-event.pbtxt b/tensorflow/tools/api/golden/tensorflow.-event.pbtxt
index 9bf8c12428..fa2f329a87 100644
--- a/tensorflow/tools/api/golden/tensorflow.-event.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-event.pbtxt
@@ -11,40 +11,40 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "FILE_VERSION_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "file_version"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "GRAPH_DEF_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "graph_def"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "LOG_MESSAGE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "log_message"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "META_GRAPH_DEF_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "meta_graph_def"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SESSION_LOG_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "session_log"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "STEP_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "step"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SUMMARY_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "summary"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TAGGED_RUN_METADATA_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "tagged_run_metadata"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "WALL_TIME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "wall_time"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-g-p-u-options.pbtxt b/tensorflow/tools/api/golden/tensorflow.-g-p-u-options.pbtxt
index 875d802a9c..5119c7fa5b 100644
--- a/tensorflow/tools/api/golden/tensorflow.-g-p-u-options.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-g-p-u-options.pbtxt
@@ -3,52 +3,52 @@ tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.config_pb2.GPUOptions\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "ALLOCATOR_TYPE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "ALLOW_GROWTH_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "Experimental"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "DEFERRED_DELETION_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "allocator_type"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "EXPERIMENTAL_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "allow_growth"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "Experimental"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "deferred_deletion_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "experimental"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "FORCE_GPU_COMPATIBLE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "force_gpu_compatible"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "PER_PROCESS_GPU_MEMORY_FRACTION_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "per_process_gpu_memory_fraction"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "POLLING_ACTIVE_DELAY_USECS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "polling_active_delay_usecs"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "POLLING_INACTIVE_DELAY_MSECS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "polling_inactive_delay_msecs"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VISIBLE_DEVICE_LIST_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "visible_device_list"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-graph-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.-graph-def.pbtxt
index 1495e847cb..318a25a092 100644
--- a/tensorflow/tools/api/golden/tensorflow.-graph-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-graph-def.pbtxt
@@ -11,20 +11,20 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "LIBRARY_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "library"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "NODE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "node"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VERSIONS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "version"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VERSION_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "versions"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-graph-options.pbtxt b/tensorflow/tools/api/golden/tensorflow.-graph-options.pbtxt
index 0844f891ca..786d831c70 100644
--- a/tensorflow/tools/api/golden/tensorflow.-graph-options.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-graph-options.pbtxt
@@ -3,48 +3,48 @@ tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.config_pb2.GraphOptions\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "BUILD_COST_MODEL_AFTER_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "BUILD_COST_MODEL_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "build_cost_model"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "ENABLE_BFLOAT16_SENDRECV_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "build_cost_model_after"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "ENABLE_RECV_SCHEDULING_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "enable_bfloat16_sendrecv"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "enable_recv_scheduling"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "INFER_SHAPES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "infer_shapes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "OPTIMIZER_OPTIONS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "optimizer_options"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "PLACE_PRUNED_GRAPH_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "place_pruned_graph"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "REWRITE_OPTIONS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "rewrite_options"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TIMELINE_STEP_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "timeline_step"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-histogram-proto.pbtxt b/tensorflow/tools/api/golden/tensorflow.-histogram-proto.pbtxt
index 2567d2fe60..3eb2d8873a 100644
--- a/tensorflow/tools/api/golden/tensorflow.-histogram-proto.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-histogram-proto.pbtxt
@@ -2,14 +2,6 @@ path: "tensorflow.HistogramProto"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.HistogramProto\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "BUCKET_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
-  member {
-    name: "BUCKET_LIMIT_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -19,24 +11,32 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "MAX_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "bucket"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "bucket_limit"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "max"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "MIN_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "min"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "NUM_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "num"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SUM_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "sum"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SUM_SQUARES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "sum_squares"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-log-message.pbtxt b/tensorflow/tools/api/golden/tensorflow.-log-message.pbtxt
index a43c5eb7e3..760739f4f3 100644
--- a/tensorflow/tools/api/golden/tensorflow.-log-message.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-log-message.pbtxt
@@ -26,18 +26,10 @@ tf_class {
     name: "INFO"
     mtype: "<type \'int\'>"
   }
-  member {
-    name: "LEVEL_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "Level"
     mtype: "<class \'google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper\'>"
   }
-  member {
-    name: "MESSAGE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "UNKNOWN"
     mtype: "<type \'int\'>"
@@ -46,6 +38,14 @@ tf_class {
     name: "WARN"
     mtype: "<type \'int\'>"
   }
+  member {
+    name: "level"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "message"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
   member_method {
     name: "ByteSize"
   }
diff --git a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-collection-def-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-collection-def-entry.pbtxt
index 3572126fbf..69bf5b31a1 100644
--- a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-collection-def-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-collection-def-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "KEY_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "key"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-meta-info-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-meta-info-def.pbtxt
index b0e9831154..8a464f1cac 100644
--- a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-meta-info-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-meta-info-def.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.MetaGraphDef.MetaInfoDef"
 tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.meta_graph_pb2.MetaInfoDef\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "ANY_INFO_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -15,28 +11,32 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "META_GRAPH_VERSION_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "any_info"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "meta_graph_version"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "STRIPPED_DEFAULT_ATTRS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "stripped_default_attrs"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "STRIPPED_OP_LIST_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "stripped_op_list"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TAGS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "tags"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TENSORFLOW_GIT_VERSION_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "tensorflow_git_version"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TENSORFLOW_VERSION_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "tensorflow_version"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-signature-def-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-signature-def-entry.pbtxt
index 48fccac99d..8c5949d067 100644
--- a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-signature-def-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-signature-def-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "KEY_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "key"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.pbtxt
index 3e683a8715..2be0432c00 100644
--- a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.pbtxt
@@ -2,14 +2,6 @@ path: "tensorflow.MetaGraphDef"
 tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.meta_graph_pb2.MetaGraphDef\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "ASSET_FILE_DEF_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
-  member {
-    name: "COLLECTION_DEF_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "CollectionDefEntry"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
@@ -23,28 +15,36 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "GRAPH_DEF_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "MetaInfoDef"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+  }
+  member {
+    name: "SignatureDefEntry"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "META_INFO_DEF_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "asset_file_def"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "MetaInfoDef"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "collection_def"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SAVER_DEF_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "graph_def"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SIGNATURE_DEF_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "meta_info_def"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SignatureDefEntry"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "saver_def"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "signature_def"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-name-attr-list.-attr-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.-name-attr-list.-attr-entry.pbtxt
index 2750bd780c..caf992f5a6 100644
--- a/tensorflow/tools/api/golden/tensorflow.-name-attr-list.-attr-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-name-attr-list.-attr-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "KEY_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "key"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-name-attr-list.pbtxt b/tensorflow/tools/api/golden/tensorflow.-name-attr-list.pbtxt
index d10faf67d0..45ddeece07 100644
--- a/tensorflow/tools/api/golden/tensorflow.-name-attr-list.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-name-attr-list.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.NameAttrList"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.attr_value_pb2.NameAttrList\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "ATTR_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "AttrEntry"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
@@ -19,8 +15,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "attr"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-node-def.-attr-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.-node-def.-attr-entry.pbtxt
index b1b62d60f1..30a9dc69f0 100644
--- a/tensorflow/tools/api/golden/tensorflow.-node-def.-attr-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-node-def.-attr-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "KEY_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "key"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-node-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.-node-def.pbtxt
index b812b4df2b..23319fdb22 100644
--- a/tensorflow/tools/api/golden/tensorflow.-node-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-node-def.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.NodeDef"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.node_def_pb2.NodeDef\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "ATTR_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "AttrEntry"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
@@ -14,25 +10,29 @@ tf_class {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
-  member {
-    name: "DEVICE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "INPUT_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "attr"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "device"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "input"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "OP_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "op"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-optimizer-options.pbtxt b/tensorflow/tools/api/golden/tensorflow.-optimizer-options.pbtxt
index 6cac5c4d99..57da2e8b55 100644
--- a/tensorflow/tools/api/golden/tensorflow.-optimizer-options.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-optimizer-options.pbtxt
@@ -10,26 +10,10 @@ tf_class {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
-  member {
-    name: "DO_COMMON_SUBEXPRESSION_ELIMINATION_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
-  member {
-    name: "DO_CONSTANT_FOLDING_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
-  member {
-    name: "DO_FUNCTION_INLINING_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
-  member {
-    name: "GLOBAL_JIT_LEVEL_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "GlobalJitLevel"
     mtype: "<class \'google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper\'>"
@@ -46,10 +30,6 @@ tf_class {
     name: "Level"
     mtype: "<class \'google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper\'>"
   }
-  member {
-    name: "MAX_FOLDED_CONSTANT_IN_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "OFF"
     mtype: "<type \'int\'>"
@@ -63,8 +43,28 @@ tf_class {
     mtype: "<type \'int\'>"
   }
   member {
-    name: "OPT_LEVEL_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "do_common_subexpression_elimination"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "do_constant_folding"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "do_function_inlining"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "global_jit_level"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "max_folded_constant_in_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "opt_level"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-run-metadata.pbtxt b/tensorflow/tools/api/golden/tensorflow.-run-metadata.pbtxt
index 808fa0fa21..17b3d88168 100644
--- a/tensorflow/tools/api/golden/tensorflow.-run-metadata.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-run-metadata.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.RunMetadata"
 tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.config_pb2.RunMetadata\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "COST_GRAPH_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -15,12 +11,16 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "PARTITION_GRAPHS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "cost_graph"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "partition_graphs"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "STEP_STATS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "step_stats"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-run-options.pbtxt b/tensorflow/tools/api/golden/tensorflow.-run-options.pbtxt
index 2f3e7f1a84..7470e4b63d 100644
--- a/tensorflow/tools/api/golden/tensorflow.-run-options.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-run-options.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.RunOptions"
 tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.config_pb2.RunOptions\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "DEBUG_OPTIONS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -23,36 +19,40 @@ tf_class {
     mtype: "<type \'int\'>"
   }
   member {
-    name: "INTER_OP_THREAD_POOL_FIELD_NUMBER"
+    name: "NO_TRACE"
     mtype: "<type \'int\'>"
   }
   member {
-    name: "NO_TRACE"
+    name: "SOFTWARE_TRACE"
     mtype: "<type \'int\'>"
   }
   member {
-    name: "OUTPUT_PARTITION_GRAPHS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "TraceLevel"
+    mtype: "<class \'google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper\'>"
   }
   member {
-    name: "REPORT_TENSOR_ALLOCATIONS_UPON_OOM_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "debug_options"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SOFTWARE_TRACE"
-    mtype: "<type \'int\'>"
+    name: "inter_op_thread_pool"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TIMEOUT_IN_MS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "output_partition_graphs"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TRACE_LEVEL_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "report_tensor_allocations_upon_oom"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TraceLevel"
-    mtype: "<class \'google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper\'>"
+    name: "timeout_in_ms"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "trace_level"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-session-log.pbtxt b/tensorflow/tools/api/golden/tensorflow.-session-log.pbtxt
index ec66d7f335..259a30546a 100644
--- a/tensorflow/tools/api/golden/tensorflow.-session-log.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-session-log.pbtxt
@@ -6,10 +6,6 @@ tf_class {
     name: "CHECKPOINT"
     mtype: "<type \'int\'>"
   }
-  member {
-    name: "CHECKPOINT_PATH_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -18,18 +14,10 @@ tf_class {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
-  member {
-    name: "MSG_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "START"
     mtype: "<type \'int\'>"
   }
-  member {
-    name: "STATUS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "STATUS_UNSPECIFIED"
     mtype: "<type \'int\'>"
@@ -42,6 +30,18 @@ tf_class {
     name: "SessionStatus"
     mtype: "<class \'google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper\'>"
   }
+  member {
+    name: "checkpoint_path"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "msg"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "status"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
   member_method {
     name: "ByteSize"
   }
diff --git a/tensorflow/tools/api/golden/tensorflow.-summary-metadata.-plugin-data.pbtxt b/tensorflow/tools/api/golden/tensorflow.-summary-metadata.-plugin-data.pbtxt
index 067f02ce8c..3d9ee9e0f2 100644
--- a/tensorflow/tools/api/golden/tensorflow.-summary-metadata.-plugin-data.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-summary-metadata.-plugin-data.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.SummaryMetadata.PluginData"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.PluginData\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "CONTENT_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -15,8 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "PLUGIN_NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "content"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "plugin_name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-summary-metadata.pbtxt b/tensorflow/tools/api/golden/tensorflow.-summary-metadata.pbtxt
index b9156521cc..9c69a2b96c 100644
--- a/tensorflow/tools/api/golden/tensorflow.-summary-metadata.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-summary-metadata.pbtxt
@@ -6,25 +6,25 @@ tf_class {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
-  member {
-    name: "DISPLAY_NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
-  member {
-    name: "PLUGIN_DATA_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "PluginData"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "SUMMARY_DESCRIPTION_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "display_name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "plugin_data"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "summary_description"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-summary.-audio.pbtxt b/tensorflow/tools/api/golden/tensorflow.-summary.-audio.pbtxt
index 781010d75e..8e761b8861 100644
--- a/tensorflow/tools/api/golden/tensorflow.-summary.-audio.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-summary.-audio.pbtxt
@@ -2,33 +2,33 @@ path: "tensorflow.Summary.Audio"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.Audio\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "CONTENT_TYPE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
-  member {
-    name: "ENCODED_AUDIO_STRING_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "LENGTH_FRAMES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "content_type"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "encoded_audio_string"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "length_frames"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "NUM_CHANNELS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "num_channels"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SAMPLE_RATE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "sample_rate"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-summary.-image.pbtxt b/tensorflow/tools/api/golden/tensorflow.-summary.-image.pbtxt
index feb9c7ee92..07b61d9e96 100644
--- a/tensorflow/tools/api/golden/tensorflow.-summary.-image.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-summary.-image.pbtxt
@@ -2,29 +2,29 @@ path: "tensorflow.Summary.Image"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.Image\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "COLORSPACE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
-  member {
-    name: "ENCODED_IMAGE_STRING_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "HEIGHT_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "colorspace"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "encoded_image_string"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "height"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "WIDTH_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "width"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-summary.-value.pbtxt b/tensorflow/tools/api/golden/tensorflow.-summary.-value.pbtxt
index ffb4f45fc5..77ba2e095e 100644
--- a/tensorflow/tools/api/golden/tensorflow.-summary.-value.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-summary.-value.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.Summary.Value"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.Value\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "AUDIO_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -15,36 +11,40 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "HISTO_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "audio"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "histo"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "IMAGE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "image"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "METADATA_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "metadata"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "NODE_NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "node_name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "OBSOLETE_OLD_STYLE_HISTOGRAM_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "obsolete_old_style_histogram"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SIMPLE_VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "simple_value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TAG_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "tag"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TENSOR_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "tensor"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-summary.pbtxt b/tensorflow/tools/api/golden/tensorflow.-summary.pbtxt
index 38de17fa9e..95263bdead 100644
--- a/tensorflow/tools/api/golden/tensorflow.-summary.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-summary.pbtxt
@@ -18,14 +18,14 @@ tf_class {
     name: "Image"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
-  member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "Value"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
+  member {
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
   member_method {
     name: "ByteSize"
   }
diff --git a/tensorflow/tools/api/golden/tensorflow.-tensor-info.-coo-sparse.pbtxt b/tensorflow/tools/api/golden/tensorflow.-tensor-info.-coo-sparse.pbtxt
index 425c35e067..b1848311cf 100644
--- a/tensorflow/tools/api/golden/tensorflow.-tensor-info.-coo-sparse.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-tensor-info.-coo-sparse.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.TensorInfo.CooSparse"
 tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.meta_graph_pb2.CooSparse\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "DENSE_SHAPE_TENSOR_NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -15,12 +11,16 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "INDICES_TENSOR_NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "dense_shape_tensor_name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "indices_tensor_name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VALUES_TENSOR_NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "values_tensor_name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-tensor-info.pbtxt b/tensorflow/tools/api/golden/tensorflow.-tensor-info.pbtxt
index 41ea393be5..9fd26d1b6c 100644
--- a/tensorflow/tools/api/golden/tensorflow.-tensor-info.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-tensor-info.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.TensorInfo"
 tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.meta_graph_pb2.TensorInfo\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "COO_SPARSE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "CooSparse"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
@@ -14,21 +10,25 @@ tf_class {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
-  member {
-    name: "DTYPE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "coo_sparse"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "dtype"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TENSOR_SHAPE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "tensor_shape"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checker.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checker.pbtxt
index bd5c36f390..925ea6df93 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checker.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checker.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "REPORTS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "reports"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checkers-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checkers-entry.pbtxt
index 7c8c68e155..e7ca821951 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checkers-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checkers-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "KEY_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "key"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.pbtxt
index 1b789f4fc9..330d6ee7be 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.profiler.AdviceProto"
 tf_class {
   is_instance: "<class \'tensorflow.core.profiler.tfprof_output_pb2.AdviceProto\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "CHECKERS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "Checker"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
@@ -22,6 +18,10 @@ tf_class {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
+  member {
+    name: "checkers"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
   member_method {
     name: "ByteSize"
   }
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.-input-shapes-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.-input-shapes-entry.pbtxt
index f0b9605bee..85aef3e8a4 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.-input-shapes-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.-input-shapes-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "KEY_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "key"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.pbtxt
index b80896a8a0..2ecfb6a971 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.pbtxt
@@ -3,124 +3,124 @@ tf_class {
   is_instance: "<class \'tensorflow.core.profiler.tfprof_output_pb2.GraphNodeProto\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "ACCELERATOR_EXEC_MICROS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "CHILDREN_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "CPU_EXEC_MICROS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "InputShapesEntry"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "accelerator_exec_micros"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "DEVICES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "children"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "EXEC_MICROS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "cpu_exec_micros"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "devices"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "FLOAT_OPS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "exec_micros"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "INPUT_SHAPES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "float_ops"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "InputShapesEntry"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "input_shapes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "OUTPUT_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "output_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "PARAMETERS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "parameters"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "PEAK_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "peak_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "REQUESTED_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "requested_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "RESIDUAL_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "residual_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "RUN_COUNT_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "run_count"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SHAPES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "shapes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TENSOR_VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "tensor_value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_ACCELERATOR_EXEC_MICROS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_accelerator_exec_micros"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_CPU_EXEC_MICROS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_cpu_exec_micros"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_DEFINITION_COUNT_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_definition_count"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_EXEC_MICROS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_exec_micros"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_FLOAT_OPS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_float_ops"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_OUTPUT_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_output_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_PARAMETERS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_parameters"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_PEAK_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_peak_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_REQUESTED_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_requested_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_RESIDUAL_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_residual_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_RUN_COUNT_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_run_count"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-multi-graph-node-proto.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-multi-graph-node-proto.pbtxt
index 33deff6497..b35d0d6e48 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-multi-graph-node-proto.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-multi-graph-node-proto.pbtxt
@@ -3,96 +3,96 @@ tf_class {
   is_instance: "<class \'tensorflow.core.profiler.tfprof_output_pb2.MultiGraphNodeProto\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "ACCELERATOR_EXEC_MICROS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "CHILDREN_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "CPU_EXEC_MICROS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "accelerator_exec_micros"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "children"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "EXEC_MICROS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "cpu_exec_micros"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "exec_micros"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "FLOAT_OPS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "float_ops"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "GRAPH_NODES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "graph_nodes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "OUTPUT_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "output_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "PARAMETERS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "parameters"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "PEAK_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "peak_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "REQUESTED_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "requested_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "RESIDUAL_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "residual_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_ACCELERATOR_EXEC_MICROS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_accelerator_exec_micros"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_CPU_EXEC_MICROS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_cpu_exec_micros"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_EXEC_MICROS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_exec_micros"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_FLOAT_OPS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_float_ops"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_OUTPUT_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_output_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_PARAMETERS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_parameters"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_PEAK_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_peak_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_REQUESTED_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_requested_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TOTAL_RESIDUAL_BYTES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "total_residual_bytes"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.-id-to-string-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.-id-to-string-entry.pbtxt
index 8c4727cf35..495a63cfeb 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.-id-to-string-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.-id-to-string-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "KEY_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "key"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.pbtxt
index 1071a82b5c..b74d7f8a55 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.pbtxt
@@ -10,17 +10,17 @@ tf_class {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
-  member {
-    name: "ID_TO_STRING_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "IdToStringEntry"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "LOG_ENTRIES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "id_to_string"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "log_entries"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-event.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-event.pbtxt
index ab3449d80f..7ac8470a7a 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-event.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-event.pbtxt
@@ -11,40 +11,40 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "FILE_VERSION_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "file_version"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "GRAPH_DEF_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "graph_def"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "LOG_MESSAGE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "log_message"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "META_GRAPH_DEF_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "meta_graph_def"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SESSION_LOG_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "session_log"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "STEP_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "step"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SUMMARY_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "summary"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TAGGED_RUN_METADATA_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "tagged_run_metadata"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "WALL_TIME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "wall_time"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-session-log.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-session-log.pbtxt
index 92ca4872ca..d1e7e9eedb 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-session-log.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-session-log.pbtxt
@@ -6,10 +6,6 @@ tf_class {
     name: "CHECKPOINT"
     mtype: "<type \'int\'>"
   }
-  member {
-    name: "CHECKPOINT_PATH_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -18,18 +14,10 @@ tf_class {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
-  member {
-    name: "MSG_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "START"
     mtype: "<type \'int\'>"
   }
-  member {
-    name: "STATUS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "STATUS_UNSPECIFIED"
     mtype: "<type \'int\'>"
@@ -42,6 +30,18 @@ tf_class {
     name: "SessionStatus"
     mtype: "<class \'google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper\'>"
   }
+  member {
+    name: "checkpoint_path"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "msg"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "status"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
   member_method {
     name: "ByteSize"
   }
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-summary-description.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-summary-description.pbtxt
index f93da2196a..6fe3c755c9 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-summary-description.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-summary-description.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "TYPE_HINT_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "type_hint"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-summary.-audio.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-summary.-audio.pbtxt
index 605e305e82..8cc8428524 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-summary.-audio.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-summary.-audio.pbtxt
@@ -2,33 +2,33 @@ path: "tensorflow.summary.Summary.Audio"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.Audio\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "CONTENT_TYPE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
-  member {
-    name: "ENCODED_AUDIO_STRING_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "LENGTH_FRAMES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "content_type"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "encoded_audio_string"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "length_frames"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "NUM_CHANNELS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "num_channels"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SAMPLE_RATE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "sample_rate"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-summary.-image.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-summary.-image.pbtxt
index 0646972196..455452b550 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-summary.-image.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-summary.-image.pbtxt
@@ -2,29 +2,29 @@ path: "tensorflow.summary.Summary.Image"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.Image\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "COLORSPACE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
-  member {
-    name: "ENCODED_IMAGE_STRING_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "HEIGHT_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "colorspace"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "encoded_image_string"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "height"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "WIDTH_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "width"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-summary.-value.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-summary.-value.pbtxt
index b319cd03d9..bc9378c75e 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-summary.-value.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-summary.-value.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.summary.Summary.Value"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.Value\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "AUDIO_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -15,36 +11,40 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "HISTO_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "audio"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "histo"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "IMAGE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "image"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "METADATA_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "metadata"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "NODE_NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "node_name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "OBSOLETE_OLD_STYLE_HISTOGRAM_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "obsolete_old_style_histogram"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SIMPLE_VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "simple_value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TAG_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "tag"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TENSOR_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "tensor"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-summary.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-summary.pbtxt
index 132ef1b7d2..c724074d8c 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-summary.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-summary.pbtxt
@@ -18,14 +18,14 @@ tf_class {
     name: "Image"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
-  member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "Value"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
+  member {
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
   member_method {
     name: "ByteSize"
   }
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-tagged-run-metadata.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-tagged-run-metadata.pbtxt
index 4dce20819d..5daec17b68 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-tagged-run-metadata.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-tagged-run-metadata.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "RUN_METADATA_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "run_metadata"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TAG_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "tag"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-bytes-list.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-bytes-list.pbtxt
index 8cf52b817f..5ca8b21ed0 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-bytes-list.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-bytes-list.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-cluster-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-cluster-def.pbtxt
index 93ff856b09..76ed034e73 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-cluster-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-cluster-def.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "JOB_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "job"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-example.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-example.pbtxt
index f7215a2037..f516cac139 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-example.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-example.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "FEATURES_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "features"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-feature-list.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-feature-list.pbtxt
index 3ad98354d6..b5b77fe3cd 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-feature-list.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-feature-list.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "FEATURE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "feature"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.-feature-list-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.-feature-list-entry.pbtxt
index cd171f4ca3..774cfc53af 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.-feature-list-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.-feature-list-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "KEY_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "key"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.pbtxt
index 3d95017d58..430f6b41b1 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.pbtxt
@@ -10,14 +10,14 @@ tf_class {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
-  member {
-    name: "FEATURE_LIST_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "FeatureListEntry"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
+  member {
+    name: "feature_list"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
   member_method {
     name: "ByteSize"
   }
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-feature.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-feature.pbtxt
index 9cca132bba..48014a90ba 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-feature.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-feature.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.train.Feature"
 tf_class {
   is_instance: "<class \'tensorflow.core.example.feature_pb2.Feature\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "BYTES_LIST_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -15,12 +11,16 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "FLOAT_LIST_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "bytes_list"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "float_list"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "INT64_LIST_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "int64_list"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-features.-feature-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-features.-feature-entry.pbtxt
index 858aee0341..8f68927d10 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-features.-feature-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-features.-feature-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "KEY_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "key"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-features.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-features.pbtxt
index 49cd12153b..94e24126f1 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-features.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-features.pbtxt
@@ -10,14 +10,14 @@ tf_class {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
-  member {
-    name: "FEATURE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "FeatureEntry"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
+  member {
+    name: "feature"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
   member_method {
     name: "ByteSize"
   }
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-float-list.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-float-list.pbtxt
index e3f01334b5..37413782a1 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-float-list.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-float-list.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-int64-list.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-int64-list.pbtxt
index 8917dc122c..0c775cf46e 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-int64-list.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-int64-list.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-job-def.-tasks-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-job-def.-tasks-entry.pbtxt
index ac6d81541a..5f0fe5c8a0 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-job-def.-tasks-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-job-def.-tasks-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "KEY_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "key"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VALUE_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "value"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-job-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-job-def.pbtxt
index ce34537fa1..20a76e517f 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-job-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-job-def.pbtxt
@@ -11,16 +11,16 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "TasksEntry"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "TASKS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TasksEntry"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "tasks"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-saver-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-saver-def.pbtxt
index 84498a64f5..24705d0558 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-saver-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-saver-def.pbtxt
@@ -15,44 +15,44 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "FILENAME_TENSOR_NAME_FIELD_NUMBER"
+    name: "LEGACY"
     mtype: "<type \'int\'>"
   }
   member {
-    name: "KEEP_CHECKPOINT_EVERY_N_HOURS_FIELD_NUMBER"
+    name: "V1"
     mtype: "<type \'int\'>"
   }
   member {
-    name: "LEGACY"
+    name: "V2"
     mtype: "<type \'int\'>"
   }
   member {
-    name: "MAX_TO_KEEP_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "filename_tensor_name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "RESTORE_OP_NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "keep_checkpoint_every_n_hours"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SAVE_TENSOR_NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "max_to_keep"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "SHARDED_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "restore_op_name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "V1"
-    mtype: "<type \'int\'>"
+    name: "save_tensor_name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "V2"
-    mtype: "<type \'int\'>"
+    name: "sharded"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "VERSION_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "version"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-sequence-example.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-sequence-example.pbtxt
index 9ab9553702..4ad3ede361 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-sequence-example.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-sequence-example.pbtxt
@@ -2,10 +2,6 @@ path: "tensorflow.train.SequenceExample"
 tf_class {
   is_instance: "<class \'tensorflow.core.example.example_pb2.SequenceExample\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "CONTEXT_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -15,8 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "FEATURE_LISTS_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "context"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "feature_lists"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-server-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-server-def.pbtxt
index af0a3b73cc..d1358cc60d 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-server-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-server-def.pbtxt
@@ -2,14 +2,6 @@ path: "tensorflow.train.ServerDef"
 tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.tensorflow_server_pb2.ServerDef\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
-  member {
-    name: "CLUSTER_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
-  member {
-    name: "DEFAULT_SESSION_CONFIG_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
-  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -19,16 +11,24 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "JOB_NAME_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "cluster"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "default_session_config"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+  }
+  member {
+    name: "job_name"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "PROTOCOL_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "protocol"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member {
-    name: "TASK_INDEX_FIELD_NUMBER"
-    mtype: "<type \'int\'>"
+    name: "task_index"
+    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
   }
   member_method {
     name: "ByteSize"
-- 
GitLab


From 10337c91efe7e3975134a7b09ea598e85877c1b0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 07:56:08 -0700
Subject: [PATCH 088/395] Preventing RemoveTrivialBinary from removing
 broadcasts.

PiperOrigin-RevId: 194937001
---
 .../toco/graph_transformations/quantize.cc    |  5 +++++
 .../remove_trivial_binary.cc                  | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc b/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc
index fa46e6bc38..347302c7a5 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc
@@ -96,6 +96,11 @@ const MinMax& GetOrComputeMinMax(Model* model, const string& array_name) {
       min = std::min(min, val);
       max = std::max(max, val);
     }
+    if (min == 0.f && max == 0.f) {
+      // Prevent downstream anger from quantized math that expects min and max
+      // to not be equal.
+      max = 1.f;
+    }
     auto& minmax = array.GetOrCreateMinMax();
     minmax.min = min;
     minmax.max = max;
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_binary.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_binary.cc
index 95a50c6179..0dfdc40e4c 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_binary.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_binary.cc
@@ -78,6 +78,25 @@ bool RemoveTrivialBinaryOperator::Run(Model* model, std::size_t op_index) {
   CHECK(is_input_constant[index_of_constant_input]);
   CHECK(!is_input_constant[index_of_variable_input]);
 
+  // If this was a broadcasting op we can't remove it as we need the broadcast.
+  // It's possible we could replace it with a cheaper op, though.
+  const auto& input_array_0 = model->GetArray(binary_op->inputs[0]);
+  const auto& input_array_1 = model->GetArray(binary_op->inputs[1]);
+  if (!input_array_0.has_shape() || !input_array_1.has_shape()) {
+    // Both input shapes must be known.
+    return false;
+  }
+  if (input_array_0.shape().dimensions_count() ==
+          input_array_1.shape().dimensions_count() &&
+      input_array_0.shape() != input_array_1.shape()) {
+    AddMessageF(
+        "Preserving %s even though it's trivial as we need to broadcast "
+        "(lhs %s, rhs %s)",
+        LogName(*binary_op), ShapeToString(input_array_0.shape()),
+        ShapeToString(input_array_1.shape()));
+    return false;
+  }
+
   // Now check if the constant operand makes this binary
   // operator trivial.
   const auto& constant_input_array =
-- 
GitLab


From 449b9e56ed8974eefdb87d7cf08ef1c1841f9d6e Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Tue, 1 May 2018 18:33:18 +0300
Subject: [PATCH 089/395] [tf.data] More debug code, since the previous 'fix'
 wasn't a fix.

---
 tensorflow/contrib/data/python/ops/BUILD             | 1 +
 tensorflow/contrib/data/python/ops/interleave_ops.py | 6 ++++++
 2 files changed, 7 insertions(+)

diff --git a/tensorflow/contrib/data/python/ops/BUILD b/tensorflow/contrib/data/python/ops/BUILD
index 7a3e42cc72..091723e0c7 100644
--- a/tensorflow/contrib/data/python/ops/BUILD
+++ b/tensorflow/contrib/data/python/ops/BUILD
@@ -184,6 +184,7 @@ py_library(
         "//tensorflow/python/data/ops:readers",
         "//tensorflow/python/data/util:nest",
         "//tensorflow/python/data/util:sparse",
+        ""//tensorflow/python:platform",
     ],
 )
 
diff --git a/tensorflow/contrib/data/python/ops/interleave_ops.py b/tensorflow/contrib/data/python/ops/interleave_ops.py
index 812a50ecbf..b3bf82ea3b 100644
--- a/tensorflow/contrib/data/python/ops/interleave_ops.py
+++ b/tensorflow/contrib/data/python/ops/interleave_ops.py
@@ -30,6 +30,7 @@ from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.util import deprecation
+from tensorflow.python.platform import tf_logging as logging
 
 
 def parallel_interleave(map_func,
@@ -239,4 +240,9 @@ def sample_from_datasets(datasets, weights=None, seed=None):
   selector_input = dataset_ops.Dataset.zip(
       (logits_ds, random_ops.RandomDataset(seed).batch(2))).map(select_dataset)
 
+  print('selector_input.output_types: ', selector_input.output_types)
+  print('selector_input.output_shapes: ', selector_input.output_shapes)
+  for i, dataset in enumerate(datasets):
+    print('dataset %i output_types: %s' % (i, dataset.output_types))
+    print('dataset %i output_shapes: %s' % (i, dataset.output_shapes))
   return DirectedInterleaveDataset(selector_input, datasets)
-- 
GitLab


From 2364000088aa95a913731c127c10f3bfffac9000 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Tue, 1 May 2018 18:35:33 +0300
Subject: [PATCH 090/395] [tf.data] More debug code, since the previous 'fix'
 wasn't a fix.

---
 tensorflow/contrib/data/python/ops/interleave_ops.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/contrib/data/python/ops/interleave_ops.py b/tensorflow/contrib/data/python/ops/interleave_ops.py
index b3bf82ea3b..140abde21c 100644
--- a/tensorflow/contrib/data/python/ops/interleave_ops.py
+++ b/tensorflow/contrib/data/python/ops/interleave_ops.py
@@ -240,9 +240,9 @@ def sample_from_datasets(datasets, weights=None, seed=None):
   selector_input = dataset_ops.Dataset.zip(
       (logits_ds, random_ops.RandomDataset(seed).batch(2))).map(select_dataset)
 
-  print('selector_input.output_types: ', selector_input.output_types)
-  print('selector_input.output_shapes: ', selector_input.output_shapes)
+  logging.warn('selector_input.output_types: ', selector_input.output_types)
+  logging.warn('selector_input.output_shapes: ', selector_input.output_shapes)
   for i, dataset in enumerate(datasets):
-    print('dataset %i output_types: %s' % (i, dataset.output_types))
-    print('dataset %i output_shapes: %s' % (i, dataset.output_shapes))
+    logging.warn('dataset %i output_types: %s' % (i, dataset.output_types))
+    logging.warn('dataset %i output_shapes: %s' % (i, dataset.output_shapes))
   return DirectedInterleaveDataset(selector_input, datasets)
-- 
GitLab


From 03cecc5eb3a0486bea54e496b000ce50d185c9dc Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Tue, 1 May 2018 18:58:55 +0300
Subject: [PATCH 091/395] [tf.data] Fix BUILD file.

---
 tensorflow/contrib/data/python/ops/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/data/python/ops/BUILD b/tensorflow/contrib/data/python/ops/BUILD
index 091723e0c7..9959ccc005 100644
--- a/tensorflow/contrib/data/python/ops/BUILD
+++ b/tensorflow/contrib/data/python/ops/BUILD
@@ -184,7 +184,7 @@ py_library(
         "//tensorflow/python/data/ops:readers",
         "//tensorflow/python/data/util:nest",
         "//tensorflow/python/data/util:sparse",
-        ""//tensorflow/python:platform",
+        "//tensorflow/python:platform",
     ],
 )
 
-- 
GitLab


From a82e0e7922d6dc657b42ef2b3a7a1a52194454c8 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 09:07:57 -0700
Subject: [PATCH 092/395] Fix crash in HloGraphDumper where it crashes on tuple
 shaped constants

The problem is that it tries to use a special logic for 0 element constants
but the logic used to check the number of elements only supports array shapes.

PiperOrigin-RevId: 194945246
---
 .../compiler/xla/service/hlo_graph_dumper.cc   |  2 +-
 .../xla/service/hlo_graph_dumper_test.cc       | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc
index 516e14b464..bb4db89f0a 100644
--- a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc
+++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc
@@ -804,7 +804,7 @@ string HloDotDumper::GetInstructionNodeInlinedOperands(
     // "{} (f32[42, 0, 10])".  The alternative, calling Literal::ToString(),
     // enumerates all of its empty dimensions (e.g.  "{ { {}, {} }, ..."), which
     // is just noise.
-    if (ShapeUtil::HasZeroElements(shape)) {
+    if (!ShapeUtil::IsTuple(shape) && ShapeUtil::HasZeroElements(shape)) {
       return Printf("{} (%s)", ShapeUtil::HumanString(constant->shape()));
     }
 
diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper_test.cc b/tensorflow/compiler/xla/service/hlo_graph_dumper_test.cc
index 4843963243..8e52d926d8 100644
--- a/tensorflow/compiler/xla/service/hlo_graph_dumper_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_graph_dumper_test.cc
@@ -131,5 +131,23 @@ TEST(HloGraphDumperTest, Constant) {
   EXPECT_THAT(graph, Not(HasSubstr("i_am_a_constant_root_instruction")));
 }
 
+TEST(HloGraphDumperTest, TupleConstant) {
+  Shape tuple_shape = ShapeUtil::MakeTupleShape(
+      {ShapeUtil::MakeShape(F32, {3, 2}), ShapeUtil::MakeShape(S32, {4, 5})});
+  HloComputation::Builder b("b");
+  auto constant = b.AddInstruction(
+      HloInstruction::CreateConstant(Literal::CreateFromShape(tuple_shape)));
+  auto gte = b.AddInstruction(HloInstruction::CreateGetTupleElement(
+      ShapeUtil::MakeShape(F32, {3, 2}), constant, 0));
+
+  HloModuleConfig config;
+  HloModule m(TestName(), config);
+  HloComputation* root_computation = m.AddEntryComputation(b.Build(gte));
+  string graph = hlo_graph_dumper::DumpGraph(
+      *root_computation, /*label=*/"tuple_constant", DebugOptions());
+  EXPECT_THAT(graph, HasSubstr("tuple_constant"));
+  EXPECT_THAT(graph, HasSubstr("constant (f32[3,2], s32[4,5])"));
+}
+
 }  // anonymous namespace
 }  // namespace xla
-- 
GitLab


From da02e19813b6d03a6ea5ff7c910d3e71644fbb34 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Tue, 1 May 2018 19:53:30 +0300
Subject: [PATCH 093/395] [tf.data] Fix debug output.

---
 tensorflow/contrib/data/python/ops/interleave_ops.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/contrib/data/python/ops/interleave_ops.py b/tensorflow/contrib/data/python/ops/interleave_ops.py
index 140abde21c..0852fc6be8 100644
--- a/tensorflow/contrib/data/python/ops/interleave_ops.py
+++ b/tensorflow/contrib/data/python/ops/interleave_ops.py
@@ -240,8 +240,8 @@ def sample_from_datasets(datasets, weights=None, seed=None):
   selector_input = dataset_ops.Dataset.zip(
       (logits_ds, random_ops.RandomDataset(seed).batch(2))).map(select_dataset)
 
-  logging.warn('selector_input.output_types: ', selector_input.output_types)
-  logging.warn('selector_input.output_shapes: ', selector_input.output_shapes)
+  logging.warn('selector_input.output_types: %s', selector_input.output_types)
+  logging.warn('selector_input.output_shapes: %s', selector_input.output_shapes)
   for i, dataset in enumerate(datasets):
     logging.warn('dataset %i output_types: %s' % (i, dataset.output_types))
     logging.warn('dataset %i output_shapes: %s' % (i, dataset.output_shapes))
-- 
GitLab


From bb8220355eda0183a3c039bef1e72c5450f58c11 Mon Sep 17 00:00:00 2001
From: Gunhan Gulsoy <gunan@google.com>
Date: Tue, 1 May 2018 11:04:26 -0700
Subject: [PATCH 094/395] Automated g4 rollback of changelist 194917415

PiperOrigin-RevId: 194962702
---
 .../tensorflow.-attr-value.-list-value.pbtxt  |  36 +++---
 .../api/golden/tensorflow.-attr-value.pbtxt   |  48 +++----
 ...ow.-config-proto.-device-count-entry.pbtxt |   8 +-
 .../api/golden/tensorflow.-config-proto.pbtxt |  72 +++++------
 .../tools/api/golden/tensorflow.-event.pbtxt  |  36 +++---
 .../golden/tensorflow.-g-p-u-options.pbtxt    |  48 +++----
 .../api/golden/tensorflow.-graph-def.pbtxt    |  16 +--
 .../golden/tensorflow.-graph-options.pbtxt    |  44 +++----
 .../golden/tensorflow.-histogram-proto.pbtxt  |  36 +++---
 .../api/golden/tensorflow.-log-message.pbtxt  |  16 +--
 ...meta-graph-def.-collection-def-entry.pbtxt |   8 +-
 ...rflow.-meta-graph-def.-meta-info-def.pbtxt |  32 ++---
 ...-meta-graph-def.-signature-def-entry.pbtxt |   8 +-
 .../golden/tensorflow.-meta-graph-def.pbtxt   |  40 +++---
 ...nsorflow.-name-attr-list.-attr-entry.pbtxt |   8 +-
 .../golden/tensorflow.-name-attr-list.pbtxt   |  12 +-
 .../tensorflow.-node-def.-attr-entry.pbtxt    |   8 +-
 .../api/golden/tensorflow.-node-def.pbtxt     |  28 ++--
 .../tensorflow.-optimizer-options.pbtxt       |  44 +++----
 .../api/golden/tensorflow.-run-metadata.pbtxt |  16 +--
 .../api/golden/tensorflow.-run-options.pbtxt  |  36 +++---
 .../api/golden/tensorflow.-session-log.pbtxt  |  24 ++--
 ...rflow.-summary-metadata.-plugin-data.pbtxt |  12 +-
 .../golden/tensorflow.-summary-metadata.pbtxt |  20 +--
 .../golden/tensorflow.-summary.-audio.pbtxt   |  28 ++--
 .../golden/tensorflow.-summary.-image.pbtxt   |  24 ++--
 .../golden/tensorflow.-summary.-value.pbtxt   |  40 +++---
 .../api/golden/tensorflow.-summary.pbtxt      |   8 +-
 .../tensorflow.-tensor-info.-coo-sparse.pbtxt |  16 +--
 .../api/golden/tensorflow.-tensor-info.pbtxt  |  24 ++--
 ...flow.profiler.-advice-proto.-checker.pbtxt |   4 +-
 ...ofiler.-advice-proto.-checkers-entry.pbtxt |   8 +-
 .../tensorflow.profiler.-advice-proto.pbtxt   |   8 +-
 ...graph-node-proto.-input-shapes-entry.pbtxt |   8 +-
 ...ensorflow.profiler.-graph-node-proto.pbtxt | 120 +++++++++---------
 ...low.profiler.-multi-graph-node-proto.pbtxt |  92 +++++++-------
 ...er.-op-log-proto.-id-to-string-entry.pbtxt |   8 +-
 .../tensorflow.profiler.-op-log-proto.pbtxt   |  12 +-
 .../golden/tensorflow.summary.-event.pbtxt    |  36 +++---
 .../tensorflow.summary.-session-log.pbtxt     |  24 ++--
 ...sorflow.summary.-summary-description.pbtxt |   4 +-
 .../tensorflow.summary.-summary.-audio.pbtxt  |  28 ++--
 .../tensorflow.summary.-summary.-image.pbtxt  |  24 ++--
 .../tensorflow.summary.-summary.-value.pbtxt  |  40 +++---
 .../golden/tensorflow.summary.-summary.pbtxt  |   8 +-
 ...sorflow.summary.-tagged-run-metadata.pbtxt |   8 +-
 .../golden/tensorflow.train.-bytes-list.pbtxt |   4 +-
 .../tensorflow.train.-cluster-def.pbtxt       |   4 +-
 .../golden/tensorflow.train.-example.pbtxt    |   4 +-
 .../tensorflow.train.-feature-list.pbtxt      |   4 +-
 ...n.-feature-lists.-feature-list-entry.pbtxt |   8 +-
 .../tensorflow.train.-feature-lists.pbtxt     |   8 +-
 .../golden/tensorflow.train.-feature.pbtxt    |  16 +--
 ...rflow.train.-features.-feature-entry.pbtxt |   8 +-
 .../golden/tensorflow.train.-features.pbtxt   |   8 +-
 .../golden/tensorflow.train.-float-list.pbtxt |   4 +-
 .../golden/tensorflow.train.-int64-list.pbtxt |   4 +-
 ...nsorflow.train.-job-def.-tasks-entry.pbtxt |   8 +-
 .../golden/tensorflow.train.-job-def.pbtxt    |  12 +-
 .../golden/tensorflow.train.-saver-def.pbtxt  |  34 ++---
 .../tensorflow.train.-sequence-example.pbtxt  |  12 +-
 .../golden/tensorflow.train.-server-def.pbtxt |  28 ++--
 62 files changed, 697 insertions(+), 697 deletions(-)

diff --git a/tensorflow/tools/api/golden/tensorflow.-attr-value.-list-value.pbtxt b/tensorflow/tools/api/golden/tensorflow.-attr-value.-list-value.pbtxt
index 004d716954..0fb1aaba28 100644
--- a/tensorflow/tools/api/golden/tensorflow.-attr-value.-list-value.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-attr-value.-list-value.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.AttrValue.ListValue"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.attr_value_pb2.ListValue\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "B_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -11,36 +15,32 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "b"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "f"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "FUNC_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "func"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "F_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "i"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "I_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "s"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SHAPE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "shape"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "S_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tensor"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TENSOR_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "type"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TYPE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-attr-value.pbtxt b/tensorflow/tools/api/golden/tensorflow.-attr-value.pbtxt
index 2996e02483..e7a3a1f02f 100644
--- a/tensorflow/tools/api/golden/tensorflow.-attr-value.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-attr-value.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.AttrValue"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.attr_value_pb2.AttrValue\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "B_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -11,48 +15,44 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "ListValue"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
-  }
-  member {
-    name: "b"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "FUNC_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "f"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "F_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "func"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "I_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "i"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "LIST_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "list"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "ListValue"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "placeholder"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "PLACEHOLDER_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "s"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SHAPE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "shape"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "S_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tensor"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TENSOR_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "type"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TYPE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-config-proto.-device-count-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.-config-proto.-device-count-entry.pbtxt
index c7022e7593..29bb3be35c 100644
--- a/tensorflow/tools/api/golden/tensorflow.-config-proto.-device-count-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-config-proto.-device-count-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "key"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "KEY_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-config-proto.pbtxt b/tensorflow/tools/api/golden/tensorflow.-config-proto.pbtxt
index ca9530de85..009d64aed0 100644
--- a/tensorflow/tools/api/golden/tensorflow.-config-proto.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-config-proto.pbtxt
@@ -3,76 +3,76 @@ tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.config_pb2.ConfigProto\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "ALLOW_SOFT_PLACEMENT_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "DeviceCountEntry"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "CLUSTER_DEF_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "allow_soft_placement"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "DEVICE_COUNT_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "cluster_def"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "DEVICE_FILTERS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "device_count"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "DeviceCountEntry"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "device_filters"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "gpu_options"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "GPU_OPTIONS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "graph_options"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "GRAPH_OPTIONS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "inter_op_parallelism_threads"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "INTER_OP_PARALLELISM_THREADS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "intra_op_parallelism_threads"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "INTRA_OP_PARALLELISM_THREADS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "isolate_session_state"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "ISOLATE_SESSION_STATE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "log_device_placement"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "LOG_DEVICE_PLACEMENT_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "operation_timeout_in_ms"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "OPERATION_TIMEOUT_IN_MS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "placement_period"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "PLACEMENT_PERIOD_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "rpc_options"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "RPC_OPTIONS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "session_inter_op_thread_pool"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SESSION_INTER_OP_THREAD_POOL_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "use_per_session_threads"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "USE_PER_SESSION_THREADS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-event.pbtxt b/tensorflow/tools/api/golden/tensorflow.-event.pbtxt
index fa2f329a87..9bf8c12428 100644
--- a/tensorflow/tools/api/golden/tensorflow.-event.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-event.pbtxt
@@ -11,40 +11,40 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "file_version"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "FILE_VERSION_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "graph_def"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "GRAPH_DEF_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "log_message"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "LOG_MESSAGE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "meta_graph_def"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "META_GRAPH_DEF_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "session_log"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SESSION_LOG_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "step"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "STEP_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "summary"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SUMMARY_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tagged_run_metadata"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TAGGED_RUN_METADATA_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "wall_time"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "WALL_TIME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-g-p-u-options.pbtxt b/tensorflow/tools/api/golden/tensorflow.-g-p-u-options.pbtxt
index 5119c7fa5b..875d802a9c 100644
--- a/tensorflow/tools/api/golden/tensorflow.-g-p-u-options.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-g-p-u-options.pbtxt
@@ -3,52 +3,52 @@ tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.config_pb2.GPUOptions\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "ALLOCATOR_TYPE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "Experimental"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "ALLOW_GROWTH_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "DEFERRED_DELETION_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "allocator_type"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "allow_growth"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "EXPERIMENTAL_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "deferred_deletion_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Experimental"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "experimental"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "force_gpu_compatible"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "FORCE_GPU_COMPATIBLE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "per_process_gpu_memory_fraction"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "PER_PROCESS_GPU_MEMORY_FRACTION_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "polling_active_delay_usecs"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "POLLING_ACTIVE_DELAY_USECS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "polling_inactive_delay_msecs"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "POLLING_INACTIVE_DELAY_MSECS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "visible_device_list"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VISIBLE_DEVICE_LIST_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-graph-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.-graph-def.pbtxt
index 318a25a092..1495e847cb 100644
--- a/tensorflow/tools/api/golden/tensorflow.-graph-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-graph-def.pbtxt
@@ -11,20 +11,20 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "library"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "LIBRARY_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "node"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "NODE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "version"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VERSIONS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "versions"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VERSION_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-graph-options.pbtxt b/tensorflow/tools/api/golden/tensorflow.-graph-options.pbtxt
index 786d831c70..0844f891ca 100644
--- a/tensorflow/tools/api/golden/tensorflow.-graph-options.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-graph-options.pbtxt
@@ -3,48 +3,48 @@ tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.config_pb2.GraphOptions\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "BUILD_COST_MODEL_AFTER_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "BUILD_COST_MODEL_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "build_cost_model"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "build_cost_model_after"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "ENABLE_BFLOAT16_SENDRECV_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "enable_bfloat16_sendrecv"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "ENABLE_RECV_SCHEDULING_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "enable_recv_scheduling"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "infer_shapes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "INFER_SHAPES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "optimizer_options"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "OPTIMIZER_OPTIONS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "place_pruned_graph"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "PLACE_PRUNED_GRAPH_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "rewrite_options"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "REWRITE_OPTIONS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "timeline_step"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TIMELINE_STEP_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-histogram-proto.pbtxt b/tensorflow/tools/api/golden/tensorflow.-histogram-proto.pbtxt
index 3eb2d8873a..2567d2fe60 100644
--- a/tensorflow/tools/api/golden/tensorflow.-histogram-proto.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-histogram-proto.pbtxt
@@ -3,40 +3,40 @@ tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.HistogramProto\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "BUCKET_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "BUCKET_LIMIT_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "bucket"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "bucket_limit"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "max"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "MAX_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "min"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "MIN_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "num"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "NUM_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "sum"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SUM_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "sum_squares"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SUM_SQUARES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-log-message.pbtxt b/tensorflow/tools/api/golden/tensorflow.-log-message.pbtxt
index 760739f4f3..a43c5eb7e3 100644
--- a/tensorflow/tools/api/golden/tensorflow.-log-message.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-log-message.pbtxt
@@ -26,25 +26,25 @@ tf_class {
     name: "INFO"
     mtype: "<type \'int\'>"
   }
+  member {
+    name: "LEVEL_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "Level"
     mtype: "<class \'google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper\'>"
   }
   member {
-    name: "UNKNOWN"
+    name: "MESSAGE_FIELD_NUMBER"
     mtype: "<type \'int\'>"
   }
   member {
-    name: "WARN"
+    name: "UNKNOWN"
     mtype: "<type \'int\'>"
   }
   member {
-    name: "level"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "message"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "WARN"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-collection-def-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-collection-def-entry.pbtxt
index 69bf5b31a1..3572126fbf 100644
--- a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-collection-def-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-collection-def-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "key"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "KEY_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-meta-info-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-meta-info-def.pbtxt
index 8a464f1cac..b0e9831154 100644
--- a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-meta-info-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-meta-info-def.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.MetaGraphDef.MetaInfoDef"
 tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.meta_graph_pb2.MetaInfoDef\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "ANY_INFO_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -11,32 +15,28 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "any_info"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "meta_graph_version"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "META_GRAPH_VERSION_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "stripped_default_attrs"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "STRIPPED_DEFAULT_ATTRS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "stripped_op_list"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "STRIPPED_OP_LIST_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tags"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TAGS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tensorflow_git_version"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TENSORFLOW_GIT_VERSION_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tensorflow_version"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TENSORFLOW_VERSION_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-signature-def-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-signature-def-entry.pbtxt
index 8c5949d067..48fccac99d 100644
--- a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-signature-def-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.-signature-def-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "key"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "KEY_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.pbtxt
index 2be0432c00..3e683a8715 100644
--- a/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-meta-graph-def.pbtxt
@@ -2,6 +2,14 @@ path: "tensorflow.MetaGraphDef"
 tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.meta_graph_pb2.MetaGraphDef\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "ASSET_FILE_DEF_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
+  member {
+    name: "COLLECTION_DEF_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "CollectionDefEntry"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
@@ -15,36 +23,28 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "MetaInfoDef"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
-  }
-  member {
-    name: "SignatureDefEntry"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
-  }
-  member {
-    name: "asset_file_def"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "GRAPH_DEF_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "collection_def"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "META_INFO_DEF_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "graph_def"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "MetaInfoDef"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "meta_info_def"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SAVER_DEF_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "saver_def"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SIGNATURE_DEF_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "signature_def"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SignatureDefEntry"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-name-attr-list.-attr-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.-name-attr-list.-attr-entry.pbtxt
index caf992f5a6..2750bd780c 100644
--- a/tensorflow/tools/api/golden/tensorflow.-name-attr-list.-attr-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-name-attr-list.-attr-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "key"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "KEY_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-name-attr-list.pbtxt b/tensorflow/tools/api/golden/tensorflow.-name-attr-list.pbtxt
index 45ddeece07..d10faf67d0 100644
--- a/tensorflow/tools/api/golden/tensorflow.-name-attr-list.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-name-attr-list.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.NameAttrList"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.attr_value_pb2.NameAttrList\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "ATTR_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "AttrEntry"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
@@ -15,12 +19,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "attr"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-node-def.-attr-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.-node-def.-attr-entry.pbtxt
index 30a9dc69f0..b1b62d60f1 100644
--- a/tensorflow/tools/api/golden/tensorflow.-node-def.-attr-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-node-def.-attr-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "key"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "KEY_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-node-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.-node-def.pbtxt
index 23319fdb22..b812b4df2b 100644
--- a/tensorflow/tools/api/golden/tensorflow.-node-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-node-def.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.NodeDef"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.node_def_pb2.NodeDef\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "ATTR_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "AttrEntry"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
@@ -11,28 +15,24 @@ tf_class {
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
-  }
-  member {
-    name: "attr"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "DEVICE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "device"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "input"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "INPUT_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "op"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "OP_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-optimizer-options.pbtxt b/tensorflow/tools/api/golden/tensorflow.-optimizer-options.pbtxt
index 57da2e8b55..6cac5c4d99 100644
--- a/tensorflow/tools/api/golden/tensorflow.-optimizer-options.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-optimizer-options.pbtxt
@@ -10,10 +10,26 @@ tf_class {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
+  member {
+    name: "DO_COMMON_SUBEXPRESSION_ELIMINATION_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
+  member {
+    name: "DO_CONSTANT_FOLDING_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
+  member {
+    name: "DO_FUNCTION_INLINING_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
+  member {
+    name: "GLOBAL_JIT_LEVEL_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "GlobalJitLevel"
     mtype: "<class \'google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper\'>"
@@ -30,6 +46,10 @@ tf_class {
     name: "Level"
     mtype: "<class \'google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper\'>"
   }
+  member {
+    name: "MAX_FOLDED_CONSTANT_IN_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "OFF"
     mtype: "<type \'int\'>"
@@ -43,28 +63,8 @@ tf_class {
     mtype: "<type \'int\'>"
   }
   member {
-    name: "do_common_subexpression_elimination"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "do_constant_folding"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "do_function_inlining"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "global_jit_level"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "max_folded_constant_in_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "opt_level"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "OPT_LEVEL_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-run-metadata.pbtxt b/tensorflow/tools/api/golden/tensorflow.-run-metadata.pbtxt
index 17b3d88168..808fa0fa21 100644
--- a/tensorflow/tools/api/golden/tensorflow.-run-metadata.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-run-metadata.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.RunMetadata"
 tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.config_pb2.RunMetadata\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "COST_GRAPH_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -11,16 +15,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "cost_graph"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "partition_graphs"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "PARTITION_GRAPHS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "step_stats"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "STEP_STATS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-run-options.pbtxt b/tensorflow/tools/api/golden/tensorflow.-run-options.pbtxt
index 7470e4b63d..2f3e7f1a84 100644
--- a/tensorflow/tools/api/golden/tensorflow.-run-options.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-run-options.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.RunOptions"
 tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.config_pb2.RunOptions\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "DEBUG_OPTIONS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -19,40 +23,36 @@ tf_class {
     mtype: "<type \'int\'>"
   }
   member {
-    name: "NO_TRACE"
+    name: "INTER_OP_THREAD_POOL_FIELD_NUMBER"
     mtype: "<type \'int\'>"
   }
   member {
-    name: "SOFTWARE_TRACE"
+    name: "NO_TRACE"
     mtype: "<type \'int\'>"
   }
   member {
-    name: "TraceLevel"
-    mtype: "<class \'google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper\'>"
-  }
-  member {
-    name: "debug_options"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "OUTPUT_PARTITION_GRAPHS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "inter_op_thread_pool"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "REPORT_TENSOR_ALLOCATIONS_UPON_OOM_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "output_partition_graphs"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SOFTWARE_TRACE"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "report_tensor_allocations_upon_oom"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TIMEOUT_IN_MS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "timeout_in_ms"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TRACE_LEVEL_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "trace_level"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TraceLevel"
+    mtype: "<class \'google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-session-log.pbtxt b/tensorflow/tools/api/golden/tensorflow.-session-log.pbtxt
index 259a30546a..ec66d7f335 100644
--- a/tensorflow/tools/api/golden/tensorflow.-session-log.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-session-log.pbtxt
@@ -6,6 +6,10 @@ tf_class {
     name: "CHECKPOINT"
     mtype: "<type \'int\'>"
   }
+  member {
+    name: "CHECKPOINT_PATH_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -14,10 +18,18 @@ tf_class {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
+  member {
+    name: "MSG_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "START"
     mtype: "<type \'int\'>"
   }
+  member {
+    name: "STATUS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "STATUS_UNSPECIFIED"
     mtype: "<type \'int\'>"
@@ -30,18 +42,6 @@ tf_class {
     name: "SessionStatus"
     mtype: "<class \'google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper\'>"
   }
-  member {
-    name: "checkpoint_path"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "msg"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "status"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
   member_method {
     name: "ByteSize"
   }
diff --git a/tensorflow/tools/api/golden/tensorflow.-summary-metadata.-plugin-data.pbtxt b/tensorflow/tools/api/golden/tensorflow.-summary-metadata.-plugin-data.pbtxt
index 3d9ee9e0f2..067f02ce8c 100644
--- a/tensorflow/tools/api/golden/tensorflow.-summary-metadata.-plugin-data.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-summary-metadata.-plugin-data.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.SummaryMetadata.PluginData"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.PluginData\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "CONTENT_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -11,12 +15,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "content"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "plugin_name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "PLUGIN_NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-summary-metadata.pbtxt b/tensorflow/tools/api/golden/tensorflow.-summary-metadata.pbtxt
index 9c69a2b96c..b9156521cc 100644
--- a/tensorflow/tools/api/golden/tensorflow.-summary-metadata.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-summary-metadata.pbtxt
@@ -7,24 +7,24 @@ tf_class {
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "DISPLAY_NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "PluginData"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "display_name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "PLUGIN_DATA_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "plugin_data"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "PluginData"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "summary_description"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SUMMARY_DESCRIPTION_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-summary.-audio.pbtxt b/tensorflow/tools/api/golden/tensorflow.-summary.-audio.pbtxt
index 8e761b8861..781010d75e 100644
--- a/tensorflow/tools/api/golden/tensorflow.-summary.-audio.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-summary.-audio.pbtxt
@@ -3,32 +3,32 @@ tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.Audio\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "CONTENT_TYPE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "content_type"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "ENCODED_AUDIO_STRING_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "encoded_audio_string"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "length_frames"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "LENGTH_FRAMES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "num_channels"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "NUM_CHANNELS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "sample_rate"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SAMPLE_RATE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-summary.-image.pbtxt b/tensorflow/tools/api/golden/tensorflow.-summary.-image.pbtxt
index 07b61d9e96..feb9c7ee92 100644
--- a/tensorflow/tools/api/golden/tensorflow.-summary.-image.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-summary.-image.pbtxt
@@ -3,28 +3,28 @@ tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.Image\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "COLORSPACE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "colorspace"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "ENCODED_IMAGE_STRING_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "encoded_image_string"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "height"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "HEIGHT_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "width"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "WIDTH_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-summary.-value.pbtxt b/tensorflow/tools/api/golden/tensorflow.-summary.-value.pbtxt
index 77ba2e095e..ffb4f45fc5 100644
--- a/tensorflow/tools/api/golden/tensorflow.-summary.-value.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-summary.-value.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.Summary.Value"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.Value\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "AUDIO_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -11,40 +15,36 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "audio"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "histo"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "HISTO_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "image"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "IMAGE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "metadata"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "METADATA_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "node_name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "NODE_NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "obsolete_old_style_histogram"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "OBSOLETE_OLD_STYLE_HISTOGRAM_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "simple_value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SIMPLE_VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tag"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TAG_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tensor"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TENSOR_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-summary.pbtxt b/tensorflow/tools/api/golden/tensorflow.-summary.pbtxt
index 95263bdead..38de17fa9e 100644
--- a/tensorflow/tools/api/golden/tensorflow.-summary.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-summary.pbtxt
@@ -19,12 +19,12 @@ tf_class {
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "Value"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Value"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-tensor-info.-coo-sparse.pbtxt b/tensorflow/tools/api/golden/tensorflow.-tensor-info.-coo-sparse.pbtxt
index b1848311cf..425c35e067 100644
--- a/tensorflow/tools/api/golden/tensorflow.-tensor-info.-coo-sparse.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-tensor-info.-coo-sparse.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.TensorInfo.CooSparse"
 tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.meta_graph_pb2.CooSparse\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "DENSE_SHAPE_TENSOR_NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -11,16 +15,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "dense_shape_tensor_name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "indices_tensor_name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "INDICES_TENSOR_NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "values_tensor_name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUES_TENSOR_NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.-tensor-info.pbtxt b/tensorflow/tools/api/golden/tensorflow.-tensor-info.pbtxt
index 9fd26d1b6c..41ea393be5 100644
--- a/tensorflow/tools/api/golden/tensorflow.-tensor-info.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.-tensor-info.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.TensorInfo"
 tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.meta_graph_pb2.TensorInfo\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "COO_SPARSE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "CooSparse"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
@@ -11,24 +15,20 @@ tf_class {
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
-  }
-  member {
-    name: "coo_sparse"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "DTYPE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "dtype"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tensor_shape"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TENSOR_SHAPE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checker.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checker.pbtxt
index 925ea6df93..bd5c36f390 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checker.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checker.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "reports"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "REPORTS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checkers-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checkers-entry.pbtxt
index e7ca821951..7c8c68e155 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checkers-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.-checkers-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "key"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "KEY_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.pbtxt
index 330d6ee7be..1b789f4fc9 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-advice-proto.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.profiler.AdviceProto"
 tf_class {
   is_instance: "<class \'tensorflow.core.profiler.tfprof_output_pb2.AdviceProto\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "CHECKERS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "Checker"
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
@@ -18,10 +22,6 @@ tf_class {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
-  member {
-    name: "checkers"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
   member_method {
     name: "ByteSize"
   }
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.-input-shapes-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.-input-shapes-entry.pbtxt
index 85aef3e8a4..f0b9605bee 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.-input-shapes-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.-input-shapes-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "key"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "KEY_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.pbtxt
index 2ecfb6a971..b80896a8a0 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-graph-node-proto.pbtxt
@@ -3,124 +3,124 @@ tf_class {
   is_instance: "<class \'tensorflow.core.profiler.tfprof_output_pb2.GraphNodeProto\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "ACCELERATOR_EXEC_MICROS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "CHILDREN_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "InputShapesEntry"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "CPU_EXEC_MICROS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "accelerator_exec_micros"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "children"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "DEVICES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "cpu_exec_micros"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "EXEC_MICROS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "devices"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "exec_micros"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "FLOAT_OPS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "float_ops"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "INPUT_SHAPES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "input_shapes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "InputShapesEntry"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "output_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "OUTPUT_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "parameters"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "PARAMETERS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "peak_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "PEAK_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "requested_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "REQUESTED_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "residual_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "RESIDUAL_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "run_count"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "RUN_COUNT_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "shapes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SHAPES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tensor_value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TENSOR_VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_accelerator_exec_micros"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_ACCELERATOR_EXEC_MICROS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_cpu_exec_micros"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_CPU_EXEC_MICROS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_definition_count"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_DEFINITION_COUNT_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_exec_micros"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_EXEC_MICROS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_float_ops"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_FLOAT_OPS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_output_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_OUTPUT_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_parameters"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_PARAMETERS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_peak_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_PEAK_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_requested_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_REQUESTED_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_residual_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_RESIDUAL_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_run_count"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_RUN_COUNT_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-multi-graph-node-proto.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-multi-graph-node-proto.pbtxt
index b35d0d6e48..33deff6497 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-multi-graph-node-proto.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-multi-graph-node-proto.pbtxt
@@ -3,96 +3,96 @@ tf_class {
   is_instance: "<class \'tensorflow.core.profiler.tfprof_output_pb2.MultiGraphNodeProto\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "ACCELERATOR_EXEC_MICROS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "CHILDREN_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "accelerator_exec_micros"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "CPU_EXEC_MICROS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "children"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "cpu_exec_micros"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "EXEC_MICROS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "exec_micros"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "float_ops"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "FLOAT_OPS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "graph_nodes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "GRAPH_NODES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "output_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "OUTPUT_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "parameters"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "PARAMETERS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "peak_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "PEAK_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "requested_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "REQUESTED_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "residual_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "RESIDUAL_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_accelerator_exec_micros"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_ACCELERATOR_EXEC_MICROS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_cpu_exec_micros"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_CPU_EXEC_MICROS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_exec_micros"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_EXEC_MICROS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_float_ops"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_FLOAT_OPS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_output_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_OUTPUT_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_parameters"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_PARAMETERS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_peak_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_PEAK_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_requested_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_REQUESTED_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "total_residual_bytes"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TOTAL_RESIDUAL_BYTES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.-id-to-string-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.-id-to-string-entry.pbtxt
index 495a63cfeb..8c4727cf35 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.-id-to-string-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.-id-to-string-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "key"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "KEY_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.pbtxt b/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.pbtxt
index b74d7f8a55..1071a82b5c 100644
--- a/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.profiler.-op-log-proto.pbtxt
@@ -11,16 +11,16 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "IdToStringEntry"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "ID_TO_STRING_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "id_to_string"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "IdToStringEntry"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "log_entries"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "LOG_ENTRIES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-event.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-event.pbtxt
index 7ac8470a7a..ab3449d80f 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-event.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-event.pbtxt
@@ -11,40 +11,40 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "file_version"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "FILE_VERSION_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "graph_def"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "GRAPH_DEF_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "log_message"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "LOG_MESSAGE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "meta_graph_def"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "META_GRAPH_DEF_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "session_log"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SESSION_LOG_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "step"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "STEP_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "summary"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SUMMARY_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tagged_run_metadata"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TAGGED_RUN_METADATA_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "wall_time"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "WALL_TIME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-session-log.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-session-log.pbtxt
index d1e7e9eedb..92ca4872ca 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-session-log.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-session-log.pbtxt
@@ -6,6 +6,10 @@ tf_class {
     name: "CHECKPOINT"
     mtype: "<type \'int\'>"
   }
+  member {
+    name: "CHECKPOINT_PATH_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -14,10 +18,18 @@ tf_class {
     name: "Extensions"
     mtype: "<type \'getset_descriptor\'>"
   }
+  member {
+    name: "MSG_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "START"
     mtype: "<type \'int\'>"
   }
+  member {
+    name: "STATUS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "STATUS_UNSPECIFIED"
     mtype: "<type \'int\'>"
@@ -30,18 +42,6 @@ tf_class {
     name: "SessionStatus"
     mtype: "<class \'google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper\'>"
   }
-  member {
-    name: "checkpoint_path"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "msg"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "status"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
   member_method {
     name: "ByteSize"
   }
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-summary-description.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-summary-description.pbtxt
index 6fe3c755c9..f93da2196a 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-summary-description.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-summary-description.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "type_hint"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TYPE_HINT_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-summary.-audio.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-summary.-audio.pbtxt
index 8cc8428524..605e305e82 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-summary.-audio.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-summary.-audio.pbtxt
@@ -3,32 +3,32 @@ tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.Audio\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "CONTENT_TYPE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "content_type"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "ENCODED_AUDIO_STRING_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "encoded_audio_string"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "length_frames"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "LENGTH_FRAMES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "num_channels"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "NUM_CHANNELS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "sample_rate"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SAMPLE_RATE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-summary.-image.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-summary.-image.pbtxt
index 455452b550..0646972196 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-summary.-image.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-summary.-image.pbtxt
@@ -3,28 +3,28 @@ tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.Image\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "COLORSPACE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "colorspace"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "ENCODED_IMAGE_STRING_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "encoded_image_string"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "height"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "HEIGHT_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "width"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "WIDTH_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-summary.-value.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-summary.-value.pbtxt
index bc9378c75e..b319cd03d9 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-summary.-value.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-summary.-value.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.summary.Summary.Value"
 tf_class {
   is_instance: "<class \'tensorflow.core.framework.summary_pb2.Value\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "AUDIO_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -11,40 +15,36 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "audio"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "histo"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "HISTO_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "image"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "IMAGE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "metadata"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "METADATA_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "node_name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "NODE_NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "obsolete_old_style_histogram"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "OBSOLETE_OLD_STYLE_HISTOGRAM_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "simple_value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SIMPLE_VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tag"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TAG_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tensor"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TENSOR_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-summary.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-summary.pbtxt
index c724074d8c..132ef1b7d2 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-summary.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-summary.pbtxt
@@ -19,12 +19,12 @@ tf_class {
     mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member {
-    name: "Value"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Value"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.summary.-tagged-run-metadata.pbtxt b/tensorflow/tools/api/golden/tensorflow.summary.-tagged-run-metadata.pbtxt
index 5daec17b68..4dce20819d 100644
--- a/tensorflow/tools/api/golden/tensorflow.summary.-tagged-run-metadata.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.summary.-tagged-run-metadata.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "run_metadata"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "RUN_METADATA_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tag"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TAG_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-bytes-list.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-bytes-list.pbtxt
index 5ca8b21ed0..8cf52b817f 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-bytes-list.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-bytes-list.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-cluster-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-cluster-def.pbtxt
index 76ed034e73..93ff856b09 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-cluster-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-cluster-def.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "job"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "JOB_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-example.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-example.pbtxt
index f516cac139..f7215a2037 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-example.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-example.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "features"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "FEATURES_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-feature-list.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-feature-list.pbtxt
index b5b77fe3cd..3ad98354d6 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-feature-list.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-feature-list.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "feature"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "FEATURE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.-feature-list-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.-feature-list-entry.pbtxt
index 774cfc53af..cd171f4ca3 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.-feature-list-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.-feature-list-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "key"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "KEY_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.pbtxt
index 430f6b41b1..3d95017d58 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-feature-lists.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "FeatureListEntry"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "FEATURE_LIST_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "feature_list"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "FeatureListEntry"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-feature.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-feature.pbtxt
index 48014a90ba..9cca132bba 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-feature.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-feature.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.train.Feature"
 tf_class {
   is_instance: "<class \'tensorflow.core.example.feature_pb2.Feature\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "BYTES_LIST_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -11,16 +15,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "bytes_list"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "float_list"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "FLOAT_LIST_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "int64_list"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "INT64_LIST_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-features.-feature-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-features.-feature-entry.pbtxt
index 8f68927d10..858aee0341 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-features.-feature-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-features.-feature-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "key"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "KEY_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-features.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-features.pbtxt
index 94e24126f1..49cd12153b 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-features.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-features.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "FeatureEntry"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "FEATURE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "feature"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "FeatureEntry"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-float-list.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-float-list.pbtxt
index 37413782a1..e3f01334b5 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-float-list.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-float-list.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-int64-list.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-int64-list.pbtxt
index 0c775cf46e..8917dc122c 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-int64-list.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-int64-list.pbtxt
@@ -11,8 +11,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-job-def.-tasks-entry.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-job-def.-tasks-entry.pbtxt
index 5f0fe5c8a0..ac6d81541a 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-job-def.-tasks-entry.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-job-def.-tasks-entry.pbtxt
@@ -11,12 +11,12 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "key"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "KEY_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "value"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VALUE_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-job-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-job-def.pbtxt
index 20a76e517f..ce34537fa1 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-job-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-job-def.pbtxt
@@ -11,16 +11,16 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "TasksEntry"
-    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
+    name: "NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TASKS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "tasks"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TasksEntry"
+    mtype: "<class \'google.protobuf.pyext.cpp_message.GeneratedProtocolMessageType\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-saver-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-saver-def.pbtxt
index 24705d0558..84498a64f5 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-saver-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-saver-def.pbtxt
@@ -15,44 +15,44 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "LEGACY"
+    name: "FILENAME_TENSOR_NAME_FIELD_NUMBER"
     mtype: "<type \'int\'>"
   }
   member {
-    name: "V1"
+    name: "KEEP_CHECKPOINT_EVERY_N_HOURS_FIELD_NUMBER"
     mtype: "<type \'int\'>"
   }
   member {
-    name: "V2"
+    name: "LEGACY"
     mtype: "<type \'int\'>"
   }
   member {
-    name: "filename_tensor_name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "MAX_TO_KEEP_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "keep_checkpoint_every_n_hours"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "RESTORE_OP_NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "max_to_keep"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SAVE_TENSOR_NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "restore_op_name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "SHARDED_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "save_tensor_name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "V1"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "sharded"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "V2"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "version"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "VERSION_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-sequence-example.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-sequence-example.pbtxt
index 4ad3ede361..9ab9553702 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-sequence-example.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-sequence-example.pbtxt
@@ -2,6 +2,10 @@ path: "tensorflow.train.SequenceExample"
 tf_class {
   is_instance: "<class \'tensorflow.core.example.example_pb2.SequenceExample\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
+  member {
+    name: "CONTEXT_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
+  }
   member {
     name: "DESCRIPTOR"
     mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
@@ -11,12 +15,8 @@ tf_class {
     mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "context"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
-  }
-  member {
-    name: "feature_lists"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "FEATURE_LISTS_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-server-def.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-server-def.pbtxt
index d1358cc60d..af0a3b73cc 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-server-def.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-server-def.pbtxt
@@ -3,32 +3,32 @@ tf_class {
   is_instance: "<class \'tensorflow.core.protobuf.tensorflow_server_pb2.ServerDef\'>"
   is_instance: "<type \'google.protobuf.pyext._message.CMessage\'>"
   member {
-    name: "DESCRIPTOR"
-    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
+    name: "CLUSTER_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "Extensions"
-    mtype: "<type \'getset_descriptor\'>"
+    name: "DEFAULT_SESSION_CONFIG_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "cluster"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "DESCRIPTOR"
+    mtype: "<type \'google.protobuf.pyext._message.MessageDescriptor\'>"
   }
   member {
-    name: "default_session_config"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "Extensions"
+    mtype: "<type \'getset_descriptor\'>"
   }
   member {
-    name: "job_name"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "JOB_NAME_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "protocol"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "PROTOCOL_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member {
-    name: "task_index"
-    mtype: "<type \'google.protobuf.pyext._message.FieldProperty\'>"
+    name: "TASK_INDEX_FIELD_NUMBER"
+    mtype: "<type \'int\'>"
   }
   member_method {
     name: "ByteSize"
-- 
GitLab


From 9477a96f88d9921020450427636db281122703fe Mon Sep 17 00:00:00 2001
From: Asim Shankar <ashankar@google.com>
Date: Tue, 1 May 2018 11:52:04 -0700
Subject: [PATCH 095/395] eager: Update sample notebooks with API changes in
 the last few releases.

Most notably:
- Avoid using tf.contrib.eager since equivalent functionality if available outside tf.contrib
- Datasets can be directly iterated on.
- Use tf.GradientTape instead of tf.contrib.eager.implicit_gradients

PiperOrigin-RevId: 194971115
---
 tensorflow/contrib/eager/README.md            |  11 +-
 .../python/examples/notebooks/1_basics.ipynb  | 364 ++++++--------
 .../examples/notebooks/2_gradients.ipynb      | 473 ++++--------------
 .../examples/notebooks/3_datasets.ipynb       |  43 +-
 4 files changed, 278 insertions(+), 613 deletions(-)

diff --git a/tensorflow/contrib/eager/README.md b/tensorflow/contrib/eager/README.md
index 9a3b780af8..762685db14 100644
--- a/tensorflow/contrib/eager/README.md
+++ b/tensorflow/contrib/eager/README.md
@@ -37,7 +37,7 @@ support for distributed and multi-GPU training and performance.
 
 ## Installation
 
-Eager execution is included in TensorFlow versions 1.7 and above.
+For eager execution, we recommend using TensorFlow version 1.8 or newer.
 Installation instructions at https://www.tensorflow.org/install/
 
 ## Documentation
@@ -48,12 +48,3 @@ For an introduction to eager execution in TensorFlow, see:
 - Notebook: [Basic Usage](python/examples/notebooks/1_basics.ipynb)
 - Notebook: [Gradients](python/examples/notebooks/2_gradients.ipynb)
 - Notebook: [Importing Data](python/examples/notebooks/3_datasets.ipynb)
-
-## Changelog
-
-- 2017/10/31: Initial preview release (in TensorFlow 1.5)
-- 2017/12/01: Example of dynamic neural network:
-  [SPINN: Stack-augmented Parser-Interpreter Neural Network](https://arxiv.org/abs/1603.06021).
-  See [README.md](python/examples/spinn/README.md) for details.
-- 2017/03: Core functionality moved out of the experimental tf.contrib namespace
-  in TensorFlow 1.7.
diff --git a/tensorflow/contrib/eager/python/examples/notebooks/1_basics.ipynb b/tensorflow/contrib/eager/python/examples/notebooks/1_basics.ipynb
index 459f2f4a7d..0279db80fa 100644
--- a/tensorflow/contrib/eager/python/examples/notebooks/1_basics.ipynb
+++ b/tensorflow/contrib/eager/python/examples/notebooks/1_basics.ipynb
@@ -1,11 +1,27 @@
 {
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "name": "Eager Execution Tutorial: Basics",
+      "version": "0.3.2",
+      "views": {},
+      "default_view": {},
+      "provenance": [
+        {
+          "file_id": "0B0kLcpwLFwKEVm9XNkFueGk4bTg",
+          "timestamp": 1504118841551
+        }
+      ]
+    }
+  },
   "cells": [
     {
-      "cell_type": "markdown",
       "metadata": {
-        "colab_type": "text",
-        "id": "U9i2Dsh-ziXr"
+        "id": "U9i2Dsh-ziXr",
+        "colab_type": "text"
       },
+      "cell_type": "markdown",
       "source": [
         "# Eager Execution Tutorial: Basics\n",
         "\n",
@@ -21,11 +37,11 @@
       ]
     },
     {
-      "cell_type": "markdown",
       "metadata": {
-        "colab_type": "text",
-        "id": "z1JcS5iBXMRO"
+        "id": "z1JcS5iBXMRO",
+        "colab_type": "text"
       },
+      "cell_type": "markdown",
       "source": [
         "# Step 1: Import Eager\n",
         "\n",
@@ -33,34 +49,34 @@
       ]
     },
     {
-      "cell_type": "code",
-      "execution_count": 0,
       "metadata": {
-        "cellView": "code",
+        "id": "RlIWhyeLoYnG",
+        "colab_type": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           }
         },
-        "colab_type": "code",
-        "id": "RlIWhyeLoYnG"
+        "cellView": "code"
       },
-      "outputs": [],
+      "cell_type": "code",
       "source": [
         "# Import TensorFlow.\n",
         "import tensorflow as tf\n",
         "\n",
         "# Import TensorFlow eager execution support (subject to future changes).\n",
-        "import tensorflow.contrib.eager as tfe"
-      ]
+        "tfe = tf.contrib.eager"
+      ],
+      "execution_count": 0,
+      "outputs": []
     },
     {
-      "cell_type": "markdown",
       "metadata": {
-        "colab_type": "text",
-        "id": "H9UySOPLXdaw"
+        "id": "H9UySOPLXdaw",
+        "colab_type": "text"
       },
+      "cell_type": "markdown",
       "source": [
         "# Step 2: Enable eager execution\n",
         "\n",
@@ -69,30 +85,30 @@
       ]
     },
     {
-      "cell_type": "code",
-      "execution_count": 0,
       "metadata": {
-        "cellView": "code",
+        "id": "WPTUfGq6kJ5w",
+        "colab_type": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           }
         },
-        "colab_type": "code",
-        "id": "WPTUfGq6kJ5w"
+        "cellView": "code"
       },
-      "outputs": [],
+      "cell_type": "code",
       "source": [
-        "tfe.enable_eager_execution()"
-      ]
+        "tf.enable_eager_execution()"
+      ],
+      "execution_count": 0,
+      "outputs": []
     },
     {
-      "cell_type": "markdown",
       "metadata": {
-        "colab_type": "text",
-        "id": "twBfWd5xyu_d"
+        "id": "twBfWd5xyu_d",
+        "colab_type": "text"
       },
+      "cell_type": "markdown",
       "source": [
         "# Step 3: Interactively Use TensorFlow!\n",
         "\n",
@@ -102,20 +118,18 @@
       ]
     },
     {
-      "cell_type": "code",
-      "execution_count": 0,
       "metadata": {
-        "cellView": "code",
+        "id": "ngUe237Wt48W",
+        "colab_type": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           }
         },
-        "colab_type": "code",
-        "id": "ngUe237Wt48W"
+        "cellView": "code"
       },
-      "outputs": [],
+      "cell_type": "code",
       "source": [
         "print(tf.add(1, 2))\n",
         "print(tf.add([1, 2], [3, 4]))\n",
@@ -131,32 +145,32 @@
         "# Most TensorFlow ops are directly usable with eager execution, giving\n",
         "# results immediately.\n",
         "print(tf.contrib.signal.hamming_window(x * y + 1))"
-      ]
+      ],
+      "execution_count": 0,
+      "outputs": []
     },
     {
-      "cell_type": "markdown",
       "metadata": {
-        "colab_type": "text",
-        "id": "IDY4WsYRhP81"
+        "id": "IDY4WsYRhP81",
+        "colab_type": "text"
       },
+      "cell_type": "markdown",
       "source": [
         "Numpy arrays are supported, too:"
       ]
     },
     {
-      "cell_type": "code",
-      "execution_count": 0,
       "metadata": {
+        "id": "lCUWzso6mbqR",
+        "colab_type": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           }
-        },
-        "colab_type": "code",
-        "id": "lCUWzso6mbqR"
+        }
       },
-      "outputs": [],
+      "cell_type": "code",
       "source": [
         "import numpy as np\n",
         "\n",
@@ -168,14 +182,16 @@
         "\n",
         "print(\"Multiplied by 42:\")\n",
         "print(tf.multiply(ones, 42))"
-      ]
+      ],
+      "execution_count": 0,
+      "outputs": []
     },
     {
-      "cell_type": "markdown",
       "metadata": {
-        "colab_type": "text",
-        "id": "PBNP8yTRfu_X"
+        "id": "PBNP8yTRfu_X",
+        "colab_type": "text"
       },
+      "cell_type": "markdown",
       "source": [
         "# Step 4: Define and Print TensorFlow Variables\n",
         "\n",
@@ -183,73 +199,66 @@
       ]
     },
     {
-      "cell_type": "code",
-      "execution_count": 0,
       "metadata": {
-        "cellView": "code",
+        "id": "3Twf_Rw-gQFM",
+        "colab_type": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           }
         },
-        "colab_type": "code",
-        "id": "3Twf_Rw-gQFM"
+        "cellView": "code"
       },
-      "outputs": [],
+      "cell_type": "code",
       "source": [
-        "x = tf.get_variable(name=\"x\", shape=[], dtype=tf.float32, initializer=tf.zeros_initializer)"
-      ]
+        "x = tfe.Variable(0.)"
+      ],
+      "execution_count": 0,
+      "outputs": []
     },
     {
-      "cell_type": "markdown",
       "metadata": {
-        "colab_type": "text",
-        "id": "45G7094TxsMb"
+        "id": "45G7094TxsMb",
+        "colab_type": "text"
       },
+      "cell_type": "markdown",
       "source": [
         "## Printing TensorFlow Variables"
       ]
     },
     {
-      "cell_type": "code",
-      "execution_count": 0,
       "metadata": {
-        "cellView": "code",
+        "id": "UJBJeZ5XxuwA",
+        "colab_type": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           }
         },
-        "colab_type": "code",
-        "id": "UJBJeZ5XxuwA"
+        "cellView": "code"
       },
-      "outputs": [],
+      "cell_type": "code",
       "source": [
         "# This does NOT print the Variable's actual value:\n",
         "print(\"Printing a TensorFlow Variable:\")\n",
         "print(x)\n",
         "print(\"\")\n",
         "\n",
-        "# A TensorFlow variable represents a reference to a tensor.\n",
-        "# The `read_value()` method provides access to the current value of the\n",
-        "# variable. Tensorflow Variables are automatically initialized according to the\n",
-        "# semantics defined in tf.get_variable().\n",
-        "print(\"Printing a TensorFlow Variable's value using .read_value():\")\n",
-        "print(x.read_value())\n",
-        "print(\"\")\n",
         "\n",
-        "print(\"Printing a TensorFlow Variable's value using .read_value().numpy():\")\n",
-        "print(x.read_value().numpy())"
-      ]
+        "print(\"Printing a TensorFlow Variable's value as a numpy array:\")\n",
+        "print(x.numpy())"
+      ],
+      "execution_count": 0,
+      "outputs": []
     },
     {
-      "cell_type": "markdown",
       "metadata": {
-        "colab_type": "text",
-        "id": "2njjWHcTpBEn"
+        "id": "2njjWHcTpBEn",
+        "colab_type": "text"
       },
+      "cell_type": "markdown",
       "source": [
         "## Changing a TensorFlow Variable's value\n",
         "\n",
@@ -257,64 +266,64 @@
       ]
     },
     {
-      "cell_type": "code",
-      "execution_count": 0,
       "metadata": {
+        "id": "v3wr6Erbo_hB",
+        "colab_type": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           }
-        },
-        "colab_type": "code",
-        "id": "v3wr6Erbo_hB"
+        }
       },
-      "outputs": [],
+      "cell_type": "code",
       "source": [
         "x.assign(42)\n",
-        "print(x.read_value())\n",
+        "print(x)\n",
         "\n",
         "x.assign_add(3)\n",
-        "print(x.read_value())"
-      ]
+        "print(x)"
+      ],
+      "execution_count": 0,
+      "outputs": []
     },
     {
-      "cell_type": "markdown",
       "metadata": {
-        "colab_type": "text",
-        "id": "uhtynjHVpTB5"
+        "id": "uhtynjHVpTB5",
+        "colab_type": "text"
       },
+      "cell_type": "markdown",
       "source": [
         "## Use a Variable just like any other Tensor"
       ]
     },
     {
-      "cell_type": "code",
-      "execution_count": 0,
       "metadata": {
+        "id": "7PbktdnHoehR",
+        "colab_type": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           }
-        },
-        "colab_type": "code",
-        "id": "7PbktdnHoehR"
+        }
       },
-      "outputs": [],
+      "cell_type": "code",
       "source": [
         "print(x + 3)\n",
         "\n",
         "# This code will broadcast the value across the list of numbers:\n",
         "print(x * [1, 2, 4])"
-      ]
+      ],
+      "execution_count": 0,
+      "outputs": []
     },
     {
-      "cell_type": "markdown",
       "metadata": {
-        "colab_type": "text",
-        "id": "GVChqwlwy1SI"
+        "id": "GVChqwlwy1SI",
+        "colab_type": "text"
       },
+      "cell_type": "markdown",
       "source": [
         "# Step 5: Debug Errors with Instant Feedback\n",
         "\n",
@@ -326,60 +335,58 @@
       ]
     },
     {
-      "cell_type": "code",
-      "execution_count": 0,
       "metadata": {
-        "cellView": "code",
+        "id": "23ap04N0v4k0",
+        "colab_type": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           }
         },
-        "colab_type": "code",
-        "id": "23ap04N0v4k0"
+        "cellView": "code"
       },
-      "outputs": [],
+      "cell_type": "code",
       "source": [
         "vector = tf.constant([10.0, 20.0, 30.0, 40.0])"
-      ]
+      ],
+      "execution_count": 0,
+      "outputs": []
     },
     {
-      "cell_type": "code",
-      "execution_count": 0,
       "metadata": {
-        "cellView": "code",
+        "id": "FCUMsIYxxRRa",
+        "colab_type": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           }
         },
-        "colab_type": "code",
-        "id": "FCUMsIYxxRRa"
+        "cellView": "code"
       },
-      "outputs": [],
+      "cell_type": "code",
       "source": [
         "# Works, because the values of `begin` and `size` (the 2nd and 3rd input\n",
         "# arguments) are within the bound of `vector`.\n",
         "print(tf.slice(vector, [1], [3]))"
-      ]
+      ],
+      "execution_count": 0,
+      "outputs": []
     },
     {
-      "cell_type": "code",
-      "execution_count": 0,
       "metadata": {
-        "cellView": "code",
+        "id": "T8me2oCNxpFp",
+        "colab_type": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           }
         },
-        "colab_type": "code",
-        "id": "T8me2oCNxpFp"
+        "cellView": "code"
       },
-      "outputs": [],
+      "cell_type": "code",
       "source": [
         "# The following does NOT work, because the value of `size` (the 3rd\n",
         "# argument) causes the indices to go out of the bounds of `vector`. The\n",
@@ -388,87 +395,86 @@
         "  print(tf.slice(vector, [1], [4]))\n",
         "except tf.OpError as e:\n",
         "  print(\"Caught error: %s\" % e)"
-      ]
+      ],
+      "execution_count": 0,
+      "outputs": []
     },
     {
-      "cell_type": "markdown",
       "metadata": {
-        "colab_type": "text",
-        "id": "irxJhAgar84v"
+        "id": "irxJhAgar84v",
+        "colab_type": "text"
       },
+      "cell_type": "markdown",
       "source": [
         "# Step 6: Using the GPU\n",
         "\n",
-        "You can place Tensors on the GPU by calling a Tensor's `.gpu()` method.\n",
+        "You can explicitly place Tensors on the GPU by calling a Tensor's `.gpu()` method. The `.device` property tells you whether the Tensor is backed by CPU or GPU memory.\n",
         "\n",
         "The first operation executing on the GPU may be slow as TensorFlow initializes. Subsequent uses will be much faster."
       ]
     },
     {
-      "cell_type": "code",
-      "execution_count": 0,
       "metadata": {
+        "id": "7J4N9baqaKCL",
+        "colab_type": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           }
-        },
-        "colab_type": "code",
-        "id": "7J4N9baqaKCL"
+        }
       },
-      "outputs": [],
+      "cell_type": "code",
       "source": [
-        "# The example code from here on will work only if your notebook\n",
-        "# is running on a machine with a functional CUDA GPU. The following\n",
-        "# line checks that.\n",
-        "is_gpu_available = tfe.num_gpus() \u003e 0\n",
-        "\n",
         "# Create some Tensors\n",
         "SIZE = 1000\n",
-        "cpu_tensor = tf.random_normal([SIZE, SIZE])\n",
+        "tensor = tf.random_normal([SIZE, SIZE])\n",
+        "print(tensor.device)\n",
         "\n",
-        "if is_gpu_available:\n",
-        "  gpu_tensor = cpu_tensor.gpu()\n",
+        "\n",
+        "if tf.test.is_gpu_available():\n",
+        "  gpu_tensor = tensor.gpu()\n",
+        "  cpu_tensor = tensor.cpu()\n",
         "else:\n",
-        "  print(\"GPU not available.\")"
-      ]
+        "  print(\"GPU not available.\")\n",
+        "  cpu_tensor = tensor"
+      ],
+      "execution_count": 0,
+      "outputs": []
     },
     {
-      "cell_type": "code",
-      "execution_count": 0,
       "metadata": {
+        "id": "4E-2n7VbzY1n",
+        "colab_type": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           }
-        },
-        "colab_type": "code",
-        "id": "4E-2n7VbzY1n"
+        }
       },
-      "outputs": [],
+      "cell_type": "code",
       "source": [
         "# Time a CPU-based matrix multiplication\n",
         "\n",
         "print(\"Time to conduct matmul on CPU:\")\n",
         "%time tf.matmul(cpu_tensor, cpu_tensor)"
-      ]
+      ],
+      "execution_count": 0,
+      "outputs": []
     },
     {
-      "cell_type": "code",
-      "execution_count": 0,
       "metadata": {
+        "id": "vbSFW-T5zhZF",
+        "colab_type": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           }
-        },
-        "colab_type": "code",
-        "id": "vbSFW-T5zhZF"
+        }
       },
-      "outputs": [],
+      "cell_type": "code",
       "source": [
         "# Time GPU-based matrix multiplications.\n",
         "\n",
@@ -481,51 +487,9 @@
         "  # Subsequent uses are much faster:\n",
         "  print(\"Time to conduct second matmul on GPU:\")\n",
         "  %time tf.matmul(gpu_tensor, gpu_tensor)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 0,
-      "metadata": {
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          }
-        },
-        "colab_type": "code",
-        "id": "E5pIOe3Rz7iW"
-      },
-      "outputs": [],
-      "source": [
-        "# Second timing demo for GPUs, after it has been used once:\n",
-        "\n",
-        "cpu_tensor = tf.random_normal([SIZE, SIZE])\n",
-        "print(\"Time to conduct CPU matmul:\")\n",
-        "%time tf.matmul(cpu_tensor, cpu_tensor)\n",
-        "print()\n",
-        "\n",
-        "if is_gpu_available:\n",
-        "  gpu_tensor = cpu_tensor.gpu()\n",
-        "  print(\"Time to conduct GPU matmul:\")\n",
-        "  %time tf.matmul(gpu_tensor, gpu_tensor)"
-      ]
-    }
-  ],
-  "metadata": {
-    "colab": {
-      "default_view": {},
-      "name": "Eager Execution Tutorial: Basics",
-      "provenance": [
-        {
-          "file_id": "0B0kLcpwLFwKEVm9XNkFueGk4bTg",
-          "timestamp": 1504118841551
-        }
       ],
-      "version": "0.3.2",
-      "views": {}
+      "execution_count": 0,
+      "outputs": []
     }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}
+  ]
+}
\ No newline at end of file
diff --git a/tensorflow/contrib/eager/python/examples/notebooks/2_gradients.ipynb b/tensorflow/contrib/eager/python/examples/notebooks/2_gradients.ipynb
index e6c7c11733..1e65b27bc8 100644
--- a/tensorflow/contrib/eager/python/examples/notebooks/2_gradients.ipynb
+++ b/tensorflow/contrib/eager/python/examples/notebooks/2_gradients.ipynb
@@ -43,11 +43,9 @@
         "# Import TensorFlow.\n",
         "import tensorflow as tf\n",
         "\n",
-        "# Import TensorFlow eager execution support (subject to future changes).\n",
-        "import tensorflow.contrib.eager as tfe\n",
         "\n",
         "# Enable eager execution.\n",
-        "tfe.enable_eager_execution()"
+        "tf.enable_eager_execution()"
       ]
     },
     {
@@ -106,7 +104,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 3,
+      "execution_count": 0,
       "metadata": {
         "cellView": "code",
         "colab": {
@@ -114,34 +112,30 @@
             "startup": false,
             "wait_interval": 0
           },
-          "height": 360,
-          "output_extras": [
-            {
-              "item_id": 1
-            }
-          ]
+          "base_uri": "https://localhost:8080/",
+          "height": 347
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 127,
+          "elapsed": 374,
           "status": "ok",
-          "timestamp": 1505502830690,
+          "timestamp": 1525154227149,
           "user": {
             "displayName": "",
             "photoUrl": "",
             "userId": ""
           },
-          "user_tz": 240
+          "user_tz": 420
         },
         "id": "O4lsC4ckAcar",
-        "outputId": "2f760690-cafb-4777-b970-91d839f99faf"
+        "outputId": "f8becb3f-498b-4cb7-9ef3-608a68cb65d0"
       },
       "outputs": [
         {
           "data": {
-            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAesAAAFXCAYAAACC+2avAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAIABJREFUeJzsnXt8VPWd99+TK7kykxtJQIebqZfaqogtrhKNa1ooEKl9\nCrpVn9ZNW6x9VWsbCi7aVUt01NZ9tq21KVZlFey2YkQNohhj3QWK2liCF5RIBCc3yEwmIZnMTOY8\nf/zmzJwzSSBAYibh+369eIU5c87vXLh8zvdu0TRNQxAEQRCEmCVurC9AEARBEISjI2ItCIIgCDGO\niLUgCIIgxDgi1oIgCIIQ44hYC4IgCEKMI2ItCIIgCDHOiIj16tWrufjii1m8eHF4269//Wvmz5/P\n0qVLWbp0Ka+//vpInEoQBEEQTjksI1Fn/eabb5KWlkZFRQWbN28GlFinpaXx7W9/+6QvUhAEQRBO\nZUbEsr7wwgvJzMwcsF36rQiCIAjCyTOqMesnn3ySsrIybr/9drq6ukbzVIIgCIIwYRk1sb722mt5\n5ZVXqK6uJicnh8rKytE6lSAIgiBMaEZNrLOysrBYLAB885vfZPfu3cc8RtzmgiAIgjCQhJFaKFpo\n29vbyc3NBeDll1+mqKjomGtYLBba2yeuuzw3N0Pubxwzke9vIt8byP2Nd06F+zsWIyLWt912Gzt3\n7sTtdnPZZZfxwx/+kJ07d/Lee+8RFxfH1KlTueuuu0biVIIgCIJwyjEiYv3ggw8O2Hb11VePxNKC\nIAiCcMojHcwEQRAEIcYRsRYEQRCEGEfEWhAEQRBiHBFrQRAEQYhxRKwFQRAEIcYRsRYEQRCEGEfE\nWhAEQRBiHBFrQRAEQYhxRKwFQRAEIcYRsRYEQRCEGEfEWhAEQRBiHBFrQRAEQYhxRKwFQRAEIcYR\nsRYEQRCEGEfEWhAEQRBiHBFrQRAEQYhxRKwFQRAEIcYRsRYEQRCEGCdhrC9AEARBOHXo6HCzcmUt\nTU2Z2O2dOBwl2GzWsb6smEfEWhAEQfjMWLmylurq6wAL9fUasJ6qqqVjfVkxj7jBBUEQhM+MpqZM\nwBL6ZAl9Fo6FiLUgCILwmWG3dwJa6JOG3e4Zy8sZN4gbXBAEQfjMcDhKgPWhmLUHh+Pysb6kcYGI\ntSAIgvCZYbNZJUZ9AogbXBAEQRBiHBFrQRAEQYhxRKwFQRAEIcYRsRYEQRCEGEfEWhAEQRBiHBFr\nQRAEQYhxRKwFQRAEIcYRsRYEQRCEGEfEWhAEQRBiHBFrQRAEQYhxRKwFQRAEIcYRsRYEQRCEGEfE\nWhAEQRBiHBFrQRAEQYhxRKwFQRAEIcYRsRYEQRCEGEfEWhAEQRBiHBFrQRAEQYhxRKwFQRAEIcYR\nsRYEQRCEGEfEWhAEQRBiHBFrQRAEQYhxRKwFQRAEIcYRsRYEQRCEGEfEWhAEQRBiHBFrQRAEQYhx\nRKwFQRAEIcYRsRYEQRCEGCdhrC9AEARBODE6OtysXFmL02mjsLADh6MEm806rGOamjKx2zuHdYww\n9oyIWK9evZrXXnuN7OxsNm/eDEBnZye33norn376KdOmTeOhhx4iIyNjJE4nCIIgACtX1lJdfR1g\nATRgPVVVS037RIuzz9dDTc33AQv19YMfI8QeI+IG//rXv866detM237/+98zb948XnrpJb70pS/x\nyCOPjMSpBEEQhBBNTZkooQawhD6b0QW9vv4qqquvZ/v27mMeI8QeIyLWF154IZmZ5j/wbdu2sXSp\neltbunQpr7zyykicShAEQQhht3eiLGoADbvdM2CfaEGH7GMeI8Qeoxaz7ujoICcnB4Dc3FxcLtdo\nnUoQBOGUxOEoAdaHYtYuHI7LAbPru61tD1AM2AAXkyY5sVr/CBxi3rwMHI5FY3cDwrCJuQSz3NyJ\nHdeW+xvfTOT7m8j3BhPz/uLi+klOTgQgOTmBnJwMsrIyuPnm5w2x7DKmTbuPgoJzaG7ew8GDt6PH\nuDMyNpKdncFNNz3Pxx+nM2NGFw8/vJCsrNhLOJuIf37Hw6iJdXZ2NocOHSInJ4f29naysrKGdVx7\ne9doXdKYk5ubIfc3jpnI9zeR7w0m7v2Vlz8XFuVduzT6+lSy2N69KRhd3zk5Z/LCC5dRWtrPwYOR\n7Xv3pnDjjYOvEUtM1D8/neG8iIxYnbWmaabPJSUlPPPMMwBs2rSJK664YqROJQiCIDB0gtlQsezB\ntg8nSU0Ye0bEsr7tttvYuXMnbrebyy67jB/+8Id897vf5Uc/+hF/+ctfKCws5D/+4z9G4lSCIAhC\nCLu9M1R+pdzauijrsWxVruUJx7JXrZrDrl2VuFzTsNkOsnr1EtaufWvQNYTYYkTE+sEHHxx0+2OP\nPTYSywuCIAiDMFSCmc1mHdSVXVn5Nk7nKsBCb6/G2rXrhxR2IbaIuQQzQRAEYXjoojxYTHewTmWD\nubyHEnYhthCxFgRBmIAYu5vpncrsdk1c3uMUEWtBEIQYYai+3SfSz3swK/rpp+cgLu/xiYi1IAhC\njDCYNVxVtXTI7UdjsOQzcXmPX0SsBUEQYoShyqhOpLxKEscmFiLWgiAIMcJQpVjm7S7a2t6ltJSw\nS3ywphojYUXLOM3YQcRaEAQhRhjKGjZub2t7F6dzFU6nconX1T1AaelU7r770mEL6XBF+ETc78Lo\nIGItCIIQIwxlDRu3l5aC0xlxibvdZ/KnPy06rjahwxVh6W4WO4xYu1FBEARh9IluGQpqPvXxCOlw\nRXg4IziFzwaxrAVBEMYRuku8ttaPx5MCLAQ0CgoODXuNoWLjQ51LktTGHhFrQRCEcYTuEr/hhv+i\npiYBeBY4xNtvu3G53ANiz4PFp4crwlLqFTuIWAuCIIxDmpsLgF5gOWChtVWjomJg7Hmo+LSI8PhC\nxFoQBGEcEG0hFxT4qK+fwrFiz5IkNjGQBDNBEITPiI4ON+Xlmygt3UZ5+TO4XO5hf69byPX1V1Fd\nfT0QoLBwN8dKAJMksYmBWNaCIAifEdEu6V27KqmtvS4cZz5aSVW0hdzcXEBt7SIqKgaOyDQiSWIT\nAxFrQRCEz4howXU6P09FRe2Qgmx0WRcUNFNf/xSQAXgoKPAcdUSmznCTxKRbWWwjbnBBEIRBOJbL\n+kQwu6RdwLts3Up4/aO7rBOBa4DFwLWhzyNHtJu9oqJ2RNcXTg6xrAVBEAZhNFptOhwl7NpVidP5\neeBdYCW9vRaqq9X6DkcJfX3r2LEjDjiMz5cWLsdqbs7B7AbPOalriUYS0WIbsawFQRAG4XjFaziW\nuM1mpbb2OsrK3KSkFA5Y32azkpychNv9bdzun1JTsyJs4UZb3QUFLeHzLVv21Elb/pKIFtuIZS0I\ngjAIw+3ypROxxDupr3+RurqXKS6OHxD71WPI5eXPhCxq8/pDvSREJ4r5fAkmy/94eoMb0WPV+/Yl\nUFhYSXZ2ETNn9kgiWowhYi0IgjAIx5tFHRHZGmABbvcWqqvT2LXrCWprrx+QrOVwlODzPcL27V1A\nNj5ffzhuPdhLQnSiWGnpNoZr+R8teczo7geNuXNlslYsImItCIIwCMfbajMisunAFvTOYk7n4kE7\ni9lsVpKSUnG7vwdYqKnRSEpaf9SXBKPotrXtAcoYjuV/PCVhEquOTUSsBUEQBmEoa3So7brI1tW1\n4HafyXAEcDChPNpLgtkKLqawsJK8vLMpKurl7ruHtvyPJsjH6+4XxgYRa0EQhEEYyhodarsusi6X\nm8svfwKnczFDCaAu+Pv3t6CSuoYnlGbRtZGXdzZbt15Bbm4GH3xwgPLyTYO6uo8myNI0ZXwgYi0I\nQswylo06IsLoBmrC9dD79iUQbaV2dLi55ZaXQiVXh5gzJ5kvfnEdzc052O0eVq26wCSkPl8PNTXf\nBzqBDVitXoqLE44plEcT3aO5uo8myDJZa3wgYi0IQswyGrXOwyXSMcwJ3Bauhy4srCTaGl65spYt\nW24Mb9u2bQNlZQG2br0CgPLyTab7sFofCO1rBa5l+vRnqaq6Ilz+NdTLiS66+/bF09HRRGNjEeXl\nz/Doo2VHdXWLII9/RKwFQYhZxjb5Se8Y9rzpGrKzi5g7V1mpBQUt+HwJvPZaErABWIgS4AyamvrD\nK0XfB2SjOphtAdJoa9uDyzXnmC8nuuhef/3TNDSswum0sHu3xo03PoHdjsSeJzAi1oIgxCxjmfwU\n6RjWhdGSnjmzJyygRotZ7fMgUAD00NbWTmkphnGWkTXmzQvyzjsP43SuwpgxPtyXE+Vuj+xXV6ex\nY8cVSOx54iJiLQhCzDKWyU+RF4WFDBVXHmgxfw5YRHLy7TidP8XptFFfr7Fgwe8oKzPex1dYtuwt\nnM7IsVu3gs023HKsQxhfIOCQuLonOCLWgiDELCMtQEdLWIv+bvXqOUReFAI4HFcOSG6LtvyhO/T7\n01Eu7nSgiwMHMnn11SVHPba3N5He3psoLKwkK6uIjo697Ntnp7z8mQGx63nz0qmp2YCawNXF/Pky\nHWuiI2ItCMIpw2Ax4fvuu5yVK2upqwvgdicDl1FfP5nBktmGEvTaWj8eTwrKCteAg4BqdgIaHR2V\nA65F9xps3Qq9vX7Uf8dv0NOTwFlnfUJDw3SczgwaGjz4fM/z+OPfCl8DJGG1eoGDzJuXwaOPXkN/\n/4BTCBMIEWtBEE4ZBosJR7fbhI3ANYPGi4dKALvhhv+ipkYD/gDk4PMlocqyAGpwuQoHWMjmHuEp\nqGQ2C273Il5//U7g1vA1bd/+AKCEuqRkfTjWDarrWVaWdch51sLEQMRaEIRThsES1gbGndOJjhfr\nFvXWrRj27WTz5oMUFf03gcCh0PbbAQuapqGywy3ActMYzGhr3eEooa7uZdzuyDX090+PuqZsQL0s\nqPGa0h70VEPEWhCEU4bBEtYqKl41CbjV+j7FxS5TIlnEot5AJLHrRYLBVSGR1YDHMQusDzWF+OjC\narNZ+fKX+9myJXINOTkHaWszZ4+D7hno5ni6ngkTAxFrQRBOGWw2azhG3dSUSUXFq1GJZB4cjuUD\nEsn27YtHucctwL1YLFPQNLMQQzvmDG0L8Klp2/vvv0lJyRFmzQqYXOIWSwD1IqASxs49N430dHP2\nOOiegSWha0mjsLABh+O6UXteQuwgYi0IQkwQnby1atUcKivfHvFWoyfSFa2jowmIxImTk9fg9Z6F\nWZyt6CIKO1BW9W3AfcDZwBG83ttoaNhCQ8P1pvM2NxcAV4XPd/jws2zYcMWA61Cegc2hZ+LG4bgO\nTYNlyzawd2/KZ96SVfjsELEWBCEmiBbRXbsqw4lUx9NqdLDyrNzcjPD3J9IVbfLkqTidG9FLsU47\nrZDZsz1s3/4A3d2ZBAKTQmumAW8CPwXeAGzAOcBiw2rpA8473OYvg5WyRbcy/SxbsgqfHSLWgiDE\nBNEi6nJN40QSqQaznJ999vrw98cSxsHEvrPzU4yW9ZEjlfzqV9excmUtjY2pHD78AR5PBt3de1Ad\nzGxEOp8ZO6C5gJ1AO3v2OLnhBicPPbT4pJq/yDzqUwMRa0EQYoJoEbXZDtLbG/mcn38oPOSioKAZ\nSAxNtTK7fo8lXqtWzWHnzntoa8sjPv4Q3d3puFzu8PGDiX12dpGp21h2dtGAkq/k5DXARcAelCir\nzmeZmR34fHfg9Z4P7AXuBiz4/VqosckLJCWlnrC7X+ZRnxqIWAuCMGocz4jLaOty9eolrF0b+ezz\n+amuVpOt1DSsaxisucn+/QHgSeBrwOQB4lVZ+TYtLbOAawgGLWzbplFREXEdDyb2M2d2snu3uT94\n9H59fRcBS1Au7/tISSmktBQcjjKWLXuL+vqrgM2mYyCD7ds/xe3+HsNxYw/2PB2OEpKTN4Zi1tIT\nfKIiYi0IwqhxPMlcg8Vjq6rs4d+Xlm4jInQZGEVv61bYtesJnM6bUC5oNYayuHjKAPFSmd1O1DSt\nLmAhdXUBGhubqKx8m/37W4gujVq1ag67dlXick3DZjvA6tVl3HnndswJZkfC1wNTuOyyI1RVqa5j\nEeu3K+qYLlQN9fDc2EM9z6efvkaaokxwRKwFQRg1oq3PfftSB8xr1jSGZX2b3b0ejKKn+mqvRu8+\nBhamTz+DqqqBGdXRmd2wAbd7Epde+if8/n9HdR4zD+6oqKgNJ7v19mosXVpJd7debuUDmoHvh86g\nAclApP+ncQ71oUP30NMzlbi4w8yblw4khLqfHduNLfHpUxcRa0EQRo3oeGpHx14aGswZ3sCg1uJQ\nfbhVL20f8ERo3URgAZFsbDia6EXHn5XYXoXfHwh9tqLizVU0NZ1BRcWrNDamYRTJSBexxSjX9lVA\nDSrT+wPgX2lufi18zqMNJHG53CQlDS+5TOLTpy4i1oIgjBrRceh9++wGoeykrq6Vvr4pKAt1IWAN\nW4tDuXxVL20Vu1ax6eXo4lVY2EBeXvCoojdz5pFQ/LkTeDG09QXgI4zdydzun1Bfr85dWLiWgS5v\njYgrezLqheFFIAd4gYKCwYV0sLjzcEutxnJkqDC2iFgLgjBqRFuU5eXP0NBgFkTzAI3lYWtxKJev\nUbCUIK4LZYV7cDiuO2YmtX58bW0LHs9PDedfh8redtHb68bvj8S0s7KmM3euOmdb27s4nStQrvh7\ngQySklbS359Mf/9cVDvQBcBfBj3/8cTxT0bYhYmFiLUgCJ8ZRqHdv99rGl6RkuKntHR92FocyuUb\n/QJgFLSKilePWfqkH19auo36eqM73A/00tX1KZr2C4wx7Vmz+sOu+VtvPURPzyaOHPkYv//HgA2f\nL5Kdrr94NDfnDCq2xxN3PpFua8LERMRaEITPDKPQKnd2RIxLSzEJ0VAu32gB7Orq5NVXf4guaD7f\nOh5/fNmAc+vH7dsXT0dHE93d8Zhd25OBa9G05zCKqdXqxeG4ElDiWVNzI2ZvwDVEZ6dDGna7e1Cx\ntdu1YcedJaFM0BGxFgRhTDhW/HWopKxoAUxMXItR0LZvjxtwzOHDxjnQG1HZ4CrrOzPTS3d3C8Hg\nTaG9zVOtiosThmy4EkloM2enT5q0i9Wrl/G9731EtNg+/XT04BBJKBOOjYi1IAhjwtEypI9GtGD2\n9+dgtpAPDzjmpptqQhncnahJWJF4dFzcM+Tnazidk0N7LwDuwGqdQXFxAqtWXRAuN2tr2wMUo2q5\nXUyatAuLxU1m5l407S7a2s5HDez4MZdc8is0bSrwGCpbXDVoOZ77loQyQUfEWhCEMWG43c2i9yso\n8Jmszby8Vlpa9PGSrfT2urDbN2GzHWDTpjJmzLDz8cdqAIfK1r4NYzwaDvPHP17OkiVr6OubgcXy\nMZdcks4f/nAlmgaXXfY4LS0/ALYA55KYuJaUlCx6erLwes8EvkZv72Ss1gdQHcwUfv+Foc9DN2g5\nFif6QiNMPESsBUEYE4abPBW934IFv+OKKx6hrs5CMHiY/v4errjiEIcPp/L++014vSo5rLfXRXHx\nL5k9+4vs21cPfA7owezG9jFvXjq//e1H9PWpnt2appGVtR5Ng5KS9bS0fAEl1KpEzO/vxu83J5Op\nuHU2Q3U0G6pBiyAMFxFrQRDGhOEmT0Xv19xcQFvbuwQCqrlKe7vGe+9VUl9/RSimq++7Ba/3Lhoa\nLMDVKFFNxyioU6Y0AqezdStE13qvXFkbcp13o4+1VEQnk6k1580LkpS0nrq6AG53K8aOZhJrFk6W\nURfrkpIS0tPTiYuLIyEhgT//+c+jfUpBEMaI4xncEZ08ZZyqZTx2sCSrDz4wj89U4zTBZjsQmtTV\nCfQxUFQvJTPzfk4/fSYdHXvp7k4JZXfrDVKeBRIpKPDQ1FRApGb6d6huZQNbnVqt71Nc7MLh+Ao2\nmxWXy80ttzzP9u1/ALKZNy+Iw/GVkXvIwinJqIu1xWJh/fr1TJ48+dg7C4IwrhnKtR0t4qtWzaG7\n20Ni4lr6+3PIzm7i7bcTaWubA3RTX78E2ExV1VJWrDiDmprb8fnsQCtvvHGE9HSLaXympn1ISclL\n+P3dJCSsIRBIAaZjdkt3A5NJTw8wa1ZPqO3p86HvazDXSa8LvSQsAZ4DJmOxrCEjYzpz5/aQlGRs\nxLLc9EJis1l5/PFvfRaPWziFGHWx1jSNYDA42qcRBCEGGMq1HS3iu3ZV4nTmoeK8GbS3dwA/wxgH\nbmrKZN++JhYufJ5gMNKk5PDhDcTH/4NJk9agaTPw+/fh9f6UhgYbEXd3MlBCxPX9D1QG94N0d/tp\nbEwNradPwUrH7GrPCZVYbWbfvgQ6OtxkZ5/HzJlHcDiWhsW5o8NNRcXwPAmCcDJ8Jpb1jTfeiMVi\nYdmyZXzzm98c7VMKgjBGDFUXHC3iym3dBhgbjAxsKnL11c8RDH4u6rsM+vtn0t9fTmFhJU7nl1FC\nrH/vB3YDS1HWsga8AawGLHg8GocP672+F6Ji1fuARabr1jOxy8s30dCwCqfTEuopHvEWqNptFdeu\nr19CX99fSE5OEvEWRpxRF+uNGzeSm5tLR0cH3/72t5k5cyYXXnjhaJ9WEIQxYKi64GgRV7HlwtBn\nN7AntIKKERcWNrBq1RIuvvgg8CEDZ0C3As/jdPaihHmx4ftEYAZKhDOIWM+6ld1FZmYuUGmYnnUd\ncB9wNoWFDTgc14Xv6WjeAn1spr7+jh1xuN3SHlQYeUZdrHNzcwHIysriyiuvZPfu3UcV69zcjNG+\npDFF7m98M5Hv70Tv7fBhNzfdVMPHH6czY0YXjz66hKwsszX56KNlrFixMbRPN2vXfotLLnmMlhYN\nFS+OuMCnTbuPd965iRUraggGVwGfALejYtDtKCs6H7gUaADsqIEaBSj394LQmkYSME7n6u6+j6lT\nz8XpXBzeIzW1kEWLjvDwwzeZrr+oqMf0olFU1EtubgZOp41ob4DF8qlpm9Np+8z+zkzkv5sw8e/v\nWIyqWPf29hIMBklLS6Onp4c33niDm2+++ajHtLd3jeYljSm5uRlyf+OYiXx/J3Nv5eXPhePRu3Zp\n9PUNZk3G8+tfLzJtOf/8PGpqNgD6HGkACzk5Z9LfH8/evSmh7XagAvgD8AVU/PkHRIu8Emz98/6o\n78wtSW222WRkfAw8hbK+D5OW9iF7987lO9+pNrmvf/zjL/DGG5W4XNOw2Q5w221ltLd3UVjYgdHi\nLyxs4ItftFJTY9zm+kz+zkzkv5twatzfsRhVsT506BA333wzFouF/v5+Fi9ezCWXXDKapxQEYYQY\nbhnWiQ6baG4uQLXhfAyj6L3zzm7OO28PZ51lrImeDEwFFpGfX09Ly2Sik8JU0xPlyk5IsBIIGL/L\nMp1j5swedu7sBH4Y3tbevoH29qvCCXB5eWdjt3fi8/nD7u7eXo21a9dTVWUfxOWvXOdJSdIeVBh5\nRlWsTzvtNKqrq0fzFIIgjBLD7TA2VFLZYOValZVvG9qGHgkd14MxvqxpBTidNxIM3kNZ2XoaG1Np\nb3+Pnh6NuLg/cs45mZx//jq2b+/A7Y4khcFeVCMSK3APA/uFb8Bq9VJcnIDDcTnnnVdLdOKa/nun\n8/M4nUuor9ewWv/IYC8jQ7UClRi1MBpIBzNBEAZluBbzUEllt9zyElu2qGzv+nqNF164g0DgrvDn\nBQvWsWDBOmpqfIbzgJpkZaGz024Yp9kTfnHYts1Fbu6DdHUB3IPFkk1c3Kf09/8EJdQagUAPkYSy\nbiwWK0uWBHA4rgx7B1SSmwvVSjQNleR2KcqKj7QKhUMYhV+6kQljwcBZcoIgCCiLWYkUgMb+/R9S\nXv4MLpfbtJ/NZuW++y7Hbvewb18ql1/+BCUlz7FtWwuqMxiAhUDAjlH8X3stgaSkROLjm4GvEmnr\n+S7gQtP2hs9lfnF4hvb2NPr7LwJmoWnX0N9fBNRgtT5KYWElKhltOSpLfDmTJ3cDsGzZW+F72LSp\njEmTfhnabwnwMzIz/5NJk+5AJao9BbiYNy+DsrL1nHfes5SVrT8h13ZHh5vy8k2Ulm4b9BkKwrEQ\ny1oQhEFxOEro61vHK68ECAS6cLu7qK6eyvbtf+Svf/22KX5tdJmDhtO5EZXBvQG4FiX6jRgt1N7e\nZKqrl2Ox3INxUIYS2Pvwem+jokJ1MTO72l1EN1BRMenFTJv2JJ98YkH913Ynykp2091dSHV1PHBZ\nKCb9MHl5ZzNp0gy83sgLRFzcJLzefwuvXVhYyUMPXXfStdLDDSkIwlCIWAuCMCg2m5Xk5CQCAWPj\nko20ta3hllseISkpNRx/bmxMY2AfbjXVCjaj6qKTgcdR86SnAN9ATbmaiu76jhxfAPyOLVuslJc/\nw+rVc9Bd7Q0NGVHJY16U21qjo6MJj2cFSvwvBHYBd4X214UdnE7V5ASexBzbzjZdR17e2SPS1ORE\nk/AEQUfEWhCEIYkWGV2Et2/vwu3+HrqlOGXKHShhzkANuuhFiV8Lqq92I5oWaRmqLG5QrmY/8L+Y\nG5skAT+jr+9RqqsTqav7G8XF8Tz99Bxuuul5tm0zCmwLaWk9/PM/r6exsQin02ilw8Dr7yISz+4h\nM/NeZs48C7vdg8/Xbyq9Gqn49FBJeIIwXESsBeEURs/YdjptFBZ2DCjPihYZFVceaIEePpyAcRBG\nQsJdBAIbUNnZk5k82YXbbRRNN/BrlKtcubaTk9fg988kGExBNTbRXd7fwe22UF2t3MdJSQA/B+ag\nLOrvEx9fFWoN+gy7dycSEWM9acyGPiHL6+3E6707fK3p6ZVs3apmTbtc7lEpvRoqCU8QhouItSCc\nwkTHmqNjqQ5HCUeOPMJrr2kEAk7i4rK45JKH2Lv3CGoaVTdwMYFAPkbxPuusc5g5s4emptcGtVhV\n1na84RgbfX3TgY+BuahxlQuAHAa6jzNRLvUl4evs6ckIX++WLY/Q16eL8SLgDmy2WcyfH4fDsZxv\nfGMnu3dH1szOLgqvM1Q51skyWusKpw4i1oJwCjBUg5NjxVJtNitPPfUvpm3l5ZtoabkFXXgtltvR\ntHOIxH5d7NnzNh99VITNtodHHilD0+CVV+7E75+NillfS2Lievx+o4A3Ar8wrZuRkYHHE+0+1qiv\nbzadr7+8jtPXAAAgAElEQVT/AKWl27DbO5k58wzee89oxV/A7NkJVFVdBsDMmUdCAzkiDVIEIdYR\nsRaEcc5wOo0NlY18tFjqcAVe07JQFnEVqnd3F8FgJb29quPX0qWVzJ07Db//34kI8wbmz8/kf/7n\nDrzeuSh39udN606ePJudO6/kllseYfv2LiAbn6+fn/98Hps3dxAMRlzdmvYL6uvVveXn/wJz0th7\nfPRRIeXlz+BwlIhLWhiXiFgLwjgnWojr6h6guDjPJNpDWdC6cKmYtcskXMMVeNU0pNLw+bemc7lc\nBQPOHxfnYc8eNxbLLJQrfSHK9R1Zd9IkJwBJSanhZLaaGo2kpPV85Svp1NQsN5wzsnZPTz6RjmgN\nwApcLls45l1VtXTYLunhtlwVhNFGxFoQxjnRQuh2n0l19SKM8eehLGg9ljrYoITIum6ghq1bCZdR\n9fWtY8eOOI4c2Y/ff6bp/Kq1Z+RcmtaI3T47dP5O4EWCwSAtLbOAr6FqoTcCC4iLu51g8MvAEVpa\nfkBFxeZBXzSefnoO77xTidM5DeVWj2SS9/a2ohLXQL0IbEHPAt+3L/64BFjqo4VYQcRaEMY5g2ds\nm+PPJ+L6LShopr7+KVRJViK9vUuorp7MSy+t4Z/+yYbb/R3gEeAAZrezF2OrT6/XRm1tVyi2nQXc\nZth3I3ANKSl9XHbZ0/zP/2Tg8fhQXczcvPiik/nzC03rt7Q0cMstzbhc01D/hX0/tE4asAO/f7ph\n//0YG6h89NEdfPnLTtzunzAcAZb6aCFWELEWhHGOLsR1dQHc7kkol7Kyns1WpMbTT885DjduIsZy\nLF1Yvd6LeP31N1EW60qUtfwEaiBHO6qLsdFFvQGP59rQPtEzoPuA5wgEGnj11UmGLO6rgY34/d+n\noeEOpky5k9bWTCCHlpZp1NT0oCzq7xPp7f0ukAt8E3gUVfaVazqf13s+Xm8iwxVgqY8WYgURa0EY\n5+iubJfLTUVFbbhcyuG4nIqKod24RiEvKurh7rsvNQl5c7O5bEpZyhpwhP5+O5GuY1ZUE5PridRG\n3wecg5o9nQc0AR+iyq6Mk7KSgCX4/Xpf8IENWDyeWSQnt2O2yO8AZgPVqBeEbuAW8vP/MzQ+MxX4\nDip2bbT6+1CW//AEWJLRhFhBxFoQJgjRtbwdHW7q6gIMZUVGx2O3blWJafooy/37WzAL3QcoUfwq\nmnY/0EYkVmxsF2pDCfWi0P7LUeJ9F8oK3xD6eRi4OXSMGo9pPp9qwKJp+4AZmIXc+HKgERdXyeLF\nm1m9+uusXbuerVuht9eC8jJsJDXVj9V6EKdzRegY87jM4T5TQRgrRKwFYYKycmUtbncyRgFsa3sX\nl2vOoCVYbvdsqqt7eeGFZwkEfoBqevI4CQkHiY930dcHyjLdiKaBGtDxBMpSNSd5KYu6m0gnsjwi\nVvi1obUnh36BalGqhFUJ//+GjrkTTZvOpEmfmu7DYslG0yLXnpmZHxbVqio75eXPhLK/rcByFi3a\nyN13XxdOWLPbzeMyBSHWEbEWhAmKEuPLiCR7fYDTuWKISVYaKuY7hUAgDlV+dROwhUDgCwQC/wvM\nAv7VsP8TKAs3A9iHxXIP8fF5BAKTQts04K+AhylTPqa11Xiut4F+LJZ7yMgoJDHxQw4f/hg4HdUi\nVG9peit9fRZaWlwUFlaSl3c2druH7m6LqT/4vHlB071Hu68ffngJ/f3xYiUL4xYRa0GYYOix6P37\nA8ALRMqjGgAL+/bFU16+iX37EigsrKS7Ow+PJxXoR8V6VwHPM3Bs5YOYXdFeIq7opSQn34HFkkIg\ncD1qulYkOe3ccx8hGFxDe7sVSEFlmF+IpnnxeBaQmPhbIuVWoCzvFoyu9by8s009vCsqjLHkr5ie\nQbT7OitrYGmaIIwnRKwFYYKgi3RdXWu4NElZqA8CU1GZ0y/S0dFEQ8Oq8PcLFqwjI8PCn/6Ui7KI\nLaj4cXTCVzbmmHIbcC9wJtCL16u3En0KJeSRY1991UJcXAoqSWwjymqPZJn7/YVRax9BxbUHTwST\nWLJwqiFiLQgThEjC2POYRfZzKMsYrFYv2dlFoVnO6vvm5hxefPEqtmypxOPJRAnkQuBhzHHoD4B7\nUOVQn6Cs9dNQ/43obvR7UWIcB/wB1VAlm2CwjWBwNsYs78j1pQE+EhPvxO+/ECXUXwX+jB7DLixs\nwOG4bljPYbCmJ7m5GcN/kIIQg4hYC8I4xihMjY3NKGs0Oqv6g9C2i0lNbaGpqdXwvYuWlgYuuiie\n1FQfHs8hIsM0koG1wNmo+dR+4N8M696DuQ57PxExVo1U4EbD9/eGfkZf35vArcyf/0fee68Bl6uA\nYPBBEhPzSEhwMW9eBg89dJ0pGexoXcgG6zr27LPXj+RjF4TPHBFrQRjHDBxxuQFlFW/AYulE0yaj\nksImAz/B6ZyDsnZ19/X7tLTcTkuLGiepYtjxeDyRrl9qNvVUVDmWbhF3hn4+jxLfhahY9FMoV/jn\nQvsaLegzQ9fXgYpPzwQ+4owzTufsszfj82XidN4aPm9f30ZgOe+8U3nU+46uH5euY8JEJG6sL0AQ\nhBMnWpiURfssYCEpCVSZlJVIzPkaVLz4Z6i4snnSVV7e2cTFTQltawLuIxA4DTW+8gPUCwGooRv/\nhnKTXwO8SE5OV+j35aiMbo9hfw34O6rL2b+EzvuvQCVnn51OVdXSIZqwdOJ0JvGlL71MefkzuFzu\nQe/bKMh2ux7rVueVrmPCREAsa0EYx+Tnt2N2KX+KsliXk529FqfT+F02A8XQQ3QS1/79h4hY6SsN\nx98R+mVHZY4b1+rC69Vbieq11A+jeocfRjVKuZWMjN/j9f4Wv/8H4WN1oR28x/mLwG243Zbw1Kyf\n/ewC3n//TZStoWq5jYIsXceEiYiItSCMA4aK0VosAZRL+xxUYtZNJCb+ioUL13PTTZdTVnYHXu8Z\nKBHXk8f0lqC7gHxgLSkpUykuDuDz+QkGe1FCrTcyIfTzDOA6VH11E+aXhAz6+oyNS14GvoDKLs9A\nxbxtBAIF5OYewOmcjHLHv0hjo5fzzvt/ZGbmUlhYSUdHPl7vx8BZKE+B2YK++urn8HrvDp970qQ7\ncDi+G35WkikuTERErAVhHDBUjLa5uQBlyR5BWco1zJ49i6qqpZSXb8LrvQtdnJOSKklIuJ2enjhg\nGiqurGqws7PvIzm5kOrqG9HHWMI+zILsDH13AGVdr0G9JAAsJCnpMIHAGjStCJVsdhvKotbLxzR6\nexPp7b2JwsJKenoScbt/gsdjwePRcDo3AuUUFq7F6bwRlQneHzpeXdP+/V48Hv2zcu9bLGdIJzJh\nwiNiLQjjgH374ol0IusKfdZdx06MYyA7OytDx6QSsUq34PPdh893H2bX9qNAKocPF1BX10JEBK8F\nqlCZ4fkogc5CDc643XD8htC+GoFAK5p2t+E7NaVLfc5An1kNVvLyzgagvn7g4I7U1Azi4n5PMHgm\nyoJ/mMREF37/atzugee12Q6OwBMWhNhGxFoQxgEdHU2ozmJKrDo6lCA7HCVs2/Ys3d0RIXc6M7nh\nho20tzehRk0aB20UYnZtu4Dv0NtrobdXL69agcoeTw/9Mo67/H3U8T5SUp6gtBS2bp0V9V1a6Pca\neXlO2toy0NuPFhR4SEpKHSRGrXHwYDvB4C8M2+8jIeE0/P7I2gkJXSQmPoHNdpBNm5aMwBMWhNhG\nxFoQxgHRjUy6u/MpLd2G3d5JWlob3d03YxS3mpofkJFRScQa34PK3Nbjyrqr20qk3MsKTCUh4X7i\n4vz4fBehksOMAtwedbwPTfuE1auXs2tXdUjw9VjyLs48Mxjq5Z3Ntm3Gmux14USwxsZUDh/eS1aW\nnVmz1g8i+oXYbAdMa3/taykSlxZOKUSsBWEcMHPmEXbvjoiVxzOJ+vqrqK/XyMy8n4Edyyx0dWWh\nOoFtQcWoVwE5KDd2GrAas8t6OZBIIHAPSsD/GZXR/RyRCVr5qASzT9AbpHi9LoqLf8mMGbPp6FiD\nxTILm62ZTZuWMWOGHYDS0m2ma2xuzglN7oL4+ATmzp2KwzEfm83Keef9P5Mwx8W9z2OPLeI3v5EM\nb+HURcRaEGIUYwZ4QcERFixYR3NzDvv3f4jbXR7ay0JcXA4DO5Y9h7Ki70fVNH+CyuZuAaaj/ukb\nBb6XSExZjzHXYIyFq45lFtRkrFzD8Vvweu/ivffUfmVl66mq+qHpXqLLsvLzD1FSsh6n8/NAN/X1\nSwA1DWzTpjKKi+/A650LHCEY/Cm/+c1msaSFUxoRa0GIUaIzwBcseCRUB52NcZqWGg+5jr/+tZfu\n7k+Bi1CW8OmoxiMbMVvRG1ATuIwCvxeoNHzuIjLUg9DPPOC7od8/aTg+zbSfPtXLWGbmcJTQ17eO\nHTvigMP8/e+dtLYas8U3huutZ8ywc+aZc0ICrpAuZMKpjoi1IMQI0bXU+/aZrd/t27twu7+HLqiZ\nmfeTnh7gwAE7s2YFuPRSqKkxCq4+0jJ6cEYGkEpy8hr6+opQJVnXAveRklLIxRd3smePO9yCNLJe\nu2Gdr6EsbTvwIcaBH8apXvX1Gjt33oPXO4nu7kwCgRRUh7PJRJLZrEAaBQVt4WcRbYlLFzLhVEfE\nWhDGEKNAt7Xtwem8CbBRX69RWFiJ2fo1dyCLi8vB6VyK07mFhgYbCQnNmEVZH2kZPTiji4yMOI4c\nyUEN2zgHZWmfTmlpAJhMS8vNqCSyDajGJMmobmcuVAw8DeU670MN67gPOJ1Jk97D5TIniLW0nAbc\nYDi/XtJ1DsrVvhyVABeplT6RLmRHG+4hCOMdEWtBGEPMgzjK0OueIZ3u7ngWLPgdzc0F2O0efL5+\namoiohsMtqISwFTcNxBIxizKHwLrAB8JCXeQmmqnt7cVvz+Prq6bQsdGyrL0TmDLlr2FuW3oY6hE\ntXZUDFwvq1qMst5/HfrcH2rCsiHqOj7GPPAjncjMaj9KvFfQ3Pxa+LmcSBeyW255iS1b1JSv+noN\nn28djz++7LjWEIRYRcRaEMaQgYM4VN0zWPB4FvHOO5U888ylVFa+zYEDqRQWVpKdXcTMmT3s2NGD\nx6N3KFPlUCrTuwg4hEokawb6uPLKqTz++DJKS7dRX38VqtXnFNO5LZaZVFS8SkGBL6r+OQnV4/t7\nqKYo0Znnt6EEWo9xL0QJsB/1wvBjIrHpDSi3ezfqBaAGZWWfvKtbxcONYQOZUyRMHESsBWEM0F22\n+/cHUMlaKlksISGDQCAiOE7n5/n615/D6VyFXtvc3d3K4cOddHbOwFwjnQccxOxyvg/4N/7+919Q\nWrqNtrY9QDHKlW22xHt7J1Fd/VVycx8gLq6SYPBclKguBJ4hMfGXTJ6sceiQUcg7iMTBdXe7FWWx\nbwQuQAm1up+MjB4uuSSN5uYUCgr+Avhpbn52hMqx9AEk+rUdPsn1BCF2ELEWhDEgeg611foAxcVT\n8PniTK5u2EN7+ySU8H0K3IbHsxGP5ybMMWA97mu2llUi1ye0tEyipUUD4oiLewzoIRj8FuamKWnA\nb2lvn40S/UuIWMQp+P134fGsJGJFd6GsZz0uvjD0nQc129ofWidyPxkZbTz00HWjEkueNy+dmprI\ntc2blz7i5xCEsULEWhDGgGj39/TpZ1BVdQUul5va2kiNMXyf/v77gVtRcd9OlGgbY8CHgZ8Dp6Hi\nwy4iItsM/BfG0q1g8FGUqNeiXNyXoBLMsgFzJzRlraeg11/7fJ9HxbHdKBd2H6rZyixUL3Er8+f3\n8NFHHQZvQCRJzelcQUXF6NRMP/TQYpKSamlq6sduD+BwLBrxcwjCWCFiLQhjwGClSR0dbm699QW8\n3lTgXVQf799hsVhD+3Whz3c210w7iSR96f29P48S4FuBNxgqLq72vxPlEg9gdqufTWLiLvx+Y1xc\nb1eqZ3FvBCJWfn7+PVRV/V+WLXsr1B5VT1J7At3CHq2aaRmNKUxkRKwF4SQYrFxI0zhmCdGqVXPY\ntasSl2saNtsBVq8uY+XKWmpqMlGCGBHI/v5VKKH7J1Ss2Si8XSir1rgtK7R/AcrCPow5lpsVtX8i\ng7ce3cP8+Vbee68y1GnsCPA14uJuJxiczWA13IcO5bFs2VuG2Lhu4SeG1tyA3R44mUcuCKckItaC\ncBJEdxnbseNOLJZEWlq+SHQbTSOVlW+H3MRq2tXatetDFmc8oAshqCztIpYsWU9dXStudwCz8LpQ\n7m/jtkmoGHRC6LNuMetx5vej9j8Ns3gfAdYQF5cNTGLTpq+wdu3bNDVl8v77/43X+wsi5VnmGu5A\noIv6+u8BZRQWqpeR3t5EdDe61erF4bhyJB69IJxSiFgLwkkQHXtubc3E7KbeyL598dxww5Ns394F\nZDNvXj8HD9pMx+lWeH19AsqtHRHA5OSPqaqqoKTkJdzuuURiyZ+gBnTEoVzZXyAu7m0SE0/Dau0h\nP9/PO++sQpVyAVyKcksfQrnN1QuFwije7cDdBIMWtm1TLxL6y4YqrzKWZx1Cud3PRDVJ0T0IFvLy\nzmbu3E6qqyO13MXFCdKoRBBOABFrQTgJomPPaqqV0UpN49Chf9DQMBNVp2yhpkajsHAtRoFsa3uX\nRx5ZwvPPr6e/H2ANMAP4kOeeUz2yOzo+QM2n/hmq3KsIVaOs1igsrKS2dkVYDE8/3YHRnR5xbx9C\nJY3prURdpKTcyec+d0FoSMh0ol8kdDIzP6S39ymUlR4EWiksTCMvz0Jb236czhWhPbVQOdbxdyIT\nBGEgItaCcBI4HCXs2mWM6YJRhAsLG+juzid6KEZW1nSCwXtCrTgP4XTm8POf/5WMjM/hdn8ntJ+b\nxMTf8KMfHaCx8QX6+rJRTU+CqCEdlqg1i/jRj14KNQc5hNc7jYHu7TtQGdr5JCSsITGxCJvtIK+/\nfiOZmVmUl3dSXR003YOxWcm5506ltdU8lzovL4etW6/A5ZpDRcVmkzBL0pcgjAwi1oJwEthsVmpr\nr6OiQh9l6QbWceCAlY6OvWRl2Wlvfx9lyUYEsKOjiba2eIwNTF555U5SUuyoEqgkQMPvn857730F\n+CbKMs4Grg8d86RpzY8+eoeGBqMlvQqze/sQytJ+ArievLxK6uuVkObmZtDe3oXDUUJX1yb++te1\n9PfnkJfXyurVXw/fb0tLtOcgKyzmgwmz9OsWhJFBxFo45TlZQRlMpMrLN9HQsCpUvuQCfoXqo51D\ncvJenM6fAq9hFD6//zz8/q8DT2F0b0cGX6SjMrvNk6+s1qmkprbgdJ6FWUhPJ9L0pBs1IUvPFrdw\n6JCV8877T7KzizjrLB93330pNpuVjAwrfv8PUUM4VMz6vvsms3JlLR988AnGF4BJk/6Ow/HdIZ9N\ndAIerBdLWxBOABFr4ZTnZAVlMLGPTjxLTEwmISEPm+0AaWmz+fBDGwOzsveimo34MIuuPviiG5X8\npR8zGYsFtmy5iNLSF4nUQOvrdaJGUBpFXwM+ADz4fB04nbfjdFrYvVtj69YHKC7OGzCas6kp0/CM\nzE1OZs8+86gvNtHPQeZSC8KJIWItnPKcrKAMJvYFBUeor9cTsRrw+1fj96syrbi421GiOR01ZcuF\nSkzTUIMyEjGKrsXyDzTtDSAXaMNYhlVSMpnKyrfxeH6KEtInAC8WSxslJalYLI/wt78l0Nv7CX7/\n5NCx/4pqQ/p703273WdSXb1owGhOu91jeEZ6k5PNwCJmzVp/1Gcjc6kFYWQQsRZOeU5WUAYT+4IC\nH2ZXduT7YLAIZeUeRDUNKUSJbyJKjL9NxH39Ppr2A5S4bgD+D/AUcXFTyM9vZu3aMr73vY9QQl2D\ncnHXk5CQQnp6DqtWzaGy8m2ami6goaGVQOBaw5V7MFvi3YCFtjYbmZn3Ehc3hXnzgjgcX6Gi4lXT\nM7Ja36e42HXM7G7JBheEkUHEWjjlOVlBGUzsm5qMiVjdmEVxHzAXZSk3oTK0dSt6DZo2GX1spGpu\nYkW5x53AfwM/Ixi04HSqeLLdrlFf/yKq8cgW4Iv4/X+jujqVHTueprVVTzp7LOo6JqFqtnWL/VpU\nYxM3Hk8u8G2SktZjs1kHeUbLhxXXl2xwQRgZRKyFU56TFZTBxN5siS5g0iR9OEcD5vnOD2O0ujVt\nKuaksNND3+k9wfVhHjVAOi+++AlPPTWX6upGlFDrDUgWAxtob3cb1r8KPclNZZtnYh7c8SAwFfg+\najZ2JCQgoisIY8uoi/Xrr7/O2rVr0TSNq6++mu9+d+jMUUGINYzJY0VFPeGM6ejv7HaNp5+eg6ZB\nRUUtH3zQR1LSSgKByWhaFgkJCeTkvMGhQzMwzneO7tsdH3+Q/v7lKOFNA/4GrEe5rI3DPJSL3e9f\nxHXX3YHqIKYnkaWH9rMQF5dNMOgCnkPVWbtRZWT7UF3QjIlsn0OJPOgxdIkxC0JsMKpiHQwGufvu\nu3nsscfIy8vjG9/4BldccQWzZs0azdMKwogRnTzW1xfJFDd/52LXrofp6cnH7U5GtQA9D11Uu7s1\nurvvZaBLvBdju87MzG5crl+i3OTdKGv6d8TFHSYYfCp0nC7cABb6+magyrgexNyx7A4CgWyUq/si\nlBv9bsP390ZdS1doTY3MzGYuv3y9xJgFIUYYVbH+xz/+gd1uZ+rUqQB87WtfY9u2bSLWwrjhaJni\n5u+2hAdzRFzKUzBbrlNR85/10qckoAKYTGbm/Vx+eT61tfmodqLGcqtzCAZ3EElYM2drJyU10tc3\nGbgg6nznh873o9Dn56K+n0JcXCWZmflcfHEQTfPT3PxsyJX/LWleIggxxKiKdWtrKwUFBeHPU6ZM\nYffu3aN5SkE4LvQZ0sYhGw899NWwUB0tU9z8XRpmIcxhYLb1p6h48JbQfsbM7CyqqpYye/bTUesk\nomZbzyYya/paVO/wmUAj55+vMWXKeurqWnC7jefrwzzCMtqqn0QwuBq3WyM9fSO//vWiE3+QgiCM\nKqMq1pqmjebygnDSRGZIR4ZsJCVFXN3G5LGiol7uvjviFjZ+19a2B6fzUpT1qgGfEB9/iJSU/Xi9\nOaSmdjB3bhJJSX/hwAErDQ31GIWzp6cJgNTUZjweo6C+jZqQFT2M42z07O3333+A555bSmNjE5dd\npiey7UG9GNQYzrMAWE1CwukEgx0Egz8I3YmFl1/uw+VyizUtCDHKqIp1fn4+Tqcz/Lm1tZW8vLyj\nHpObmzGalzTmyP3FFk6nMdlL/Xz5Zbj55s08/PBCiopO49lnrx/02I8//pitWz/C651OUpKb5OT7\n6euLCGt//4P097t5//2vMmuWPXzcsmUbaGiwY8z61rQscnMzyM+fTUuLMRu8ELOl3YNyg98U3max\n5JKbm8HNN+8OCfUSYD5KqA9jjImXlc3i2Wf/lWXLnuJPf5ocWkPD5UpizZo3ePrpa07mccY04+3v\n5vEi9zexGVWxPvfcc/nkk0/49NNPyc3N5YUXXuCXv/zlUY9pb+866vfjGX1YwkTls7y/kRoQUVjY\ngbI8jVauxp/+dA11dXdywQWn09ycg93eyaOPltHfHx8+trj4v/F6VUJXX9/AMiz4HL29izjnnDWc\nddaF4evcuzcF1RAl0go0Le1+PvjgAG1tjcBqIpb0PZhd1wdRLvGI0H75ywHa27tC6+qubiuwHKv1\nAdzuVeFrfuGFRzj33Cc57bROMjPvx+M5K3TMQvbufW3C/v2Uf3vjm1Ph/o7FqIp1fHw8a9as4Tvf\n+Q6apvGNb3xDksuEESE6S9vne4SkpNTjFm+Ho4QdO35Pa2ukhSf4AQutrZnU1NwYPseKFea4rsrC\nNoqzLvzmjmB9fRdRX78k3IpUNTHRO5KloHqEZ1BS8gRO57+gLO40EhPfRNM6CQSM1xYAesOJYfPm\nBXnooa8Aegx9Sfj4KVPexGJJQLnmu4EFBAIZNDRYaGhYQWHhWjyeReHrLShoobx8k0zIEoQYZNTr\nrOfPn8/8+fNH+zTCKUZ0lvb27V243SrufDzDOGw2K7m5X6S19RuGrZtRYmseB/nxx+mmYxMT38Pn\n08up9gNWLJZVaNoUVCb4wtA6R8JrNDVl8vTTc/D5nmf79k85csSH378aj8cSilXrE7bgnHOCfPCB\nJ6pF6BPAdSxePPD+VAxdnyftxuc7Pfyyoa7j58A09KSzrKzpzJ0bicd3dSXIhCxBiFGkg5kwLonO\n0lZzno8+jGMo13lHxweYLeJ/oCxRv2n71KkdpvXmzTuNurprUAKryq1UUuUToWP+GlrrJlQzkhfZ\nv99LRcWr3HnnpVRWvs3WreD361neVlRWOeiZ521tB+jtjVxDYuJHLFwYqX8+WjigtHQbZsv/QmAR\nqu5aY9as/rAY5+ZmcP75zx7zGQqCMDaIWAvjkugWnz5fPzU1Rx/GMdQozKys6TidxqQuG5BGUtLf\n8fkiLmhN8wMRgfzb3zKJuLKNomhDJXlp5OfXc/75f2H7dhdu909wuy1UV2vs2lUZVZetZ3nvITPz\nAOnpnTQ2FnHWWekEg/fQ2WnHZjvIpk3fZMaMSLLa0cZ75ucbx2lG3PLJyZlkZ1fS2FhEefkzOBwl\n5OZmyIQsQYhhRKyFcUl0r2qXy01SkhLv/PxD+Hx+Skqeo6OjiezsImbOPGKY0+wGati6FcrLn+G0\n047Q0BA993k+gUAzxlpop3MzYBbIwTuBvQtYsFrfp67u/2KzWSkt3UZ9fUTQW1ryMQu8L3TeFfT2\n/gaPZzVOp1qvrGxod/TRmrZYLAHMDViUWz47243TuSo8xxrW8+yz18uELEGIYUSshQmBUbzLyzdR\nXX0jSvwiohSZ01wDLKe3V1m5Cxaso6xsPXV1AdzuScDngQcJBmcCT6JaeU5mxoxuYKBAKsv7DlQ8\nuAOYDnhITvawbNlb2O2dFBT4TFZrMNiIWeCdwCpAw++fynDd0UezhpubC1DDO9TLSUrKc5SWQmNj\nUagP95kAAB2VSURBVOhFwLy+DOsQhNhFxFqYcETE1Ni9y0J2dhFz565n61bo7Y1s37LFD8SRlLSP\nSy9NYceO9/H7jT221wCn8cYbbXz8cdMQ8fJ/Ae4HvoxyN/8Tra1NtLbGU1+vYbM1oAR9BtCIGk/5\nKOACckhI8HHmmU9y8KATtzsXo5C3tb2Ly6WGhETHp49mDUeuU5VxlZYqC728/JmQRS3ubkEYL4hY\nCxOOiEh1YU4QcwNJJCe3mJK21Pzoa+nr0/jb39bQ3z8Ts+V8EbAEp1Nj6dJKamuvo69vHVu39hMM\ndqFqnp/D3GnsTuDfw59drmYiPb9dwC9DP28DLAQCGtOmraOjw4/bnYkS9rMAC07nCioqlAteud87\nqa9/kc2bXyQ//xCbNpWZ4tg6Qwm5uLsFYfwhYi1MOHQx2rcvno6OylDMugefzx9yj3cCG7Bavbjd\nzcC3ULHddPr6JqHKsIyWc6T0yuWahs1mJTk5iWBQj1u7gD9hFvjZUZ+Nru0tqOlYz5v22bEjLtTA\nxAIsxVjGFXGFW1Bu/GsIBi3hF4j6+h8OeA5DubXF3S0I44+4sb4AQTgROjrclJdvorR0G+Xlz+By\nucPf6WL05z/PZ+7cacTHJwAaBw7o7nErcC3Tp2eRnNwN/A8qE/tS1HCM6cDtKOt3FZAMPAW4sNkO\nAkZXuxvVuSwdJewQGdox1Gd96EdX1D6HMQu8uYzLbu8M7Wd277tc007gCQqCMJ4Qy1oYM06mZejR\nSpaG2ieSYBaJ1WZm5vL66z6MFmvEoq4M/VKfU1LuZNOmbwJGV3sNKiFtPpFe3/XAdeidxPLz/8E5\n56Tw1lsPEAza6OvbT1/fYlR2trLwi4sT8PnSTOVn8CbgJjHxQ1avXoamESr5ysKY+Ka/QAiCMHER\nsRbGjOEI7lAcrWRpqH2ys4v4whfWsWNHHHAYny8Nl+t0Iv20zRYrmMurZs/+AmvXvk1T00cUFPhY\nsOB3vPZaGr293ai49TWhdeopK3s93EnM4bjB9BLicrmpqNBjxgEcjiux2azh8jOVld4K3ArY8Ps1\n1q5dD2CqzY6LqyQ/HzZtWjKsZyYIwvhFxFoYM4YjuDC4BR6dkd3W9i6NjbOprHw7FKtuort7CkYL\n9PDhvRw4kIjb/RNAjcMsLFwL5KFi1p+iOnzplu1HGC3xDz98h927fwxsob5+Cvn573DxxbBt23J0\nK1qNpnRTVXXLkPetu+n1+9LLuxyOEqqqluJyufnSl17G7Y5MBDPHrNXPL3zhbLZuveL4HrogCOMS\nEWthzBhux6zBLHCHoyTkEv48cASncwVf//rDIctT1Vfr61qtD5Ca6sfpXAG8gVHwsrKm09PTh9t9\nLSr+vBHoBeJJT7fS3R3pbOb1TkElhy0HOmlp6aalxQU4UAlk76EakMwIdwY7mls/+r5eeukOzjjj\ni8yceYR587yDdGTTpMOYIJyiiFgLY8ZwSog6OtzU1bWiMqe7gIXU1QUAyMs7G6cz4gJWiVadKAs5\nsj9kk5WVHJpdbSznctHR0YT6Z2BM9Ipj0qQP+dKXckNWs3EQhje0dgORUiy9i9mtqBj2tVRXR9z6\nQ8Xmoz0LXu9cdu9ewu7dkUYtA5+NlFwJwqmIiLUwZgynhGjlytqw21qJ4gbc7klUVNSGRk1GLE2b\n7QC9vS+i1y4b909N3R/6HEnqSk1tCVnincCDoTOqY71eDYvlEcrK9CYqicBpgHGKlTG+fQ7K6s4I\nb9Nd10PF5gc2V4mUiG3fHsfOnZcPsMyl5EoQTk2kdEuIaQa29vQBC2lqysThKKGsbD3nnfcsZWXr\n2bSpDKvVO+j+2dlFoX1fo6wswM6dV5KXdzaRUq5CoMh07JtvJlFVtZTSUg3l+p5i+F5PSoOI0Kah\nLHe1TXUecw8am+/ocOPz9ZCYeCeqocq9wFfDx+ovJIIgCCCWtRCj6K7j/ftbMDcoSQYmY7d7BrXM\ni4vfCrmgzfvPnNkzYF+zZbsA1S50cfjYI0eacbnchvi4RiQBbQEqLn4xSqi/Sn7+b9C0Plpbn0OP\no1dUbB7gAbDbPaxcWUtNzfdRVv2LZGZm0Nv7K/z+81Gu9oU0Nb02sg9VEIRxi4i1EJNEXMeq21hm\nppf09BaysuzMmrWeVasuoLx804A4sB4Hb2xM5fDhvUPuv2LFGezc+Qlxcb9H09rQtGyU5RwZien3\n9/GlL71McXE8mzYtYfHi/6Kt7V5UMtmnXHppOllZ7tCam3E4bmDZsrdobdXj6CrePm1aIYWFkU5q\nDsflLFv2FsYGLTNnPovdnkF19VVE9wQfbu25IAgTFxFrIWYwJmIpi7oTo5ht3fp/wvuqyVqROPCO\nHXdywQX/v717D66yvvM4/s4dSAI5QIBEuiGAEay2TC11YVxCsY0SwKBopXWkRZuV0sEx7Qw3124t\n3VBTrbZDhyJip1AqWNYkUAhVA4RWKcvWTTEqZYg0CLmS5DQJhlzI2T8eTs41yUlyDufJyef1jyR5\n8jy/x4if/G7f379QVTWelBQb+/bdicVyj5frjbra+/f/DZttKvZtXUbFskSMFd23AGeBWVitVyks\nXAgcYO7cmRQUrMAepnFxO/rorR/qPsMabMye7VhwVlv7IcYsVAuw8PqCMc8V7mvXHtA8tYgorMU8\nPM+Jfg3jPGnPbUru88A1NaMpKjIWf5WW2igpeZ709AleVl4bVcpsNvszXgVGYdTyjsGo8x2O8yEc\nsIeKitFERUW4PLOqarzHO2zYcAenTm2msXEyHR2tdHZ67iNft+6oS3GT5OTN5OU9isWS4LHCvbfj\nMUVk+NACMzEN9wBOSLjavXjMfZuSo0421/853uV7rdYZFBau6F6k1VNdbSOclwMP4Dhw4zxGr95+\nTSy1tR9y7twZl2c6/wJhr1V+773/Q2VlCq2t99HZGef1evf3nDDh1u6hbvf30l5qEQH1rMVE3Lcy\n/eu/dhET00RFxWjWrj3iUmTEfcjY4LywrAXn3qx9LrukpBqr1blKWTze64I7evUjRpyisvJx4G3g\nBSIj4/nqVyPIy3MMs3uOCuwBFpGQ8DyTJ6fS0HCW8vIUsrPfICmpvcfiJjq+UkS8UViLaTiOthxF\nQ8NZ3n13NE1NI4H5lJaOwbl2uMWSwNGjj/LUUwc5caKZrq6RjBr1X3z6adL178nEOQjtK8dd63I3\n0dLSRXGxZ4979OirhIe/CtTT1TWRq1dPYN9j3dlp429/20xj4z9Zu9Y+x96Ja489DhhDevpE4FPK\nyjZQWRlGWZmNhQt/1UPBEx1fKSLeKazFNOxBlZ2dT1mZY07Xfq6z+/ytxZJAdPQorNYngDCamowg\njI6OoqLimNeeqc3m8hG5uf9Gbu4ujh6toqnJ0eP+9NOP6ezcdP3j3TiOtQQIo7LyNpYuzae6ehoQ\nAdTg3LNPSDhDenqj28pv43urqpJU01tE+kVhLUHR2/GYnoVQjLlfb/O37td6C0LnZ9XWfkBl5SPA\nCUpLLZw6Vcirr36ZoqIyjCpmxtx3Z2eM030XERb2PDabYw82NFJdzfW2NQNfJyrqP/nsZ79w/ZeE\n5S7z0KrpLSKDobCWoOjteEz3cHPupdo5iqZ0Ar/AmKOezJkzZzl/fjqpqSlenwVZwHPAOowe8hKW\nLv0B7e3P4dqTHwf8DmNOu4nY2Gu0tPwEo6zoFYzKaP/h8j2xsVO89pg1Dy0ig6WwlqDo7XhMz3Bb\n7lIYpKHByoIFu64vLmsB/oF9q9XVqzbuv38zpaVrenyWUVrU8XFbW6rb12MxVoR/B8ee6k20tKzC\nqP8dS2Rkk8u2LIhlzhz7QjdXmocWkcFSWEvA+XIetfPQcF/h5r5PGTbjHLaNjZNdnlldfRqoAyYB\nTURHv097u+PZMTEfc/Wq4+Pw8L8QE3Mzra2OeyYm3sq8eYc5e3YkKSlW2tvDXY6wTE4u46WXHvXr\nvyNVLhMRO4W1BFxP51E79557Kh/qjWtP+Z/ANYzDMIxqYBbLRS9D369h1P22MW9eM7Gxjmd/97uZ\nfOtbRiETi+Ui+fnfIDfXtcb41KmfsnfvCurqjIM6GhutREc79/4fHVS49jYtICKisJaA8zbk7d57\ndi8f2ltYTZpUh2Pl9SGc545HjPgB+fkP88QT53Ad2o4HrEAR77wziowMG3v3Oupul5be7vKMvDxj\nq1hP88z+HtrubVpAREQVzCTgfKnK5R5Wb74J2dlv0Nho7b7GXiXs3XerMHrKBzAWejm+b8aMO0hN\nTfFS4awZo/DJclpbV7hUN3O/f0ZG8fUiLF9mz547AFi27CSf+cxmFizY79Euf1DlMhHpjXrWEnB9\nrYb2drBFa2sUhYXLgV0899yXWbfuKCUlnVitMcDNGNXGwFix7Riurq39kIwMSEq6wsKFO6iqGk9S\n0mWgg2PHYl3mod17r84nfZWWHqKk5C1Gjap2mR+/eHEPZWUr8PcwtVaMi0hvFNYyIN4WRCUmxvf6\n9Z7mdD0XjD0HrMIeqJ6lPH+CUdP7MAAjRjzDzTfPor7+LJWV36Gy0kJpqY2srF28+ebd3W2Jiemk\ntXU39pO2ej4cxCg9arWGYbXux3PPt/+HqbViXER6o7CWAfG2IMo4PrLnr/cURp5bq27FOBrTGA72\n/PoM4JcYx1oa27UmT95BRMStVFZauq9zPuXKOewTEp4nPX2i18NBjLY6lx5twbPmuIapReTGUljL\ngHhbEFVfbyU7e7+X86h774m6b+OaNOk0V69eBuppb48lKanNrUjKOVpaEl32OZ84EU56uvftYO5t\nnTLlZrZv77l4iethHwtJTt7MuHFpNDaeIyHhM0yb5jgFTFuuRORGUFjLgHjbJ716dZHP51E7y8tb\nQFvbDv7yl3CgHputDat1JRBGUZG3gy+Wc+edr2G1Ovd468nLM+a43ed9fS332dNhH/ZtWYmJD3Zv\n3bLTlisRuREU1jIg3vZJL1z4v7ifRz1lSkGfC6YslgRiYqKxWpdgzEN3YAR9JpDgtd73nDlxFBW9\nhrElq5k5c+KwWBJYv/4LLFu2n7//PYk//nEbqak3M2WKY7GZL4u3+jN/rC1XInIjKKxlQLztk25s\njMJ5fjc9PdLrcLOd8xCyMWy+H1iBo7e8B1jutSf80ktLiI4+SkXFNVJSOsnLWwzAsmX7XRarffTR\nHj76aEX3YrPBcB7m96USm4iIvyisxS+MHuV8jIAdQVTU/1FefgvZ2W/0OI9rDCHbe9MzgCqce6kj\nR3aQkbGrx+pm3nq/jY2TXe7hz9XbzsP8PVVi05YrEQkEhbX4hdHDHIOx//l3dHQ8S1lZGGVlrvO4\nnr3p/wYex3FutKOXmpFB9/nWvs4LWyyf0NoamNXb5887rxL3XolNRCQQFNbiF3l5CwgL28mxY9do\nauqgq8v7PK7nnukXcD43Oioql8jIz2CxXGTjxvsAKC8fhXNIfvzxKI/n238JGDNmOg0Nz2CzJREW\nVk1q6nTS0nb5pcebmtrMqVMa8haRG09hLT5raLDy1FN/vL5q+zJz5sTx0ktLsFgSsFgSiI6Oxmpd\njrE4zHuouS/IioyMp7PTfu0YOjpS6ej4Bq2tNnJzd7F9ewoNDX93uV99/VngHpe2uf4S8DWysnax\nffsK/Gnr1kza2jTkLSI3nsJafLZu3VEOH7YPWdsoKnqN6Oij3cPAjmHiTGDP9TlnXELNfUHWqFEN\nxMUZ+5g/+eQ8Vms29gM39u/v5NSpXxAbm4QxFx4HtDB2bIpH227EquyxYzXkLSLBobAWn3lWEoun\nouJa99cdw8QJwHIyMjznlh2FRzqxWkfQ1PQdmprGMHv2LqZOnUBh4Rjsq8BttjAqK22MGPEMsAl7\nwE+btsujbVqVLSKhTGEtPjMC0V6TOxb4gKQkxypvX4aJ7QuyMjKKKS1d2v35iorR7N17B7CLwsJm\nHD3pZq5ds7gVRfG8r1Zli0goU1iLz/LyFnDy5C+prjZqcsMSYEf31/szTOzeE66t/ZCHH4aUFBsx\nMVW0ta3u/lpExA/Yvv3fe72fVmWLSChTWIvPLJYEJk26jepqx1B4VdX4Ad3LuSdcW/uhy2lZ8fHb\naWtzPCM19TZ/NF9EZMhSWIvPvJ07PdC5YeeecEYGLqdlRURYcV79nZbWNui2i4gMZQpr8Zn7udPJ\nyZvJy3t00Pd1HxKfMyee6GjNP4uI2CmsxWfuq8EnTLgViyWhuyBJZaWF5OSGfh8T6bk4bLGOmRQR\ncaKwHqYGcg5zT9ujPKuS9e+YyIEsDtM50iIynCish6mBnMPc0/aoG3VMpHNA19Z+QGXlasCic6RF\nJOQprIcJ957oxx/H0t+A7akH3FOP29+9X9cefBbGXuyv+9x+EZGhSmE9TLj3pJOTc+mpfndf3EN4\n40ajmIkxZ93Y3eMeSO+9N54V1GKv/1kVy0QktAUsrLds2cLrr7/OuHHjAMjJyWHevHmBepz0wT3o\nxo6dwuzZA1tx3VMIJybGU1fX3OMzB9v7de/BJyeXMWFCl1aMi0jIC2jPeuXKlaxcuTKQjxAfuQfd\ntGnXBtzL9TWE+1uvu69hc88580e1qExEhoWAhrXNZgvk7aUf/Fk729cQ7u8z+xo2V0lRERmuAhrW\nu3fvprCwkNtuu43169cTHx8fyMdJL/wZdL6GcH+feaNWlYuIDDVhtkF0f1euXMnly5c9Pp+Tk8Os\nWbOwWCyEhYXx4osvUldXR25u7qAaKwNXX29l9eoizp+PIzW1ma1bMxk71lxDyA8//Dtef91Y3Q02\nvva1Pezd+/VgN0tEJOgGFda+unTpEqtWreLAgQN9Xuu8QCnUuC/AupGys/NdCpdkZfl/X/Jg36+x\n0cratUddeuxmmpMO5s8v0EL53UDvN9QNh/frS8CGwevq6khMTATgrbfeIi0tLVCPEh84hpitQBFv\nvgnZ2W+YqvKX5qRFRLwLWFj/9Kc/5aOPPiI8PJybbrqJH/3oR4F6lPjAsSisCFhOa2sYhYUD3/vs\nbeW2L78diohI/wUsrPPy8gJ1axmADRvu4NSpzVRVTcJmG/wiLm8rtwsKVviruSIi4kQVzIaJzZvf\nu3685Wv0VrnM3mMuL4+goaGCcePSmDr1isdwuVZui4jcOArrYcIRrpnAHkaO7CAjA49tV44e8x5g\nA5WVYbz/vudwube91vX1VrKz9+skLBERP1NYDxOOcE0AlpOR4Qhf5/nnf/yjEyOA43DuOZeXjyI7\nO9+jHrh95faGDV9g1qxfcfHiOvpbC1zHXYqI9E5hPUz0VsjE9TSr3RjD5M04D5dfvnyGsrKnsQdx\ne/sOfvObh7vvkZ2dz8WLtzKQoXF/H/ghIhJqFNbDRG/bolznnxeRkPA8kycn09Cw+fqc9accPZqA\ncxCfOBHu5R4tDOQkL81/i4j0TmEtbvPPY0hPn8j27fe5XJOWthXnIIZ6L/e4D2OuO5bk5DLy8h4d\nwPN13KWIiDuFtfhU63vOnDiKil4D4oFm5syJ87hHTMxhzp4dSUqKtV8nYvnzkBERkVB0Q8qN9keo\nl5Qbqu/nSynQofx+vgjl9wvldwO931A3HN6vL+pZB0ioVfhSKVARkeBRWAeIKnyJiIi/hPd9iQyE\nVjiLiIi/KKwDJCXlnxirpkErnEVEZDA0DB4gWuGsymQiIv6isA6Q/i7ICsVgU2UyERH/UFibRCgG\nm+btRUT8Q3PWJhGKwaZ5exER/1DP2iRCseSm5u1FRPxDYW0SoRhsKqQiIuIfCmuTULCJiEhPNGct\nIiJicgprERERk1NYi4iImJzCWkRExOQU1iIiIiansBYRETE5hbWIiIjJKaxFRERMTmEtIiJicgpr\nERERk1NYi4iImJzCWkRExOQU1iIiIiansBYRETE5hbWIiIjJKaxFRERMTmEtIiJicgprERERk1NY\ni4iImJzCWkRExOQU1iIiIiansBYRETE5hbWIiIjJKaxFRERMTmEtIiJicgprERERk1NYi4iImJzC\nWkRExOQU1iIiIiansBYRETE5hbWIiIjJKaxFRERMTmEtIiJicoMK68OHD7N48WJmzpzJBx984PK1\nbdu2kZGRwcKFC/nzn/88qEaKiIgMZ4MK67S0NLZs2cLs2bNdPl9eXk5RURGHDh1i+/btPPvss9hs\ntkE1VEREZLgaVFhPnTqVKVOmeARxcXExmZmZREZGMnnyZFJSUjh9+vSgGioiIjJcBWTOuqamhqSk\npO6PJ06cSE1NTSAeJSIiEvIi+7pg5cqVXL582ePzOTk5LFiwwOv3eBvyDgsLG0DzREREpM+w/vWv\nf93vm06aNImqqqruj6urq5kwYYJP35uYGN/v5w0ler+hLZTfL5TfDfR+Q12ov19f/DYM7tybXrBg\nAYcOHaK9vZ1PPvmECxcu8LnPfc5fjxIRERlWwmyDWKb99ttvs2nTJhobGxk9ejQzZszglVdeAYyt\nW/v27SMyMpKnn36au+66y2+NFhERGU4GFdYiIiISeKpgJiIiYnIKaxEREZNTWIuIiJicacN6x44d\nzJgxA6vVGuym+NXPf/5z7rvvPpYuXcrjjz9OXV1dsJvkV3l5eSxcuJCsrCzWrFlDS0tLsJvkN73V\nwh/Kjh8/zr333ss999zDyy+/HOzm+NXGjRuZO3cuS5YsCXZTAqK6upoVK1aQmZnJkiVL2LlzZ7Cb\n5Dft7e089NBDLF26lCVLlrBly5ZgNykgurq6uP/++1m1alWv15kyrKurq3n33XdJTk4OdlP87tvf\n/jb79++noKCA+fPnh9x/gHfddRcHDx6ksLCQlJQUtm3bFuwm+U1PtfCHsq6uLjZt2sSOHTv4wx/+\nwMGDBykvLw92s/zmgQceYMeOHcFuRsBERESwYcMGDh06xJ49e9i9e3fI/Pyio6PZuXMnBQUFFBQU\ncPz48ZAsW71z506mTZvW53WmDOvc3FzWrl0b7GYERGxsbPefW1tbCQ835Y9gwObOndv9TrNmzaK6\nujrILfKfnmrhD2WnT58mJSWFm266iaioKBYtWkRxcXGwm+U3X/ziFxk9enSwmxEwiYmJzJw5EzD+\n3zJt2jRqa2uD3Cr/GTlyJGD0sjs7O4PcGv+rrq6mpKSEhx56qM9r+6xgdqMdOXKEpKQkbrnllmA3\nJWBefPFFCgsLiY+PD6lhK3f79u1j0aJFwW6G9MJbHf/3338/iC2Sgbp48SJnzpwJqQJUXV1dPPDA\nA1y4cIFHHnkkpN4NHB3T5ubmPq8NSlj3VG/8qaeeYtu2bbz66qvdnxuKvZi+6qnn5OSQk5PDyy+/\nzG9/+1vWrFkThFYOnC/14rdu3UpUVNSQmyscSC38oWwo/v0ST1euXOHJJ59k48aNLqN3Q114eDgF\nBQW0tLSwevVqzp07x/Tp04PdLL84duwY48ePZ+bMmZw8ebLP64MS1j3VGz979iyXLl0iKysLm81G\nTU0Ny5Yt4/e//z3jxo27wa0cOF/rqS9evJgnnnhiyIV1X++Xn59PSUnJkBw1GEgt/KFs0qRJVFZW\ndn9cU1Pjcx1/MYfOzk6efPJJsrKy+MpXvhLs5gREXFwcX/rSl/jTn/4UMmH93nvvceTIEUpKSmhr\na+PKlSusXbuWvLw8r9ebasI0LS2Nd955h+LiYo4cOcLEiRPJz88fUkHdl4qKiu4/FxcXM3Xq1CC2\nxv+OHz/OK6+8wtatW4mOjg52cwImVHqkt99+OxcuXODSpUu0t7dz8OBB7r777mA3y69C5WfVk40b\nNzJ9+nS++c1vBrspftXQ0NA9PHz16lVOnDgRUv+//N73vsexY8coLi7mZz/7GXfeeWePQQ0mnLN2\nFhYWFnJ/0V544QXOnz9PeHg4ycnJPPvss8Fukl/9+Mc/pqOjg8ceewyAz3/+8/zwhz8MbqP8xLkW\n/qpVq1xq4Q9VERERPPPMMzz22GPYbDYefPBBn1amDhXf//73OXnyJFarlfnz57NmzRqWLVsW7Gb5\nzV//+lcOHDhAWloaS5cuJSwsjJycHObNmxfspg1aXV0d69evp6uri66uLjIzM0lPTw92s4JGtcFF\nRERMzlTD4CIiIuJJYS0iImJyCmsRERGTU1iLiIiYnMJaRETE5BTWIiIiJqewFhERMTmFtYiIiMn9\nPyQ+uNKCpR6MAAAAAElFTkSuQmCC\n",
+            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAecAAAFKCAYAAAAnj5dkAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMS4yLCBo\ndHRwOi8vbWF0cGxvdGxpYi5vcmcvNQv5yAAAIABJREFUeJzs3Xt8VPWdP/7X3M5MkpkkM8mEAAER\nQoICgUBALkUEQ7FucekDEeWL3VZXu121dler39pu1Vbb77b+2m1/3277qNXa2kUptGttt/tDEWqp\nyDWBiC6ES8slXDJJJpfJ3C+/P8JM5nLOmTOTmWQm83r+RebMnJyTAO/z+Xzen/dbFQqFQiAiIqKc\noR7rCyAiIqJYDM5EREQ5hsGZiIgoxzA4ExER5RgGZyIiohzD4ExERJRjtGN9AWE220DWzm02F8Nu\nd2bt/LmukO+/kO8d4P0X8v0X8r0D+XH/VqtJ8lhBjJy1Ws1YX8KYKuT7L+R7B3j/hXz/hXzvQP7f\nf0EEZyIionzC4ExERJRjGJyJiIhyDIMzERFRjmFwJiIiyjEMzkRERDmGwZmIiCjHMDgTERHlGAZn\nIiKiJDy+ADrtTnh8gVH5fjlTvpOIiCjXBIJBbNt9Gq3tNvT0e2Ap1aOxzopNq2uhUWdvfMvgTERE\nJGHb7tPYdfhi5Ovufk/k683NdVn7vpzWJiIiEuHxBdDabhM91treldUpbgZnIiIiEX0OD3r6PaLH\n7ANu9DnEj2UCgzMREZGIMqMellK96DGzyYAyo/ixTGBwJiIiEqHXadBYZxU91lhXCb0ue20pmRBG\nREQkYdPqWgBDa8z2ATfMJgMa6yojr2cLgzMREZEEjVqNzc112LByBvocHpQZ9VkdMYcxOBMRESWh\n12lQZS4ete/HNWciIsqa0a6sNV5w5ExERBk3VpW1xgsGZyIiyrixqqw1XvDxhYiIMmosK2uNFwzO\nRESUUWNZWWu8YHAmIqKMGsvKWuMFgzMREWXUWFbWGi+YEEZERBk3VpW1xgsGZyIiyrixqqw1XjA4\nExFR1ox2Za3xgmvORESUMawIlhmKRs7t7e34x3/8R3zmM5/Bli1bcPnyZTzxxBMIBAKwWq34zne+\nA0EQYj7zzW9+E8eOHYNKpcJTTz2FhoaGrNwAERGNPVYEy6ykPzGn04lvfOMbWLp0aeS1H/zgB9i8\neTO2bt2K6667Djt27Ij5zMGDB3Hu3Dls27YNzz//PJ5//vnMXzkREeWMcEWw7n4PQhiuCLZt9+mx\nvrS8lDQ4C4KAF198EVVVVZHXDhw4gFtvvRUAsGrVKrz//vsxn3n//ffR3NwMAJgxYwb6+vrgcDgy\ned1ERJQjlFQE43R3apJOa2u1Wmi1sW9zuVyRaeyKigrYbLG/lK6uLsyePTvytcVigc1mg9FozMQ1\nExFRCjy+QFYzppNVBHt150mcPG/ndHcKRpytHQqFMvIes7kYWm320uytVlPWzp0PCvn+C/neAd5/\nId+/xVKCl3/3IfYfvwxbrwvW8iIsmTMR962bDY0mc4HRVFYEq7kInXZXwjG9oMW+41ciX4enu4uL\nBDywfm7GrkFMPv/u0wrOxcXFcLvdMBgMuHr1asyUNwBUVVWhq6sr8nVnZyesVvFqMWF2uzOdS1HE\najXBZhvI2vlzXSHffyHfO8D7L+T7t1pN+L+/ao3pDNVpd+HNvWfhdHkz3hmqYUZFzPcKC4WCou9/\n79glfGLxlKztfc6H373cw0Naj07Lli3Dzp07AQBvvfUWVqxYEXN8+fLlkeMffvghqqqqOKVNRDSK\n3F7/qHaG2rS6Fs1NNagoNUCtAipKDVg+pxpur3hwZgMMeUlHzsePH8e//uu/oqOjA1qtFjt37sQL\nL7yA//2//ze2bduGSZMmYf369QCAf/qnf8K3vvUtLFiwALNnz8bdd98NlUqFp59+Ous3QkREw+z9\nyTtDZaI4SPR6dnxFMAA4cd6ObpHrYAMMeUmD85w5c/Dqq68mvP6zn/0s4bXvfe97kT8//vjjI7w0\nIiJKl7l0qDNUssCYLFlM6rjcvubooN9YZxWd7mYDDHks30lENA54fAHYel1AKASruRhWQSsbGLUa\nFbbuapcsGpKsqEh4X3NYONELQMx6tlgDjIYZFqxqnAyPL8AALYHBmYgojwWCQbz+zim898EVuL1D\n68gGQY3mxdfhzlumAxDvDJUsuMod37Byhsx6tg0bVs6IBN3oBhg9/W7sOnIRbae78MfWS9xWJYPB\nmYgoj23bfRrvHOmIec3tDeL3f/4L3G6faGeoZEVD1i2bJnv85nmTJNezu/s9eHXnSXz29lkxAVev\n02BPawf2tHTEvFdstE1sfEFElLfkgiwAtJy0RaaOq8zFkdFssqIhFzsdsscRCsFSKp3Mte/4lYSy\nnUqqiNEwBmciojE0krKWckEWAOwDHtHtSmVGvWRwNZsMqKkySh4XdBpYyorQWCdfuyI+4CZ7IOC2\nqlic1iYiGgPpdHGKz5wOB1mxjGwAMJv0otuV9DqNbLKYqViQPO72BvDG3rPYtLoWLrcf70VV/4oW\nv11L7lq5rSoRgzMR0RhQmu0MyAdyqSAKAAvqrZLZ0GJZ1OFkMQBYv+J6/LntciTJLFprexc2rJyB\nLWvr8T/netAz4E14T3zATfZAwKztWAzORESjLNn6a3S2MyAfyDetrkUoFIrL1tagefFU/O2y6ySv\nITqLWmwfs8Ppg0ckMAOxo+IF9VWKA26yBwIaxuBMRDTKlKy/hqeDlQTy/7WmHnfeUhuzz7lmUrmi\n2tLhZLF4SqehUwm4yR4IaBiDMxHRKEtl/VVpINfrNKixKu9hkKwymNJp6HQCrtQDAQ1jcCYiGmWp\nrL9mOpHK6fFh69uncOJcD+wDXlhK9WiorUTzwhpYSg0x3zuVUTEDbmYxOBMRjQGlgS9TiVThpLL4\nJK/ufg/2tAwVB6mIyxgXGxUDQHefO/JnTk9nB4MzEdEYSGU6eP2K6+F0+3HinB29Do9oIE82TR2f\nVCZGKmNcr9OgoswQyRjv7vfAIKgBqODxBliGMwsYnImIxlB4v7LSzk9LZ1fjnjV1KNZrJd/TWGfF\nw3c1Rs6TrJJYPCUZ49F9mlmGM/MYnImIRlH0CFerUaXc+em941fg8vrxd7fNgqlYkNxmVVwkYP3y\naQCSVxKLl0rGeDSxoE7pYXAmIsqAZNPKYiPcYoMOFzodkfco7fzU0t6F1vY/Y7K1BE63T/Q9+49f\nxicWT1FUSSye2aSH1xeI1OVWGtzjgzqlj8GZiGgElJbhFBvhSgXLZJ2fACAE4KJtUPJ4V68LZzv6\nMH1ymWxSmZhBtw9Pv3woci/rV0xXFNxZhjNzGJyJiEZASRnOVNd8ozs/KR3tJlAB33n9aCQDO7G3\nsx71U83QaVU4ftYO+4Abgk4DtzcQWU+OvhclwZ1lODOHwZmIKE1Ky3CmuuZbbtQDKhUaaitj+h+n\nIngtXyv+YUEsO9zjC8DW68K//eqoZC3tZ+9fhEAgiNZTXeh1eGEQhj7r9QVYhjMLGJyJiNKktHpX\nqmu+To8fT790EGaTgImWYlzucY74WsMPC2L0Og0ErRp2kQYWwNC9bH37FE6et6PP4YXZqMf8ukps\nWDkdDqeP+5yzgMGZiChNSqt3KV3zFbQqeP2hyOh1qNuTF3qdGh5fUPazydgH3Hh150mcPG8XXRuX\nuxdBp8G+qNaQdsdQ4RKNWsWtU1nC3eJERCny+ALotA+NZhvrrKLviV9/3bS6FsvmVMue1+sPib6u\nUqV5oVEEnRr7jl9Bd78HIQxPd7/+zikAww8Q4sSvq7W9Cx6feOcqGhmOnImIFBLLzJ4/sxKrF07G\nsVPdsmU4NWo17l1bj5Pn7Sknebm9QSyfU40j7TbRNWElfH7xkfd7H1zBnbfUQq/TYNPq2si6cp/D\nC0upAbOmluO9qFFzNG6dyh4GZyIihcQys9850oHmpho898BNSctwprqlKUytGirh+T/netIKztWW\nIlzpcYkec3uHksEmVhRj2+7TaDvTjT6HF+VGPRpqK7Bh5QyckHig4Nap7OG0NhGNC+Gp5mxNsybL\nzAYQad0oJRAMIhQKRTKdlQqGgE67SzJhKxmXxy//hlAo8uARnvYOryu/sfes4ql7yhyOnIkor8kV\nAckkpZnZcrbtPo13jqS+Ncpi0qOmypj2vue+QR/0WjU8IlPbBmGogpjcg8ez9y+K/DlZ60jKDAZn\nIsprckVAHr1nYca+z0j7KqdaiCRaSZEOGo0K9VPNMVnTSpUbBcybWYl3Wy8lHFs2txouj1/2wcPh\n9CnuoEWZweBMRHkr2VSz25tkOjcFep0GDTMqsEckwCmZ3k21EEm0C50OPP7DffB4AzAIGoRCoZS2\nVjXOrMTmNXXQadRoOWlDz4AHZSU6LKiz4p5bZ8IfkK5GFr8ljMlfoyPt4Lx9+3a8+eabka+PHz+O\n1tbWyNezZ8/GggULIl+/8sor0Gj4pEVEmZNsqtne78nICCQ8dd52phvAUIJWMDQ03bygXnoKPboZ\nRqqFSOKFE8FSTQibaCnG5jV10KjVQ9nYwRCOtneh1+FB25luaDSnsWl1rWSiGteVx0baf283btyI\njRs3AgAOHjyI//7v/445bjQa8eqrr47s6oiIZCSbajaX6jHQJ56lrEQ4uO48dCGmjGbw2rbfeTMr\nRYtwSK2Dz59Zmdaas5jwA4Icg6DGV/6uKdKAY9vu0zH3Eb0EEH7A4LpybsjItPYPf/hDvPDCC5k4\nFRGRYnJbkxrrKmEQtBgQ+Vwq7R27+z1QSxQBaTvdDc+qQMI5tr7dHjP9HQ6CtyyYhBpriWw3KRWk\nSn7EShaYAeBjDZNQrB/6b37A6cXh/+kUfV+4tCfXlXPHiINzW1sbJk6cCKs1NtXe6/XiscceQ0dH\nB9auXYvPfvazsucxm4uh1WbvL4LVasraufNBId9/Id87MP7v/+G7GlFcJGD/8cuw2V0wl+qxZM5E\nPLh+LoDY+w8Egnj5dx8OvbfXBWt5EZbMmYj71s2GRjO8s/TFNz6ICfhSgdA+4IZG0MFaWRI5/0/e\n+ADvHktclwaAAx9ehcsjPy2tJDADgLXcgAX1VXj70PlIk4toRXotHlg/F3pBi5++eRxvHzwHj1d8\nnTr+PmoUXkOuy+e/+yMOzjt27MCnPvWphNefeOIJ3HHHHVCpVNiyZQuampowd+5cyfPY7SMv7C7F\najXBZhN7fi4MhXz/hXzvQOHc/7qlUzEw6MFRXxfs/R4cOH4ZXq8fD9/ViJ6e4VHq1l3tMUG30+7C\nm3vPwunyxrR3fO+Ysqlns8mAgNcX+RnHnz9essAMDK1jz5tZibbT3TFtHOPNq63E8jnVeOvAedHz\neLx+/OWCHbuOXExa9CT+PsaDfPi7L/fwMOLgfODAAXz1q19NeP2ee+6J/HnJkiVob2+XDc5EROmS\nWkstLhKwfvk0ANlp7xidLDWSrVLRFtRbsbm5Dp5VQ1PvxmIBb+w9G7MWPG9mBUKhEP7tV0clR9qV\n5UUo0mvRclJ8KlvqPig3jKhC2NWrV1FSUgJBEGJeP3v2LB577DGEQiH4/X60tLRg5syZI7pQIiIx\nckFx//HLkYphSoqIAMNJZmLUqqEmFBWlBjQ31WDT6tpIZTKb3Zk0qMtVBrOY9JFzAsPblor1Wmxu\nrsNzD9yEbz64BF/7TBM8ngDeOdJxrWuVuCVzJg7tX05SVWzZnGomfeWgEY2cbTYbLBZL5Ouf/OQn\nWLRoERobG1FdXY0777wTarUaq1evRkNDw4gvlojGv2TJWvHkgq7N7sLZjj5Mn1yWkfaOK+dPwtrF\nU1Fm1EOrUSVkZAs6FTw+8bGsXqfGktlV+GPr5YRjy+dUY8vaetn71WpU2HXkIlpOdsoG3IprmeH3\nrZuNy1f7YTEJku+3lOpx79r6SDY35Y4RBec5c+bgpz/9aeTrBx98MPLnL33pSyM5NREVGLkynHLB\nQy7oqtTAC68fjZxr3sxK7BbZyiTW3hEQ31YUvpb49eVk+5c9viBuXVgDrUYje14p8ZXQxKgAPHpn\nA2qqTNBo1NDrNFhQXyX5uQV1Vk5n5yhWCCOinCBXhlNsL3GY3Eg3nMUcPtetCyejualGci9v9Kg9\nflsRAHT3uSN/Tmd9efeRDty7dlbS7UrxswceX0DR2rGl1ABrXAWvTatrEQyFsO+DK5HEMoOgwfK5\nnM7OZQzORDTmlCZrSYke6fb0u6GSKNBx9FQ3nnvgpoTgGAgGsXVXu+iovaLMkDCir59qTqsUZ9uZ\nHnh8AckymFKzB6saJyddOwbEE7s0ajW2rKnHxltqYbM7AZUK1vIijphzHIMzEY25kXZ80qjVkZHu\n2Y4+vPD6UdH39Qy4I2vQ0eeTG7UHAsGEgiL7jl+BRF0SWcnuRap4idcfkK0IJuhUWNEwSXYkrNdp\nUFOVv/t+Cw2DMxGNuZF2fArT6zSYPrlMeg0awHdePxpJmtq0uhb+QEhy1L637RK8EoU7lBYLiSZ1\nL0Mj91N496h48ZK2092yFcG8vhBUKhUTu8YR/iaJaMyF143FpLoHV+5c4QAXHpFu231adtTu8QbT\nCsJSpO4lvE9bKgD3ObwoNwriB69pbe+KbBsDALfXj067M+Y1yh8cORNRTshk44XwZ9rOdMPW64IK\n4lPCre1dWLds2oi6RSlRbhTQNKtK9F6UFC+xlBrQUFsRU2glXnjKPLxG3namGza7S3HWO+UWBmci\nygnR68YjbbwQPtfnNhTh4LEOfEdiDdo+4IbL45fM9s6EcqOAZ+9bDFOx+MhXSUWy6IeUd1vFR9jh\nKfN0s94pt/AxiohySjiTOZXAHK7SJTaFayrWoUKi4lc4oG1aXYtbF05GNgaWbq8fb+w9i8vdg6LX\nl6wi2c3zqrGqcTL8gRDuWlWLxTdMEH1vY10lAOktXvHT3pTbOHImorwltfXozlumY8cfz0amdvWC\neNSNXgMOBkOi3Z1SMVTeU4VA1NDW7R3K9t7TeikmES08xSy3T3tSZQk+/Isde49dgV7QAAjB7Q3C\nIKgBqOD1BWKm/7v73CPKeqfcweBMRHlLagr35PleXOh0RF53X8u41qiBwLUAbBA0CIVCCASDQxnb\np7pGdC2L6q24Z00dnvv5Ick9yVJTzGLr7cUGbdw9RCd7Dd3EsjnVuDeq7Gemst5p7DE4E1Fekkuk\n6rA5RF8PRI2M3d6h5hHBENBUZ0WvI3mRDznGEgFeXwB2BcVC4gurxK+3F+m1+Porh5Ke5+T53piv\n5Ubh7DyVXxiciSjjUm1ekY6efrdkhrXcnuB477Z2yGZBp3Ieh9MLs0yjiTD7gBu2XhcErTrmZxRe\nb+9U0OEKGPoZxE9VR2eqd/W6RpT1TmOHwZmIMibd5hXp2HVEOrtarppWvFQCebLzHDphg0bBbQo6\nDb63rRV2hw8Wk4AF9VUxPyO56eloekGTMFUdnal+5q/dWX1AouxhtjYRZUx4Dbi734MQYot9ZJLH\nF0Dbaek14urKsUt6CihIKnN7A7A7fACAngEvdh2+iNfeORXJOgcgWUhFKYOgTTnrnXIHR85ElBEj\nbV6RCrkpbQCwlhlwyebMyPfKJLUKQAgQi99/bO3A0XYb7ANeWEr1mDezErcunIyWk12wO8Tv1Xtt\n+YAZ2OMPR85ElBFKmldkytuHz0seU6uAY6d7Mva9MikoEZiBofaWPQPeyIzD7iMdUKlUeOa+RZKl\nO5mBPX4xOBNRRsgV05ALIh5fABc7B3DR5lBUJMPjC2D/h9K9jTO1hpwLWtu7IOg0aJpVJXqcGdjj\nF6e1iQjAyDOsU93GEwgG8do7p7Dvg8uRfbsGQYPlc6tx960zJRPIbL2umD2/41l4xiGTdccpPzA4\nExW4TGZYpxJEtu0+jd1HYrcwhfceq1QqbG6uE39gCOXe0DiV7PBUhGccMll3nPIDgzNRgctkowSl\nQcTjC6DlpPTUdGu7DYFAEG1nuhMeGKzmYhgEdWS0nQuyNZUeP+MQ3gdN4x/XnIkKWLIMa6WNEuIb\nTyRrXtHn8MgW6uju92BP6yXRLVl6nQY33Sje/GG8ELRqNDfVcNq6gHHkTFTAlGRYy43U0p0SL9Jr\nIWhU8AbEh5xS08QtJ20IBEM4fjY3s7GTUauGZuXNJj2cHr/o2rleq8a3/mEpypmFXdAYnIkKWCqN\nEsTWf1OdEo8O5lKBGZCeJu4Z8GSk1OZYWXzDBKxfcT3KjHr8+t0zoslzK+ZPYmAmBmeiQqYkw1pq\ndLx+xfUpFx2JD+ZSykp06Bv0JbyercSr0WAQNNiyth7F+qH/dpmBTXIYnIkKXLIgITU6drr9KU2J\ny61vR6soNWDuDDP+2Ho54Vi+BmYA+FjDxEhgBpQnz1FhYnAmKnByQUIuoJ44Z1c0JR6eDvf6Aoo6\nLTXWVSKYg9ul0lVRKj8iZgY2iWFwJiIA4kFCLmGs1+HB4hsmoPujqwnH6qeWIxAMYeuu9pjpcL3M\nFihTkRY3Xl+B25dch+d/cXjkN5QDyo0CvvaZJpiKxctvEklJOzgfOHAAjz76KGbOnAkAqKurw7/8\ny79Eju/btw/f/e53odFocPPNN+Ohhx4a+dUS0aiSSxgTdBq0X+wFMLwWrNepgJAK7x+/gpb2zphA\nnKz94YDLjwMfXR1KFvPnzh7mkegf9MLl8TM4U8pGNHJevHgxfvCDH4gee+655/DSSy9hwoQJ2LJl\nC9auXYvaWiY6EOWD6MxsqYQxtzcQ2QoUXgv2+EIAQteOiwdYg6BGKAR4fOLHMxGY9Tq15PnV6qEm\nE+kyCBrF5UPZmILSlZVp7QsXLqCsrAwTJ04EAKxcuRLvv/8+gzNRjhPLzJ4/sxKrF07GsVPdsA+4\nUW6U3qOrRCYre5mNAubXWdF2ujsmmW353Go8+zPxqfFQEBC0Knj9qa1rq1XAyvmToFar8M4RZdu5\nGmZYmORFaRlRcD59+jT+4R/+AX19fXj44YexfPlyAIDNZoPFYom8z2Kx4MKFCyO7UiLKOrHM7HeO\ndKC5qQbPPXDTUGKXP4inXzo4hlc5zHXtAeHhDXOgUathLS+CVqPC1rfbJbddhQD4ZfZYSwmGgOam\nKagyF0GlUqG1vQs9A27ZUt/NTVNS/j5EwAiC87Rp0/Dwww/jE5/4BC5cuIBPf/rTeOuttyAI6a2t\nmM3F0Gqz94RptZqydu58UMj3X8j3Dii/f6fLiz+3JW5fAoDWU134zLo5qJlUDrfXD6u5CJ12VyYv\nMy1ubwB7Wjqwp6UDVnMR5s6ohKBTY0/rJdnPpbsl670Pr+LzG+bh0XsWwu3140r3IL7+0/2w9boT\n3qtWD73/wfVzodGMTaVk/t3P3/tPOzhPmDABt99+OwBg6tSpqKysxNWrVzFlyhRUVVWhq6sr8t6r\nV6+iqkq8H2mY3e5M91KSslpNsNkGsnb+XFfI91/I9w4ov/9AMIiv/fSg5FR1d58bD39nN5pmVWHT\n6lo0zKhQVExkNNnsLuw+nN0Zuv0fXMa6pddFpqpLtGrMq60U/VkEg8Af9v0VXq8/5QYimcC/+7l/\n/3IPD2k/zr355pt46aWXAAxNY3d3d2PChKFi9DU1NXA4HLh48SL8fj/27NkTmfImotwQ3axi665T\nuNwj/4Dc6/BGmk9sWl2L5qYaVJQaoFYNJUllmirjZxy5ngEP+hyxWeebVtdiVeMkqCUuOJUGIkRh\naY+cV69ejccffxzvvPMOfD4fnnnmGfz+97+HyWTCmjVr8Mwzz+Cxxx4DANx+++24/vrrM3bRRJS+\n+KQvs0nAoMuv+PPh0pzRhUuMxTq8sfcvOHyiE70O6W5TqcjFMiRq1VDTjmgatRprF0/FHyWm0pU0\nECGKl3ZwNhqN+PGPfyx5fNGiRdi2bVu6pyeiLIlP+pJr3SgmOthEFy7Z3FyHdcum4emXD2YsQOea\nYAii+5ZTaSBCpAT7ORMVEKX1reWIdasKT4+bigUUG8Zv4cGKUr1ooA03EBETbiBClIrx+6+IiBLI\nleNUSq5bVZFei0td2UvuHGuNdVbJQMsuU5RJDM5EBaTMqIfZJIhOZet1ahiLdOgZ8KC8RI95MysQ\nCoVw7HQ3+hxeWEqTd6sCRhb4c9nK+ZNkAy27TFEmMTgTFRC9ToOSIvHgXGUuxlP3LoxJ8Gptt6HP\n4UW5UY+G2gpsWl0LjVqdkenxfPOJm6ZCo06+EsguU5QJDM5EBcTjC8Dp9okec7p9sNmdsJqL8et3\nz8SMiu0OD/a0dECjVmFzc11GpsfzidmoY1IXjSoGZ6ICIhdUu/s9+NrLh2AxCXB6xPfltrZ3Yf2K\n6fjDgXNQqSBbunI8MRbLT1FHNwrhVDZlAoMzUQGR2/ITJre1yj7gxvM/P5y0YMl4M+jyweMLJARe\nsaS4xjprZPqfKF3820NUQOS2/Cih1agKLjADQK8jsTIYMJwU193vQQhDsw/hKmpEI8HgTDTORO87\nFhNdelOVYo1Mf7odI/KcWCERuaQ4luykkeK0NlEOS2UtU2yKdfm8yVi3NDbLOLzlZ/2K6/HLnSdx\n4KNOxaUyg5lrxZxXxAqJyK3fs2QnjRSDM1EOSmctU2zf8Zt7z8Lp8op2RXpj71+w/6POrN1DPlOr\nhmp7W2QKibBkJ2UTgzNRDhILtOGvxQJtsinWDStnxIz8CnGfcipWzp+EtYunxsxYDDi9uNjpQE2V\nEaZiIbJ+L9YukiU7aaQYnIlyTKqBFkg+xWrrdUHQqiPBps/hkc3YLiRTqoxwuv0JJTfDMxRevx/P\n/6IFHTYHgqGhUfVkqxFf+fQCluykrGFwJsox6axlyk2xCjoN/u1XR2Ef8Eamx29fMhVq1VCXpZEQ\ntGp4/bm/EF1jLcGDd8zGntYOtJ3uTgik/kBIcm3/+V+04EKnI/J1MARc6HTg+V+04Nn7FrNkJ2UF\ngzNRjklnLVNuitXtDcDtHcocDk+PO93+EQdmAPjnTQ048FEn3j16KSPny4abGybi3tvqoVGrce/H\n6+FZlZhkp1FDNHlrwOlFh82R8DoAdNgcGHB6I1PcTP6iTOJWKqIck077wUAwiFAoBIMwfMwgqGEQ\nxP+Jnzhnh8UkiB5LxSt/OIlFCLMMAAAgAElEQVSb508as0phgjb5XjCNJvY94UCqZIR7sdMh+dAR\nDA0dJ8oGBmeiHBS9F1mtAipKDWhuqpFcy9y2+zTeOdIRGSEDgNsbhNsrPuXc6/DghussI77OK3YX\n/s8vWyDoUtwwnQHlJTo01FZCp5X/b2xP66W0i4LUVBmhlrg1tWroOFE2cFqbKAeF9yKvWzYtJkM4\nnscXgM3uTDnzWtBpcM+aOpzvdMSsp6bD4xubNefeQR8On1B231KJdMmYigVMthpFf0aTreK/E6JM\nYHAmyiHhoiPRLRvF9jlH74NOJ+s6FArhavcgBl3SdbTHk5EUBfnKpxdIZmsTZQuDM1EOiC86ohc0\nMVPU8fuc4/dBp8rjC+Ibvzgy4uvOFyMpCiJotXj2vsUJ+5yJsolrzkQ5YOvb7TENFKIDc7TW9i4M\nOL0sIJKiTBQFMRULuGGahYGZRgVHzkRjKBAMYuuuU3j36CVF7+/pd+Nip0NyH3Suq7YUodPuyvq2\nK/W1XtNWcxEaZlSwKAjlHQZnojHi8QXwy50n8d7xK4o/oxc0qKkyJu3JnItunj8RaxdNRZFei1f/\nvxM4cd4OlzeYkWIo8UIAHr97PhbPm4yBPldmT040ChiciUZZeH35yImrsDt8KX46hDf2nsWgO9XP\njb332i7jT0cvJ7yejVF0eYke0yeXwSBoMZD50xNlHdeciUZZOJkr9cA8tHd5T+ulhP3LGjWwsnEi\nKkpztxNSIMmOK71O+X9HS26sgtmokzw+n40nKM8xOBMl4fEF0Gl3wuMLiH6d6rmykcwVCAJqlVqy\nslg+ULJf2iBo0NxUg/s/eSMWzpog+p4pVUZsbp6Z6csjGlWc1iaSEL+9yWwSUFIkwOn2Ke6xHC+b\n3aBaT9owt9aSN80o4pmKdBB0atmfT4lBiw0rZ0CjVsd0hOrpd6PMKKBxZiU2r6lT/PsgylUMzkQS\n4vcS9wx40TMwXLQjWY/leIFgEDsPXchKAhQA9A56sfeY8uSyXNNYXwlBq5Hdv20f8ESKiYSrqLEj\nFI1HIwrO3/72t3HkyBH4/X587nOfw8c//vHIsdWrV6O6uhoazdA/lhdeeAETJohPQxHlmlSmn5WW\nhty2+zT2tHSM6LoErQpef462f7rm5saJOHuxHxdtg4o/o9WocO/H6wEAgWAI77Z2iD7AiBUTYUco\nGo/SDs779+/HqVOnsG3bNtjtdnzqU5+KCc4A8OKLL6KkpGTEF0k02uR6KsdTUhoyU2vNC+qtOH6m\nBw63f8TnyhatSoWnP7sIW99uR+upLvQ5vBB0atk15dJiHfyBELQaFTRqFXRa8fdnopgIUT5IOzgv\nWrQIDQ0NAIDS0lK4XC4EAoHISJkon8n1VI6npDSkXLBXqQBTkYB+p3yda4OggVarzunADADvHb+C\njatm4t61s3DX6qFa4V5fAF97+ZDkZ+wOL/ocHuw6clF0WnsoG30yi4lQwUg7OGs0GhQXD40UduzY\ngZtvvjkhMD/99NPo6OjAwoUL8dhjj0Glkm4rZzYXQ6vNXmC3Wk1ZO3c+KOT7T/fel8+bjDf3nlXw\nvkmomVQu+x5TWRGs5qHqWAnXV16EuuvK8WeRPcDRqiuKse+D3F9T9niD8KtUqCwrwmC3EyUmA2pM\nBljLDbD1ukU/Yy0vQs2kcrT96pjo8UAQMOh1qJ5QlvL18O9+4crn+x9xQtiuXbuwY8cOvPzyyzGv\nf+ELX8CKFStQVlaGhx56CDt37sRtt90meR673TnSS5FktZpgsxVuKYJCvv+R3Pu6pVPhdHnR2t4F\n+4Ab5UY9Sop0cLp9sA94YDYZ0FhXiXVLp8JmG4h0lJJKTGqYUSE6Kuwf9CQNzBMtxfjr5fz5Hf7i\nDx/hg9Ndkf3YBkGDynKD5PsbZlTg4qVe0YeXsPfbLmPd0utSmtbm3/3CvHcgP+5f7uFhRMF57969\n+PGPf4yf/vSnMJliv8n69esjf7755pvR3t4uG5yJco1GrcaGlTNwc8NEQKWCtbwIep0mIQgP1cdu\nl2zvGHbnLdNx8nxvpPVgWHxBkXhmowCvP/U91WPp0EedMV+7vQFc7BxEtaUI9gFPZD3ZIGiwfG41\nNq2uhT8QQrlRQK9DfHq/d9CTdttHonyTdnAeGBjAt7/9bbzyyisoLy9POPbFL34RP/rRjyAIAg4d\nOoS1a9eO+GKJRkv8HufogBufHRy/5Sp+i1U4mO88eB4XOh0pX8sN0yx4P4X627nsSo8LZpMejTPL\nsPam61BtKY6MhDVqoHFmJfa0ijcBsYyg7SNRvkk7OP/hD3+A3W7HF7/4xchrN910E+rr67FmzRrc\nfPPN2LRpE/R6PW688UaOmimvSAXcQDAU2fIDAANOLw6f6BQ7BVrbbQgEgmg7042efg9kUi4kLZ9T\njXvWzMTJ8/a8a3QhxT7gwf6POqFRq7FlbX3Msc1r6nC6o1/0IYaZ2lRIVKFQKCc2TWZzbSAf1h6y\nqZDvP5179/gC+OqL+0WDoVoFLLphAjavmYk39v4FR050ot+ZnSYUpmItHr2zAZOtppS7V+ULi0nA\ngvqqmCWAcBvNo+1d6B30wHJtbT+VSmxh/LtfmPcO5Mf9Z23NmSjfhaeci/RauDx+lBn1stuegiHg\nwEdXceCjq0nPrbrWUzhdA04/nvtFCwyCGotvqIJBUCddn843PQPehCprGrUad62qxarGyUAoBKu5\nmCNmKjgMzlSQwmvKLSc70TPgjZTUrCjVY870CpSVCOgdlN93nEym5qTc3iD+dOwKplQZ01qzzgfh\nKmtajUpyrZ/1sqmQMDhTQYpfUw5nT3f3e/DuUfGEpLE2MDg+1pzFhKusxRchSbV+OdF4wUdRKjjZ\natuYbb2D2VnbzgVmkwFFeq3k76W1vSutFp1E+YrBmcalcM9ltzex1GUqdbMzqdwoYGXjJOi16f2z\ny9VlV4Mw8v9GGusq4fL4JX8v4ZE1UaHgtDaNK/H7k63mIjTMqIhZs0ylbnaYkh7J1RVF6OxxiXZT\nUqmAL909H5ayIrSf68XlntQr4uXqwLGi1ICOLvn7MQgaeH0BmE16FBt0GHT50OsYrrIWLkIi9XtR\nUr+caDxhcKZxJX4tudPuSliz1Os0aKyzyvYNjmYx6THnegsOfHQVnmsBWqMeanPo8YWgAhAC4HL7\nJfs0h0LAsz8/BBVUst2Z8pHD5cOqBZPRdrob3f3uoZkBNeDzBSPBd/2K6XA4vZGqamKlTjVqSP5e\nuMeZCg2DM40bcmvJ8T2Xw92NWk7a0DPgiWRriykp0uFPbbG1rwNBYIKlCJe6nAh/rC/JmrDXFwKQ\nE2UFMqp/0Ie1i6bgrlW1kYALICH4FuuH/7uR6sEc/r2E65lHj6yJCgmDM40bcmvJ8T2XNWo1NjfX\nYcPKGZF9zg6XD7sOX0DbmZ5IYGiYYUHbmW7Rc15KMpVbKCylhkgQjg646dTAjv+9SDURIRrvGJxp\n3JBbS5Zas4wOKKZiAfeunRUz5WqzOyVrPcspLdZlrXJYrsnGlLPUyJqoUDBbm8aN8FqymPgAEs7m\njt+eEw7MxmIBv373DL63vS3l67CY9Nh068yUP5ePjEVaTjkTZQFHzjSuxK9ZVpYPZ2sD0t2m7rxl\nOnb88Sxa223o7vcoys6W4nB58dPffZSxe8plTrcfTrcfpmJhrC+FaFxhcKZxJX7Ncsa0Cgz0uSLH\npbpNnTzfG1MaM93APPTZ8Zf0JSUYAi52OnDDNMtYXwrRuMJpbRqXwmuWBmH4+VMum3u81qzONrUK\nqKkyjvVlEI07DM5UMMaqMthYaaqvhEFQlqhVbhSgVg01/qixlkCvG/6vwSBoYCwSn2SbbDVySpso\nCzitTQWjzKiH2SSgZ2Bk3abyhVarwavP3ob/OW2D1+fHd147KloAxSBo8Ox9iyMtM8NFQmx2J6BS\nwVpeBJUqhOd/0YIOmwPB0NCIebLViK98esEY3BnR+MfgTAVDr9Ng1nUW7Dt+ZawvZVS0n+8FANRY\njRhwehGSqrICQNBpYkbAep0GNVWxjeCfvW8xBpxeXOx0oKaKI2aibGJwpoKyec1MtLTb4PbmaKHq\nDOp1eNDV68L2XSfx52OX4Q2IB2ePNxBToEWOqVhg8hfRKOCaM41LUl2p9DoNKssMY3RVo8tsMuB3\ne89i95EO2ezzcIUvIsodHDmTLLEGBbkm+hq1GpVsV6ptu0/jom0w49dQXiKgpEiLrj53zjS2mD3d\njP0fJK9uJlbhKx9+70TjGYMziZIq1hHdenGsiV1jsUEXsy0quivVhpUzJLdSjVTvoBcurz9nArOx\nSIu2093odcgnvy2bUx1T4Ssffu9EhYDBmURJFesAhlsvjjWxa5Tq0dza3oWlN07I6laqXAnMAOBw\n+ZO/CYBOp4r5Oh9+70SFgI/ClCBZ68X4etRjQe4axXT3u/GDX7eNw4aNI/Nu62Vs230aQH783okK\nBYMzJVDSenGspVNQJFm/5ULVctIWWWPO9d87UaFgcKYE4daLYqRaL442uWuk1NgHPJHkr1z/vRMV\nCgZnSpBK68WxIneNwFDVK7UKqCiQbVMjYTbpI1nZuf57JyoUDM4katPqWjQ31aCi1HCt5rIBzU01\nOdW7d/2K6yVrR5cYtHjms4vw/X++BRUcYctaUG+NBN58+L0TFYK0s7W/+c1v4tixY1CpVHjqqafQ\n0NAQObZv3z5897vfhUajwc0334yHHnooIxdLoye+9WIu7nd1OH3wSFT6Cq+dlhn1mDXVjPcKpGQn\nANRPLcPJ831J32cQNFg2N3YrVT783okKQVrB+eDBgzh37hy2bduGM2fO4KmnnsK2bdsix5977jm8\n9NJLmDBhArZs2YK1a9eitpZP3vko3HpxrMgVwwivkYptnwoB+P6ONiw/1Y31N08vmOCsVgGf/cQN\n+MqL+xEQ2dmlF9T40j2NEDRqWM3FkoF3rH/vRIUureD8/vvvo7m5GQAwY8YM9PX1weFwwGg04sKF\nCygrK8PEiRMBACtXrsT777/P4EwpUVIMI7xGGr0vN1p3vwdv7j2LQx8WRmAGhjpFVZmLcUvjZLxz\npCPh+MfmTsT0iWVjcGVElIq0gnNXVxdmz54d+dpiscBms8FoNMJms8FiscQcu3DhQtJzms3F0Gqz\nN31mtZqSv2kcy7f7f/GND0SLYRQXCXhg/dzI6w/f1YjiIgHvf3AJtl636Lku9zizfr1jTa0GplWX\n4juPrIAgaPHIpgUoKdZj//HLsPW6YC0vwpI5E3HfutnQaAor1STf/u5nUiHfO5Df95+RCmGh0MhL\nO9jt2fsP1Go1wWYbyNr5c12u33/81LXT48NbB86Jvve9Y5fwicVTYqZjP7F4CqZXG/Fv29tG65Jz\nyv23z0JDbSVMxQL6+lyR19cvn4Z7b78BZ/7aHfnZ9vRkvq54Lsv1v/vZVMj3DuTH/cs9PKQVnKuq\nqtDV1RX5urOzE1arVfTY1atXUVVVlc63oXFOaura4fZJtnQMF8OoMhcnfL4QVZQa0HTDBMm1Y4Og\n5doxUR5Ka35r+fLl2LlzJwDgww8/RFVVFYxGIwCgpqYGDocDFy9ehN/vx549e7B8+fLMXTGNG+E6\nzt39HoQwPHXderJT8jPlJj28vgA8vkDC5wsR9x8TjU9pjZwXLFiA2bNn4+6774ZKpcLTTz+N3/zm\nNzCZTFizZg2eeeYZPPbYYwCA22+/Hddff31GL5ryn1wdZ49POtQ6nD48/fIhmE0CnJ7Cq/WsUSOS\nhW0Q1AiGQggEg+wYRTTOpL3m/Pjjj8d8PWvWrMifFy1aFLO1igqT3DaodGpjA4DXPxSZegbkWyGO\nFwZBA68vALPJgCK9JqYXtdsbxO4jHVCrVOwYRTTOsGUkZZySbVBye5QNgkZyzXm8U6uG9mhbTAY0\n1lVi/Yrp6HN4sPPQefz52GXRz7S2d2HDyhmc3iYaRxicKeOU9ASW26NcWW5AV687EqD1WjU8/uz3\nSv7nuxrQM+DBG386g95BZf2QM23l/ElYu3hqzGzDG3vP4k9HxQMzEJskR0TjA4MzZZTcWvKREzas\nWzYNpmIBACJlI1vbu2AfcMNsMqDYoMWFTkfsOf1BGAQ13N7EAG0QNCjSa2Af4TS3SgW88t8n0DPg\nRaa3AWvUgFajhscn/4BhEDTYcEstivXD/yyV9K1mxyii8YfBmTJKtieww4OnXz6IpllVkSnu6DrO\nRXotvv7KIYkzq0RftZYXYVJlMQ58JJ3hrUQoNLyOLVb2MhUq1dD5wl2xnti8AIJWjadfPoheh/RD\nhNcXgMPpjQnOStbmmbFNNP4wxZMyKlmf5V6HF7sOX8R/vH0y8lq4jrPL45cMRF5fANWWooTXL3Q6\ncPRUl8gnxk64Jk8wBNh63fiXn+7H7/b9FQvrpVtcAuIjYLmfp1oFrGqcxI5RROMQgzNlVLI+y2F/\nbL2MV986iUBweJhqLNZBL9ECUqdV42qPS/RYsuniseb2BrHr8EWEADQ31Ui2uWyYYUGfwwOPbzgZ\nTu7nubJxMu5dO4vbqIjGIU5rU0Z5fAGsapyMQDCElnYb+mSmcfe0dECjHt4G9Js/nZXM0s71AKzE\nsVPdeO6Bm7B+xfXY+vYpnDhnR6/Dg3KjHiVFOrSd6cYfWy8lZLeLrc031lVyxEw0jjE4U0ZEb5/q\n7vdA0Krg9Sev2xXeBgQA+z6QzkiWky9br6Kzqv/+kzdG9oHvPHQBe1qGO0jFZ7ezxzJR4eF8GGVE\ndClNAIoCMzAcsGx2p2g2thJL50xAjbUkrc+Opvg1Zb1OgzKjHm2nxdfMW9u7Eqa4q2R6MBPR+MHg\nTCOmZLuPlEjAUolnYyvh8wfh8ozNvuRUiGVVy2a3X3twIaLCw+BMafH4Aui0OyNTs+l2hQoHLGt5\nETRq8QCdLGy/13ZFtNKYEjoNUG4U0vqsHL1ODUE7fOUGQYPQtTrY0eSysbl/mahwcc2ZUiJWmrOh\nthJmk6Co3nV8ecropCadVoWAN3E6XC+oMb/Wiv0fXRU950g6UvkCkN17nK4qc3FMMRW3N4B3jnRA\nFVcHW6tRodigE3244P5losLF4EwpESvNuaelA1OqjIqC88caqnHTDdWoqTJGKoUBQ9O7UmvObm8Q\nK+dPwoGPruZ0a0iVauiho6G2AsdOiU/zRyfA9Tk82HnwfEJFNACYUmVkNjZRAWNwJsXk1padbh9W\nNU7C+x9eFc2cVquBSRUl+PAvduw9diVhu5CxWCebdf3j3x5XFJjD1blGm8WkxxfvmoeyEgEXOx0x\n2dfRevrd+OXOkzhx3o6efo/kUrvT7Yc/EMp4KVEiyg8MzqSYfPKSB2sXT8WGW2rx2tvtQ8FnwIOy\nEgGzppZDr9fi3dZLkffHbxd6Y+9fZLdD9Q36FF3jWARmAJhfV4k/HbsUme5Xq4YqhMXT6dR47/iV\nyNdS18tmFkSFjcGZFJNr8xhOXtLrNLj/2h5eW68LCIVQZtRL1sxube/CumXT0HJyZLWxR5v62gjd\nUjq0dh4KhWKm+6WCrldhMRUmgxEVNgZnUkyuzWN08lIgGMSv3z0TGUWWGQXJpCv7gBsXOx2K1qtz\nycrGyVg1fxKgUqGsRJB8+JAaQSfDZDCiwsbgTEmFt0uVGfWKSknGJ43JZUObTQZUmYvSDmIjoVYD\nQZmBrFoFlJYMPViEr89i0mN+XSVUAL6/oy3pw0cIQGmxgH6n/MNH/EicyWBEhY3BmSSJbZsKJ3FJ\nlZJMtSBJY10lXB7/qAdmYCgw11SVoMM2KDoNfcuCydh4S22knaXL40eZUY9fv3tG8cOHqViH/sHk\nswIrGydj7aIpLM1JRAAYnEmG2Lap6CQusWSlPodHtiCI2ahH36AHZpMBc6ab4XT78b3tbZm/eIVc\n7gD+n4eW4Ve7T+PDv3RjwBWA2ShgYVTP6fB9moqFlB8++gd9srMCFpMeC+qHs9aJiAAGZ5IgF4TC\ne3XjR3iBYBA7D12QDEYVpQZ87TNNcLh82HXkIt4/fjntetqZYh9ww+UJwFgsQNBpoXIFoJaoVAbI\nZ6wDQw8f9riSm1KBefmcamxZW8+RMhEl4KM6iVJa8zm6jOe23aexp6VDMhg11lXCVCxgT2sH9rR0\njHlgBoCyEj12HjofadoRwvAMwbbdp2PeGwgGsfPgecm9yRWlBjx17wLJcqBq1VAp0opSA5qbavCZ\n22cxMBORKI6cCUBs0le4W5LctiljsYCtu9qHM7JLBLi80s0nJltLsGl17YiaZGSD3eHBn4+Jt6qM\nnyHYtvs09kTt1Y7XWFeJQDAk2cM6BODxu+dj+uQyBmUiksXgXODkkr6ktk0V6TX4j7dO4v0Ph2td\n9yZJehp0+eDxBbD17VNpN6nIFqmRfnQhELmHCrVqKKFr0+pa+AMhyYcai8nAwExEijA457H40W46\n5JK+Nq2uxcnzvQm1ny/aBnHRNpjS9+lzeLH17VPYF1UdK9eVG/WRQiBy0/yhELB20RRo1Gpo1FC0\nF5yISA6Dcx6SG+2mkvGbLOlr3bJpcLqVlc1Mxlyqx4lzPSM6h0EYCmweXwAqZH9fdEmRLhJM5ab5\nLaWx1byU7AUnIpLD4JyHkm1xUipZ0tfFTkfafZrjOZw+eP3KE8AMggZeXwDma12emhfWwFJqiFz3\nHw6cw5+Oiq8VZ4rTPTQVr9dpFFdHAwCNWo3NzXWSe8GJiJJJKzj7/X585Stfwfnz5xEIBPDEE0+g\nqakp5j2zZ8/GggULIl+/8sor0Gj4H9RIpbPFSUqypK+aKqPkcTla9VBWcnQZaaWBWaUCbmmcjA0r\nZ8Dh9IoGtipzMdYumpr14Gwf8MQ0n0h1RKzXadi4gojSklZw/u1vf4uioiK89tprOHXqFL785S9j\nx44dMe8xGo149dVXM3KRNEzJFielASHZaFDQaVA/1ZzyOnEKA+QEEyxFuPfj9QCAYr30X09LqQEG\nQZ32dixBq4I/EJKdGo9vPpFsRJyJHAAiIiDN4HzHHXfgk5/8JADAYrGgt7c3oxdF0pR0hkqF2Ghw\n/swKBEMhfPXF/ejp90TWeuVaOmaKvd+DAac3UipTPshJFwvR69TwyHSA8vmTL1hLJXDFj4gzlQNA\nRBSWVnDW6XSRP//85z+PBOpoXq8Xjz32GDo6OrB27Vp89rOfTf8qKSKVtU8lxEaDv373DN6JOn84\nKC+bUw29To22Mz2wD7gh6DIftD2+IJ5+6SD6Br2yQa7P4YFH4vuqVEBTfVVM3+R4ZpMeKhVEH3LU\nKmDl/EmKE7gylQNARBSWNDhv374d27dvj3ntkUcewYoVK/Af//Ef+PDDD/HjH/844XNPPPEE7rjj\nDqhUKmzZsgVNTU2YO3eu5Pcxm4uh1WZvKtBqNWXt3KPt4bsaUVwkYP/xy+jqdaGyvAhL5kzEfetm\nQ6MRH6kpuf8aAG6vH21nukWPn+7oww+fWA0AuNLtRCAYwH/vO4cDx6+g15G5vcvhPdPhIFdcJOCB\n9bF/d0xlRbCai9BpdyV83lpehEfubkTFzpN4++A5uDyJQfxj8ycDAN7cezbh2G1Lp+HzG+Ypula5\nn1fbmW58bkMRDMLY5l2Op7/76Sjk+y/kewfy+/6T/q+xceNGbNy4MeH17du3Y/fu3fj3f//3mJF0\n2D333BP585IlS9De3i4bnO12p9JrTpnVaoLNNpC18482jy+AZTdW4dbGSTHTvz094nuPU7n/TrsT\nNpGABwBdvS60n+3CntYOtLbb0i4mkmp7yPeOXcInFk9JmBWYO92Cd450JLx/7nQLnA4P1i+fhs1r\n6/H/vt6KE+ftsA94Iklc65ZOBQA4Xd6EBK9PfWxaxn5eZ/7aPaZJYePt736qCvn+C/negfy4f7mH\nh7Qe6S9cuIDXX38dv/zlL6HXJ65xnj17Fj/84Q/xwgsvIBAIoKWlBbfddls634qiiK1tNsyoQHPT\nFFhKDRlJQkq2pr3r8AXZEpZKTLYaEwqbyOnpF090k4rv/mAQnXYnyox6WIsE3P/JGyWTtUa65SnT\nOQBERECawXn79u3o7e3Fgw8+GHntpZdewiuvvIJFixahsbER1dXVuPPOO6FWq7F69Wo0NDRk7KIL\nldja5p7WS9jTegkVGUpC0us0aKitxJ4WkRHpDEtMyc50LJ9TjXtvq8P2PWfw3gdXIuvVekENny8o\nOqLWC5qEIOfxBXBUYkvZ3qOX8afWy7CU6rF83mSsWzpVdlvTSLY8ZToHgIgIAFShkFib+dGXzemH\nfJjeAOS34nh8AXz1xf1Jp5Kbm2oSkpCU3n94ZN5yshM9A97I9HNFqR6zpprhDwZx4KPOpOdRqYZK\nWsazmPR4/sElkXvz+AKw9boQCASx52iH5L5lg6DB9x75WORzgWAQL/3+I+xXcC2A+M8kk4ZnNBL3\nP491tna+/N3PlkK+/0K+dyA/7j/j09qUWUq24iTrIxyWaiGSaPEj8/Ao1uHy4r3jV2Q2LsWqNhfj\nck9iDsGCemvMdel1GtRYjdi6q122oIj32kNLlbkYgWAQX3/lcErT4iP5mSjBimBElGnchJkDwkFR\nrp9weG0zmehey6mQqzzm8Q1F6WRTLAZBA4OgwZUeZ+TP4f7FqxZMxqrGyfD4YjOnlbSQjF673fp2\ne0qBGUj/Z5Kq8PQ4AzMRjRSD8xhLVo4zHMzCa5vJpJuEpHRkLkavU+OmG6vg9gbg9gYQAiJ/Xjqn\nGg0zLGg73YWvvngAX31xP7buakcgGFT8fcNrtx5fAK2nulK+PiZmEVG+4bT2GEulHGd0Na/ufrfo\nZ9JNQpLLOk5m2Zxqyb2+Le22mCIl8QU65L5vdJ9kYOhn1euQ7xstholZRJRvOHIeY3LT1VK1nZ97\n4CY8/8BNWLVgMipKDVCrhqaOm5tq0m5LqHRkDgwFTVXU92xumiL5gCFVPSw8KyD3fVfOn4R7P14f\nWXcv0mtRbhQUXWNYkXhq+nIAABFXSURBVF6D9Sump/QZIqKxxpHzGEtnK45ep8HEihLc+/F6eFZl\nrtlCfJ1tQacRDa4r50/C2sVTI9/T4wukPOruiZoVSNbtKTphLtWRs8cbgMPplW2iQUSUa/g/Vg5I\ntRVhtEy2JYzOOu7pd+Otw+dx4MOrkc5PBkGDJbOr0Nw0JeFhYNZUs2wt63gqADsPnsfmNXVJs53j\ns8hTUVlexPVmIso7DM45YKRbcTLdqlCv02BPawfebY3d3uT2BrD/w068e63Ax/yZlQgBOHaqC939\nHhgENQAVvL6A5Kg7LBgC9rRegkajjuxBFnvQkEuY0+vUKDFo0evwSn6/JXMmcr2ZiPIOg3MOSXUU\nnK1WhXIBMRwAu/s9CXWtwyPsJbMnoP28XVG3qmR7kOUS5nz+IL5413wIWjWMxQLe2Hs2YfbhvnWz\nJWuOExHlKgbnPJatVoUj2VYFACfP9cKucG04PiM9nnztaj2s5UWRwC42+yDVpYuIKJfxf648pXR/\ntNznO+1O0fcpLXgixZ5CwY9ke5D1Og2KDYldzwCg2KBLGHGzEAgRjQccOecpudFtd78bPf1uTKwo\nSTgWPxVebtRjfl0lNjfPjEyFy2WQK5FKS8hke5A9vgAGXeKj8EGXL7Idi4hoPOHIOU8lG93uOnxB\n9PX4UqF2hwd7Wjrw9VcOR6p2AUMZ5M1NNZF91AZBeQBUEpjVKmDVgslJM9L7HB7YB8SDc6/DMypl\nOYmIRhtHzjkqWQa2XGtHAGg705MwqpSbCr/Q6cDWt9tx79pZAGIzyMOdo/7Udhltp7sjCVfzZ1Zc\ny9Yefq2htgLHTtnQIxFQw0IhYO2iKUkT19gvmYgKEYNzjkklA7t5YY1kcBZLtEqW6PXe8SvYcEtt\npGBHIBjEr989E3MtDTMq0Nw0BZZSQyTwb7wl9kFCo1YlnRK3lIoH1vBDSZFeC5fHjzKjnv2Siajg\nMDjnmFQysC2lBlSkMKosM+pRbtRLJmx5fUG89nY77v/kjZLXEr83GUjcApZODXC5XtLzZ1Zi9cLJ\nMSN0pUVaiIjyEYNzDpGbdj58ohPrlk2DqXi4tnSqpT/1Og3m10lPhQPAifP2SAa3XDa43N7k+Epj\nuw5fQNuZHtnAKtVLOryfurmpBs89cBP7JRNRQWBwziFy0869Di+eefkQFs6KneJOtfTn5uaZOPFX\nOy73OEWP2weGk6yUdsuSEqkBvnaW7Bq6kp7O4QeCTJUqJSLKZQzOOSRZ20a7I3GKO9XSnxq1Gl/5\nu4V47P++B48vmHA8PB0eCIagF9SRql9i70mlbKhc9TMlRU+UPhAQEY0HDM45ROn+4tb2LqxbNi2S\nMKXXaVIq/Vms12HFvEmS0+EAsPXtdtHADADzZlYkJIqNpGyokl7SzMwmokLC4JxjwtPRh090SrZH\n7O534+mXD6LP4U07MIpNh8+bWYFQKISvvrhfMlAaBA2CwRB2tw6vW4+0bKiShxJmZhNRIWFwzjHh\naep1y6bhmZcPSWZWhwN3KoExfho6fjr81++eSTpq93gDOHaqW/RYskQxOeGHhZaTNvQMeGKytcMP\nH0REhYLBOUeZigUsnKW8hGZ8YIwOxIFAEFt3tYtOQ4enw5UkZQFAmVFAr8QDw0jWhePXzqP3OXPE\nTESFhsE5h8VPPZeVSO9RDgfGijJDQhGTMqMeZy/1R94rNtpW2omqcWYl2s50Z61iV/TaefS2MSKi\nQsLgnMPERpNff+WQbGAUKxwitX4cPdpOlpRlMemxoP7a2rbmNCt2ERFlEYNzHogeTcoVHQGkC4eI\niZ6GlkvKWj6nGlvW1kcCb6p7q4mIKDUMznkkEAzCHwhCr1XD4x/a5mQQNFg2txqbVteiu8+taGo6\nLH4aWi7oRmeCp7q3moiIUsPgnMOik7q0GhW+/sphXOh0xLzH7Q3A6fLjctegov3C0eKnoVMNuqns\nrSYiIuXSCs6/+c1v8P3vfx9Tp04FACxbtgyf//znY97z5ptv4uc//znUajXuuusubNy4ceRXWyDE\nOlMZ9Fp02AZF37//o6vY/9FVGAQ1KsuKACQG5ylVRjjdfkXT0Ay6RERjK+2R8+23344nn3xS9JjT\n6cQPf/hD7NixAzqdDnfeeSfWrFmD8vLytC+0kIgldYkF3HhubxAXbYMJgXj5vElYt3Qq/IEQp6GJ\niPJAVqa1jx07hrlz58JkMgEAFixYgJaWFqxevTob3y4vKK1D7fEF0HKyc0Tfa9Dlw9OfXRTZJ1wz\nqRw22wA0anBETESUB9IOzgcPHsT9998Pv9+PJ598EjfeeGPkWFdXFywWS+Rri8UCm00+i9hsLoZW\nm73RnNVqysp53V4/7P0emEv1MAiJP85AIIiXf/ch9h+/DFuvC9byIiyZMxH3rZsNjUad8N4f/Ooo\negbEy3YqZR/woKjEgOnXlURey9b954NCvneA91/I91/I9w7k9/0nDc7bt2/H9u3bY177m7/5Gzzy\nyCO45ZZb0NraiieffBK/+93vJM8RCoWSXojdLt7CMBOsVhNstoGMnlNsXVisxvXWXe0xU9Sddhfe\n3HsWTpc3odzm1l3t2K2wIpgcs0mPgNcXueds3H++KOR7B3j/hXz/hXzvQH7cv9zDQ9LgvHHjRtlk\nrsbGRvT09CAQCECjGRr5VlVVoaurK/Kezs5OzJ8/P5Vrznli68LxVbfkSmKKldtMtkc5vJbc0++G\nTqeGV6TlIwAsqLdyTZmIKI+l3t8PwIsvvojf//73AID29nZYLJZIYAaAefPm4YMPPkB/fz8GBwfR\n0tKCpqamzFxxDkgWdD2+AAD5kpjhAiBhycpnLptTja99pgnPPXATvvW5Jfjuwx/DrQsnwyAM/9wN\nggarF05mMRAiojyX1przunXr8KUvfQmvv/46/H4/nn/+eQDAT37yEyxatAiNjY147LHHcP/990Ol\nUuGhhx6KJIeNB0qCbpW5WHbfcXwBELn3VpTqce/aemjU6pikrv+1ph533lILW68LCIVgvVbpi4iI\n8ltawbm6uhqvvvpqwusPPvhg5M+33XYbbrvttvSvLIcpDbpyJTHjC4DIv1d6mlqv06DGakz3VoiI\nKAexQlgaUgm6qdShZs1qIiICGJzTpjSQplISkzWriYgIYHBOWzbrULN8JhFRYUsrW5uGhQNpvoxw\nPb4AOu3OSEY5ERHlHo6cC4TSoilERDT2GJwLhJKiKURElBvG7ZCJ07fD3F6/oqIpRESUG8bdyFls\n+nb5vMlYt3Rqzk/fKu1clSp7v7KiKURElBvGXXAWm76VajSRK7K9HmwuVV6pjIiIxl5uDyVTpLTm\nda4JP1B093sQwvB68LbdpzNyfoOgRWOdVfRYfNEUIiIae+MqOKfSaCJXjNYDxabVtWhuqkFFqQFq\nFVBRakBzUw2rjxER5aBxNa2dSqOJXKG0icZIsfoYEVH+GFcj53DNazG5On0bfqAQk40HinwrmkJE\nVIjGVXAGxKdv71gxPWenb/PxgYKIiLJrXE1rA+LTtzWTymGzDYz1pUliNyoiIoo27oJzWD41j8jm\nerDHF8DlrkEEfAGOwomI8sS4Dc75KJMPFDF7pwc8sJhYS5uIKF8wOI9TrKVNRJS/OIQah/K1GAsR\nEQ1hcB6H8rEYCxERDWNwHodGe+80ERFlFoPzOMS900RE+Y0JYRmWrbaPqeLeaSKi/MXgnCHZbvuY\nqui90xpBh4DXxxEzEVGe4LR2hmS77WO69DoNJlaWMDATEeURBucM4NYlIiLKJAbnDODWJSIiyqS0\n1px/9KMfYd++fQCAYDCIrq4u7Ny5M3L84sWLWLduHebMmQMAMJvN+MEPfpCBy81N+dhHmoiIclda\nwfnzn/88Pv/5zwMA/vM//xPd3d0J77n++uvx6quvjuzq8kR461J0ucwwbl0iIqJUjShb2+/347XX\nXsMvfvGLTF1P3uLWJSIiypQRBee33noLH/vYx2AwGBKOdXV14Qtf+AI6OzuxefNm3HHHHSP5Vjkv\nm20fiYiosKhCoVBI7g3bt2/H9u3bY1575JFHsGLFCtx///149tlnUVNTE3Pc4XBg586duOOOOzAw\nMICNGzfitddeQ1VVleT38fsD0GoZzIiIiJIGZylOpxMbN27Ef/3XfyV976OPPop77rkHS5YskXyP\nzTaQzmUoYrWasnr+XFfI91/I9w7w/gv5/gv53oH8uH+r1SR5LO2tVCdOnMD06dNFj+3fvx/f+ta3\nAAwF8RMnTuD6669P91sREREVlLSDs81mg8ViiXnt+eefx4ULF9DU1IS+vj5s2rQJn/70p/Hggw9i\nwoQJI75YIiKiQpD2tHamcVo7ewr5/gv53gHefyHffyHfO5Af95+VaW0iIiLKDgZnIiKiHMPgTERE\nlGMYnImIiHJMziSEERER0RCOnImIiHIMgzMREVGOYXAmIiLKMQzOREREOYbBmYiIKMcwOBMREeWY\nggjO3d3d+Pu//3vce++9uPvuu3Hs2LGxvqRR4/f78eSTT+Kee+7BXXfdhcOHD4/1JY26gwcPYunS\npdizZ89YX8qo+uY3v4lNmzbh7rvvRltb21hfzqhrb29Hc3MzfvnLX471pYy6b3/729i0aRM2bNiA\nt956a6wvZ1S5XC48+uij2LJlCzZu3Ji3/+61Y30Bo+HNN9/E3/7t32LdunU4ePAgvv/97+Pll18e\n68saFb/97W9RVFSE1157DadOncKXv/xl7NixY6wva9ScP38eP/vZz7BgwYKxvpRRdfDgQZw7dw7b\ntm3DmTNn8NRTT2Hbtm1jfVmjxul04hvf+AaWLl061pcy6vbv349Tp05h27ZtsNvt+NSnPoWPf/zj\nY31Zo2bPnj3/f3v3D5JaFIAB/BNvRtHfK9ewLVqKIlqaoqJoimgTWguChhqL4g7NRrQooZiDQ2Bo\nBEFDEVE0BOGoREtLiFEXScqSQHhDcHnCe5EP3j3q+X7TuWf6DlzOxz2IB/39/VhYWEA6ncb8/DzG\nx8dFxyqbFOU8NzdnjjOZjFTXV87MzGB6ehoAoKoqXl5eBCeylqZp8Pv90HVddBRLXV9fY3JyEgDQ\n3d2NXC6Ht7c3NDU1CU5mDYfDgVAohFAoJDqK5YaGhjAwMAAAaGlpwcfHB4rFIux2u+Bk1piamjLH\n1bzfS1HOwNf904uLi8jn84hEIqLjWKaurs4cRyIRs6hl0dDQIDqCEIZhoK+vz3xWVRXPz8/SlLOi\nKFAUaba3Ena7HY2NjQCAeDyO0dFRaYr5d7Ozs3h8fEQgEBAd5Z/U3Nsbi8UQi8VK5paXlzEyMoKD\ngwNcXl5ifX29Jo+1v1v73t4eUqlU1b6oP/Hd+mXHf+mVz9nZGeLxeE3udT8RjUZxe3uLlZUVHB0d\nwWaziY5UlporZ4/HA4/HUzJ3c3ODXC6H1tZWjI2NYXV1VVC6/+tPawe+Suv8/Bw7OzslX9K15m/r\nl5HL5YJhGObz09MTNE0TmIisdHV1hUAggN3dXTQ3N4uOY6lkMgmn0wm3243e3l4Ui0Vks1k4nU7R\n0coixa+1T09PcXh4CAC4u7uD2+0WnMg6Dw8PiEaj8Pv9qK+vFx2HLDI8PIyTkxMAQCqVgsvlkuZI\nW3avr6/Y3NxEMBhEW1ub6DiWSyQS5mmBYRh4f39He3u74FTlk+JWqmw2i7W1NeTzeXx+fkLXdQwO\nDoqOZYnt7W0cHx+js7PTnAuHw3A4HAJTWefi4gLhcBj39/dQVRWapklzzLe1tYVEIgGbzYaNjQ30\n9PSIjmSZZDIJr9eLdDoNRVHQ0dEBn88nRVnt7+/D5/Ohq6vLnPN6vSV7QC0rFArQdR2ZTAaFQgFL\nS0uYmJgQHatsUpQzERFRNZHiWJuIiKiasJyJiIgqDMuZiIiowrCciYiIKgzLmYiIqMKwnImIiCoM\ny5mIiKjCsJyJiIgqzC8iivHPF8qqogAAAABJRU5ErkJggg==\n",
             "text/plain": [
-              "\u003cmatplotlib.figure.Figure at 0xa813090\u003e"
+              "\u003cmatplotlib.figure.Figure at 0x7f7a18dfb8d0\u003e"
             ]
           },
           "metadata": {
@@ -155,7 +149,7 @@
         "\n",
         "import matplotlib.pyplot as plt\n",
         "\n",
-        "plt.scatter(inputs.numpy(), labels.numpy())\n",
+        "plt.scatter(inputs, labels)\n",
         "plt.show()"
       ]
     },
@@ -168,14 +162,12 @@
       "source": [
         "## Step 2: Define our TensorFlow variables\n",
         "\n",
-        "We'll use Keras's object-oriented [`Dense`](https://www.tensorflow.org/api_docs/python/tf/contrib/keras/layers/Dense) layer to create our variables. In this case, we'll create a `Dense` layer with a single weight and bias.\n",
-        "\n",
-        "(**Note**: We're using the implementation of `Dense` found in `tf.layers.Dense` though the documentation link is for `tf.contrib.keras.layers.Dense`. When TensorFlow 1.4 is released, the documentation will also be in `tf.layers.Dense`) "
+        "We'll use Keras's object-oriented [`Dense`](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Dense) layer to create our variables. In this case, we'll create a `Dense` layer with a single weight and bias."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 4,
+      "execution_count": 0,
       "metadata": {
         "cellView": "code",
         "colab": {
@@ -183,27 +175,23 @@
             "startup": false,
             "wait_interval": 0
           },
-          "height": 34,
-          "output_extras": [
-            {
-              "item_id": 1
-            }
-          ]
+          "base_uri": "https://localhost:8080/",
+          "height": 34
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 22,
+          "elapsed": 332,
           "status": "ok",
-          "timestamp": 1505502830753,
+          "timestamp": 1525154229931,
           "user": {
             "displayName": "",
             "photoUrl": "",
             "userId": ""
           },
-          "user_tz": 240
+          "user_tz": 420
         },
         "id": "z9r-ZeyrXu3A",
-        "outputId": "6230a7a3-29fe-4d08-f101-da80425bad82"
+        "outputId": "e19a698e-5892-4fcd-80d3-1394605ee72c"
       },
       "outputs": [
         {
@@ -212,7 +200,7 @@
               "[]"
             ]
           },
-          "execution_count": 4,
+          "execution_count": 48,
           "metadata": {
             "tags": []
           },
@@ -222,7 +210,7 @@
       "source": [
         "# Create TensorFlow Variables using Keras's Dense layer.\n",
         "\n",
-        "wb = tf.layers.Dense(units=1, use_bias=True)\n",
+        "wb = tf.keras.layers.Dense(units=1, use_bias=True)\n",
         "\n",
         "# We can access the underlying TensorFlow variables using wb.variables.\n",
         "# However, the variables won't exist until the dimensions of the input\n",
@@ -240,7 +228,7 @@
         "id": "docKLUaonYG_"
       },
       "source": [
-        "## Step 3: Define our loss function\n",
+        "## Step 3: *Define the loss function*\n",
         "\n",
         "Our loss function is the standard L2 loss (where we reduce the loss to its mean across its inputs)."
       ]
@@ -261,15 +249,14 @@
       },
       "outputs": [],
       "source": [
-        "def loss_fn(inputs, labels, wb):\n",
+        "def loss_fn(predictions, labels):\n",
         "  \"\"\"Calculates the mean L2 loss for our linear model.\"\"\"\n",
-        "  predictions = wb(inputs)\n",
         "  return tf.reduce_mean(tf.square(predictions - labels))"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 6,
+      "execution_count": 0,
       "metadata": {
         "cellView": "code",
         "colab": {
@@ -277,36 +264,32 @@
             "startup": false,
             "wait_interval": 0
           },
-          "height": 34,
-          "output_extras": [
-            {
-              "item_id": 1
-            }
-          ]
+          "base_uri": "https://localhost:8080/",
+          "height": 34
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 24,
+          "elapsed": 348,
           "status": "ok",
-          "timestamp": 1505502830875,
+          "timestamp": 1525154234538,
           "user": {
             "displayName": "",
             "photoUrl": "",
             "userId": ""
           },
-          "user_tz": 240
+          "user_tz": 420
         },
         "id": "RkNbXoXkpjVH",
-        "outputId": "c36fc98d-3a57-4074-901d-c10ae017ae3f"
+        "outputId": "e4688f3c-e29f-416d-f541-6d81953b5660"
       },
       "outputs": [
         {
           "data": {
             "text/plain": [
-              "\u003ctf.Tensor: id=40, shape=(), dtype=float32, numpy=7.3549819\u003e"
+              "\u003ctf.Tensor: id=1252, shape=(), dtype=float32, numpy=16.979801\u003e"
             ]
           },
-          "execution_count": 6,
+          "execution_count": 50,
           "metadata": {
             "tags": []
           },
@@ -316,47 +299,43 @@
       "source": [
         "# Test loss function (optional).\n",
         "\n",
-        "loss_fn(inputs, labels, wb)"
+        "loss_fn(wb(inputs), labels)"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 7,
+      "execution_count": 0,
       "metadata": {
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           },
-          "height": 51,
-          "output_extras": [
-            {
-              "item_id": 1
-            }
-          ]
+          "base_uri": "https://localhost:8080/",
+          "height": 51
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 57,
+          "elapsed": 418,
           "status": "ok",
-          "timestamp": 1505502830981,
+          "timestamp": 1525154260083,
           "user": {
             "displayName": "",
             "photoUrl": "",
             "userId": ""
           },
-          "user_tz": 240
+          "user_tz": 420
         },
         "id": "K_7beXoHOU7t",
-        "outputId": "1ad0856a-02ec-4117-a6c0-b41030981d87"
+        "outputId": "8f55c028-fe2b-4edb-ad68-a849afc60623"
       },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "w: tf.Tensor([[ 1.56891453]], shape=(1, 1), dtype=float32)\n",
-            "b: tf.Tensor([ 0.], shape=(1,), dtype=float32)\n"
+            "w: -0.311619\n",
+            "b: 0.000000\n"
           ]
         }
       ],
@@ -364,31 +343,20 @@
         "# At this point, the variables exist, and can now be queried:\n",
         "\n",
         "w, b = wb.variables\n",
-        "print(\"w: \" + str(w.read_value()))\n",
-        "print(\"b: \" + str(b.read_value()))"
+        "print(\"w: %f\" % w.numpy())\n",
+        "print(\"b: %f\" % b.numpy())"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
         "colab_type": "text",
-        "id": "YIlebeb_qYtC"
+        "id": "JVDWpL9VYWdP"
       },
       "source": [
-        "## Step 4: Create our gradients function using `implicit_value_and_gradients()`\n",
-        "\n",
-        "With a loss function defined, we can calculate gradients and apply them to our variables to update them.\n",
+        "## Step 4: Create an optimizer\n",
         "\n",
-        "To calculate the gradients, we wrap our loss function using the `implicit_value_and_gradients()` function.\n",
-        "\n",
-        "`implicit_value_and_gradients()` returns a function that accepts the same inputs as the function passed in, and returns a tuple consisting of:\n",
-        "\n",
-        "1. the value returned by the function passed in (in this case, the loss calculated by `loss_fn()`), and\n",
-        "1. a list of tuples consisting of:\n",
-        "  1. The value of the gradient (a `tf.Tensor`) with respect to a given variable\n",
-        "  1. The corresponding variable (`tf.Variable`)\n",
-        "\n",
-        "Test it out below to get a feel for what it does. Notice how the first value of the returned tuple (the loss) is the same as the value returned in the cell above that tests our loss function."
+        "We'll use a `GradientDescentOptimizer` to fit our model."
       ]
     },
     {
@@ -403,87 +371,29 @@
           }
         },
         "colab_type": "code",
-        "id": "v1spZQ4NwW1U"
+        "id": "DudNEebMKDWN"
       },
       "outputs": [],
       "source": [
-        "# Produce our gradients function. See description above for details about\n",
-        "# the returned function's signature.\n",
-        "\n",
-        "value_and_gradients_fn = tfe.implicit_value_and_gradients(loss_fn)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 9,
-      "metadata": {
-        "cellView": "code",
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          },
-          "height": 153,
-          "output_extras": [
-            {
-              "item_id": 1
-            }
-          ]
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 46,
-          "status": "ok",
-          "timestamp": 1505502831114,
-          "user": {
-            "displayName": "",
-            "photoUrl": "",
-            "userId": ""
-          },
-          "user_tz": 240
-        },
-        "id": "21WMcpsmFFLd",
-        "outputId": "f51b3171-33f5-4f87-8bf7-0be2dc8edc8a"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "Outputs of value_and_gradients_fn:\n",
-            "Loss: tf.Tensor(7.35498, shape=(), dtype=float32)\n",
-            "\n",
-            "Gradient: tf.Tensor([[-3.00773573]], shape=(1, 1), dtype=float32)\n",
-            "Variable: \u003ctf.Variable 'dense/kernel:0' shape=(1, 1) dtype=float32\u003e\n",
-            "\n",
-            "Gradient: tf.Tensor([-4.06519032], shape=(1,), dtype=float32)\n",
-            "Variable: \u003ctf.Variable 'dense/bias:0' shape=(1,) dtype=float32\u003e\n"
-          ]
-        }
-      ],
-      "source": [
-        "# Show outputs of value_and_gradients_fn.\n",
-        "\n",
-        "print(\"Outputs of value_and_gradients_fn:\")\n",
-        "\n",
-        "value, grads_and_vars = value_and_gradients_fn(inputs, labels, wb)\n",
-        "\n",
-        "print('Loss: {}'.format(value))\n",
-        "for (grad, var) in grads_and_vars:\n",
-        "  print(\"\")\n",
-        "  print('Gradient: {}\\nVariable: {}'.format(grad, var))"
+        "optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.1)"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
         "colab_type": "text",
-        "id": "JVDWpL9VYWdP"
+        "id": "YBeJYxY8YaiO"
       },
       "source": [
-        "## Step 5: Create an optimizer\n",
+        "### Step 5: Define a training step\n",
         "\n",
-        "We'll use a `GradientDescentOptimizer` to fit our model."
+        "To fit model variables to the data we'll need to:\n",
+        "\n",
+        "1. Calculate the gradients of the loss with respect to the model variables.\n",
+        "2. Use `optimizer` to compute updates to the variable values based on those gradients.\n",
+        "\n",
+        "To calculate the gradients, we use the [`tf.GradientTape`](https://www.tensorflow.org/api_docs/python/tf/GradientTape) context manager\n",
+        "and its `gradient` function to compute gradients through computation conducted within its context:\n"
       ]
     },
     {
@@ -498,94 +408,72 @@
           }
         },
         "colab_type": "code",
-        "id": "DudNEebMKDWN"
+        "id": "diDZfrMJM3OC"
       },
       "outputs": [],
       "source": [
-        "optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.1)"
+        "def run_step(inputs, labels):\n",
+        "  with tf.GradientTape() as g:\n",
+        "    loss = loss_fn(wb(inputs), labels)\n",
+        "  # Compute the partial derivatives of loss with respect to the variables\n",
+        "  grads = g.gradient(loss, wb.variables)\n",
+        "  optimizer.apply_gradients(zip(grads, wb.variables))\n",
+        "  return loss"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
         "colab_type": "text",
-        "id": "YBeJYxY8YaiO"
+        "id": "1WWepgmJQOzc"
       },
       "source": [
-        "### Step 5a: Test Our Optimizer\n",
-        "\n",
-        "Now we have everything needed to start fitting our variables to the data!\n",
-        "\n",
-        "In the next cell, we'll demo these capabilities. We'll:\n",
-        "\n",
-        "1. Print the current values of `w` and `b`\n",
-        "1. Calculate the loss and gradients\n",
-        "1. Apply the gradients\n",
-        "1. Print out the new values of `w` and `b`\n",
-        "\n",
-        "You can run the cell multiple times. Each time, you should see the values of `w` and `b` get closer to their true values of 3 and 2."
+        "Repeatedly running the training step will nudge the variables towards the values that best fit the data (i.e., \"w\" will move closer to 3.0, while \"b\" will tend to 2.0):\n",
+        "\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 11,
+      "execution_count": 0,
       "metadata": {
-        "cellView": "code",
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           },
-          "height": 102,
-          "output_extras": [
-            {
-              "item_id": 1
-            }
-          ]
+          "base_uri": "https://localhost:8080/",
+          "height": 51
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 103,
+          "elapsed": 380,
           "status": "ok",
-          "timestamp": 1505502831285,
+          "timestamp": 1525154412590,
           "user": {
             "displayName": "",
             "photoUrl": "",
             "userId": ""
           },
-          "user_tz": 240
+          "user_tz": 420
         },
-        "id": "diDZfrMJM3OC",
-        "outputId": "d585fff0-ecb3-4e98-9b33-bbae07a95d8c"
+        "id": "ya5Qxz5XQlhU",
+        "outputId": "8dd47155-a6c1-44c5-c279-617c803f1723"
       },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "Values of w, b, BEFORE applying gradients:\n",
-            "(array([[ 1.56891453]], dtype=float32), array([ 0.], dtype=float32))\n",
-            "()\n",
-            "Values of w, b, AFTER applying gradients:\n",
-            "(array([[ 1.86968815]], dtype=float32), array([ 0.40651903], dtype=float32))\n"
+            "Values of w, b BEFORE applying gradients: 2.725763, 1.894334\n",
+            "Values of w, b AFTER  applying gradients: 2.774932, 1.922555\n"
           ]
         }
       ],
       "source": [
-        "# Test the optimizer.\n",
-        "\n",
-        "print(\"Values of w, b, BEFORE applying gradients:\")\n",
         "w, b = wb.variables\n",
-        "print(w.read_value().numpy(), b.read_value().numpy())\n",
-        "print()\n",
-        "\n",
-        "# Calculate the gradients:\n",
-        "empirical_loss, gradients_and_variables = value_and_gradients_fn(\n",
-        "    inputs, labels, wb)\n",
-        "optimizer.apply_gradients(gradients_and_variables)\n",
-        "\n",
-        "print(\"Values of w, b, AFTER applying gradients:\")\n",
-        "print(w.read_value().numpy(), b.read_value().numpy())"
+        "print(\"Values of w, b BEFORE applying gradients: %f, %f\" % (w.numpy(), b.numpy()))\n",
+        "run_step(inputs, labels)\n",
+        "print(\"Values of w, b AFTER  applying gradients: %f, %f\" % (w.numpy(), b.numpy()))\n"
       ]
     },
     {
@@ -602,51 +490,44 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 12,
+      "execution_count": 0,
       "metadata": {
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           },
-          "height": 397,
-          "output_extras": [
-            {
-              "item_id": 1
-            },
-            {
-              "item_id": 2
-            }
-          ]
+          "base_uri": "https://localhost:8080/",
+          "height": 364
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 225,
+          "elapsed": 580,
           "status": "ok",
-          "timestamp": 1505502831550,
+          "timestamp": 1525154278709,
           "user": {
             "displayName": "",
             "photoUrl": "",
             "userId": ""
           },
-          "user_tz": 240
+          "user_tz": 420
         },
         "id": "VukGe-huNaJ4",
-        "outputId": "f0a8d665-1910-477c-d8ab-c94ccdc4afcd"
+        "outputId": "c79c8e63-c781-451e-f74f-20815d8da49f"
       },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "[2.111051321029663, 2.3047544956207275, 2.4602210521698, 2.5850086212158203, 2.6851789951324463, 2.7655951976776123, 2.830157995223999, 2.8819968700408936, 2.9236228466033936, 2.9570505619049072]\n"
+            "[0.9409681558609009, 1.3733772039413452, 1.7128530740737915, 1.9793939590454102, 2.188689708709717, 2.3530514240264893, 2.4821391105651855, 2.583533763885498, 2.6631851196289062, 2.7257626056671143]\n"
           ]
         },
         {
           "data": {
-            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAd0AAAFXCAYAAADnFpTQAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAIABJREFUeJzs3Xd4FFUbBfAzu+m9koSShBQCSC+igIAgRRGkChJEiggo\nHURAEBQBQeADRcWCha50ULFLk6IivYRQQwskhPS6O/P9sckmm4Rkk2x2difn9zz7bLuZvC8JHO7M\n7FxBkiQJREREVOlUchdARERUVTB0iYiIzIShS0REZCYMXSIiIjNh6BIREZkJQ5eIiMhMjArdlJQU\njB8/Hk8//TS6d++OkydPVnZdREREiiMY8znd6dOno2XLlujbty80Gg0yMzPh4uJijvqIiIgUo9TQ\nTU1NRa9evfDbb7+ZqyYiIiJFKnX38s2bN+Hp6YkZM2agd+/emD17NjIzM81RGxERkaKUGroajQbn\nzp3DoEGDsH37djg4OOCzzz4zR21ERESKUmro+vv7w9/fHw0bNgQAdO3aFefOnSvxa3g5ZyIioqJs\nShvg4+ODgIAAXL16FbVr18aRI0cQGhpa4tcIgoC4uBSTFSkHX19Xq+8BUEYfSugBYB+WRAk9AMro\nQwk9ALo+jFFq6ALArFmzMHXqVGg0GtSqVQsLFy6sUHFERERVkVGhW7duXWzdurWyayEiIlI0XpGK\niIjITBi6REREZsLQJSIiMhOGLhERkZkwdImIiMyEoUtERCbRuXM7uUuweAxdIiIyCUEQ5C7B4hn1\nOV0iIqKy+OijFTh69BAEQYUhQ4ajU6fOuH8/HnPmzER6ehq0Wi2mTJmOJ59sgwUL3kZU1HkAArp3\n74nnn39B7vIrDUOXiEhh5s6dhd27d5h0mz169MLcue8aNXbv3t9x+XI01qz5Fg8eJODll4egadNm\n+PXXn9Cq1eN48cVhkCQJmZmZOH/+POLi7uGbbzYBANLSUk1at6Xh7mUiIjKp06dP4qmnugIAPD29\n0LRpc5w/fw716j2CH37Yha+++hyXLkXD0dERtWrVwp07t7F8+RIcPXoYTk7OMldfuTjTJSJSmLlz\n3zV6VloZCq80l/e8ceOm+Oijz3H48EEsWDAXAwcOxuDBA/D11xtx9Ohh7Ny5DX/88StmzHhLjrLN\ngjNdIiIyifxwbYbff/8VoijiwYMHOHXqBOrXfwSxsbHw8PDEs8/2wrPP9sLFixeQmJgIUdSiffsn\n8fLLoxEdHSVzF5WLM10iIjKJvLOX27d/EmfPnsbQoS9AEFR49dXx8PT0wp4932PjxrWwsbGBk5Mz\nZs16G7GxsXj99TcgSSIEQcDo0eNk7qJyCVIlrThv7esjKmmNR2vvQwk9AOzDkiihB0AZfSihB8D4\n9XS5e5mIiMhMGLpERERmwtAlIiIyE4YuERGRmTB0iYiIzIShS0REZCYMXSIismjHjx/DmTOn9M93\n7NiKn3/+0STbXrv2K5Nsx1gMXSIismjHjx/D6dP5odurV1907fqMSba9Zo15Q5dXpCIiogrbsGEN\n7O3t0bfvAHzwwVJcvnwJK1Z8gmPH/sGPP+7C7NnzDMZHRV3Ahx8ug0aTDWdnN7z55hx4eXlj8+ZN\n2LlzG2xsbBAcXBujR4/Fzp1boVbb4Ndf92DixNfx779/w8nJCQMHDsa4caNQp04ETp48gczMTMya\nNRdr136FK1cuo2PHzhg5cgwAYMaMqYiLu4fs7Cz07/8CevTohVWrViI7OwvDh0eidu0QzJ49D7/8\nsgebN2+CVqtB/foNMGXKdJOuE8zQJSJSGOe5s2Bv4qX9snr0QloJiyg0btwM3367Hn37DkBU1AXk\n5ORAq9Xi1KkTaNy4mcFYjUaD5csX4733liEsrBY2bdqGTz/9CDNmvIX167/Bli27YWNjg7S0VDg7\nu+C55/rqQxYA/v33b4Pt2dra4Ysv1mDz5k2YPn0KvvpqPVxcXDFgQC8MGBAJNzc3zJw5B66ursjK\nysLIkUPQvn1HjB49Ftu2bcaXX64HAFy/fg2///4LVq36Emq1GkuXLsIvv+wx2awaYOgSEZEJRETU\nRVTUeaSnp8PW1hYREXVx/vw5nDx5HJMmTTMYGxNzHVeuXMakSa9BrVYhO1sDHx9fAEBYWDjmzn0T\n7dp1wBNPdDDqe7dt2w4AEBoahpCQUHh6egEAqlevgXv37sLNzQ3ffbcBBw7sAwDcu3cPN2/GoH79\nBgYrIv3779+4eDEKI0cOgSRJyM7OhpeXV0X/aAwwdImIFCZt7rslzkorg42NDfz9A/Djj7vQsGFj\nhIWF4/jxf3H79i0EBQUXGi0hJCQUn3zyZZFrL7///gqcOPEfDh7cjzVrvsSaNd+W+r1tbe0A6BZc\nsLW11b8uCAK0Wi2OHz+G//77F5999jXs7OwwbtwoZGdnF7MlCd26dceoUa+V40/AODyRioiITKJx\n46bYuHEdmjRphkaNmmDHjq0ID69TZFxgYDAePEjEmTOnAeh2N1+9egUAcPduLJo2bY4xY8YhLS0N\nGRnpcHJyQlpaWrnrSktLhaurK+zs7HD9+jWcPXtG/56trS20Wi0AoHnzR7F37+948OABACA5ORmx\nsbHl/r7F4UyXiIhMonHjpli79is0aNAQ9vYOsLe3L3I8F9DNit99dxGWL38fy5cvQnZ2Dp5//gXU\nqhWId96ZnRuwEvr3HwhnZxe0adMOs2a9gb/+2o+JE183OLGppJOc8t5r1ao1duzYisGDn0dgYBAa\nNGioH9OzZ2+89NJARETUxezZ8/Dyy2MwefJrEEUJtra2mDx5Gvz9/U32Z8Sl/R5CSctNWXsfSugB\nYB+WRAk9AMroQwk9AFzaj4iIyOIwdImIiMyEoUtERGQmDF0iIiIzYegSERGZCUOXiIjITBi6RERk\ndt99txFZWVlyl2F2DF0iIjK7zZs3Iisrs9j3RFE0czXmw9AlIqIK27BhDbZu1V0n+YMPlmLCBN2S\neseO/YN582YbjN2yZRPi4+MwbtxovPTSSwCAzp3bYeXK5Rg2bBDOnDmF/v17Ijk5CQBw4cJ5jBs3\nCgCQmZmJhQvfwciRL2H48ME4eHC/uVo0CV4GkohIgbyaNyj29YRjZ4p9vazjCyvL0n79+g3Et99u\nxIcfforQ0BqIi0tBZmYGGjRoiLFjJ+aOMry8Y94lHb/5ZjWaN38UM2a8hdTUVIwcOQQtWz4Ke3sH\no+qUG0OXiIgqrCxL++lIuTcdtVqN9u07Fnq/qH/+OYpDhw5g48Y1AHSLJdy9G4vAwGCT9VKZGLpE\nRApk7Ay1vOMLK9vSfkXZ2dkbLF6gVqshirrgzc7OP+FKkiS8++5i1KoVWKF65cJjukREZBLGLu0H\nAE5OzgbL9RVeeycgoDqios4DAPbt+0P/+qOPPoYtWzbpn0dHR5myhUpn1Ey3Y8eOcHFxgUqlgo2N\nDbZs2VLZdRERkZUxdmk/AOjZsxemTh2PgAB/LFmyssgSfUOHjsR7770DFxcXNG3avMDrL+ODD5bi\npZcGAgD8/QOwaNH/Kq8pEzNqab9OnTph27ZtcHd3N2qjFy9ehKdnQIWLk5OSlpuy9j6U0APAPiyJ\nEnoAlNGHEnoATLy0nyRJZfrc1IABA5CTk2P0eCIioqrAqNAVBAEjRoxA37598d1335U6/sSJE/jw\nQ+uZ7hMREZmDUcd0N23aBF9fXyQkJGDYsGEICQlBixYtHjq+Ro0aWLp0Ebp164769R8xWbFERETW\nzKhjugWtXLkSzs7OGDZs2EPH/PDDD3j22WfRvHlzHDlyBDY2/GQSERFRqWmYkZEBURTh7OyM9PR0\nHDx4EGPHji3xa7p3747nn38B3323EXPnvosJE6aYrGBzUdLBfWvvQwk9AOzDkiihB0AZfSihB8D4\nE6lKDd34+HiMHTsWgiBAq9WiR48eaNu2bakbfvfd97Bv3594//2F6NatOyIi6hpVEBERkVKVeiJV\nrVq1sHPnTuzYsQO7d+/GK6+8YtSGPTw88f77y5GdnY0JE8ZAo9FUuFgiIrJMsbF3MGTIAJNuMzr6\nIg4f/kv//ODB/Vi//huTbFuupQUr9YpU3bo9g759n8d//x3DqlUfVea3IiIimRW+wEVFXbp0EUeO\n5Idu27btEBn5kkm2XdLSgpWp0s9wmj9/Efbv34tFi95F165PP/SSYEREZN00Gg3eeWc2Ll68gNq1\nQzFr1tuwt7c3GHPr1k0sW7YYSUmJcHBwwHvvLYCLiw/++OM3fP3151Cr1XB2dsHy5R/jiy9WITs7\nG6dPn8TgwcOQlZWJCxfOYdKkaViw4G3Y2dkjOjoKiYkPMGPGW9iz53ucPXsa9es3wMyZcwAAS5a8\nh6ioc8jKykKHDp0wfPgrBksLenh4YMWKT/D330fw5ZefIScnBzVq1MTMmXPg4GD6lYsqPXS9vLyx\nePH/MGxYJCZMeBW7d/8MtVpd2d+WiKjKmjvXHrt3m/af9x49NJg7t+TdsTEx1zFjxhw0aNAQCxe+\ng+3bN2PgwMEGYxYvXoBp02aiRo2aOHfuDObOnYslS1bim2++wLJlH8HHxwdpaamwsbHByy+PRlTU\neUyc+DoAYM+e7w1m06mpKfj0069w8OA+vPHGJKxa9RVq1w7BiBEv4tKlaISFhWPUqNfg6uoKURQx\nYcIYXLlyyWBpQTc3NyQlJWLNmi+xYsXHsLd3wPr132DTpnUYOvRlk/4ZAmZaZah79x7o1asPduzY\nhs8//wSjR5d89jMREVkfPz9/NGjQEADQtesz2LLlW4PQzcjIwJkzJzF79hsFFjjQ3Tds2Bjz589B\nx46d0b79k0Z9vzZtngAAhISEwcvLG7VrhwAAatcOQWzsbYSFheP333/Grl07oNVqkZBwH1evXkVI\nSBgKLi149uwZXLt2BWPGjIAkSdBoNGjQoFHF/0CKYbYP0C5YsAQHD+7HggXvoEuXbrlNExGRqc2d\nm1XqrLQyFD6mW/gQrySJcHV1w5dfrte/lveRoalTZ+D8+bM4dOggRox4EatXryv1+9nZ2QEAVCqV\n/nHec61Wizt3bmPTpvVYvXotnJ1dsGDB2wbLBObXJaFly8cwZ867ZWm3XMy2tJ+Pjw/ee28pMjMz\nMWHCa2W6ljMREVm+2Ng7OHtWty7vr7/+jEaNmhi87+TkjICA6vjzz9/0r124cAGA7lhvvXqPYMSI\nUfDw8MS9e3fh5ORksPxfSYq7zlNaWhocHR3h5OSMhIT7OHLkkEEtedt+5JGGOH36JG7dugkAyMrK\nxI0bMWXo3HhmvVRUz5690aPHduzevQOrV3+KkSPHmPPbExFRJQoKCsa2bd9h4cK3ERwcgl69+hUZ\nM2fOu3j//YX45psvodVq0LNnD/Tv/yI+/ngFbt68AQBo3rwlwsLCUa2aH9at+xrDh0di8OCHXwUR\nKP7M6bCwcISHRyAysh+qVfNDo0aN9e/lLS3o4+OLFSs+wcyZczB37kxkZ+dAEASMHDkGtWoFVvBP\npJg6y3oZSGM97AojcXFxeOKJlsjMzMSffx7S74O3NEq6Soq196GEHgD2YUmU0AOgjD6U0ANg4qX9\nTMnX1xcLFy5Beno6Jk0ay93MRERUZZg9dAGgV6++ePrpZ3Ho0EF8/fVqOUogIiIyO1lCVxAELF78\nP3h4eOCdd97C9evX5CiDiIjIrGQJXQDw8/PD/PmLkZ6ehsmTxxV75hkREZGSyBa6ANCv3wB06dIN\nBw7sw5o1X8lZChERUaWTNXQFQcCSJSvg7u6BuXNnVdrnooiIiCyBrKELAP7+AZg3byHS0lK5m5mI\nyEoZu7Tfnj3f4/79eDNUZJlkD10AGDBgEDp16ox9+/7Ehg1r5S6HiIjKwZil/X78cTfi4uKKfa8q\nfITUIkJXEAQsXfoBXF3d8NZbM3H79i25SyIiojLKW9pv8OD+mD17epFF4vfu/R0XLpzHvHmzMXx4\nJLKystCxY0d88smHGDHiRfz5528YN24UoqJ0l4ZMSkpE//49AegC+eOPV2DkyJcwdOgg7Nq13ez9\nmYJFhC4AVK9eA++8swApKcmYMmU8dzMTEVVA8+bOxd5MNb44MTHX0afP81i3bjOcnJywfftmg/c7\ndOiEevXqY86cd/Hll+v1a+26u3tg9eq16NSpSzFb1c2ev/9+J1xcXPH559/g88+/wa5d2xEbe6dM\n9VkCiwldABg06EV06NARv//+K779doPc5RARURkUXtrv1KmTRcZIkoTCc6pOnTqXuu2//z6Cn376\nAcOGDcIrr7yE5OQkqzz51qwLHpRGEAQsW/Yh2rV7DLNnz0CHDh3h7x8gd1lERFbn2DHjVucp7/ji\nlLa038M4OjrqH6vVakiS7thudnZ2gVESJk16HS1bPlbRMmVlUTNdAKhZsxbmzJmHpKRETJ06gbuZ\niYisRGlL+wGAs7Mz0tJSH7qNgIAauHDhHAAYLAH46KOPY9u2LdBoNACAGzdikJWVacryzcLiQhcA\nhgwZhieeaI9ffvkJW7Z8K3c5RERkhLyl/QYP7o+UlORil/Z7+ulnsWTJQv2JVIVnxy+8EInt27di\n+PDBSE5O1r/eo0cvBAfXxogRgzFkyAAsWbIQWq220nsyNbMv7WesmJjraNfuMdjZ2eLAgX/g5+dn\nosqMo6Tlpqy9DyX0ALAPS6KEHgBl9KGEHgALXtrPWIGBQZg9+20kJiZi2rRJ3M1MRERWz2JDFwCG\nDXsZrVu3xZ4932PHjq1yl0NERFQhFh26KpUK//vfSjg5OWHGjKm4d++e3CURERGVm0WHLgDUrh2C\nN9+cg4SEBMyYMVXucoiIiMrN4kMXAEaMGIVWrR7H7t07rPbSX0RERFYRuiqVCitWfAQHBwdMnz4F\n8fFVd4UKIiKyXlYRugAQEhKGGTPeQnx8PGbO5G5mIiKyPlYTugDwyitj0KLFo9ixYxt++GG33OUQ\nERGViVWFrlqtxooVH8Pe3h7Tpk1CQsJ9uUsiIiIymlWFLgCEh9fBG2/MQlzcPbz55htyl0NERGQ0\nqwtdABgzZiyaNWuOrVu/w08//Sh3OUREREaxytDV7Wb+BHZ2dnj99YlITHwgd0lERESlssrQBYCI\niLp4/fUZuHs3FrNnz5C7HCIiolJZbegCwGuvTUDjxk3x7bcb8OuvP8ldDhERUYmsOnRtbGzwwQef\nwNbWFlOnTkRSUqLcJRERET2UVYcuANSrVx+TJ0/DnTu3MWfOm3KXQ0RE9FBWH7oAMH78ZDRo0Agb\nNqzFH3/8Jnc5RERExVJE6Nra2uKDDz6BjY0NJk8eh5SUZLlLIiIiKkIRoQsADRo0xMSJU3H79i3M\nnTtb7nKIiIiKUEzoAsDEiVNRv34DrF37Ffbt+1PucoiIiAwYHbqiKKJ3794YPXp0ZdZTIXZ2dvjg\ng4+hVqsxefI4pKamyF0SERGRntGhu2bNGoSGhlZmLSbRqFETjB8/CTduxGDevDlyl0NERKRnVOjG\nxsZi37596N+/f2XXYxKTJ7+BunXr4auvvsDBg/vlLoeIiAiAkaG7YMECTJs2DYIgVHY9JmFvb48V\nKz6GSqXCxIljkZaWJndJREREsCltwN69e+Hj44N69erh6NGjRm/Y19e1QoVVVJcuHTBt2jS89957\nWLZsAT744IMyb0PuHkxFCX0ooQeAfVgSJfQAKKMPJfRgLEGSJKmkAcuWLcOuXbugVquRlZWFtLQ0\ndO7cGYsXLy5xw3Fx8p/ElJmZiaeeegIXL0Zh5849ePzxNkZ/ra+vq0X0UFFK6EMJPQDsw5IooQdA\nGX0ooQfA+P84lLp7efLkydi7dy9+//13LFu2DK1atSo1cC2Fg4MDli//CCqVChMmvIr09HS5SyIi\noipMUZ/TLU6LFo9i9OixuHbtKhYunCd3OUREVIWVKXQfffRRrFq1qrJqqTRvvPEmQkPD8NlnH+Po\n0SNyl0NERFWU4me6AODo6Ijlyz8GAEyc+CoyMjJkroiIiKqiKhG6ANCq1WN45ZUxuHz5EhYtmi93\nOUREVAVVmdAFgBkz3kJwcG2sWrUS//77t9zlEBFRFVOlQtfJyQkrVnwMURQxYcKryMzMlLskIiKq\nQqpU6ALA44+3wcsvj0J09EUsWfKe3OUQEVEVUuVCFwDefHMuAgODsXLlchw/fkzucoiIqIqokqHr\n7OyM5ctX6nczZ2VlyV0SERFVAVUydAGgbdt2GDp0BC5cOI///c86rrBFRETWrcqGLgC89dY7qFUr\nECtWLMOpUyfkLoeIiBSuSoeui4srli37EFqtFuPHv4rs7Gy5SyIiIgWr0qELAO3bP4kXXxyGc+fO\nYPnyJXKXQ0REClblQxcA5s6dhxo1amL58iU4c+a03OUQEZFCMXQBuLq6YenSD6DRaDB+/Bjk5OTI\nXRIRESkQQzdXx45PYdCgF3HmzCl8+OH/5C6HiIgUiKFbwNtvz4e/fwCWLl2E06e5m5mIiEyLoVuA\nu7sHli5dgZycHAwdOhSpqalyl0RERArC0C2kc+duiIwcgv/++w8DBvRGcnKS3CUREZFCMHSL8f77\nyzFo0CD8889R9O3bEwkJ9+UuiYiIFIChWwwbGxusWbMGgwa9iJMnj6N372cRFxcnd1lERGTlGLoP\noVarsWzZhxg+fCTOnz+LXr2exp07t+Uui4iIrBhDtwQqlQoLFy7Bq6+OR3T0RfTs2Q03bsTIXRYR\nEVkphm4pBEHAnDnzMGXKG7h+/Rp69uyGK1cuy10WERFZIYauEQRBwBtvvIlZs+bi1q2beO65pxEV\ndUHusoiIyMowdMtg/PjJmD9/Ee7ejUWvXk/j9OlTcpdERERWhKFbRiNHjsGSJSuQkJCAPn2exfHj\nx+QuiYiIrARDtxyGDBmGDz9chZSUZPTt2xNHjhyWuyQiIrICDN1yev75F/DZZ18hMzMDAwf2xoED\n++QuiYiILBxDtwJ69uyNr75aD41Gg0GD+uG3336WuyQiIrJgDN0K6tr1aaxb9x1UKhVeemkQfvhh\nt9wlERGRhWLomkCHDh2xceNW2NnZ4+WXh2Dbts1yl0RERBaIoWsirVu3xebNO+Ds7IIxY17Ghg1r\n5S6JiIgsDEPXhFq0eBTbtu2Gp6cnJk58DatXfyZ3SUREZEEYuibWqFETbN/+I3x9q2HGjKn4+OMP\n5S6JiIgsBEO3EtSrVx87d+5BQEB1zJ37JpYuXQRJkuQui4iIZMbQrSRhYeHYuXMPAgODsGjRfCxY\n8A6Dl4ioimPoVqLg4NrYuXMPQkJCsWLFUsyePZ3BS0RUhTF0K1mNGjWxc+dPqFu3Hj777BNMnToR\noijKXRYREcmAoWsGfn5+2L79RzRo0Ahr136FceNGQ6PRyF0WERGZGUPXTLy9vbFt2240b94Cmzdv\nwujRI5CTkyN3WUREZEYMXTPy8PDE5s078fjjbbBr13YMHz4YmZmZcpdFRERmwtA1MxcXV2zcuBXt\n2z+Jn3/egyFDBiI9PV3usoiIyAwYujJwcnLC2rXfokuXbti79w8MGtQPqakpcpdFRESVrNTQzc7O\nRv/+/dGrVy/06NEDK1euNEddiufg4IAvv1yHHj164dChg+jfvxeSkhLlLouIiCqRTWkD7OzssGbN\nGjg6OkKr1eKFF15Au3bt0KhRI3PUp2h2dnb49NMvYW9vjy1bvkWfPj3w3Xc74O3tLXdpRERUCYza\nvezo6AhAN+vlR11My8bGBitXfooXXxyK06dPok+f7rh7967cZRERUSUodaYLAKIook+fPoiJiUFk\nZGTps9zgYHiJRa+8lHDsTLHDvZo3KPZ1WcerhCI9VGY9XwFwGDkan3++Cr16PY2tW3ejevUaFd9+\ngT6s6s+/oNweLKaeco5HzHWLqofjOd4SxisiL4CH/v0uzKjQValU2LFjB1JTU/Hqq6/i0qVLCAsL\nK/Fr1CqhyGu+vq4P+QZFx1rC+MI9VHY9n376Mby83LFo0SL07v0M/vjjDwQHB1d4+3l9yP3nWZHx\napVgUfWUZ/xDv8ZK6i843uBrLaCe8ozXP7eQeso7vrh/a+Wsp8zjoYy8MJYglfFiwCtXroSzszOG\nDRtW4ri4OOs+G9fX11WWHiRJwtKli7B48QJUr14D27btRkhIyf/BKYlcfZiSEnoA2IclUUIPgDL6\nsPgeRBHIzISQmQEh9x4ZmRCyMiFkZgKZGRAyMuE+dJBRmyt1ppuQkABbW1u4uroiMzMThw8fxiuv\nvFLhPqh4giBg6tTpcHBwxDvvzEbPnk9jy5ZdqFu3ntylERHJy8gAzHtfN7bg89z3szLzt5M7Hpm5\n28nIHZf3fna2cbWZKnTj4uIwffp0iKIIURTxzDPPoH379sYVQeU2duwEODo6YMaM19G79zP47rsd\naNiwsdxlEREVJUlAdjaE9DQI6em5tzT9PdLTIaQ95D1JA9fEFMMAzMrSPy9XAJa1fEEAHB0hOThA\ncnCE5OICyccXkoM9JAdHIO91BwdIjo6Avb3hcwcHuBj5vUoN3YiICGzfvr2CLVF5jBgxCg4Ojpg8\neRz69OmBTZu2onnzlnKXRUTWSJKAjIwioWd4nw4UfC2taEgWHZf7ulZb7tIcCpZZXAB6+0BydCga\ngA4ORQPRwQGSfe57jo4FxjoCDvaGz/O2aWsLCGU7NluYyUKX5BUZOQQODg4YO3YU+vV7Dhs2bMbj\nj7eRuywiqkyiqAuylBQIqakQUpINH6emQJWaCkg5cI5/UCQQ9Y/T8h8jIx2CCdbzllQqSE7OkJyc\nACcniN4+kJyc9K9JTk6QnAs8dnIGDN43fM+rpi/i00VdANo7AHZ2FQ5AS8bQtQJ9+z4POzt7jB49\nHAMH9sGaNZvQvv2TcpdFRAVJku64YEpKbiimFB+aqbrHKv17KbrX8h6npEBISzU6IJ2KK8XWVh9u\nors7pIDqucH38PArGJYlhSTs7U0bir6ukCz5RCoTY+haiR49noODw3oMH/4iBg9+HqtXr0GXLk/L\nXRaR9cvJyZ095oeeKi0lPwALhmZaam5gFgzRlPyvL+fFgyQ7O0iurpBcXCEGBUN0dc197gLJxS3/\nsasrJFfc0+i4AAAgAElEQVQ3iC4ukFxc4FHTDwlZAJxzw9HRUReMtram/TMik2HoWpHOnbth/frN\nGDJkIIYOjcSnn36JHj16yV0WkbxEEUJyEoTERKiSEiEkJkJISoQqMbH415ISgdRkeCcl6YKynMtr\nSmo1JBddOIoB1SE560JRdHXLD0gXV/2Y/OB0g+iS/1hycdHNHsvD1xXaKjRLVAKGrpVp164DNm3a\nhkGD+mPkyKH48MNV6N9/oNxlEVWMkcGpSnxQJECF5KQyHauUnJwAd3eInl6QagUWmUnqQzMvLAuG\npqsrRGfdPRwdFX3skSoHQ9cKPfZYa2zZshMDBvTB2LGjkJWVhcGDX5K7LKrqSgzOB/qQzAvSigan\n6O4BsXp1iPXqQ/LwgOTuAbHQveThAdHDE5KHJ0R3D0ju7oC9PXx9XfGAM0SSAUPXSjVr1gLbtn2P\n559/DpMnj0NmZgZefnm03GWRUmRkQBUfB9X9eKjux0OIj4fq/n2o7scDmalwi40zTXB6eEKsXgNi\n/UfyQ1Iflh6FXvPUvwc7u0psnqjyMHStWMOGjbBjxx707dsDM2dOQ0ZGJsaNmyh3WWSJ0tL0AaoP\n0fgCz+/H54bsfaji43UXLShB3hFIyckZoocHg5PISAxdKxcRURe7du1B3749MW/eW8jISMfrr8+A\nwGNNyiVJhiEaHwchNyzzn+cFqm52KqSnl75Ze3uI3j7QhIVD8vaG6O2ju/n4QPLxzX3uDc/QWojX\n2up21TI4icqEoasAISFh2LlTN+NdsuQ9ZGZmYvbstxm81kKSdB9FiYszDMr4eMNdvPfv658bc8at\n5OCgC9HwiEIh6gvJx0cfonnPJWcX404MqmKfqyQyJYauQgQGBmHXrp/Qt28PrFy5HBkZ6Zg/f7Hc\nZVVdkqQ7eSg2Fqo7t6G6GwukJcL5+q1Cx0lzH2dllb5JR0eIPr7Q1K2nuwpQboDqZ6O5AZoXrnB2\n5tm1RBaGoasgAQHVsWPHHvTv/xxWr/4MWVlZ+Prr1XKXpTzp6VDF3oH6bm6g6oP1DtR37kAVeweq\nu7HFzkYLXj1IcnKG6OMDTf1HdCFaIDAfGqJEZNUYugpTrVo1bN/+PQYM6IN1677BvXt3MG/eYtSu\nHSJ3aZZPo4Hq3l1daBYIT/Wd27rHsXd0AZuU+NBNSCoVRN9qutmof4D+pg2oDrewIDywdc4PUafi\nLuBHRErG0FUgLy9vbN26C6+8Mgy//PIL9u/fjwkTpmDs2ImwL++Vb6yZJEF4kKAL0oKz0dhYqGIL\nzFTj7pX4kRfRwwNiQAA0TZvlBmkARL8AiAHVIfr76+59fAGbh/y18nWFhsdCiao0hq5Cubm5Y+PG\nrfjzzz2YMGEiFi2ajy1bvsWiRcvQrl0HucsznbQ0qO8WmJnmBqsqNm+GGgvV3TslHjOVHBwg+gcg\np9XjEIsJUq2fP0T/AN0ViIiIKoChq2CCIGDAgAFo0aINFi2aj9WrP0O/fj3Rp08/vP32Qvj5+cld\n4sPlnoikvhEDJMXB4eKV/BlqgWBVJSc9fBMqFUQ/f90xU78AXaDm7uoV/fz1wSq5e/CEIyIyC4Zu\nFeDm5o758xdjwIBBmDZtErZt24Jff/0FM2fOxtChL0OtVstTWFoa1DdioI65BlXMdaivX4c6RndT\nxVyHKiVZP9S10JeKnp4Qa9SEpnkLaP0Dip2hij6+gFy9EREVg6FbhTRq1AQ//PAb1q79GvPnv40Z\nM17Hpk0b8P77/0OTJs1M/w2zs6G6dVMfpLowvaZ7fP06VPFxxX6Z5OQEbWAQcgJbQxsYBKd6dZDs\n6gWtf26g+gcADg6mr5eIqJIxdKsYtVqNoUNH4JlneuDtt2dh8+ZN6Nr1SQwdOgIzZ74Fd3cP4zcm\nirqPzsRch+r6NYNZqjrmOlR3bkMQxSJfJtnaQluzFjSPNIA2MAjawCCIuffawGBIPj4Gu3udfF2R\nxROQiEgBGLpVVLVq1fDRR59h0KAXMW3aJHz11Rf4/vtdePvt+ejb93nd1awkCUJCAtS5s1OVfvdv\n7u7gmzcgZGcX2bYkCBADqiPn0ccKhGkQxKBg3b1/AHf7ElGVxNCt4to2boIDH32OXz/7GCd3bkPW\nqyNxcdZ0NPX0hGNsLFRpqcV+nejjkztTDS4UrEHQ1qhV/kW5iYgUjKGrdFlZUF+OLjBLzdv9mzt7\nTUgAAAzOvQEAEu4jOeE+Yn184dGmLVA7JDdYdTNVba1AwMVFro6IiKwWQ1cJJAlCfDxsoqOgvhgF\ndXQUbC5GQX0pGrh9C17FXPBBsreHtlYgNI2b5odpkC5Qf4m+iCnz38btO7cReOEC3hs6Ak891VWG\nxoiIlIWha01EEapbN2Fz8QLUFy/mh2t0FFQPHhQZrq1eA2jfHhkBNQ1OVBKDgiBW8wNUqmK/Taem\nzXHwmR5YunQRPv30Iwwa1B/du/fEu+++hxo1alZ2l0REisXQtUQ5OVBfvQL1xagCs9eLsLl0sci6\nqJJKBW1wbeS0ehza8Aho6kRAWycC2vA6kFxc4evritRynPnr4uKCOXPm4fnnX8C0aZPwww+78Oef\nv2PatJkYOXI0bG1tTdUtEVGVwdCVU1oabC5dzA/V3Fmr+uoVCBqNwVDJwQHa0HBo6tTJD9fwCGhD\nQiv1pKV69epj5849+PbbDXj77VmYO/dNfPvtBixe/D+0avVYpX1fIiIlYuiagXD/ftHjrdEXob55\no8hY0d0DmibN8kO1Th1owiMg1gqU7WM2KpUKL7wwGF27Po13352Ldeu+QY8eXRAZOQSzZ78NLy9v\nWeoiIrI2DF1TkSSobt8qsEv4ItQXL8AmOgqq+/eLDNf6+SP7iQ76UNXWiYAmPAJStWoWex1gLy9v\nLFv2IQYMiMS0aZOwfv0a7NnzPd56ax4GDoyE6iHHiImISIehW1YaDdTXrhaatUZBHR1d5DOtkkoF\nMTAIWc1bFtglXAfaOhGQ3NxlaqDiWrV6DL/9th9ffPEpFi2aj4kTX8OGDWuxePH/UL/+I3KXR0Rk\nsRi6D5OeDpvTJwuEq+5sYfWVyxBycgyGSnZ20IaGI7tAqGrCI6ANDVPsNYJtbW0xZsxYPPdcb8ya\nNR3ff78TnTq1xahRr2Hq1Olw4ed4iYiKYOgCEFKSYXPqJGxOnoDNqeOwOXkCuHIZnoU+3yq6uELT\nsBG0deoW2CVcB2JQcJW9rGH16jXw5Zdr8dtvP2P69Nfx8ccfYMeOrZg/fzGeeeZZ3eUkiYgIQBUM\nXSE5CTanT+kC9uR/uvsrlw3GiG7uQLt2yKgdVuCEpgjdNYMZIsV66qmuOHCgHVasWIIPP1yOYcMi\n0blzVyxY8D6CgoLlLo+IyCIoOnSF5KQiM9giAevugewn2kPTqAk0TZoip1ETiMG14VvNrVyfb63K\nHB0dMX36bPTtOwBvvDEZv/76Mw4e3I9Jk17Hq6+Oh52dndwlEhHJSjGhW6aAbdwUmsZN9AHL2atp\nhYfXwdatu7F163eYM+dNLFjwDjZv3oRFi5ahbdt2cpdHRCQbqwzdIgF74jhsrl4xGKML2A7QNG7C\ngJWBIAjo128AOnfuioUL5+Grr75Anz7Pol+/AZg7dz6qVasmd4lERGZn8aGrD9gTx/NnsKUFbOOm\nupObGLCyc3f3wHvvLcWAAYMwbdpkbNnyLX755Se8+eYcDBkyDOoqegIaEVVNFhW6QlJi0V3EJQRs\nTpOm0DRqwoC1Ak2bNsdPP/2Br79ejQUL3sEbb0zGpk3r8P77y9GoURO5yyMiMgvZQteogPXwQHa7\nJ3Nnr00YsFZOrVZjxIhX8OyzPTFnzkxs27YFXbp0wPDhIzF9+iy4WfEFQ4iIjGGW0DUI2JPHYXvy\nONTXrhqMYcBWHX5+/li16ku88MKLmD59Cr744lPs2rUD8+YtRK9effnZXiJSrMoJ3T/+gOPev2Bz\n6oRxAdu4KcTAIAZsFdO+/ZPYu/cwVq5cjuXLl2DUqOFYv34tFi1agtDQcLnLIyIyucoJ3U6dkHcR\nQIOAzTsGy4ClXPb29pgy5Q306dMfM2ZMxR9//Ib27R/HuHGTMGHCFDgo9DKaRFQ1VU7oTp+OpPD6\nDFgyWu3aIdi4cSu+/34XZs16A0uXLsLWrd/lnvncW+7yiIhMotS12GJjYzFkyBA888wz6NGjB9as\nWVP6VhcuRHaPXjwmS2UiCAJ69HgOf/31D0aNeg03bsRg4MA+6NevH44d+wdSoWthExFZm1JDV61W\nY8aMGfjxxx+xadMmrF+/HpcvXy7ty4jKzcXFFfPmLcSvv+5HixaPYuvWrXj66U7o0OFxfPbZx0hI\nKLo+MRGRNSg1dH19fVGvXj0AgLOzM0JDQ3Hv3r1KL4yoQYOG+P77X/DTTz+hZ8/euHQpGrNmTUej\nRhEYNWoY9u/fC1EU5S6TiMhoZTqme/PmTVy4cAGNGjWqrHqIDKhUKnTt2hXNmrVGfHw8Nm/ehHXr\nvsb27VuxfftWBAUFIzJyCAYOjIS/f4Dc5RIRlUiQjDxQlpaWhhdffBGvvvoqnnrqqRLHBgej2BnI\nsWNpxY5v3ty52NflHK9SqYr0YE315ynYhyXUU57xeT3kjZckCX//fRTr13+DXbu2Iz39LADAwcER\nLi4ucHBwgCAIFlN/npgYFeKKWbnK0v/8C4/39XU16EPuesozvmAPllBPecf7+roiMLD4vT3WUD8A\ntGzpavV5Aej+fhvDqJmuRqPB+PHj8dxzz5UauHlUqqIF+Pq6PmRs8duQe3zhHuSup7zj8/qwlHrK\nM16lUhmMf/bZznj22c5ISkpCSIgKqakpyMzMQGZmBtRqNVxcXJCUdB9hYWEWUX9JX2MNf/6Fxxd8\nbAn1lGd83nNLqaf844v/AmupX/c11p8XxjJqpjtt2jR4enpixowZRm+4uP/RW5PC/5u3Vkrow9ge\nTp8+hQ0b1mDLlu+QlJQIAGjbth0iI4ege/eesn/mVwk/C0AZfSihB0AZfSihB6Dk/1QUVGpGHzt2\nDLt378aRI0fQq1cv9O7dG/v3769wgUSm1rBhIyxcuASnTkXh448/R5s2T+Dgwf0YM+ZlNGpUBzNn\nvo6zZ8/IXSYRVWFGH9MtK2v/n4uS/vdl7X1UpIcrVy5hw4Z12LhxHeLidGfdN2vWHJGRL6F3775w\ncTHuf6emoISfBaCMPpTQA6CMPpTQA2DCmS6RNQsJCcOsWXNx4sR5fPPNRnTp0g0nThzHlCnj0aBB\nHUyc+Br++ecoL7xBRGbB0KUqwdbWFk8/3R3r1n2H48fPYcaM2fDx8cWGDWvRvXtntGvXCqtWrcT9\n+7zwBhFVHoYuVTkBAdUxadLr+PvvE9i8eSd69eqDq1ev4K23ZqJRozoYOXIo9u79gxfeICKTk20R\neyK5qVQqtG//JNq3fxL379/Hli2bsH79GuzcuQ07d25DYGAQXnhhMF54YTCqV68hd7lEVMmys4H0\ndCA9XShwr3uclpb/WkZG0TEbNxr3PXgi1UMo6eC+tfdhzh4kScKxY/9g/fo12L59K9LT06BSqdCx\n41OIjHwJXbp0g62tbbm2rYSfBaCMPpTQA6CMPsrSgyiimMDLv8/IKDksC79XOEA1mvIv0GNsknKm\nS1SAIAho0eJRtGjxKObNW4gdO7Zh/fpv8Ntvv+C3336Br281DBgwCIMHD0FISNELbxCRjiQBaWlA\naqqAlBQBKSmGj9PSdI+1WiA+3v6hQVg4LE1BpZLg5AQ4OenuvbzEAs91rzk7S3B0zB9jeF/0NehX\nkS8ZZ7oPoYT/QQLK6MMSejh37iw2bFiDzZs34cGDBwCA1q3bIjJyCJ599jk4OjqWug1L6MMUlNCH\nEnoATN+HJAFZWXnhmB+SqanIDUvd4/zwzH+vuK+RpPKHpL19ySFX8N7R0XCMs/PDw9LREbC3N/2q\ns8Z+ZIih+xD8S2k5LKmHzMxM7NnzPdatW4MDB/YCANzc3NGv3/OIjHwJDRs+fDEQS+qjIpTQhxJ6\nAPL70GhQKAx1j4ubZRYOybzHea/n5JQvjezsJLi6SnBxAVxcdI9dXQFXVwnOzvmPdWN0z11cJNSs\n6YTs7LQiM0y12sR/WJWMoVtBSvtLac0stYdr165i48a12LhxPWJj7wAAGjduisjIIejTpx/c3NwN\nxltqH2WlhD4srQdJ0h2rTEwUkJgoICkp7x548CD/eeH309JUSE6Wyr3bVRAMw9DZuWAwokBA5j/P\nC1NdkOaHp719+Xq3tJ9FeTF0K0hJvwjW3oel96DRaPD7779i/fpv8OuvP0Or1cLR0RE9e/ZGZORL\naNXqMQiCYPF9GEsJfVRWD5mZKBSQMAjJwqFZ8P3sbOOD09ZWgru7BC8vFRwdtUVmjwXDMO/1ggGa\n956Tk+l3s5aVEn6fAIZuhSnpF8Ha+7CmHmJj7+Dbbzdg/fo1uHbtKgAgLCwckZEvYeTIobCzc5O5\nwoqzpp/Hw5TUQ04ODELz4YFZ9P3MTOMTTK2W4OEhwd0d8PCQ9Dd3d6nQ86Lv54Wl0n8W1oShW0FK\n+kWw9j6ssQdRFHHo0EGsW/cNfvhhF7KysgAAoaFhaN26rf4WEFBd5krLzlp+Hnlnz8bHC7h/X0B8\nvID4eBXu3xeQnm6PO3dy9KFZcBduWXbVCoIuFN3dJXh65gem4fOi73t46HbXVnSWaS0/i5IooQeA\noVthSvpFsPY+rL2HBw8SsG3bZhw48Cf27z+A1NT8XkJCQvUB3KbNE1YRwnL+PDIzDUM0Li7vsapQ\nuOoeZ2QYl2pubg+bZepC82GzUFfXsq+nakrW/ncDUEYPAEO3wpT0i2DtfSihB0DXx507D3DmzCn8\n9ddBHDp0AEeOHEZKSrJ+TO3aIWjT5gk8/ngbtGnzhEVeCcuUP4+cHCAhQReexYWmLlhV+sepqaWH\nqL29BB8fw5u3twQfH1H/PDTUCZKUCg8PCW5ugI2VXrFACX83lNADwNCtMCX9Ilh7H0roASi+D61W\naxDChw8fMgjh4ODaBiFco0ZNc5ddREk/D61Wd7Zt4QDNn5EWfE+FxMTSQ9TGJi808wPU17domOa9\n7uxc+m5bJf9OWRsl9AAwdCtMSb8I1t6HEnoAjOtDq9Xi7NnTBiGcnJykfz8oKNgghGvWrFXZZUOj\nAe7dE3DnjoDYWBUyMx1x7VpWkRlpfLyAhAQBolhy4qlUEry8Cs9Ciz7OC1N398q5kEFV+Z2ydEro\nAWDoVpiSfhGsvQ8l9ACUrw+tVotz587gr78O4NChgzh8+BCSkhL17wcGBqNNm/wTs2rVCizT9tPS\ngNhYAbdvq/SheueOgNu38x/fu1d6kHp46EKy5Bmp7ubpKcl+4YOq/DtlaZTQA8DQrTAl/SJYex9K\n6AEwTR+6ED6LQ4cO4K+/DuLw4b8KhXAQWrdui8cea4v69dtDrQ7MDVEVYmMF3LmjC1LdTYXk5IeH\nqZ2dBH9/CQEBIgICJAQESPD3FxEW5gBb23T4+OhC1ctLQjnXgJANf6cshxJ6AIwPXSs9fYCoalKr\n1ahTpxHc3BqjcePx6NVLwokT93DqVDyuXMnErVu22LTJD5s21QBg99DtuLtLqFFDRPPmulD195dQ\nvXr+44AA3ey0uN26vr4OiIvTVl6TRArG0CWyEJIEJCejwK5e3Wy04K7e2FjdCUiGaufedMdLfXyy\n4eAQj5yca0hMPI2srCsAbgG4CT8/EW3ahKB9+1Zo3botAgODIMh9SSKiKoShS2QGWi1w6xZw5oyq\nwK7egrt7da+VdGEGJyfdDLRuXU3uzFTM3eWrm6FWr67b3as7XuoKoCFE8RGcP38Ohw8fxF9/peDw\n4YPYtu0Atm37BgBQo0ZN/WeEW7dui6CgYIYwUSXiMd2HUNJxBmvvw1p6SEkBrl9X4do1Fa5fF3D9\nukr//ObNkldv8fExPG4aEKAL1bxdvQEBItzcKn4WryiKuHDhfG4IH8Thwwdx//59/fs1atTUnxnd\nunVbBAfXLhLC1vLzKIkSegCU0YcSegB4IlWFKekXwdr7sJQetFrdmb7Fher16wLu3y/+0kQ+PiKC\ngiSEhqrh5ZVtcGJSQIAIP7/yr9BSUaIoIirqAg4dOoBDh/7CoUMHDEK4evUaBiFcu3YIqlVzs4if\nR0VYyu9URSmhDyX0ADB0K0xJvwjW3oc5e0hNhT5M84JVF6oq3LhR/EowtrYSAgMlBAWJ+ltwcP5z\nFxfz91FekiQhKuoC/vrrAA4f1oVwfHy8/n1//wA0bdoEgYEhqFMnAuHhEQgPrwNvb28Zqy47a/hZ\nGEMJfSihB4BnLxMVSxR1s9W8UL12LT9Ur18v7iQlHW9vEQ0aFAxV3ew1KEg3a5X7c6emIggC6tat\nh7p162HEiFcgSRIuXozSh/CRI4ewZ8+eIl/n7e2tD+D8WwRq1qwFlZwXJyayMAxdUpz0dBjMVAvu\nAo6JUSErq+hs1cZGQq1aEho00BQJ1aAg3fHUqkgQBERE1EVERF0MHz4SAGBjo8GRI/8hOvoiLl6M\nwqVLuvu//z6CI0cOGXy9o6MjQkPDUadOnQKhHIGQkFDYy7VPnUhGDF2yOpIE3L1reGy14Gz13r3i\nZ1aenhLq1Ss6Uw0K0p35a60XvTc3T09PtGjxKFq0eNTg9czMTFy9egXR0VEFwvgiLl+OxpkzpwzG\nqlQqBAUFo06dCISF1cndVa2bIbu7e5izHSKz4j8zZJEkSbcb+MIFFWJjgTNn7PWhGhOjKnbJNrVa\nQs2aEtq10+hDVXevu7m7y9BIFeLg4IB69eqjXr36Bq+LooibN2/khvFF/cw4OjoKP/+8Bz//bLi7\nulo1v9wwDjc4bhwQUJ0fZyKrx9AlWUmS7mL6Fy6oEBWlu124oMbFiyokJRX8B1Z3dSU3Nwnh4WKB\nMJX0M9caNThbtUQqlQqBgUEIDAxCp05dDN67f/8+oqOj9Luqo6OjcOlSNA4e3I+DB/cbjHV2dkF4\neDjCwyMMZsjBwbVha23XoaQqi/9EkdnExeWHa37Iqoss76ZWSwgJEfHEEyIiIkS0bGkPb+80BAWJ\n8OCeR0Xx9vaGt3drPPZYa4PX09PTcflydIEw1s2Qz507ixMnjhuMtbGxQe3aIQXCOFx/7+Ji3Bml\nRObC0CWTu39fKBSsulvhz7GqVBJq15bQurUGdevqAjYiQkRoqGjwuVVfX3vExYlm7oLk5OTkhIYN\nG6Nhw8YGr2s0GsTEXEN0dLTBSVzR0RcRHX0RP/6422B89eo1DM6mzpsh+/i4mLMdIj2GLpXbgwdA\nVJS60K5hVZGP3QiChKAgCS1b5hiEa1iYCAcHmYonq2RjY4OQkDCEhISha9en9a9LkoR79+4VOYkr\nOjoK+/b9iX37/jTYjouLC/z9A+DvHwA/P//cez+D1/z8/OHk5GTuFknhGLpUqqQk4MIFtUGwRkWp\nij1LODBQRJcuGkREaBERIaJuXV248t8uqkyCIMDPzw9+fn5o27adwXupqSkFPt6kmyHfvHkdt2/f\nxqVL0SVu193dA/7+/vDzC4C/v39uKPvnhnL+Y378iYzF0CW9lBTkBqraIFxjY4uGa61aIp56SpM7\na9Wibl0R4eEinJ1lKJyoBC4urmjatDmaNm2ufy3vKkhZWVm4d+8u7t6NRWxsLO7evYPY2FjExt5B\nbOyd3NfvICrqQonfw8vLq5hgDjAI6WrV/HjCFzF0q6LUVODixfwzhfPC9fbtouFao4aIjh01ubNW\n3ey1Tp38SxsSWTN7e3vUqhWIWrUCSxyXkZGBe/fuFgjm/HDOC+Zbt27i/PmzD92GIAjw9vbRB3HB\nXdsFw9nHxxc2PA1fsfiTVbD0dODff4HDh230s9eoKBVu3CgargEBIjp00Oh3CeftHnblyZ9EcHR0\nRFBQMIKCgkscl5aWhrt3Y/VBnB/M+Y+vXLlc5GIhBalUKvj6Vis0Yy46g7a2612TDkNXITQa3a7h\n48fVOH5chf/+081gRREAHPXjqlUT8cQT+WcL581eeeEIoopzdnZGSEgoQkJCSxyXmppisBu74K7t\n/F3a53Hy5PGHbsPGxgZeXl5wc3OHu7s73N09Ctx76J97eHjAza3ovVopFwy3MgxdKyRJQEyMgOPH\n1fjvP13InjqlNrhKk5OThJYttWjZ0gaBgZn62aunp4yFExEA3XHmsDBXhIWFP3SMJElITk4q9hjz\n3bt3ERt7B8nJiUhIeICYmOvIzs4uUw2urm7FhLW7QVjnv+ZpEOCOjo68Olg5MXStQEICcOJEXsDq\nQrbgx3JUKgl164po1kyLZs1ENG2qm73a2OSdMJIjY/VEVB6CIOhnrBERdYsdk3dCmCRJyMjIQHJy\nEhITE5GUlISkpAe597rniYmJ+vcL3sfEXEdKSnKZarOzs9PPmjnLLhuGroXJyABOn87bTawL2mvX\nDI/BBgaKeO65HDRtqgvZhg21PGuYqAoTBAFOTk5wcnKCv39Amb9eq9UiOTnJIKSTkhILBHhigZth\nkF+/fg05OWX7j72rq5s+gH18vGBraw9HR139jo6OcHJy1t87ORV87lRgXP69s7Pu3hrCnKErI60W\niI5W4b//VPpZ7PnzKmg0+bttPD0ldOyoyQ1YLZo0EeHrK8lYNREpjVqthqenFzw9vcr8tWWdZRcM\n7piY6zh79rTJ+rC3ty8S2oXDOu/2sJA3DHvDcLezs6vwbnWGrplIEnD7tqA/Bnv8uBonTqiRlpb/\nA7S3l9CkiW43cdOmulvt2hJ46ISILFVFZ9ne3s6IibmH9PR0ZGSkF3ufd8vIyEB6elqh+4LjdK+l\npaUjOTkZsbGxyMhIhyia5jKyarW6SFjnzcT3799r1DZKDd2ZM2di79698Pb2xu7du0sbTrmSkqDf\nRZx3NnHBKzgJgoSICBFNm4r6WWzduiLs7GQsmojIzFQqFZydneFcScfIJElCVlZWgSDXBXZ6enEB\nXtmry4cAAAsRSURBVFyQFwx9w/vExESkp6eVafd6qaHbp08fvPjii5g2bVqFGleyrCzg7FmVwdnE\nly4ZHluoXl1E9+45aNpUN5Nt3FjLz8ASEVUyQRDg4OAABweHcu0+N4ZJQ7dFixa4detWhQpSElEE\nLl/WHYfNm8meOaNCTk7+PmBXV91C6rrdxLqZrL8/j8MSESlRWS7vyWO6pbh7V3ccNu9kpxMn1EhJ\nyQ9YW1sJDRqI+mOwzZrplqZTFb3oExERVXEM3QIkCTh/XoUDB9Q4fhw4csS5yPWIw8K06NYt/2Sn\nRx4xXPuViIjoYSotdH19reOA5Y0bwG+/6W6//w7cvZv/np+fCj17Ao8+CrRqBbRoAXh4qAGoAVjP\naiHW8rMoiRJ6ANiHJVFCD4Ay+lBCD8YyKnQlqezHI+PiUsr8NeaQlAQcPGiD/fvV2L/fBpcv589k\nq1UT0a+fFu3aadCzpyMcHVMMPq6TkwPExclQdAXkXbHGmimhB4B9WBIl9AAoow8l9AAY/x+HUkN3\nypQpOHr0KBITE9GhQweMGzcOffv2rXCB5pKVBfzzj1ofsidOqCCKuiR1dpbQpYsG7dpp0K6d7tKJ\neSHr62t9AUtERJat1NBdunSpOeowGVHUfXxn3z5dyB49mr8QgI2NbhGAdu10t2bNtOCa0kREZC6K\nOJHq+nUB+/frdhkfOKBGQkL+LuN69fJCVoPHH9dy8XUiIpKNVYZuQoLuuGzebPb69fyQrV5dxMCB\nOWjXToMnntDCz4+fjyUiIstgFaGbkQEcPZp/XPb0aRUkSbfL2M1NwjPP5Ohns6GhvFYxERFZJosM\nXa0WOHVKpd9l/PffamRl6ZLUzk5Cmzb5u4wbNdKtG0tERGTpLCKuJAm4elXAvn26kD140AZJSfnT\n1YYN80O2VSstnJxkLJaIiKicZAvde/cEHDyYv8v45s3847KBgSJ69tTtMm7TRgsfHx6XJSIi62e2\n0E1NBY4cUetns+fP56/C4+kp6UO2XTsNgoMZskREpDyVFro5OcDx4/nHZf/9Vw2NRrfL2MFBQvv2\nugtStG+vQYMGXCCAiIiUr1JCt2dP4M8/XZCaqgtZQZDQpImov/JTy5ZaODhUxncmIiKyXJUSurt3\nAyEhEvr10+0ybttWAw+PyvhORERE1qNSQvfaNcDJKa0yNk1ERGS1KuVIalBQZWyViIjIuvH0JSIi\nIjNh6BIREZkJQ5eIiMhMGLpERERmwtAlIiIyE4YuERGRmTB0iYiIzIShS0REZCYMXSIiIjNh6BIR\nEZkJQ5eIiMhMGLpERERmwtAlIiIyE4YuERGRmTB0iYiIzIShS0REZCYMXSIiIjNh6BIREZkJQ5eI\niMhMGLpERERmwtAlIiIyE4YuERGRmTB0iYiIzIShS0REZCYMXSIiIjNh6BIREZkJQ5eIiMhMGLpE\nRERmwtAlIiIyE4YuERGRmRgVuvv370e3bt3QtWtXfPbZZ5VdExERkSKVGrqiKGLevHlYvXo1vv/+\ne/zwww+4fPmyOWojIiJSlFJD99SpUwgKCkKNGjVga2uL7t274/fffzdHbURERIpSaujevXsXAQEB\n+ud+fn64d+9epRZFRESkRKWGriRJ5qiDiIhI8WxKG+Dv74/bt2/rn9+9exfVqlUrdcO+vq4Vq8wC\nKKEHQBl9KKEHgH1YEiX0ACijDyX0YKxSZ7oNGzZETEwMbt26hezsbPzwww/o1KmTOWojIiJSlFJn\numq1GrNnz8bw4cMhSRL69euH0NBQc9RGRESkKILEg7ZERERmwStSERERmQlDl4iIyEwYukRERGZS\n6olUZbF//34sWLAAkiShb9++eOWVV0y5ebOYOXMm9u7dC29vb+zevVvucsolNjYW06ZNQ3x8PNRq\nNfr3748hQ4bIXVaZZWdnIzIyEjk5OdBqtejatSvGjh0rd1nlIooi+vbtCz8/P6xatUrucsqlY8eO\ncHFxgUqlgo2NDbZs2SJ3SeWSkpKCN998E9HR0VCpVFiwYAEaN24sd1lGu3r1KiZNmgRBECBJEm7c\nuIEJEyZY5d/xr7/+Glu2bIEgCKhTpw4W/r+9u3mJag8DOP6dHKRQexElCyzIjCySFr1AEyamSTXV\nxGCLNiVRbdIow14oghYJLfoHWkREEBEaRG1EszGmQiuGYIgwIhhMKkRT5yXPnOcu4l64G+89x7nz\na7rPZz1n+A6HmYcznHmmo4P8/HzTWY7cunXrr/fCv/qslQxJp9NSX18vsVhMfvz4IXv37pWhoaFM\nPX3WDAwMSDQaFb/fbzrFtS9fvkg0GhURkcnJSdmxY0dOngsRkXg8LiIilmVJU1OTRCIRw0Xu3Lx5\nU9ra2uT48eOmU1yrq6uTsbEx0xmzdvbsWbl//76IiExPT8vExIThIvfS6bT4fD4ZHh42neLYyMiI\n1NXVSSqVEhGRkydPSldXl+EqZ96/fy9+v19SqZRYliWHDx+WT58+zXhMxr5e/l12NG/YsIH58+eb\nzpiV0tJSqqqqACgoKKCioiJnV3fOmzcP+HnVa1mW4Rp3RkZGePr0KU1NTaZTZkVEsG3bdMasTE5O\nMjg4SDAYBMDr9VJYWGi4yr1wOMyyZcv+tqo3l9i2TSKRwLIsksnkv1q89Cv58OED69evJz8/n7y8\nPDZu3Eh3d/eMx2Rs6OqO5l9TLBbj3bt3VFdXm05xxbZtAoEAPp8Pn8+Xk6/j6tWrtLe34/F4TKfM\nisfj4ciRIwSDQe7du2c6x5VYLMaiRYs4f/48+/fv59KlSySTSdNZrj1+/Jjdu3ebznBl8eLFNDc3\nU1tbS01NDUVFRWzZssV0liOVlZUMDAwwPj5OIpEgFArx+fPnGY/J2NAV/bnvL2dqaorW1lYuXLhA\nQUGB6RxX5syZw4MHDwiFQkQiEYaGhkwnOdLX10dJSQlVVVU5/x65e/cunZ2d3Lhxgzt37jA4OGg6\nyTHLsohGoxw8eJCuri7mzp2bs/8RPj09TW9vLzt37jSd4sr379/p6enhyZMn9Pf3E4/Hc+4+moqK\nCo4ePUpzczPHjh1j9erVeL0z3yqVsaHrdkez+m9YlkVrayv79u2jvr7edM6sFRYWsmnTJvr7+02n\nOPL69Wt6e3vZvn07bW1tvHz5kvb2dtNZrpSWlgJQXFxMQ0MDb9++NVzkXFlZGWVlZaxbtw6AxsZG\notGo4Sp3QqEQa9eupbi42HSKK+FwmPLychYuXEheXh4NDQ28efPGdJZjwWCQzs5Obt++zYIFC1i+\nfPmMj8/Y0P2ddjTn+hUJ/LwLe+XKlRw6dMh0imujo6NMTEwAkEwmef78OStWrDBc5czp06fp6+uj\np6eH69evs3nzZq5du2Y6y7FEIsHU1BQA8XicZ8+eUVlZabjKuZKSEpYsWcLHjx8BePHiRc6utX30\n6BF+v990hmtLly4lEomQSqUQkZw9F6OjowAMDw/T3d39j+ckYz8Z+l12NP95NTI2NkZtbS0tLS1/\n3XSRK169esXDhw9ZtWoVgUAAj8fDqVOnqKmpMZ3myNevXzl37hy2bWPbNrt27WLbtm2ms/6Xvn37\nxokTJ/B4PKTTafbs2cPWrVtNZ7ly8eJFzpw5g2VZlJeX09HRYTrJsWQySTgc5sqVK6ZTXKuurqax\nsZFAIIDX62XNmjUcOHDAdJZjLS0tjI+P4/V6uXz5MkVFM/9jku5eVkoppbJEN1IppZRSWaJDVyml\nlMoSHbpKKaVUlujQVUoppbJEh65SSimVJTp0lVJKqSzRoauUUkpliQ5dpZRSKkv+AO2e4yf8wTuC\nAAAAAElFTkSuQmCC\n",
+            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAd8AAAFKCAYAAABcq1WoAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMS4yLCBo\ndHRwOi8vbWF0cGxvdGxpYi5vcmcvNQv5yAAAIABJREFUeJzs3Xd4U2X/BvD7ZLRpumlLS6EDgbKh\niIggU7aAgPhDRKsIUoYgiK++ioAguBARXmZBEARFUBGhiChIEQcqe+/RMlpGd9KRcX5/nDZtaFra\nkuY07f25rlw5zXmSfPMk5OY5Oec8giiKIoiIiMhhFHIXQEREVN0wfImIiByM4UtERORgDF8iIiIH\nY/gSERE5GMOXiIjIwVSOeJJbtzLs/pi+vlqkpOjt/rhkjf3sGOxnx2A/Owb7WRIQ4FnsOqcd+apU\nSrlLqBbYz47BfnYM9rNjsJ/vzWnDl4iIyFkxfImIiByM4UtERORgDF8iIiIHY/gSERE5GMOXiIjI\nwRi+REREDsbwJSIih/vxx61YtGi+3GXIhuFLRETkYA45vSQREZEtGzeux65dPwMAOnbsjOeeG45/\n/tmHFSuWwNVVA1/fGnjnndk4eHB/kdtUKueNMKesPDZ2C7p37wSNxkfuUoiInN6MGVOxdetmuz2e\nQiGgb98BmDFjdontbty4hgMH/sGKFV8AAKKjX0DXrt3x3XcbMH78q2jZshX27PkVaWmpNm/z8/O3\nW82O5nSbnTMy0jFixHN48cUX5S6FiIjuw9mzZ9G0aXOoVCqoVCo0b94S58+fRdeu3fHxxx/giy9W\noUGDhvDz87d5mzNzupGvp6cXOnTohF27duHkyRNo0qSp3CURETm1GTNm33OUWhYBAZ6lms1OEABR\nFC1/GwwGCIICvXv3Rdu27fDbb3H4739fxezZc2zeFhYWbreaHc3pRr4AEB09DgCwYsVSmSshIqLy\niohoiOPHj8FoNMJoNOLkyROIiGiI1as/g1KpwoABT6Jbt564fPmizducmdONfAGgR49eqFevHr79\ndgPefnsG/P2de/MDEVF1FBQUjFatHsKECdEwm0X07z8AQUG1EBgYhEmTxsHT0wuenp4YOvQ56PX6\nIrc5M0EsPOavIKXZ/FBW69d/jokTJ+LNN6di8uQ37P74JCnt5iO6P+xnx2A/Owb7WRIQ4FnsOqfc\n7AwAL774Ijw9vbBq1Qrk5ubKXQ4REVGpOW34enp6YtiwKNy8mYQfftgkdzlERESl5rThCwAvvTQa\nCoUCMTFL4ICt50RERHbh1OEbFhaO3r374ujRw/j7731yl0NERFQqTh2+ADB6tHTY0fLlS2SuhIiI\nqHScPnwfeaQ9mjdviR9/3Ir4+Ctyl0NERHRPTh++giAgOnoszGYzVq5cLnc5REQkk/Pnz1kGYe+8\n8xZycrLL/ViHDx9ESkqyvUorwunDFwAGDhyMgICa+PLLL5CZmSl3OUREJIM9e35FQkI8AGDmzA/g\n6qop92Nt27alQsPXKc9wdTdXV1e8+OJLmDPnfWzY8BVGjoyWuyQiIrqHYcMGY+3ajRBFEX36PIaF\nC5ehUaMmmDx5PN54420EBdWCyWTCnDnv4fr1azAajXjppTFo3boNtm+PxaZNG6FSqVG/fgQGDhyM\nH37YhD17foWvry+mT38LX3yxAZ9+Oge+vr44c+Y0UlNT8OyzL2Dbtq1IS0vFokXLIQjAzJlTkZWV\nhezsbLz66uvQ6TKxd28cLl26iNmz5+DMmZP4+ut1UCpVaNiwMSZMePW+X3uVCF8AeOGFkZg/fy5W\nrFiKF198CQpFlRjUExFVOPcZU+FqxykFoRDg3ncAdPeYrKFhw8a4ePECjEYDGjVqjOPHjyIiohGS\nk5MRFFQLAPDLLz/Bz88fb701HampqZg4cQzWrPkaX3+9DnPmzEdgYBC2bduCOnXqoG3bdujSpRua\nNGlm9TxKpQoLFizFzJlTcezYUSxYsASzZk3DwYP7ER5eF/36DUSnTl1w4MC/+PLLNXjvvY9Rv34E\nJk9+A15eXlizZiWWLfscLi4umDbtTRw9ehgtWkTeVxdVmfANCAjA4MFDsH79Ouza9TN69Ogtd0lE\nRFSCyMgHceLEMeTm5uCpp57Gnj270bLleURENLS0OX78KI4cOYSjRw8DAHJycmAwGNC9ey9MmfI6\nevXqg+7de5W4iblxY2n2Oz8/f8tMSL6+ftDpMlGjhh/WrPkM69evhcFggEZj/TiXLl1EUlIiJk8e\nDwDQ6TKRmJiIFi3u77VXmfAFgFGjxmL9+nWIiVnK8CUiKiXdjNn3HKWWRUCAJ3SlOLdzq1atsW7d\nauTkZKNfvwHYtm0rjh07ggcffMjSRqVS4/nnRxT5To+KehE9evRBXNxOvPLKWCxeXPwOt0ql0uay\nKIrYuPEr+PvXxLRps3D69EksWjTf6r5qtbSped68Rfd8PWVRpbbNNmvWHB06dMJvv+3GqVMn5S6H\niIhKEBoahqSkJGRm6qDVusPPzw9798ZZhW+TJs3w++97AAApKcmIiVkMs9mMmJjF8Pf3x9Chz6FZ\ns+ZITEyEIAgwmUxlqiEtLRW1a9cBAOzZsxtGoxEAoFAoYDKZEBoajsuXL1l2vlq5Mga3bt2879de\nqvA9e/YsunfvjnXr1gEAbty4gaioKAwbNgwTJ06sVBMbcK5fIiLn4evri6CgIABS0N64cQM1awZa\n1j/2WHe4uWkxZswIvPHGq2jRIhIKhQJarTtGj34REyeOhSAIaNAgAi1btsL8+R9j//5/Sv38vXv3\nxYYNX+LVV19G06bNcOfOHWzbtgWRkQ9i6tT/4vr1a5g48TX85z8TMXbsCKSlpcLfP+C+X/c9pxTU\n6/UYPXo0wsPD0bBhQzz33HN466230KlTJ/Tp0wfz5s1DUFAQhg0bVuxjVMTUUsVNWWUymdCu3YO4\nceM6Dh06xbl+7xOnBnMM9rNjsJ8dg/0sua8pBV1cXLBixQrUrFnTctvff/+Nbt26AQC6du2Kv/76\nyw5l2odSqcSoUWOQk5ODtWs/l7scIiKiIu4ZviqVqsjeX1lZWXBxcQEA+Pn54datWxVTXTk988xz\nnOuXiIgqrfve27k0U/n5+mqhUinv2a6sihvSBwR44qWXRuLTTz9FXNxPePbZZ+3+3NVJSZtOyH7Y\nz47BfnYM9nPJyhW+Wq0W2dnZ0Gg0SEpKstokbUtKir5cxZXkXr8pDBv2IhYsWIC5cz9Bjx79IQiC\n3WuoDvjbjWOwnx2D/ewY7GfJff3ma0v79u2xY8cOAMDPP/+Mjh07lq+yChQWFo5evR7H4cOH8M8/\nf8tdDhERkcU9w/f48eOIiorC999/jy+++AJRUVEYP348Nm/ejGHDhiE1NRUDBw50RK1lxrl+iYio\nMrrnoUb24MhDjQoTRRHdunXEyZPH8e+/RxESEmr3Oqo6bj5yDPazY7CfHcPe/RwXtwtdunSz2+M5\nit03OzsLzvVLROTcbty4jp07d8hdht1V6fAFgEGDnoK/fwDWrVvDuX6JiCqRYcMGw2QywWg0okeP\nTjh9Wjot8OTJ45GYeAMAMG/eRzh8+CA+/3wFVq6MwaxZ0zFu3EvYv/8fTJ36huWx+vaVRsaXLl3E\nK6+MwcSJY/HWW68hI6Nybumo8uGbP9dvenoaNmz4Su5yiIgqpRqtm9m8aAptNfQcN8pmG8/o4ZY2\nmrWrgfDwUj1n/pSC586dsUwpaDabraYUfOaZKERGPogXXxwFADAaDViy5LNip42dP/9jvP76FCxY\nsBRt2jyCTZs2lqs/KlqVD19AmutXOlPXUpjNZrnLISIiFEwpeOzYETz11NM4efIELlywnlLwbvnT\nAxbn5MkT+Oij2Rg/Pho7dvxomRChsqlSUwoWp2bNmnjyyf/D119/ybl+iYhsSD5w/J5tMpasuGeb\n7Kjh8Jw8AbDTlIJ3U6vVAFDk3A35sxFpNBosXBhT6c/tUC1GvoA01y8AxMRwtiMiosqgNFMK5k/t\ndzd3d3fcuXMbAHD+/Dno9dLJnOrXb4B9+/4EAOzcuaNMMxw5UrUJ3+bNW+DRRztyrl8iokrkXlMK\nhoXVxZkzp/G//31idb/69SOg0bhhzJgR2LHjRwQFBQMAJk78D9au/Rzjx0fjxx9jS9yELacqfZzv\n3bZv34YXXngGzz33AubNW2j3mqoiHhfpGOxnx2A/Owb7WVJtj/O9W8+evREWFo5vvvkat2/flrsc\nIiKqpqpV+HKuXyIiqgyqVfgC0ly/Hh6enOuXiIhkU+3C19PTC88+G4WkpERs2fK93OUQEVE1VO3C\nFwBGjhwNQRCwfPkSOGB/MyIiIivVMnzDw+uid+++OHz4EP79t3IeA0ZERFVXtQxfgHP9EhHJ6ccf\nt2LRovl2eSydLhP//LMPALB27WocP3603I+VmJiIkyfvfbav+1Vtw7ddu0fRrFkLxMb+gISEeLnL\nISKicjpz5rQlfKOihqNZsxblfqyDB//FqVMn7FVasarFuZ1tyZ/r95VXxmLVqhV4551ZcpdERFSt\n3LhxDf/5zyu4eTMJQ4YMQ79+A6zWf/fdRuzc+RMEQYGOHbvgmWeew9mzp/HJJx9BrVbDxcUFM2d+\ngHnz5kCv1yEkJBTHjx9Fly7dkJaWisOHDyI1NRWXLl1EdPRY7Ny5A5cvX8L06bPRtGkzLFw4DydP\nnkBubi4GDhyMDh06Y9Wq5VCpVAgMDELt2iH49NM5EAQBWq0WU6bMgKdn8SfOKItqG76ANNfvu+9O\nx7p1a/Daa/+Fh4eH3CURETncjBmu2LrVfnGgUAB9+7pixoycEtslJMRj1aovodNlYvjwYejb9wnL\nhAjXr19DXNwuLFmyEgAwduxIdO3aHT/+uBWDBj2F3r374sCBf5GcfAfDhkXh4sULGDDgSatNzgkJ\n8Viy5DNs3boZ69atxqpVX2L79q3YuXMH6tdvgKCgYEyYMBk5OdkYMmQg+vcfiD59+sHHxwcdOnTG\nxIlj8frrUxASEopNm77Bpk0b8cILI+3SR9U6fPPn+v344w+wceN6jBgxSu6SiIiqjRYtIqFSqeDt\n7QN3d3ekpaXBx8cHAHDq1AlcvZqACRNGAwD0eh0SE6+jQ4fOmDv3QyQkxKNbtx4ICwvHiRPHbD5+\no0ZNIAgC/Pz8Ua9eAyiVSvj6+kGnOwJXV1ekp6dhzJgRUKlUSE1NKXL//OkJAcBgMKBx4yZ2e+3V\nOnwBaa7fBQs+wYoVSzF8+MhiJ2gmIqqqZszIuecotSykczuX5vGsp/0rPAugSqVGu3aP4o033i5y\nr88++wJ//rkXs2fPwPjxk4p9dKVSaXNZFEUcOnQABw/ux6JF0mbmHj06Frl/RU5PWO2TJn+u3wsX\nzuPXX3+RuxwiomrjxImjMJlMSElJQVZWFry8vC3rGjZsjIMHDyA7OxuiKGL+/LnIycnGd99tQHp6\nGnr27IOnnx6Gs2dPQxAEm9MOliQtLRU1awZCpVLh99/3wGQyw2AwWE1hWJHTE1b7kS8gzfX79ddf\nIiZmCbp37yV3OURE1UJoaDimTXsT164lIDp6nNUIMygoCEOGPIOXXx4FhUKBTp26wNVVg9q1QzBt\n2pvw8PCAWq3GlCnvIDU1BcuWLURAQM1SP/dDD7XFl1+uwfjx0ejYsTPat++AuXM/QPfuPTF79gz4\n+Phi4sT/YM6c9/Dll2vg4uKKGTNm2+21V6spBUsyaFBf/PHHXvz2299o1Kix3R7X2XFqMMdgPzsG\n+9kx2M8STilYCtHR0kk3VqxYKnMlRERU1TF88xSe6/fOnTtyl0NERFUYwzdP/ly/2dnZnOuXiIgq\nFMO3EM71S0REjsDwLSR/rt/ExBvYunWz3OUQEVEVxfC9S/5cvzExiznXLxERVQiG71041y8RUcUr\nzZSCu3fvdFA1jsfwtYFz/RIRyW/dujVyl1BhGL42cK5fIqKKlz+l4PPPP43Y2B+s1n311Rc4f/4s\npkx5HQcP7scbb0zC+PHROH36FPr27WZpN3XqGzh4cD/0eh2mTn0DEyeOxfjx0Th//pyjX06ZMHxt\nyJ/r12w2Y9WqFXKXQ0RU4Vq3drd5WblSbWkzbpzGZpvoaI2lzdq1aoSHl+45ExLi8eGH87BwYQxW\nroyx2s9m2LDn4eHhgfff/xgAcOHCecybt6jYMxBu3Lgebdu2x4IFS/Haa29i0aJPy94JDsTwLcbA\ngYPh7x+AdevWIDMzU+5yiIiqHFtTChanfv0GcHFxKXb9sWNHsXnzdxg/PhqffPIhdLrK/b3NiRWK\nodFoMHz4SMyd+yHn+iWiKu/AAd092yxZkn3PNlFRBkyerMGtW6V51uKnFLybWq22ebvRaMxbr8Kr\nr76OZs1alOaJZceRbwleeGEkXFxcsGLFUpjNZrnLISKqUkqaUhAAzGbbh3sKgoDs7GxkZ2fj7Nkz\nAIAmTZrht9/iAACXLl3E11+vq9Da7xfDtwSBgYEYNOgpzvVLRFQB8qcUnDRpbJEpBQEgIqIhRo16\nvsj9Bg58CtHRL+D992eiYUPpN+Cnnnoa164lYNy4l/DRR7MRGfmgQ15DeXFKwXs4duwIunXriM6d\nu+Kbb3649x2qGE4N5hjsZ8dgPzsG+1nCKQXvQ/PmLdG+fQfs2bMbp0+fkrscIiKqAhi+pcC5fomI\nyJ4YvqXQq1cfhIZyrl8iIrIPhm8pSHP9jkZ2djbWrVstdzlEROTkGL6lNGxYFDw8PLFy5XIYDAa5\nyyEiIifG8C0lT08vDBv2HOf6JSKi+8bwLQPO9UtERPbA8C2DunUfQK9ej+PQoYPYv59z/RIRUfmU\nK3x1Oh3Gjx+PqKgoDB06FHv37rV3XZVWwVy/POyIiIjKp1zh+/3336Nu3bpYu3YtFixYgPfee8/e\ndVVa7dt3QNOmzREb+wOuXk2QuxwiInJC5QpfX19fpKamAgDS09Ph6+tr16IqM0EQMHr0OJhMJs71\nS0RE5VLuczuPHDkS8fHxSE9PR0xMDCIjI4ttazSaoFIpy11kZZOdnY2wsDDk5ubi6tWrcHd3l7sk\nIiJyIuWaz/eHH35AcHAwVq5cidOnT2PKlCnYtGlTse1TUvTlLrA4cp+4+/nnR2Du3A+xePFyvPji\nS7LVUdHk7ufqgv3sGOxnx2A/S+w+scLBgwfRoUMHAECjRo1w8+ZNmEym8lXnpDjXLxERlVe5wjcs\nLAxHjhwBAFy7dg3u7u5QKqvOZuXSyJ/r9/z5c9i9e6fc5RARkRMpV/g+/fTTuHbtGp577jm89tpr\nmDFjhp3Lcg7R0WMBADExS2SuhIiInEm5fvN1d3fHggUL7F2L08mf6zcu7lecPn0KjRo1lrskIiJy\nAjzD1X0qmOt3mcyVEBGRs2D43qeCuX7XIzmZc/0SEdG9MXzvU+G5fteuXS13OURE5AQYvnbAuX6J\niKgsGL52wLl+iYioLBi+dsK5fomIqLQYvnbCuX6JiKi0GL52xLl+iYioNBi+dsS5fomIqDQYvnbE\nuX6JiKg0GL52NnDgYPj7B2Dt2tXQ6XRyl0NERJUQw9fONBoNhg8fibS0VGzcuF7ucoiIqBJi+FYA\nzvVLREQlYfhWgMDAQAwcOBjnz59DXNwuucshIqJKhuFbQTjXLxERFYfhW0FatIhEu3aPYvfuXThz\n5rTc5RARUSXC8K1AnOuXiIhsYfhWoN69H0doaBjn+iUiIisM3wqkVCrx0kujkZWVhXXr1shdDhER\nVRIM3wo2bFgU3N09ONcvERFZMHwrmJeXN4YNew43blxHbOwPcpdDRESVAMPXAV56aQwEQcAHH8xC\nZmam3OUQEZHMGL4OULfuA3j55Ym4fPkSpk79r9zlEBGRzBi+DvLmm1PRokUkvvpqLbZu3Sx3OURE\nJCOGr4O4uLhg2bKVcHNzw+TJr+Datatyl0RERDJh+DpQ/foNMGvWh0hLS8X48aNhMpnkLomIiGTA\n8HWwqKjh6NOnH/74Yy8WL/6f3OUQEZEMGL4OJggC5s1biMDAIHz44SwcPnxQ7pKIiMjBGL4y8PPz\nw6JFMTAajRgzZiR0Op3cJRERkQMxfGXSuXNXjBv3Ci5evIBp096UuxwiInIghq+M3nprGpo1a4F1\n69Zg61ae/YqIqLpg+MrI1dXVcvjRa69NwPXr1+QuiYiIHIDhK7OIiIaYOfN9pKamYsKEMTCbzXKX\nREREFYzhWwm88MII9O79OPbu3YMlSxbKXQ4REVUwhm8lIB1+tAg1awbigw/exdGjh+UuiYiIKhDD\nt5Lw9/fHwoXLYDAYePgREVEVx/CtRLp27YbRo1/G+fPnMH36FLnLISKiCsLwrWSmTp2Bpk2bY+3a\nz/Hjj7Fyl0NERBWA4VvJ5B9+pNFoMHnyeCQm3pC7JCIisjOGbyXUsGEjzJjxHpKTk/Hyy6N5+BER\nURXD8K2kXnzxJfTs2Rt798Zh2bLFcpdDRER2xPCtpARBwKefLkZAQE28994MHDt2RO6SiIjIThi+\nlVhAQAAWLlxqOfxIr9fLXRIREdkBw7eSe+yxHoiOHotz587inXfelrscIiKyA4avE5g6dSYaN26K\nNWtW4qeffpS7HCIiuk/lDt8tW7bgiSeewJNPPom4uDg7lkR302g0WLZsJVxdXfHqqy8jKSlR7pKI\niOg+lCt8U1JSsHjxYnz11VdYtmwZdu3aZe+66C6NGzfBjBmzcefOHc5+RETk5MoVvn/99RfatWsH\nDw8P1KxZE7NmzbJ3XWTDiBHR6N69J+LifsXy5UvkLoeIiMqpXOF79epVZGdnY8yYMRg2bBj++usv\ne9dFNgiCgAULlsLfPwCzZ8/AsWNH5S6JiIjKQRBFUSzrnZYvX46DBw9i0aJFuH79Op5//nns3r0b\ngiDYbG80mqBSKe+7WJJs374djz/+OBo3boz9+/dDq9XKXRIREZWBqjx38vPzQ6tWraBSqRAaGgp3\nd3ckJyfDz8/PZvuUFPsfnxoQ4IlbtzLs/rjO4KGHOuCll0bjs89iMH78RHz00bwKe67q3M+OxH52\nDPazY7CfJQEBnsWuK9dm5w4dOmDfvn0wm81ISUmBXq+Hr69vuQuksps+fRYaN26Czz//DD//vF3u\ncoiIqAzKFb6BgYHo1asXhgwZglGjRmHq1KlQKHjIsCNpNBosXSodfjRx4jgkJSXJXRIREZVSuX7z\nLauK2PzAzRqSFSuW4u23/4uuXbth/frv7P6fIPazY7CfHYP97BjsZ4ndNztT5fHSS2Pw2GPdsXv3\nLnz22TK5yyEiolJg+Dq5gsOP/PHuu9Nx4sRxuUsiIqJ7YPhWAYGBgZg/fzFyc3MxduxIZGVlyV0S\nERGVgOFbRfTs2QcjRozC6dOn8O670+Quh4iISsDwrULeeWc2GjZshJUrl2Pnzh1yl0NERMVg+FYh\nbm5uWLZsFVxcXPDKK+Nw8+ZNuUsiIiIbGL5VTNOmzTBt2kzcvn0LEyeOhQOOJCMiojJi+FZBo0aN\nRZcuj2HXrl+wcmWM3OUQEdFdGL5VkEKhwMKFy+Dn54eZM6fh1KmTcpdERESFMHyrqMDAIHz66WLk\n5ORgzJgRyM7OlrskIiLKw/Ctwnr3fhzDh4/EqVMnMWvWdLnLISKiPAzfKm7GjPcQEdEQK1Ysw65d\nP8tdDhERgeFb5Wm1WixdutJy+NGtW7fkLomIqNpj+FYDzZu3wNtvz8CtWzcxadI4Hn5ERCQzhm81\nMXr0OHTu3BW//LIDq1atkLscIqJqjeFbTeQfflSjRg3MnDkVp0+fkrskIqJqi+FbjQQF1cKnny5G\ndnY2xowZycOPiIhkwvCtZvr06Yvnnx+BkyeP4733ZspdDhFRtcTwrYZmznwP9es3QEzMYuzevUvu\ncoiIqh2GbzXk7u6OZctWQq1WY8KEMbh9+7bcJRERVSsM32qqRYtIvPXWdNy8mYRXX32Zhx8RETkQ\nw7caGzduAjp27IIdO7Zj9eqVcpdDRFRtMHyrMYVCgUWLlsHX1xfvvDMFZ8+ekbskIqJqgeFbzdWq\nFYx58xYhOzsbo0ePQE5OjtwlERFVeQxfQt++/REVNRwnThzD+++/K3c5RERVHsOXAADvvvsB6tWr\nj6VLFyIu7le5yyEiqtIYvgSg4PAjlUqFCRPG4M6dO3KXRERUZTF8yaJly1Z4881pSEpKxKuvjufh\nR0REFYThS1bGj5+IDh064aeftuGLLz6XuxwioiqJ4UtWpMOPYuDj44Pp09/C6dOn5S6JiKjKYfhS\nEcHBtfHJJwuRlZWF/v374+LFC3KXRERUpTB8yab+/Qdg8uTXcf78eTz+eDf8/fc+uUsiIqoyGL5U\nrDffnIbly5cjLS0Ngwf3w/fffyt3SUREVQLDl0o0atQorF//HVxdNRg9egQ+/fRj7gVNRHSfGL50\nT126PIbY2J9Rp04IPvhgFiZNehm5ublyl0VE5LQYvlQqjRs3wfbtuxAZ2Qrr16/DM88MRlpaqtxl\nERE5JYYvlVpgYBC+//5H9O7dF3v37kHfvj1w5cplucsiInI6DF8qE3d3d3z++TqMGTMeZ8+eQZ8+\nj+HAgX/lLouIyKkwfKnMlEol3n33fXz44SdITk7GoEF9sXXrD3KXRUTkNBi+VG4jRozCunUboFSq\nMHJkFBYtWsA9oYmISoHhS/ele/de2LLlJ9SqFYx3352G//xnEgwGg9xlERFVagxfum/Nm7fATz/9\nimbNWmDt2s/x7LP/h/T0NLnLIiKqtBi+ZBe1agVjy5af0KNHL8TF/Yr+/Xvh6tUEucsiIqqUGL5k\nNx4eHlizZj1GjozGqVMn0bv3Yzhy5JDcZRERVToMX7IrlUqFDz6Yi9mzP8StWzcxYEAfbN++Te6y\niIgqFYYvVYjo6HFYvforAMDw4cMQE7OYe0ITEeW5r/DNzs5G9+7dsWnTJnvVQ1VInz598cMP2xEQ\nUBPTpr2FKVNeh9FolLssIiLZ3Vf4Ll26FN7e3vaqhaqgli1b4aeffkXjxk2xcuVyvPDCM8jMzJS7\nLCIiWZU7fC9cuIDz58+jS5cudiyHqqI6dUIQG7sDXbo8hl9+2YEnnuiNGzeuy10WEZFsBLGcP8RF\nR0dj2rRp2Lx5M2rXro0nn3xblRT0AAAgAElEQVSy2LZGowkqlbLcRVLVYDAYMH78eCxfvhy1a9dG\nbGwsIiMj5S6LiMjhVOW50+bNmxEZGYmQkJBStU9J0ZfnaUoUEOCJW7cy7P64ZM3e/Txr1seoVSsU\nM2dOxaOPdsBnn61G9+697Pb4zoqfZ8dgPzsG+1kSEOBZ7LpyhW9cXBwSEhIQFxeHxMREuLi4ICgo\nCO3bty93kVQ9CIKAl19+BaGhYXj55VF47rmn8f77H2PEiFFyl0ZE5DDlCt/58+dblhcuXIjatWsz\neKlM+vcfgODgYERFDcWbb76GS5cuYsaM2VAq+fMEEVV9PM6XZNO6dRts374LERENEROzGC+++Bx0\nOp3cZRERVbj7Dt8JEyaUuLMVUUnCwsKxbdsv6NixM376aRsGDnwcSUmJcpdFRFShOPIl2Xl7+2D9\n+u/wzDPP4ciRQ+jTpxtOnTopd1lERBWG4UuVgouLC+bPX4wpU6bj6tUE9OvXE7t375K7LCKiCsHw\npUpDEARMmvQfxMSsQm5uDoYNewpr166WuywiIrtj+FKlM2jQU/j2263w9vbGa6+9glmz3oHZbJa7\nLCIiu2H4UqXUtu0j+PHHXahXrz4WLvwUo0YNR1ZWltxlERHZBcOXKq0HHqiHH3/ciXbtHsXWrZvx\n5JN9cevWLbnLIiK6bwxfqtR8fWtg48bNeOqpp3HgwH706dMNZ8+ekbssIqL7wvClSs/V1RWLFy/H\n66+/hfj4y+jbtwd+//03ucsiIio3hi85BUEQ8Prrb2HRohjo9ToMGTIQX3/9pdxlERGVC8OXnMqQ\nIc/gm29+gIeHB155ZSw+/HAWyjkrJhGRbBi+5HTat++AH3/chbCwcMyb9zHGjh2J7OxsucsiIio1\nhi85pfr1G2D79l/Rpk1bbNr0Lf7v/wbgzp07cpdFRFQqDF9yWv7+/vjuu60YOPBJ/P33X3j88W64\nePG83GUREd0Tw5ecmkajwbJlqzBp0n9w6dJF9OnTjXtCE1Glx/Alp6dQKDBlynTMn78YGRkZePLJ\nfnjhhWE4ffqU3KUREdnE8KUqY9iwKGzZ8hPatGmL7dtj0aVLO0yYMAbx8VfkLo2IyArDl6qUhx56\nGLGxP2Pdug1o1KgJNmz4Cu3aPYgpU17HzZs35S6PiAgAw5eqIEEQ0LNnH/z66+9YuvQzBAfXxmef\nxeDhh1vigw/eRVpaqtwlElE1x/ClKkuhUGDw4CH4888DmDPnU3h6euLTT+eiTZsWWLhwPvR6vdwl\nElE1xfClKk+tVmP48JH4++/DmDp1JgBg1qzpaNs2EqtXr4TBYJC5QiKqbgTRAefmu3Urw+6PGdCm\nOUzmoqXrx72C7JHRAADPcaOg/vuvIm0MrR9CxvLVAADN2tXQzp9r8zmS/zoIuLhAee4svIc+abNN\nxryFMHTuCgDw6dUFitu3i7TJHvIM9P99GwDg/s7bcI39oUgbU2gY0r7fBgBw2b4NHlP/a/P5Urfu\ngDm4NoTUFPh262izjW7KdOQMHgIA8Hr2/6CysddvbtfuyJw7HwDgtnA+3FZ/VqSNqNVCdfoUbt3K\ngGr/P/AaPcLm86WvWgtjy1YAAN+2kRCMxiJtsqLHImv0ywAAj0kvw2XvniJtjM1bIn21dL5m16+/\nhPvHH9h8vuQ9+wAPDyguX4LP4P4222TOmYfcbj0BAD79ekJx47plndlsRkZ6OlZl6fG60Yjw8Lr4\nrmEjtDxxHBAEq8cx1wpGauzPAACXXT/D443JNp8v9butMIfXBTIzUaPzIzbb6F5/CzlDnwUAeA1/\nFqpjRyzrlAoBJrOI3I6dkTl/MQDALWYx3JYvLfI4okqFlL8PAwBURw7Ba0SUzedLj1kF40MPAwB8\nOz4MwcZIP2v4S8iaMAkA4PGfSXDZvbNIG2Ojxkj/8hsAgOt3G+H+/rs2ny9l116IPr5QXL8Gn/69\nbLbJnP0Rcvv0BQB4D+oLpY2d4XL6DYBu5nsAAO1H70GzcX2RNmZ/f6TuiAMAqPfshufkCTafL+3r\nTTA1iAByc1Gj3YOWfi5MP+k/yI4aDgDwjB4O9YH9RR7H0LYdMpasAABoVi6Hdsn/bD5f8oHjAADl\nyRPwjnraZpuMRTEwtHsUAODb9VEI6WlF2mQ/+zz0k98AALhPeR2uO7YXaWOqVx9pGzcDAFy2bobH\njKk2ny9l+68Qa9aEcPMmfPs8ZrNN5ozZyO0/EADgPWQglBeKHi+f06sPdO9/DADQzpsDzZdfFGkj\nenkjZfcfCAjwROqWn+A5frTN50tbuwGmJk0BADVaN7PZRs7vcnsJCPAsdp3Krs9E5AQUCgW8fXzw\nwpBncBoivvjic+y4fAmBajW8vX3g5uYmd4lEVMU578g3wLNCHpesVYd+vnLlMj7++AN8883XEEUR\nDz/8CKZOnYFHHmnvsBqqQz9XBuxnx2A/S0oa+fI3X6r2wsLCsWhRDPbs2Yc+ffrhn3/24YknemPo\n0CdxrNCmYSIie2H4EuVp1Kgx1qz5Ctu370KHDp3w66870a1bR0RHD+c5o4nIrhi+RHdp3boNvvtu\nKzZu3IzIyFbYvHkTHn20DV577RVcv35N7vKIqApg+BLZIAgCunR5DDt2xGHlyrV44IF6WLt2Ndq2\njcQ777yN5GROX0hE5cfwJSqBIAjo338A9uzZhwULliAgoCaWLl2Ihx5qgblzP0RmJncqIXJqZjOg\n0wGZmQ59Wu7tTCViP1vLycnBmjUrMX/+XNy+fRv+/v6YOPE1vPDCSGg0mnI/LvvZMdjPjnHf/SyK\ngMEAITsLQlYWoNdDyM6GkKWHkJUFITsL0GdJfxe6HdlZEPRZBW2yCrXRF2pT+PbsbOkpFQqkffUt\nDI91t1MvlLy3M8OXSsR+ti0zMwMxMUuwZMlCZGSko3btOnj99bcwZMgzUKnKfvg8+9kx2M92IoqA\nTgdBp4Ogy4Sg00Ghy4SgywR0OngrTMhISraEoJCVBdwVggXhWNBG0OuB/DA1mexbsloN0U0L0c0N\n0GggarUQNRrLbaKPL3RvvwNznRC7PSfDl8qN/Vyy5OQ7+N//PsWqVcuRnZ2NBg0i8Oab09Cv3xMQ\n7jpbVknYz45RLftZFKWQKxSUQmZmwbLuruXM/GUdhMyMguXC6/Q6CHaIDlEQADc3KfzcCsIQbm4Q\nNW4QtXnrNG557Qq3KRScGqkd7g5UjRugzbsux3+K7xfDl8qN/Vw6169fwyeffISvvloLk8mEyMhW\nmDLlHXTu3LVUIcx+dgyn6Oe8sFSkp0FIS4OQnlZ8GBYXmlZ/Z0Iwm8tfjiBAdPeA6O4uXTw8Cy17\nWK9zl9Z5BvkhzaQoCNG84LQEZn6AuroWOaVrVcLwpXJjP5fNhQvn8NFH72Hz5k0AgA4dOuHtt99B\n69ZtSrwf+9kxHNLPZjOEjPS84Ey3CtGC5XTp70LLVuttnB+9tEStuyUMzR6eQKHl/Nvh7gGzR35o\neuSFqDtErUeRUIVWW+aA5OdZwvClcmM/l8+xY0fw/vvvYteuXwAAvXv3xVtvTUPjxk1stmc/O0ap\n+jk3F0JaGhQZxQVlqnVopqVBUWhZyMwo8yZZUauF2csbopcXRC9vmL3zl32k68JBefeIs/Cy1h1Q\nyH8QCz/PEoYvlRv7+f7s2/cnZs+egX/+2QdBEPDUU0/jjTemICws3Kod+9nOjEYIKSlQJN+BIvkO\nhDvStWeuDvrrN4sP0Yx0aeegMhAFAaKnF0Rvbyk4ve5a9vKG6C2FqNnLO29d/rJ0O9TqCuoIefDz\nLGH4Urmxn++fKIrYuXMH3nvvXZw8eRxqtRpRUcPx6qtvIDAwEAD7uURmM4TUFCiSky0hWjhQFcl3\nICTfgeJO3nVyMhRpqaV+eFGthujtDbMlQPOCMi9ArZel8LSMUr29IXp4VorRZmXCz7OE4Uvlxn62\nH7PZjM2bv8OHH87G5cuXoNVqER09Di+//Arq1w+pHv0sitIo0yo4k6G4Yx2iluWUZAgpKaXaYUhU\nqWCu4QfRz0+6riFdm/1qWJa9wmsjRVRLI868wIWbW5Xe6UcO/N6QMHyp3NjP9mcwGPDVV2sxd+6H\nSEpKhI+PDyZNmoQePfqhQYMIucsrPVGU9qS1MRq1CtSUZGn9HWm5NDsTiQoFxBo1pPDMD1K/QoFa\no4YUsr41LIErenrdM0T5eXYM9rOE4Uvlxn6uOHq9HitXLsfChfOQmiptJm3YsBH69n0C/foNQNOm\nzcp0rLDdmM0Qbt+GMukGFEmJUCQlQZFYaDnphnR96yaE3NzSPaSvr1WIWo9M85Z9a0D0ywtTb58K\n2ZTLz7NjsJ8lDF8qN/ZzxcvISMcff/yKr77agLi4XcjOO91deHhd9Os3AP36PYFWrVrffxAbjVDc\nvpUXoolQJCbeFah5t926WeLZhUQXF5gDg2AOCIDZz79oiNbwsx61+vjIcoIDW/h5dgz2s4ThS+XG\nfnaM/H7OzMzErl0/IzZ2C375ZQf0eh0AIDi4Nvr27Y9+/Qbg4YcfgVKpLLizwQDFrZt5o9OkvBC9\nAcXNJOuQvX2rxN9ORY0G5ppBMAcGwhxUC6bAQClk8y9BtWAODIToW8NpfyPl59kx2M8Shi+VG/vZ\nMWz1c1ZqKvbH/oCD27bg8l9/wFuvRy0AdTUaNK/hh1C1Gp46HRR3bpd4XKmo1cJcMxCmoFp5IRpk\nFbJSuAZKm3qdNFRLi59nx2A/S0oK38qxLYiousnKgjIhHsqEK1DExwOpt+B58Yr1ZuDkZIQCePLu\n+2ZnA9evIQPARYUCBv8AuNVvAP9mzSEE15HC1TJaDZIOhanioUrkbBi+RBUhOxvKawlQXLmSF7Lx\nUMRflpbj46G4dbPIXfInJDR7ecMcGAhj0+Yw1wwsGK3mBaohoCb+jr+CH3b9jG3btuLGjevArZvw\nOHYUPXr0RL/QAXisVWu4u7s79jUTUalxszOViP1cjNxcKK4m5IXpFSjyrqWQvQJlUqLNu4lqNcy1\n68AUGg5TaCjMIaEwhYbBq2kE7rh6wRwYJJ1Lt5TMZjMOHtyP2NgtiI3dgvj4ywAANzc3dO3aHX37\n9kevXn3g5eVtj1ft9Ph5dgz2s6RCfvOdM2cODhw4AKPRiNGjR6Nnz57FtmX4Oq9q288GAxTXrlqP\nWuPzlhPiobhx3ebvrKJSCXPtEJhCpVA1h4TCFBIKU2g4zKGhUrgW3lkqjz36WRRFHD9+DNu2/YDY\n2C04e/YMAECtVqNTpy7o128AevfuCz8/v/t6HmdWbT/PDsZ+ltg9fPft24eVK1dixYoVSElJwaBB\ngxAXF1dse4av86qy/Ww0QnHjuvWoNX85IR6K69ds7hksKhTSyDUktFCwhsEcGibdViu4XIfVVEQ/\nnz17BrGxUhAfP34UAKBUKtG+fQf07fsEHn+8H4KCatn1OSu7Kvt5rmTYzxK7h6/JZEJOTg60Wi1M\nJhPat2+PP//80/rwh0IYvs7LafvZZIIi8YYUpFcuW0asls3E167aPJZVFASYawVbNgebQkKlYM1f\nDq5dISfBr+h+vnz5ErZt24rY2B9w4MC/AABBEPDQQw+jX78B6Nu3P0JDwyrs+SsLp/08Oxn2s6RC\nDzXasGED9u/fj48//rjYNhXxJrRp4wmzjZHJuHG5GDnSkLeswd9/F/0PQevWJixfLp3IYO1aNebP\nd7H5HH/9pYOLC3DunAJDh7rZbDNvXjY6d5a+xHv10uL27aJ7lQ4ZYsB//yudCeidd1wRG1t0ZBQa\nasb330uzqWzfrsLUqa42n2/rVj2Cg0WkpgLdutneoWbKlBwMHiydwu/ZZ91w+nTRMwV17WrE3Lk5\nAICFC12wenXRQNFqRZw+rcStWxnYv1+B0aNt98GqVVlo2VJ6L9q2dYetswdGR+di9GjpfZk0yRV7\n9xbtg+bNTVi9Wnpfvv5ahY8/tt0He/bo4OEBXL4sYPBADWA0QDAYAEPetdGIJRiLx02xAIAO2Iur\nqFPwAAoloFLh/8L3YUbfP2EOCcP033rhu31hgEpptWdwrVpmxMZK78uuXUq88YYGtnz3nR7h4SIy\nM4HOnW2/L6+/noOhQ6XOGT5cg2PHCj6bCoUCZrMZHTsaMX++9L7ExKixfHnRz6ZKBfz9t3T875Ej\nCowYYft9iYnJwkMPSe9Lx45a6PXS6zIaTcjK0iMrS4/c3PkQxTkAAD+/b2A0doebmxvUhf6D0aiR\nGV9+mZX3OlV4/33b78uuXTr4+ADXrwvo39/279azZ+egTx+pDwYNckN8fNHPZr9+RsycKfXBRx+5\nYOPGop9Nf38RO3boAQB79igxebLt9+Xrr7PQoIEZublAu3buln4ubNKkXERFSZ/N6GgNDhwo+p3R\ntq0JS5ZIn82VK9VYssT2d8aBA9L7cvKkAlFRtt+XRYuy0a6d9J3RtasW6elFvzOefdaAyZOl74wp\nU1yxY0fRfy/16pmxcaP0vmzdqsKMGbbfl+3b9ahZU8TNmwL69LH9vsyYkYP+/aX3ZcgQN1y4UPR9\n6dXLiPffl96XefNc8OWXRd8XLy8Ru3frERDgiS1b9Bg/3vb7snZtFpo0kd6H1q1t/3uR87vcXirs\nUKOdO3fi22+/xapVq0ps5+urhUple1R8PxQ2Tj/n6alBQID0hms0ts9Q5+qqQECAOq998WexCwjw\nhIsLcOdO8W18fLQICJCWVSrb7dzdXREQIP3D0Gptt1GrFZY3ytu7+Ofz8/NAQEDxzwUAXl5ulppc\nXGy3c3NzQUCA9EH18LDdJn9DRkCAJ3x9i38+X193y/MplYCt8zh4eNzP+yICJhOQKwVswKyp8Dh/\nGBnH9VCkfFv0gRQKCA3qA62GAuHhwDf1gMy8sywpVZZwVT05CO4fDJJqugkoDhV9qLK+L25uxbfx\n9Cx4X1xdi7ZTKBTQaEr3vuTXVJb3Jb+di4sCLi7e8Pb2xvPPT0OdOvWwadMm/PxzMkQxFWlpqVCr\n1dBqtXB3d4eLi9ryfF5exT+fv7/0OcnJKb6Nt3dBH6jVtttptQV94F7M9LQqVUEf+JRwJsoaNaQ+\nyM0taHP390bh7wxb7wsAaDSl/86Qntd+3xnFfaZcXBSlfF+kz6bZbL/vjNK9L9p7vi9ASf9e5Psu\nd4Ryj3z37t2LBQsW4LPPPoOPj0+JbbnZ2Xk5tJ+NRigvX4Ly7Bkoz52BKu9aee4cFLpMq6aiQgFz\naBiMEQ1himgkXTeIgKlBBEQn3LO3MnyeU1NTsGPHdmzbtgW7d+9CTo40yqlb9wHLaS4jIx+U53zT\ndlIZ+rk6qMh+FkXAaATyNnbBYBBgMEj/wTIagdxcoci6/PWF/zYYhEL3kf7TMXSoAV5e9qvV7pud\nMzIyMGzYMKxevbpUe04yfJ1XhfRzVhaU589Bde4MlGfPQHXurBSyFy8UOVG/6OICU70GBeEa0RDG\nBg1hqldf+u9wFVHZPs+ZmRnYuVM6zeXOnT9bTnNZu3Ydy2ku27RpW+x+HpVVZevnyshkks7jkpsr\nBVlODpCTI11b3ybdLt0GZGdLyzk5AtRqV6Sm5tgIwMLhd3cAFl1XEKrW7SrKypVZls3v9mD38N2w\nYQMWLlyIunXrWm776KOPEBwcbLM9w9d53U8/C2mpBeFaaDSriL9S5DAds4cnTBERMDXIC9eIhtLf\noeGV5qT8Fakyf56zsrIQF/crYmN/wI4d25GengYA8PX1RYsWkWjZshVatoxEixaRCA0Nq9Qj48rc\nz2YzkJUlhVh2trSclSXkBaFgFXZSEBaEoK1l27cVF6QFj2syyfP+qdUi1GppBKpSiXnXgIuLdHvB\npaBd/nL+xcVFzLtP0XVqtfVjFrST1nl5AW3amGwdCVhuPLczlds9+1kUobiZJIXr2TPSaDY/bG2c\naMLsH5A3ipXC1digIUwNG8EcVKtanwLRWT7Pubm5+P33PYiN3YK9e/fgypXLVut9fHzQvHkkWraM\ntARyeHjdShPIZelnUcwf0UmBqNcXBGN2tmAVjkWvC9rq9db3KQhX679zciq+jwRBhKurFDouLiI0\nGmnZ1VXMuw3QaETL+vy2rq62lotfX7OmFnq9/q4ALByu1gGoUlXNf/4MXyo3Sz+bzVDEX8nbVHy2\n0G+yZ6FISy1yP1NIKEwNCsLVmBe2om8NGV5F5eesn+fU1BQcPXoER44cxtGj0uXSpYtWbby9fdCi\nRUs0b97SEsrh4Q/Y3GGyLERRGqnpdAJ0uruvbd8mii5ISTFYQq9wgNoKUbPZ/ong5iaFnkYjXbu5\niXBzk/7Ovy58e36wFQ486/C71/qCx3BUyDnr59neGL5UJkJKMlRHDkN15BA8Lp6F4dgJqM6fhZA3\nz2w+UaWCqe4DeTs8ReSNZhvCWK+BtDsklVpV+jynpaXi2LGjeYF8CEePHsGFCxcAuAPwAOABrTYI\nDzzQEqGhTREcHIGaNetBqw1EVpaiSGjq9cWHqz02kapURcPPVgjmL+cHp5tb6f8uuL1qjvDuVpU+\nz/eDsxpRsYSMdKiOHoHq8CGojhyE+tBBKO/alKjSamGMaGS9w1NEQ5jqPlAhJ5ygysFgANLTBaSn\nAxkZAtLSBMvf6ekCMjKKjjCloNRCpwuGTtfHchtgnTh6PXD8uHQpLa1WhLu7CHd3oEYNs2XZ+rrk\n22rXdkdWViY0GunxNJpqsUsBVUL82FUnej1Ux49BfeQgVIcOQnXkEJTnz1nt/GSuUQO5XbvB0OpB\nGFs+CO9Oj+C2WwkHk1KlZDYDmZn54WkdmmlpUnCmp8OynB+sGRkFt+WflKOslEoRHh5S2Pn6iqhT\n5+5QlJbV6hxkZNzAnTuXkZh4Htevn8aNG+cgiukAMgFkws1NRNOmddGqVSO0bNkSLVu2Qv36Dcq9\nl3VAAHDrVoVv7CO6J4ZvVZWbC9XJ49KI9vBBqA8fgvLMKatTKpo9vWB4tCOMkQ/CENkKxsgHYQ4J\ntd4uFuAJcPORQ4mi9Ptj4dC0DkncFZgC0tJQaFkKUVEsW3hKe3yK8PQEgoLM8PIS8y4otFxwm6en\nCA8PEVqtdai6uJRl02qtvEs7AIBOp8Px48dw9Oghy+/IBw/GYf/+Xy330Gq1aNq0uWWHrpYtW6FB\ngwioOIQlJ8LffKsCoxHKM6ehPnIob0R7EKqTJ6yOmRXd3GBs3jJvRCsFremBevcc0bKf7092NpCc\nLFhd7tyRrlNSCv7OzFQhOdlsGZ0aDGULTkGQQlMKTxHe3sWHZuG/vb0L7uPmVjl/j9Tr9Thx4hiO\nHj2MI0eky9mzp2Eq9B9JNzc3NGnSLG+HrlZo0SISDRs2KhLI/Dw7BvtZwh2uqhKzGcoL56E6fNAy\nolUdPwohK8vSRHRxgbFps7wRrRS2poiGlWa2HWeVk1NykOYvF/67tJtu3dwALy9zMSPNoiHq7Y1C\nISsWeyrKqiorKwsnTx63jI6PHDmMM2dOwVjoxOIajQZNmzbL28taCuSOHR9Gamp2CY9M9sDvDQnD\n11mJIhRXLhca0R6C6shhKDILXreoVMLUqIlls7ExshWMjZtK2/7soKr2c04OLCPPwkFaeDR6d5Dq\ndKULUq1WRI0a0sXXV4SfX/F/+/lJt4WEVM1+dqTs7GycOnXCKpBPnz4Jg8FgaaNUKlG7dh3UqRNi\nuYSEhFqua9euA1dX2xMUUOlV1e+NsuLezs5AFKG4cd0SsurD0rUiJaWgiSDA1CACuS1bFWw+btZC\nGjZVczodkJQkIClJgdu3bQdp4dFpZmbpglSjkQLygQfMRYLz7kt+kPLtkIdGo0GrVq3RqlVry205\nOTk4ffqkZXP1hQtncOnSZfz11x8obtwRGBiUF8YhCAkJsyzXqSOFtIeHh6NeElVhHPnKRLh1C+rD\nB6x2iFLcumnVxhReN29E21oa0bZoCdGj+P9JVQS5+zkzsyBUExMFJCUJSExU5N0mWNZlZNw7TF1d\nC8KzpCAt3EZrewY2u5O7n6uL/H7Ozc3FtWtXcfVqAq5eTUBCQrxlOT4+HtevX7XahF2Yr6/vXaEs\nBbMU1qHw8fGtNGf0kgs/zxKOfOWWlQX1/n+gOrgf6vxDfK5dtWpiql0HOY/3LxjRtoyssmeDEkUp\nVPNDND9Uk5IKQjV/3b029fr7mxESYkZgoIigIBGBgWYEBBQfpNX8O5HyuLi4oG7dB1C37gM215tM\nJiQlJSIhIQFXr8YjISHesnz1agLOnTuDo0cP27yvu7tHoVCWRs/5f4eEhCIgoOZ9n92LnB/DtyIY\njVAdPgiXvXug3rsH6n//hpA3PRsgnd84p0cvy2+0hpYPQqxZU8aC7UMUgfR0WI1MExMVuHlTsBq1\n3rxZ8o5IgiCFZt26+aEqXdesWRCwQUEiAgJEe/20TWRFqVQiOLg2goNro23bR4qsF0URd+7cQULC\nlbyRc0Ewx8dL16dPn7L52K6urggOrm0VyvnBHBISilq1gnnYVDXAd9geRBHKUyfhsjdOCts//7Da\nKcrYtDlyO3aG4eFHYGz1IMzBtZ1qCCaKQGoqrDb9Wo9SC/7Ozi7+dSkUIvz9RdSrZ7aEaGCgaBWw\ngYFSqPLEWVSZCYIAf39/+Pv7W/3GXFh6elpeKCcgIeGKZVkaSSfgt99227yfUqlErVrBVjuF1ahR\nA76+NQpd+6FGjRrw8vLmKNpJ8TffclJcvpQ3so2Dy++/QXH7tmWdse4DMHTsgtxOnWFo3xGiv79s\ndZaG0Qhcvy4gPl6BhAQBV64oLMtJSSrcuCGWOOOKQiGNSvM3/dasab0ZWLqWgpf/obdN7s9zdVGZ\n+lmv1+Patat3/d58xbKcmHgDZrO5xMdQKBTw9fWFr68Uyn5+fpbl/KDOX65RI3+dL1wqeJNRZepn\nOfE3XzsQkpLg8ru0Gdnl99+gjL9iWWcKDEL2U08jt1MXGDp0grlOiIyVFmU2SzstXbkiBWp+sMbH\nSyF77Zpg8wT1CoWIWolNSOEAAAkTSURBVLWAJk3Md41SrUet/v6iXefAJKoOtFotGjSIQIMGETbX\nGwwGXL9+DdevX0NycjJSUpILXd+x+jslJRmXLl20OvFISTw8PG2MpouGduEwd3d3r/Y7ktkTw7cY\nQloq1H/+IY1s9+6B6sxpyzqztw9yHu8vbUru1AWm+g1k3YwsisDt24JVoMbHFyxfvSogN9d2fUFB\nZjz4oBmhoWaEhZkREiIiNFT6OzhYRHCwJ27d0jv4FRGRWq1GWFg4wsLCS9XebDYjPT3NKpALL9+5\nU/T2s2dPI6vQCXpK4uLiYmMUbTu869ULQW6uAHd3d2i17uU+F3dVxvDNl5UF9T/7LJuSVUcOQ8jb\n5CO6uSG3y2PI7dgFhk6dpWNrHfxhSk0FEhIUVqPXwiPY4nZg8vc3o2lTsyVQ88M1LMyM2rWlWV2I\nyPkpFAr4+PjCx8cXQL1S30+v1xcJ6sIj7Ltvv379Ok6dOlmm2jQaDdzd3eHu7gGtVpsXyh5511q4\nuxddzg/u/PvZauvMv3dX3/A1GqE6dMB6j+S8cyGLKhWMDz1sGdkaHnxImp26AmVmSuEaHy9YQjZ/\nOT5egfR02+Hq5SWdACI/WMPCCpZDQszg+QCIqCRarRZarRa1a9cp9X2MRiNSUlKKDe2srAwkJ6dC\np9PlXTKh1+uh0+mQlJQInU6H3ELnnr+/2u8O6sIhbyvIrf8TkL/s6+sLb2+f+66ptKpP+JrN1nsk\n//WnZY9kURBgbNYChg6dYOjUGblt28PeqZWTgyKbhfODNT5ewJ07tv8Hp9VKI9VHHpHCVBrBFmwa\n9va2a5lERPekUqkQEBCAgIAAm+tLs8OVwWCAXq+zBHTBcmbe33rLsvV66zDPb5OamorMzIxS/+59\nN4VCga+++haPPda9XPcvq6obvqJYsEfy73uK7pFcrz5yBg+R9kh+tCPEGn52eVq9Hjh/XoEzZxQ4\nezb/WonLlwWYzUVHry4uIkJCRDRvbiwSrKGh0vGu3MeBiKoatVoNb28fu442RVFEbm6uzXC2DvOi\n6wEgIqKh3Wq5lyoVvoqkRGlUm79HckK8ZZ2pVjCyhzyD3A6dYOjYGeYybGKxJTMTOHs2P2CVlqBN\nSBCKzKNao4YZbdqYUK+eFKjSCFbaRFyzplitZqMhIqoogiDA1dUVrq6uqGGnAVVFcerwFdJSof7j\nd2lT8u+/We+R7OuLnH4DpLDt1AWmevXLtUdyaipw5owS584VjGbPnlXg2rWiiVmzphkdOpgQEWFG\nRIQZDRtK1/7+FX4oNRERORHnC19RhNvC+cCOWPgdOFCwR7JWi9zHukt7JHfsJO2RXIYh5e3bQqHN\nxAWbjG/eLPoYwcFmdOlitISrdDHB19dur5KIiKow5wtfnQ7un3wIGI0wPPwIDB07S5cHH7rnHLai\nCNy8Kdz1e6x0sbXDU2ioGd27G/NGsdKItkEDM7y8KurFERFRdeB84evhgTv7j8M/LBBpetunXhNF\n6XSJ1qNY6XfZtDTrTc+CICI8XESbNgarzcX165vh7u6IF0RERNWN84UvADEgAHB3hzkzAwkJgtVe\nxfnLd09Fp1RKx8N26GC22lxcr56Zk58TEZFDOV34iiIwfbor/v0XOHXKA1lZ1iGrVouoX99cZKen\nBx4wc/o5IiKqFJwufPV6YMMGNbKzYQnZ/IBt2NCE8HDOnENERJWb08WUuztw4kQmAgM9kZzME/4T\nEZHzccrTO6jVDp/XgIiIyG6cMnyJiIicGcOXiIjIwRi+REREDsbwJSIicjCGLxERkYMxfImIiByM\n4UtERORgDF8iIiIHY/gSERE5GMOXiIjIwRi+REREDiaIoijKXQQREVF1wpEvERGRgzF8iYiIHIzh\nS0RE5GAMXyIiIgdj+BIRETkYw5eIiMjBnC5833//fTz99NMYOnQojh49Knc5VdqcOXPw9NNPY/Dg\nwfj555/lLqdKy87ORvfu3bFp0ya5S6mytmzZgieeeAJPPvkk4uLi5C6nStLpdBg/fjyioqIwdOhQ\n7N27V+6SKi2V3AWUxT///IMrV65gw4YNuHDhAqZMmYINGzbIXVaVtG/fPpw7dw4bNmxASkoKBg0a\nhJ49e8pdVpW1dOlSeHv/f3v398r6H8Bx/LkzubBxzDJaIblRSigXWHJBLlz7kRa3cqVc0FKUq7lS\nKAp/gLZwI0pZuZgr5UJRXGExy8evxgU6d6fOt9x8a3vbp9fjbrt61i5ee38+n7bfpjNsy7IslpaW\niEajpNNpFhYW6OjoMJ1lO5ubm1RXVzM+Ps7d3R3Dw8Ps7u6azvqRcmp84/E4nZ2dANTU1PD09MTr\n6ytut9twmf00NzdTX18PQFFREW9vb3x+fuJ0Og2X2c/l5SUXFxcagwyKx+O0tLTgdrtxu93Mzs6a\nTrIlj8fD+fk5AM/Pz3g8HsNFP1dOXXZOpVL/fJglJSXc398bLLIvp9NJQUEBAJFIhPb2dg1vhoTD\nYSYnJ01n2Nr19TXv7++MjIwwODhIPB43nWRLPT09JBIJurq6CAaDTExMmE76sXLq5Ptf+mXMzNvf\n3ycSibC+vm46xZa2trZoaGigoqLCdIrtPT4+sri4SCKRYGhoiIODAxwOh+ksW9ne3sbv97O2tsbZ\n2RmhUEjPMXwjp8bX5/ORSqX+vk4mk5SWlhossrfDw0OWl5dZXV2lsLDQdI4txWIxrq6uiMVi3N7e\nkp+fT3l5Oa2trabTbMXr9dLY2EheXh6VlZW4XC4eHh7wer2m02zl+PiYQCAAQG1tLclkUrervpFT\nl53b2trY29sD4PT0FJ/Pp/u9GfLy8sLc3BwrKysUFxebzrGt+fl5otEoGxsb9Pb2Mjo6quHNgEAg\nwNHREV9fX1iWRTqd1v3IDKiqquLk5ASAm5sbXC6XhvcbOXXybWpqoq6ujoGBARwOB9PT06aTbGtn\nZwfLshgbG/v7Xjgcxu/3G6wS+X/Kysro7u6mr68PgKmpKX79yqmzR07o7+8nFAoRDAb5+PhgZmbG\ndNKPpb8UFBERyTJ99RMREckyja+IiEiWaXxFRESyTOMrIiKSZRpfERGRLNP4ioiIZJnGV0REJMs0\nviIiIln2BzQKNGAGnBgwAAAAAElFTkSuQmCC\n",
             "text/plain": [
-              "\u003cmatplotlib.figure.Figure at 0xc1dc310\u003e"
+              "\u003cmatplotlib.figure.Figure at 0x7f7a18df6b50\u003e"
             ]
           },
           "metadata": {
@@ -668,13 +549,10 @@
         "  w_at_step = []\n",
         "  b_at_step = []\n",
         "  for step_num in range(num_training_steps):\n",
-        "    loss, gradients_and_variables = value_and_gradients_fn(inputs, labels, wb)\n",
-        "    loss_at_step.append(np.asscalar(loss.numpy()))\n",
-        "    \n",
-        "    optimizer.apply_gradients(gradients_and_variables)\n",
+        "    loss_at_step.append(run_step(inputs, labels))\n",
         "    w, b = wb.variables\n",
-        "    w_at_step.append(np.asscalar(w.read_value().numpy()))\n",
-        "    b_at_step.append(np.asscalar(b.read_value().numpy()))\n",
+        "    w_at_step.append(np.asscalar(w.numpy()))\n",
+        "    b_at_step.append(np.asscalar(b.numpy()))\n",
         "\n",
         "  print(w_at_step)\n",
         "  t = range(0, num_training_steps)\n",
@@ -688,171 +566,12 @@
         "\n",
         "train_model(inputs, labels, wb, optimizer, num_training_steps)"
       ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "colab_type": "text",
-        "id": "UNurY9VJ-hpH"
-      },
-      "source": [
-        "## Other Ways to Compute Gradients\n",
-        "\n",
-        "Using our loss function as an example (`loss_fn()`), there are several other ways we could compute gradients:\n",
-        "\n",
-        "1. `tfe.implicit_gradients()`\n",
-        "1. `tfe.gradients_function()`\n",
-        "1. `tfe.implicit_value_and_gradients()`\n",
-        "1. `tfe.value_and_gradients_function()`\n",
-        "\n",
-        "Each of these functions does the following:\n",
-        "* Wraps a function.\n",
-        "* Returns a function with the same input signature as the wrapped function.\n",
-        "\n",
-        "They differ only in what information they return.\n",
-        "\n",
-        "### Gradients-only functions\n",
-        "\n",
-        "The following two functions return a function that returns only the variables' gradients:\n",
-        "\n",
-        "1. `tfe.gradients_function()`: Returns the partial derivatives of the function `f()` with respect to the parameters of `f()`.\n",
-        "1. `tfe.implicit_gradients()`: Returns the partial derivatives of the function `f()` with respect to the trainable parameters (`tf.Variable`) used by `f()`.\n",
-        "\n",
-        "In our example above, the `tf.layers.Dense` object encapsulates the trainable parameters.\n",
-        "\n",
-        "### Value and gradients functions\n",
-        "\n",
-        "The following two functions are identical to their counterparts above, except that they also return the value of the wrapped function.\n",
-        "\n",
-        "1. `tfe.implicit_value_and_gradients()`\n",
-        "1. `tfe.value_and_gradients_function()`\n",
-        "\n",
-        "### Gradient demos\n",
-        "\n",
-        "In the demos below, we show examples for the `implicit_*` functions, since our existing loss function works seamlessly with these versions. (The other versions require that your parameters are tensors and tensors only; in our example, we're using a `Dense` layer.)\n"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 13,
-      "metadata": {
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          },
-          "height": 85,
-          "output_extras": [
-            {
-              "item_id": 1
-            }
-          ]
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 100,
-          "status": "ok",
-          "timestamp": 1505502831671,
-          "user": {
-            "displayName": "",
-            "photoUrl": "",
-            "userId": ""
-          },
-          "user_tz": 240
-        },
-        "id": "aEoCftnfAIH5",
-        "outputId": "72f1c1dc-a574-463f-f860-c4e5f48fcdaa"
-      },
-      "outputs": [
-        {
-          "data": {
-            "text/plain": [
-              "[(\u003ctf.Tensor: id=673, shape=(1, 1), dtype=float32, numpy=array([[-0.26846504]], dtype=float32)\u003e,\n",
-              "  \u003ctf.Variable 'dense/kernel:0' shape=(1, 1) dtype=float32\u003e),\n",
-              " (\u003ctf.Tensor: id=671, shape=(1,), dtype=float32, numpy=array([-0.32890949], dtype=float32)\u003e,\n",
-              "  \u003ctf.Variable 'dense/bias:0' shape=(1,) dtype=float32\u003e)]"
-            ]
-          },
-          "execution_count": 13,
-          "metadata": {
-            "tags": []
-          },
-          "output_type": "execute_result"
-        }
-      ],
-      "source": [
-        "# tfe.implicit_gradients() demo\n",
-        "gradients_fn = tfe.implicit_gradients(loss_fn)\n",
-        "\n",
-        "# Returns only gradients and variables:\n",
-        "gradients_fn(inputs, labels, wb)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 14,
-      "metadata": {
-        "colab": {
-          "autoexec": {
-            "startup": false,
-            "wait_interval": 0
-          },
-          "height": 102,
-          "output_extras": [
-            {
-              "item_id": 1
-            }
-          ]
-        },
-        "colab_type": "code",
-        "executionInfo": {
-          "elapsed": 88,
-          "status": "ok",
-          "timestamp": 1505502831785,
-          "user": {
-            "displayName": "",
-            "photoUrl": "",
-            "userId": ""
-          },
-          "user_tz": 240
-        },
-        "id": "bbgCUdCzAVhH",
-        "outputId": "152aa9b6-9e42-4b7e-848a-9423c0b1929c"
-      },
-      "outputs": [
-        {
-          "data": {
-            "text/plain": [
-              "(\u003ctf.Tensor: id=688, shape=(), dtype=float32, numpy=1.0623235\u003e,\n",
-              " [(\u003ctf.Tensor: id=720, shape=(1, 1), dtype=float32, numpy=array([[-0.26846504]], dtype=float32)\u003e,\n",
-              "   \u003ctf.Variable 'dense/kernel:0' shape=(1, 1) dtype=float32\u003e),\n",
-              "  (\u003ctf.Tensor: id=718, shape=(1,), dtype=float32, numpy=array([-0.32890949], dtype=float32)\u003e,\n",
-              "   \u003ctf.Variable 'dense/bias:0' shape=(1,) dtype=float32\u003e)])"
-            ]
-          },
-          "execution_count": 14,
-          "metadata": {
-            "tags": []
-          },
-          "output_type": "execute_result"
-        }
-      ],
-      "source": [
-        "# tfe.implicit_value_and_gradients() demo\n",
-        "value_gradients_fn = tfe.implicit_value_and_gradients(loss_fn)\n",
-        "\n",
-        "# Returns the value returned by the function passed in, gradients, and variables:\n",
-        "value_gradients_fn(inputs, labels, wb)"
-      ]
     }
   ],
   "metadata": {
     "colab": {
+      "collapsed_sections": [],
       "default_view": {},
-      "last_runtime": {
-        "build_target": "",
-        "kind": "local"
-      },
       "name": "Eager Execution Tutorial: Working with Gradients",
       "provenance": [],
       "version": "0.3.2",
diff --git a/tensorflow/contrib/eager/python/examples/notebooks/3_datasets.ipynb b/tensorflow/contrib/eager/python/examples/notebooks/3_datasets.ipynb
index 0088da5c4b..bfcc7feb07 100644
--- a/tensorflow/contrib/eager/python/examples/notebooks/3_datasets.ipynb
+++ b/tensorflow/contrib/eager/python/examples/notebooks/3_datasets.ipynb
@@ -16,7 +16,9 @@
         "\n",
         "We recommend using the `Dataset`s API for building performant, complex input pipelines from simple, re-usable pieces that will feed your model's training or evaluation loops.\n",
         "\n",
-        "If you're familiar with TensorFlow graphs, the API for constructing the `Dataset` object remains exactly the same when eager execution is enabled, but the process of iterating over elements of the dataset is slightly different.  You will use a Pythonic `Iterator()` class instead of using `make_one_shot_iterator()` and `get_next()`. As a result, the discussion on iterators in the [Programmer's Guide](https://www.tensorflow.org/programmers_guide/datasets) is not relevant when eager execution is enabled."
+        "If you're familiar with TensorFlow graphs, the API for constructing the `Dataset` object remains exactly the same when eager execution is enabled, but the process of iterating over elements of the dataset is slightly simpler.\n",
+        "You can use Python iteration over the `tf.data.Dataset` object and do not need to explicitly create an `tf.data.Iterator` object.\n",
+        "As a result, the discussion on iterators in the [Programmer's Guide](https://www.tensorflow.org/programmers_guide/datasets) is not relevant when eager execution is enabled."
       ]
     },
     {
@@ -48,11 +50,8 @@
         "# Import TensorFlow.\n",
         "import tensorflow as tf\n",
         "\n",
-        "# Import TensorFlow eager execution support (subject to future changes).\n",
-        "import tensorflow.contrib.eager as tfe\n",
-        "\n",
         "# Enable eager execution\n",
-        "tfe.enable_eager_execution()"
+        "tf.enable_eager_execution()"
       ]
     },
     {
@@ -137,32 +136,27 @@
       "source": [
         "# Step 3: Iterate\n",
         "\n",
-        "Use `tfe.Iterator` on the `Dataset` object to get a Python iterator over the contents of the dataset.\n",
-        "\n",
-        "If you're familiar with the use of `Dataset`s in TensorFlow graphs, note that this process of iteration is different. Here there are no calls to `Dataset.make_one_shot_iterator()` and no `get_next()` calls."
+        "When eager execution is enabled `Dataset` objects support iteration.\n",
+        "If you're familiar with the use of `Dataset`s in TensorFlow graphs, note that there is no need for calls to `Dataset.make_one_shot_iterator()` or `get_next()` calls."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": 0,
       "metadata": {
         "colab": {
           "autoexec": {
             "startup": false,
             "wait_interval": 0
           },
-          "height": 153,
-          "output_extras": [
-            {
-              "item_id": 1
-            }
-          ]
+          "base_uri": "https://localhost:8080/",
+          "height": 153
         },
         "colab_type": "code",
         "executionInfo": {
-          "elapsed": 201,
+          "elapsed": 388,
           "status": "ok",
-          "timestamp": 1505952405928,
+          "timestamp": 1525154629129,
           "user": {
             "displayName": "",
             "photoUrl": "",
@@ -171,7 +165,7 @@
           "user_tz": 420
         },
         "id": "lCUWzso6mbqR",
-        "outputId": "ec027d30-96c6-4ea4-9ee1-ef74ec1ae29a"
+        "outputId": "8e4b0298-d27d-4ac7-e26a-ef94af0594ec"
       },
       "outputs": [
         {
@@ -179,9 +173,9 @@
           "output_type": "stream",
           "text": [
             "Elements of ds_tensors:\n",
-            "tf.Tensor([4 9], shape=(2,), dtype=int32)\n",
+            "tf.Tensor([1 9], shape=(2,), dtype=int32)\n",
             "tf.Tensor([16 25], shape=(2,), dtype=int32)\n",
-            "tf.Tensor([36  1], shape=(2,), dtype=int32)\n",
+            "tf.Tensor([ 4 36], shape=(2,), dtype=int32)\n",
             "\n",
             "Elements in ds_file:\n",
             "tf.Tensor(['Line 1' 'Line 2'], shape=(2,), dtype=string)\n",
@@ -191,22 +185,19 @@
       ],
       "source": [
         "print('Elements of ds_tensors:')\n",
-        "for x in tfe.Iterator(ds_tensors):\n",
+        "for x in ds_tensors:\n",
         "  print(x)\n",
         "\n",
         "print('\\nElements in ds_file:')\n",
-        "for x in tfe.Iterator(ds_file):\n",
+        "for x in ds_file:\n",
         "  print(x)"
       ]
     }
   ],
   "metadata": {
     "colab": {
+      "collapsed_sections": [],
       "default_view": {},
-      "last_runtime": {
-        "build_target": "",
-        "kind": "local"
-      },
       "name": "Eager Execution Tutorial: Importing Data",
       "provenance": [],
       "version": "0.3.2",
-- 
GitLab


From 07c58859c2ec62757f110dc56da9946d415b72ee Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 11:52:43 -0700
Subject: [PATCH 096/395] Boosted trees: support indicator column.

PiperOrigin-RevId: 194971229
---
 .../python/estimator/canned/boosted_trees.py  | 292 ++++++++++++------
 .../estimator/canned/boosted_trees_test.py    |  58 ++++
 .../boosted_trees/stats_ops_test.py           |  55 +++-
 3 files changed, 315 insertions(+), 90 deletions(-)

diff --git a/tensorflow/python/estimator/canned/boosted_trees.py b/tensorflow/python/estimator/canned/boosted_trees.py
index 085dace1b3..d281fd90ea 100644
--- a/tensorflow/python/estimator/canned/boosted_trees.py
+++ b/tensorflow/python/estimator/canned/boosted_trees.py
@@ -49,35 +49,10 @@ _TreeHParams = collections.namedtuple('TreeHParams', [
 
 _HOLD_FOR_MULTI_CLASS_SUPPORT = object()
 _HOLD_FOR_MULTI_DIM_SUPPORT = object()
+_DUMMY_NUM_BUCKETS = -1
 
 
-def _get_max_buckets(feature_columns):
-  """Gets the maximum number of buckets from feature_columns.
-
-  Args:
-    feature_columns: a list/set of tf.feature_column.
-
-  Returns:
-    max_buckets: the maximum number of buckets among bucketized_columns.
-
-  Raises:
-    ValueError: when unsupported feature_columns are given.
-  """
-  if not feature_columns:
-    raise ValueError('feature_columns must be a non-empty list/set of '
-                     'tf.feature_column.')
-  max_buckets = 1
-  for fc in feature_columns:
-    if isinstance(fc, feature_column_lib._BucketizedColumn):  # pylint:disable=protected-access
-      # N boundaries creates (N+1) buckets.
-      max_buckets = max(max_buckets, len(fc.boundaries) + 1)
-    else:
-      raise ValueError('For now, only bucketized_column is supported but '
-                       'got: {}'.format(fc))
-  return max_buckets
-
-
-def _get_transformed_features(features, feature_columns):
+def _get_transformed_features(features, sorted_feature_columns):
   """Gets the transformed features from features/feature_columns pair.
 
   Args:
@@ -91,22 +66,33 @@ def _get_transformed_features(features, feature_columns):
     ValueError: when unsupported features/columns are tried.
   """
   # pylint:disable=protected-access
-  for fc in feature_columns:
-    if not isinstance(fc, feature_column_lib._BucketizedColumn):
-      raise ValueError('For now, only bucketized_column is supported but '
-                       'got: {}'.format(fc))
   transformed_features = feature_column_lib._transform_features(
-      features, feature_columns)
-  # pylint:enable=protected-access
+      features, sorted_feature_columns)
   result_features = []
-  for column in sorted(transformed_features, key=lambda tc: tc.name):
-    source_name = column.source_column.name
-    squeezed_tensor = array_ops.squeeze(transformed_features[column], axis=1)
-    if len(squeezed_tensor.shape) > 1:
-      raise ValueError('For now, only supports features equivalent to rank 1 '
-                       'but column `{}` got: {}'.format(
-                           source_name, features[source_name].shape))
-    result_features.append(squeezed_tensor)
+  for column in sorted_feature_columns:
+    if isinstance(column, feature_column_lib._BucketizedColumn):
+      source_name = column.source_column.name
+      squeezed_tensor = array_ops.squeeze(transformed_features[column], axis=1)
+      if len(squeezed_tensor.shape) > 1:
+        raise ValueError('For now, only supports features equivalent to rank 1 '
+                         'but column `{}` got: {}'.format(
+                             source_name, features[source_name].shape))
+      result_features.append(squeezed_tensor)
+    elif isinstance(column, feature_column_lib._IndicatorColumn):
+      source_name = column.categorical_column.name
+      tensor = math_ops.to_int32(transformed_features[column])
+      if len(tensor.shape) > 2:
+        raise ValueError('Rank of indicator column must be no more than 2, '
+                         'but column `{}` got: {}'.format(
+                             source_name, features[source_name].shape))
+      unstacked = array_ops.unstack(tensor, axis=1)
+      result_features.extend(unstacked)
+    else:
+      raise ValueError(
+          'For now, only bucketized_column and indicator_column is supported '
+          'but got: {}'.format(column))
+    # pylint:enable=protected-access
+
   return result_features
 
 
@@ -120,9 +106,87 @@ def _local_variable(tensor, name=None):
       name=name)
 
 
-def _cache_transformed_features(features, feature_columns, batch_size):
+def _group_features_by_num_buckets(sorted_feature_columns):
+  """Groups feature ids by the number of buckets.
+
+  Derives the feature ids based on iterating through ordered feature columns
+  and groups them by the number of buckets each feature require. Returns a
+  sorted list of buckets and a list of lists of feature ids for each of those
+  buckets.
+
+  Args:
+    sorted_feature_columns: a list/set of tf.feature_column sorted by name.
+
+  Returns:
+    bucket_size_list: a list of required bucket sizes.
+    feature_ids_list: a list of lists of feature ids for each bucket size.
+
+  Raises:
+    ValueError: when unsupported features columns are provided.
+  """
+  bucket_size_to_feature_ids_dict = collections.OrderedDict()
+
+  # TODO(nponomareva) for now we preserve the previous functionality and bucket
+  # all numeric into the same num of buckets. Can be easily changed to using
+  # each numeric's real buckets num, but we need to test that it does not cause
+  # a performance hit.
+
+  # We will replace this dummy key with the real max after we calculate it.
+  bucket_size_to_feature_ids_dict[_DUMMY_NUM_BUCKETS] = []
+
+  max_buckets_for_bucketized = 2
+  max_buckets_for_indicator = 2
+
+  feature_idx = 0
+  # pylint:disable=protected-access
+
+  for column in sorted_feature_columns:
+    if isinstance(column, feature_column_lib._IndicatorColumn):
+      num_categorical_features = column.categorical_column._num_buckets
+      if max_buckets_for_indicator not in bucket_size_to_feature_ids_dict:
+        bucket_size_to_feature_ids_dict[max_buckets_for_indicator] = []
+
+      for _ in range(num_categorical_features):
+        # We use bucket size of 2 for categorical.
+        bucket_size_to_feature_ids_dict[max_buckets_for_indicator].append(
+            feature_idx)
+        feature_idx += 1
+    elif isinstance(column, feature_column_lib._BucketizedColumn):
+      max_buckets_for_bucketized = max(max_buckets_for_bucketized,
+                                       len(column.boundaries) + 1)
+      bucket_size_to_feature_ids_dict[_DUMMY_NUM_BUCKETS].append(feature_idx)
+      feature_idx += 1
+    elif not isinstance(column, feature_column_lib._IndicatorColumn):  # pylint:disable=protected-access
+      raise ValueError(
+          'For now, only bucketized_column and indicator column are supported '
+          'but got: {}'.format(column))
+
+  # pylint:enable=protected-access
+  # Replace the dummy key with the real max num of buckets for all bucketized
+  # columns.
+  bucket_size_to_feature_ids_dict[
+      max_buckets_for_bucketized] = bucket_size_to_feature_ids_dict[
+          _DUMMY_NUM_BUCKETS]
+  del bucket_size_to_feature_ids_dict[_DUMMY_NUM_BUCKETS]
+
+  feature_ids_list = list(bucket_size_to_feature_ids_dict.values())
+  bucket_size_list = list(bucket_size_to_feature_ids_dict.keys())
+  return bucket_size_list, feature_ids_list
+
+
+def _calculate_num_features(sorted_feature_columns):
+  num_features = 0
+  for column in sorted_feature_columns:
+    if isinstance(column, feature_column_lib._IndicatorColumn):  # pylint:disable=protected-access
+      num_features += column.categorical_column._num_buckets  # pylint:disable=protected-access
+    else:
+      num_features += 1
+  return num_features
+
+
+def _cache_transformed_features(features, sorted_feature_columns, batch_size):
   """Transform features and cache, then returns (cached_features, cache_op)."""
-  num_features = len(feature_columns)
+  num_features = _calculate_num_features(sorted_feature_columns)
   cached_features = [
       _local_variable(
           array_ops.zeros([batch_size], dtype=dtypes.int32),
@@ -132,7 +196,7 @@ def _cache_transformed_features(features, feature_columns, batch_size):
   are_features_cached = _local_variable(False, name='are_features_cached')
 
   def cache_features_and_return():
-    """Caches transoformed features.
+    """Caches transformed features.
 
     The intention is to hide get_transformed_features() from the graph by
     caching the result except the first step, since bucketize operation
@@ -144,7 +208,8 @@ def _cache_transformed_features(features, feature_columns, batch_size):
           the graph.
     """
 
-    transformed_features = _get_transformed_features(features, feature_columns)
+    transformed_features = _get_transformed_features(features,
+                                                     sorted_feature_columns)
     cached = [
         state_ops.assign(cached_features[i], transformed_features[i])
         for i in range(num_features)
@@ -349,6 +414,8 @@ def _bt_model_fn(
     ValueError: mode or params are invalid, or features has the wrong type.
   """
   is_single_machine = (config.num_worker_replicas <= 1)
+
+  sorted_feature_columns = sorted(feature_columns, key=lambda tc: tc.name)
   if train_in_memory:
     assert n_batches_per_layer == 1, (
         'When train_in_memory is enabled, input_fn should return the entire '
@@ -364,24 +431,26 @@ def _bt_model_fn(
   # the dimension max_splits_per_layer, instead of max_splits (for the entire
   # tree).
   max_splits = (1 << tree_hparams.max_depth) - 1
-  max_buckets = _get_max_buckets(feature_columns)
   train_op = []
   with ops.name_scope(name) as name:
     # Prepare.
     global_step = training_util.get_or_create_global_step()
-    num_features = len(feature_columns)
+    bucket_size_list, feature_ids_list = _group_features_by_num_buckets(
+        sorted_feature_columns)
     # Extract input features and set up cache for training.
     training_state_cache = None
     if mode == model_fn.ModeKeys.TRAIN and train_in_memory:
       # cache transformed features as well for in-memory training.
       batch_size = array_ops.shape(labels)[0]
-      input_feature_list, input_cache_op = _cache_transformed_features(
-          features, feature_columns, batch_size)
+      input_feature_list, input_cache_op = (
+          _cache_transformed_features(features, sorted_feature_columns,
+                                      batch_size))
       train_op.append(input_cache_op)
       training_state_cache = _CacheTrainingStatesUsingVariables(
           batch_size, head.logits_dimension)
     else:
-      input_feature_list = _get_transformed_features(features, feature_columns)
+      input_feature_list = _get_transformed_features(features,
+                                                     sorted_feature_columns)
       if mode == model_fn.ModeKeys.TRAIN and example_id_column_name:
         example_ids = features[example_id_column_name]
         training_state_cache = _CacheTrainingStatesUsingHashTable(
@@ -446,34 +515,61 @@ def _bt_model_fn(
         gradients = gradients_impl.gradients(loss, logits, name='Gradients')[0]
         hessians = gradients_impl.gradients(
             gradients, logits, name='Hessians')[0]
-      stats_summary_list = [
-          array_ops.squeeze(
-              boosted_trees_ops.make_stats_summary(
-                  node_ids=node_ids,
-                  gradients=gradients,
-                  hessians=hessians,
-                  bucketized_features_list=[input_feature_list[f]],
-                  max_splits=max_splits,
-                  num_buckets=max_buckets),
-              axis=0) for f in range(num_features)
-      ]
-
-      def grow_tree_from_stats_summaries(stats_summary_list):
+
+      stats_summaries_list = []
+      for i, feature_ids in enumerate(feature_ids_list):
+        num_buckets = bucket_size_list[i]
+        summaries = [
+            array_ops.squeeze(
+                boosted_trees_ops.make_stats_summary(
+                    node_ids=node_ids,
+                    gradients=gradients,
+                    hessians=hessians,
+                    bucketized_features_list=[input_feature_list[f]],
+                    max_splits=max_splits,
+                    num_buckets=num_buckets),
+                axis=0) for f in feature_ids
+        ]
+        stats_summaries_list.append(summaries)
+
+      accumulators = []
+
+      def grow_tree_from_stats_summaries(stats_summaries_list,
+                                         feature_ids_list):
         """Updates ensemble based on the best gains from stats summaries."""
-        (node_ids_per_feature, gains_list, thresholds_list,
-         left_node_contribs_list, right_node_contribs_list) = (
-             boosted_trees_ops.calculate_best_gains_per_feature(
-                 node_id_range=last_layer_nodes_range,
-                 stats_summary_list=stats_summary_list,
-                 l1=tree_hparams.l1,
-                 l2=tree_hparams.l2,
-                 tree_complexity=tree_hparams.tree_complexity,
-                 min_node_weight=tree_hparams.min_node_weight,
-                 max_splits=max_splits))
+        node_ids_per_feature = []
+        gains_list = []
+        thresholds_list = []
+        left_node_contribs_list = []
+        right_node_contribs_list = []
+        all_feature_ids = []
+
+        assert len(stats_summaries_list) == len(feature_ids_list)
+
+        for i, feature_ids in enumerate(feature_ids_list):
+          (numeric_node_ids_per_feature, numeric_gains_list,
+           numeric_thresholds_list, numeric_left_node_contribs_list,
+           numeric_right_node_contribs_list) = (
+               boosted_trees_ops.calculate_best_gains_per_feature(
+                   node_id_range=last_layer_nodes_range,
+                   stats_summary_list=stats_summaries_list[i],
+                   l1=tree_hparams.l1,
+                   l2=tree_hparams.l2,
+                   tree_complexity=tree_hparams.tree_complexity,
+                   min_node_weight=tree_hparams.min_node_weight,
+                   max_splits=max_splits))
+
+          all_feature_ids += feature_ids
+          node_ids_per_feature += numeric_node_ids_per_feature
+          gains_list += numeric_gains_list
+          thresholds_list += numeric_thresholds_list
+          left_node_contribs_list += numeric_left_node_contribs_list
+          right_node_contribs_list += numeric_right_node_contribs_list
+
         grow_op = boosted_trees_ops.update_ensemble(
             # Confirm if local_tree_ensemble or tree_ensemble should be used.
             tree_ensemble.resource_handle,
-            feature_ids=math_ops.range(0, num_features, dtype=dtypes.int32),
+            feature_ids=all_feature_ids,
             node_ids=node_ids_per_feature,
             gains=gains_list,
             thresholds=thresholds_list,
@@ -486,32 +582,50 @@ def _bt_model_fn(
 
       if train_in_memory and is_single_machine:
         train_op.append(distribute_lib.increment_var(global_step))
-        train_op.append(grow_tree_from_stats_summaries(stats_summary_list))
+        train_op.append(
+            grow_tree_from_stats_summaries(stats_summaries_list,
+                                           feature_ids_list))
       else:
-        summary_accumulator = data_flow_ops.ConditionalAccumulator(
-            dtype=dtypes.float32,
-            # The stats consist of gradients and hessians (the last dimension).
-            shape=[num_features, max_splits, max_buckets, 2],
-            shared_name='stats_summary_accumulator')
-        apply_grad = summary_accumulator.apply_grad(
-            array_ops.stack(stats_summary_list, axis=0), stamp_token)
+        dependencies = []
+
+        for i, feature_ids in enumerate(feature_ids_list):
+          stats_summaries = stats_summaries_list[i]
+          accumulator = data_flow_ops.ConditionalAccumulator(
+              dtype=dtypes.float32,
+              # The stats consist of grads and hessians (the last dimension).
+              shape=[len(feature_ids), max_splits, bucket_size_list[i], 2],
+              shared_name='numeric_stats_summary_accumulator_' + str(i))
+          accumulators.append(accumulator)
+
+          apply_grad = accumulator.apply_grad(
+              array_ops.stack(stats_summaries, axis=0), stamp_token)
+          dependencies.append(apply_grad)
 
         def grow_tree_from_accumulated_summaries_fn():
           """Updates the tree with the best layer from accumulated summaries."""
           # Take out the accumulated summaries from the accumulator and grow.
-          stats_summary_list = array_ops.unstack(
-              summary_accumulator.take_grad(1), axis=0)
-          grow_op = grow_tree_from_stats_summaries(stats_summary_list)
+          stats_summaries_list = []
+
+          stats_summaries_list = [
+              array_ops.unstack(accumulator.take_grad(1), axis=0)
+              for accumulator in accumulators
+          ]
+
+          grow_op = grow_tree_from_stats_summaries(stats_summaries_list,
+                                                   feature_ids_list)
           return grow_op
 
-        with ops.control_dependencies([apply_grad]):
+        with ops.control_dependencies(dependencies):
           train_op.append(distribute_lib.increment_var(global_step))
           if config.is_chief:
+            min_accumulated = math_ops.reduce_min(
+                array_ops.stack(
+                    [acc.num_accumulated() for acc in accumulators]))
+
             train_op.append(
                 control_flow_ops.cond(
-                    math_ops.greater_equal(
-                        summary_accumulator.num_accumulated(),
-                        n_batches_per_layer),
+                    math_ops.greater_equal(min_accumulated,
+                                           n_batches_per_layer),
                     grow_tree_from_accumulated_summaries_fn,
                     control_flow_ops.no_op,
                     name='wait_until_n_batches_accumulated'))
diff --git a/tensorflow/python/estimator/canned/boosted_trees_test.py b/tensorflow/python/estimator/canned/boosted_trees_test.py
index c8c52d3bc6..95bb9b5a3b 100644
--- a/tensorflow/python/estimator/canned/boosted_trees_test.py
+++ b/tensorflow/python/estimator/canned/boosted_trees_test.py
@@ -46,6 +46,7 @@ INPUT_FEATURES = np.array(
         [3.0, 20.0, 50.0, -100.0, 102.75],     # feature_2 quantized:[2,3,3,0,3]
     ],
     dtype=np.float32)
+
 CLASSIFICATION_LABELS = [[0.], [1.], [1.], [0.], [0.]]
 REGRESSION_LABELS = [[1.5], [0.3], [0.2], [2.], [5.]]
 FEATURES_DICT = {'f_%d' % i: INPUT_FEATURES[i] for i in range(NUM_FEATURES)}
@@ -101,17 +102,25 @@ class BoostedTreesEstimatorTest(test_util.TensorFlowTestCase):
 
   def _assert_checkpoint(self, model_dir, global_step, finalized_trees,
                          attempted_layers):
+    self._assert_checkpoint_and_return_model(model_dir, global_step,
+                                             finalized_trees, attempted_layers)
+
+  def _assert_checkpoint_and_return_model(self, model_dir, global_step,
+                                          finalized_trees, attempted_layers):
     reader = checkpoint_utils.load_checkpoint(model_dir)
     self.assertEqual(global_step, reader.get_tensor(ops.GraphKeys.GLOBAL_STEP))
     serialized = reader.get_tensor('boosted_trees:0_serialized')
     ensemble_proto = boosted_trees_pb2.TreeEnsemble()
     ensemble_proto.ParseFromString(serialized)
+
     self.assertEqual(
         finalized_trees,
         sum([1 for t in ensemble_proto.tree_metadata if t.is_finalized]))
     self.assertEqual(attempted_layers,
                      ensemble_proto.growing_metadata.num_layers_attempted)
 
+    return ensemble_proto
+
   def testTrainAndEvaluateBinaryClassifier(self):
     input_fn = _make_train_input_fn(is_classification=True)
 
@@ -325,6 +334,55 @@ class BoostedTreesEstimatorTest(test_util.TensorFlowTestCase):
         [[0.353850], [0.254100], [0.106850], [0.712100], [1.012100]],
         [pred['predictions'] for pred in predictions])
 
+  def testTrainEvaluateAndPredictWithIndicatorColumn(self):
+    categorical = feature_column.categorical_column_with_vocabulary_list(
+        key='categorical', vocabulary_list=('bad', 'good', 'ok'))
+    feature_indicator = feature_column.indicator_column(categorical)
+    bucketized_col = feature_column.bucketized_column(
+        feature_column.numeric_column(
+            'an_uninformative_feature', dtype=dtypes.float32),
+        BUCKET_BOUNDARIES)
+
+    labels = np.array([[0.], [5.7], [5.7], [0.], [0.]], dtype=np.float32)
+    # Our categorical feature defines the labels perfectly
+    input_fn = numpy_io.numpy_input_fn(
+        x={
+            'an_uninformative_feature': np.array([1, 1, 1, 1, 1]),
+            'categorical': np.array(['bad', 'good', 'good', 'ok', 'bad']),
+        },
+        y=labels,
+        batch_size=5,
+        shuffle=False)
+
+    # Train depth 1 tree.
+    est = boosted_trees.BoostedTreesRegressor(
+        feature_columns=[bucketized_col, feature_indicator],
+        n_batches_per_layer=1,
+        n_trees=1,
+        learning_rate=1.0,
+        max_depth=1)
+
+    num_steps = 1
+    est.train(input_fn, steps=num_steps)
+    ensemble = self._assert_checkpoint_and_return_model(
+        est.model_dir, global_step=1, finalized_trees=1, attempted_layers=1)
+
+    # We learnt perfectly.
+    eval_res = est.evaluate(input_fn=input_fn, steps=1)
+    self.assertAllClose(eval_res['loss'], 0)
+
+    predictions = list(est.predict(input_fn))
+    self.assertAllClose(
+        labels,
+        [pred['predictions'] for pred in predictions])
+
+    self.assertEqual(3, len(ensemble.trees[0].nodes))
+
+    # Check that the split happened on 'good' value, which will be encoded as
+    # feature with index 2 (0-numeric, 1 - 'bad')
+    self.assertEqual(2, ensemble.trees[0].nodes[0].bucketized_split.feature_id)
+    self.assertEqual(0, ensemble.trees[0].nodes[0].bucketized_split.threshold)
+
 
 class ModelFnTests(test_util.TensorFlowTestCase):
   """Tests bt_model_fn including unexposed internal functionalities."""
diff --git a/tensorflow/python/kernel_tests/boosted_trees/stats_ops_test.py b/tensorflow/python/kernel_tests/boosted_trees/stats_ops_test.py
index f0bb84e69a..5cceb98cff 100644
--- a/tensorflow/python/kernel_tests/boosted_trees/stats_ops_test.py
+++ b/tensorflow/python/kernel_tests/boosted_trees/stats_ops_test.py
@@ -224,7 +224,7 @@ class StatsOpsTest(test_util.TensorFlowTestCase):
       self.assertAllClose([[[-.424658], [-.6]], [[-.043478], [.485294]]],
                           sess.run(right_node_contribs_list))
 
-  def testCalculateBestGainsWithMinNodeWEight(self):
+  def testCalculateBestGainsWithMinNodeWeight(self):
     """Testing Gain calculation without any regularization."""
     with self.test_session() as sess:
       max_splits = 7
@@ -271,6 +271,59 @@ class StatsOpsTest(test_util.TensorFlowTestCase):
       self.assertAllClose([[[-0.75]], [[-0.014925]]],
                           sess.run(right_node_contribs_list))
 
+  def testCalculateBestGainsWithMinNodeWeightNoSplitOnFeturePossible(self):
+    """Testing Gain calculation without any regularization."""
+    with self.test_session() as sess:
+      max_splits = 7
+      node_id_range = [1, 3]  # node 1 through 2 will be processed.
+      stats_summary_list = [
+          [
+              [[0., 0.], [.08, .09], [0., 0.], [0., 0.]],  # node 0; ignored
+              [[0., 0.], [.15, .0036], [.06, .007], [.1, .2]],  # node 1
+              [[0., 0.], [-.33, .068], [0., 0.], [.3, .04]],  # node 2
+              [[0., 0.], [0., 0.], [0., 0.], [0., 0.]],  # node 3; ignored
+              [[0., 0.], [0., 0.], [0., 0.], [0., 0.]],  # node 4; ignored
+              [[0., 0.], [0., 0.], [0., 0.], [0., 0.]],  # node 5; ignored
+              [[0., 0.], [0., 0.], [0., 0.], [0., 0.]],  # node 6; ignored
+          ],  # feature 0
+          [
+              [[0., 0.], [0., 0.], [.08, .09], [0., 0.]],  # node 0; ignored
+              [[0., 0.], [.3, .5], [-.05, .6], [.06, .07]],  # node 1
+              [[.1, .1], [.2, .03], [-.4, .05], [.07, .08]],  # node 2
+              [[0., 0.], [0., 0.], [0., 0.], [0., 0.]],  # node 3; ignored
+              [[0., 0.], [0., 0.], [0., 0.], [0., 0.]],  # node 4; ignored
+              [[0., 0.], [0., 0.], [0., 0.], [0., 0.]],  # node 5; ignored
+              [[0., 0.], [0., 0.], [0., 0.], [0., 0.]],  # node 6; ignored
+          ],  # feature 1
+      ]  # num_features * shape=[max_splits, num_buckets, 2]
+
+      (node_ids_list, _, _, _,
+       _) = boosted_trees_ops.calculate_best_gains_per_feature(
+           node_id_range,
+           stats_summary_list,
+           l1=0.0,
+           l2=0.0,
+           tree_complexity=0.0,
+           min_node_weight=1,
+           max_splits=max_splits)
+
+      # We can't split either of the nodes on the first feature
+      self.assertEqual(2, len(sess.run(node_ids_list)))
+      self.assertAllEqual([], sess.run(node_ids_list)[0])
+      self.assertAllEqual([1], sess.run(node_ids_list)[1])
+
+      # Now check when we can't split on any feature
+      (node_ids_list, _, _, _,
+       _) = boosted_trees_ops.calculate_best_gains_per_feature(
+           node_id_range,
+           stats_summary_list,
+           l1=0.0,
+           l2=0.0,
+           tree_complexity=0.0,
+           min_node_weight=10,
+           max_splits=max_splits)
+      self.assertAllEqual([[], []], sess.run(node_ids_list))
+
   def testMakeStatsSummarySimple(self):
     """Simple test for MakeStatsSummary."""
     with self.test_session():
-- 
GitLab


From 8e918c3d202bb0eed6b423eb78a6ef45629f952e Mon Sep 17 00:00:00 2001
From: RJ Ryan <rjryan@google.com>
Date: Tue, 1 May 2018 12:02:59 -0700
Subject: [PATCH 097/395] Improve shape inference for tf.contrib.signal.frame.

PiperOrigin-RevId: 194972934
---
 .../signal/python/kernel_tests/shape_ops_test.py   |  2 +-
 tensorflow/contrib/signal/python/ops/shape_ops.py  | 14 +++++++-------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensorflow/contrib/signal/python/kernel_tests/shape_ops_test.py b/tensorflow/contrib/signal/python/kernel_tests/shape_ops_test.py
index 64cc8c7ea5..f132050153 100644
--- a/tensorflow/contrib/signal/python/kernel_tests/shape_ops_test.py
+++ b/tensorflow/contrib/signal/python/kernel_tests/shape_ops_test.py
@@ -119,7 +119,7 @@ class FrameTest(test.TestCase):
     frame_step = 1
     result = shape_ops.frame(signal, frame_length, frame_step,
                              pad_end=True, pad_value=99, axis=1)
-    self.assertEqual([1, None, None, 3, 4], result.shape.as_list())
+    self.assertEqual([1, 2, None, 3, 4], result.shape.as_list())
 
     result = shape_ops.frame(signal, frame_length, frame_step,
                              pad_end=False, axis=1)
diff --git a/tensorflow/contrib/signal/python/ops/shape_ops.py b/tensorflow/contrib/signal/python/ops/shape_ops.py
index 1ddc2941ec..91862f0cc0 100644
--- a/tensorflow/contrib/signal/python/ops/shape_ops.py
+++ b/tensorflow/contrib/signal/python/ops/shape_ops.py
@@ -43,13 +43,13 @@ def _infer_frame_shape(signal, frame_length, frame_step, pad_end, axis):
   outer_dimensions = signal_shape[:axis]
   inner_dimensions = signal_shape[axis:][1:]
   if signal_shape and frame_axis is not None:
-    if frame_step and frame_length is not None:
-      if pad_end:
-        # Double negative is so that we round up.
-        num_frames = -(-frame_axis // frame_step)
-      else:
-        num_frames = (frame_axis - frame_length + frame_step) // frame_step
-      num_frames = max(0, num_frames)
+    if frame_step is not None and pad_end:
+      # Double negative is so that we round up.
+      num_frames = max(0, -(-frame_axis // frame_step))
+    elif frame_step is not None and frame_length is not None:
+      assert not pad_end
+      num_frames = max(
+          0, (frame_axis - frame_length + frame_step) // frame_step)
   return outer_dimensions + [num_frames, frame_length] + inner_dimensions
 
 
-- 
GitLab


From 5c18dc63d752af4a810ed70c6aa18d4f7dd2601a Mon Sep 17 00:00:00 2001
From: Benoit Steiner <bsteiner@google.com>
Date: Tue, 1 May 2018 12:17:48 -0700
Subject: [PATCH 098/395] Simplified shape inference.

PiperOrigin-RevId: 194975603
---
 .../core/grappler/costs/graph_properties.cc   | 354 ++++++++----------
 .../core/grappler/costs/graph_properties.h    |  34 +-
 .../grappler/costs/graph_properties_test.cc   |  15 +-
 tensorflow/core/grappler/op_types.cc          |   4 +
 tensorflow/core/grappler/op_types.h           |   2 +-
 tensorflow/core/grappler/utils.cc             |  24 +-
 tensorflow/core/grappler/utils.h              |  19 +-
 .../core/grappler/utils/topological_sort.cc   |  18 +-
 .../core/grappler/utils/topological_sort.h    |   4 +-
 .../grappler/utils/topological_sort_test.cc   |  34 +-
 10 files changed, 251 insertions(+), 257 deletions(-)

diff --git a/tensorflow/core/grappler/costs/graph_properties.cc b/tensorflow/core/grappler/costs/graph_properties.cc
index 313f63149d..a12d9b932b 100644
--- a/tensorflow/core/grappler/costs/graph_properties.cc
+++ b/tensorflow/core/grappler/costs/graph_properties.cc
@@ -256,18 +256,14 @@ typename DisjointSet<Handle>::Rep* DisjointSet<Handle>::Find(Handle value) {
   return root;
 }
 
-bool IsQueue(const NodeDef& node) {
-  return str_util::EndsWith(node.op(), "QueueV2");
+bool IsEnqueue(const NodeDef& n) {
+  return (n.op().find("Enqueue") != std::string::npos &&
+          n.op().find("EnqueueMany") == std::string::npos);
 }
 
-// Returns true if the node is an Enter op AND its input is a Queue.
-bool IsEnterWithQueue(const NodeDef& node, const GraphView& graph) {
-  if (IsEnter(node)) {
-    GraphView::InputPort input(&node, 0);
-    GraphView::OutputPort fanin = graph.GetRegularFanin(input);
-    return IsQueue(*fanin.node);
-  }
-  return false;
+bool IsDequeue(const NodeDef& n) {
+  return (n.op().find("Dequeue") != std::string::npos &&
+          n.op().find("DequeueMany") == std::string::npos);
 }
 
 bool HasAnyUnknownDimensions(const TensorShapeProto& proto) {
@@ -428,7 +424,8 @@ class SymbolicShapeRefiner {
     }
     return it->second.inference_context.get();
   }
-  Status UpdateNode(const NodeDef* node, bool relax, bool* refined) {
+
+  Status UpdateNode(const NodeDef* node, bool* refined) {
     NodeContext* node_context = GetNodeContext(node);
     if (node_context == nullptr) {
       TF_RETURN_IF_ERROR(AddNode(node));
@@ -519,8 +516,12 @@ class SymbolicShapeRefiner {
       }
     }
 
+    // Make sure we schedule the fanout of resources (which have no input)
+    // whenever the resources are updated.
+    *refined |= inference_context->num_inputs() == 0;
+
     if (!*refined) {
-      // No input shape has changed, we're done
+      // No input shape has changed, we're done.
       return Status::OK();
     }
 
@@ -573,51 +574,6 @@ class SymbolicShapeRefiner {
     }
   };
 
-  // Compute the shape of the tensors outputed by node 'node' at output port
-  // 'port_index' as the intersection of shape1 and shape2.
-  ShapeHandle OutputAsIntersection(const NodeDef* node, int port_index,
-                                   ShapeHandle shape1, ShapeHandle shape2) {
-    if (shape1.SameHandle(shape2)) {
-      return shape1;
-    }
-    InferenceContext* ctx = GetContext(node);
-    ShapeHandle merged = shape1;
-    if (!ctx->RankKnown(shape2) && !ctx->RankKnown(shape1)) {
-      // Return either one since they're expected to represent the same value.
-      return shape1;
-    } else if (!ctx->RankKnown(shape2) && ctx->RankKnown(shape1)) {
-      return shape1;
-    } else if (ctx->RankKnown(shape2) && !ctx->RankKnown(shape1)) {
-      return shape2;
-    } else {
-      const int rank = ctx->Rank(shape1);
-      if (ctx->Rank(shape2) != rank) {
-        // We detected an inconsistency, return an unknown shape. This can
-        // happen in the fanout of a merge node since during the initial
-        // propagation we optimistically assume that all the inputs to the merge
-        // node have the same shape.
-        return GetUnknownOutputShape(node, port_index);
-      }
-      for (int d = 0; d < rank; ++d) {
-        if (!ctx->Dim(shape1, d).SameHandle(ctx->Dim(shape2, d))) {
-          if (ctx->Value(ctx->Dim(shape1, d)) !=
-              ctx->Value(ctx->Dim(shape2, d))) {
-            DimensionHandle new_dim;
-            if (ctx->Value(ctx->Dim(shape1, d)) < 0) {
-              new_dim = ctx->Dim(shape2, d);
-            } else if (ctx->Value(ctx->Dim(shape2, d)) < 0) {
-              new_dim = ctx->Dim(shape1, d);
-            } else {
-              new_dim = GetUnknownOutputDim(node, port_index, d);
-            }
-            TF_CHECK_OK(ctx->ReplaceDim(merged, d, new_dim, &merged));
-          }
-        }
-      }
-    }
-    return merged;
-  }
-
   // Compute the shape of the tensors outputed by node 'node' at output port
   // 'port_index' as the union of shape1 and shape2.
   ShapeHandle OutputAsUnion(const NodeDef* node, int port_index,
@@ -822,6 +778,7 @@ class SymbolicShapeRefiner {
         status.Update(SetUnknownShape(&node, output_port));
       }
     }
+
     return status;
   }
 
@@ -884,29 +841,6 @@ class SymbolicShapeManager {
   DisjointSet<shape_inference::DimensionHandle> dims_;
 };
 
-Status GraphProperties::MergeEnqueueShapesAndTypes(
-    SymbolicShapeRefiner* shape_refiner, const NodeDef* qnode,
-    const std::vector<ShapeAndType>& shapes_and_types,
-    std::vector<ShapeAndType>* queue_shapes_and_types) {
-  if (shapes_and_types.size() != queue_shapes_and_types->size()) {
-    return errors::InvalidArgument(
-        "Enqueue nodes mixed number of tensors: ", shapes_and_types.size(),
-        "  vs ", queue_shapes_and_types->size());
-  }
-  for (size_t i = 0; i < shapes_and_types.size(); ++i) {
-    const ShapeAndType& a = shapes_and_types[i];
-    ShapeAndType& b = (*queue_shapes_and_types)[i];
-    if (a.dtype != b.dtype) {
-      return errors::InvalidArgument("Enqueue nodes mixed dtypes for tensor ",
-                                     i, ": ", DataTypeString(a.dtype), " vs ",
-                                     DataTypeString(b.dtype));
-    }
-
-    b.shape = shape_refiner->OutputAsIntersection(qnode, i, a.shape, b.shape);
-  }
-  return Status::OK();
-}
-
 Status GraphProperties::RelaxEnqueueShapesAndMergeTypes(
     SymbolicShapeRefiner* shape_refiner, const NodeDef* qnode,
     const std::vector<ShapeAndType>& shapes_and_types,
@@ -936,7 +870,7 @@ Status GraphProperties::RelaxEnqueueShapesAndMergeTypes(
 // inputs are UnknownShapes. So we need to ignore the input from NextIteration
 // nodes to propagate any known shape from the Merge node.
 Status GraphProperties::UpdateMergeNode(SymbolicShapeRefiner* shape_refiner,
-                                        const NodeDef* node, bool relax,
+                                        const NodeDef* node,
                                         bool* new_shapes) const {
   InferenceContext* c = shape_refiner->GetContext(node);
   if (!c) {
@@ -955,15 +889,8 @@ Status GraphProperties::UpdateMergeNode(SymbolicShapeRefiner* shape_refiner,
   bool out_initialized = false;
   for (const GraphView::Edge fanin :
        shape_refiner->graph().GetFaninEdges(*node, false)) {
-    // Skip back edges during the initial propagation phase. This is equivalent
-    // to assuming that all the inputs to the merge nodes are fed by the same
-    // shape, and will be corrected as needed in the relaxation phase.
-    if (!relax && IsNextIteration(*fanin.src.node)) {
-      continue;
-    }
-
     InferenceContext* in = shape_refiner->GetContext(fanin.src.node);
-    if (!relax && !in) {
+    if (!in) {
       // Handling a loop for the first time, the back edge won't have any shape
       // info.
       continue;
@@ -976,11 +903,7 @@ Status GraphProperties::UpdateMergeNode(SymbolicShapeRefiner* shape_refiner,
       out = input;
       continue;
     }
-    if (relax) {
-      out = shape_refiner->OutputAsUnion(node, 0, input, out);
-    } else {
-      out = shape_refiner->OutputAsIntersection(node, 0, input, out);
-    }
+    out = shape_refiner->OutputAsUnion(node, 0, input, out);
   }
 
   if (*new_shapes || !shape_refiner->EquivalentShapes(out, c->output(0))) {
@@ -994,11 +917,10 @@ Status GraphProperties::UpdateMergeNode(SymbolicShapeRefiner* shape_refiner,
 // Manually propagate the input shape for Enter nodes and update any Merge node
 // outputs.
 Status GraphProperties::UpdateEnter(SymbolicShapeRefiner* shape_refiner,
-                                    const NodeDef* node, bool relax,
-                                    bool* new_shapes) {
+                                    const NodeDef* node, bool* new_shapes) {
   auto enter_ctx = shape_refiner->GetContext(node);
   if (!enter_ctx) {
-    TF_RETURN_IF_ERROR(shape_refiner->UpdateNode(node, relax, new_shapes));
+    TF_RETURN_IF_ERROR(shape_refiner->UpdateNode(node, new_shapes));
     enter_ctx = shape_refiner->GetContext(node);
   }
 
@@ -1012,53 +934,54 @@ Status GraphProperties::UpdateEnter(SymbolicShapeRefiner* shape_refiner,
     enter_ctx->set_output(0, input);
     *new_shapes = true;
   }
+  auto* outputs = in->output_handle_shapes_and_types(fanin.port_id);
+  if (outputs) {
+    enter_ctx->set_input_handle_shapes_and_types(0, *outputs);
+    enter_ctx->set_output_handle_shapes_and_types(0, *outputs);
+    *new_shapes = true;
+  }
   return Status::OK();
 }
 
-Status GraphProperties::UpdateShapes(SymbolicShapeRefiner* shape_refiner,
-                                     bool relax, const NodeDef* n,
-                                     bool* new_shapes) const {
+Status GraphProperties::UpdateShapes(
+    SymbolicShapeRefiner* shape_refiner,
+    const std::unordered_map<const NodeDef*, const NodeDef*>& resource_handles,
+    const NodeDef* n, bool* new_shapes) const {
   if (IsEnter(*n)) {
     // The Enter shape function always forwards an UnknownShape, so do the right
     // thing here.
-    TF_RETURN_IF_ERROR(UpdateEnter(shape_refiner, n, relax, new_shapes));
+    TF_RETURN_IF_ERROR(UpdateEnter(shape_refiner, n, new_shapes));
   } else if (IsMerge(*n)) {
     // Properly handle merge nodes.
-    TF_RETURN_IF_ERROR(UpdateMergeNode(shape_refiner, n, relax, new_shapes));
+    TF_RETURN_IF_ERROR(UpdateMergeNode(shape_refiner, n, new_shapes));
+  } else if (IsEnqueue(*n)) {
+    TF_RETURN_IF_ERROR(
+        UpdateEnqueue(n, resource_handles, shape_refiner, new_shapes));
   } else {
     // Rely on regular TF shape refinement for all the other nodes.
-    bool updated = false;
-    TF_RETURN_IF_ERROR(shape_refiner->UpdateNode(n, relax, &updated));
-    if (updated) {
-      // We want to avoid propagating through loops on the merge pass because
-      // the shapes are not guaranteed to converge.
-      if (relax || !IsNextIteration(*n)) {
-        *new_shapes = true;
-      }
-    }
+    TF_RETURN_IF_ERROR(shape_refiner->UpdateNode(n, new_shapes));
   }
   return Status::OK();
 }
 
 // Propagates the shapes in the transitive fan-out of <new_shapes>.
 Status GraphProperties::PropagateShapes(
-    SymbolicShapeRefiner* shape_refiner, bool relax, TopoQueue* new_shapes,
-    const std::unordered_map<const NodeDef*,
-                             std::unordered_set<const NodeDef*>>& resources,
+    SymbolicShapeRefiner* shape_refiner, TopoQueue* new_shapes,
+    const std::unordered_map<const NodeDef*, const NodeDef*>& resource_handles,
     int num_loops) const {
   // Limit the number of iterations to prevent infinite loops in the presence of
   // incorrect shape functions. The algoritm should converge in at most
   // num_nested_loops^2 * max_rank. We approximate max_rank with the constant 4.
   // The same applies to resources.
-  VLOG(1) << "Propagating (relax=" << relax << ") " << new_shapes->size()
-          << " new shapes through " << num_loops << " loops and "
-          << resources.size() << " resources" << std::endl;
+  VLOG(1) << "Propagating " << new_shapes->size() << " new shapes through "
+          << num_loops << " loops and " << resource_handles.size()
+          << " resources" << std::endl;
 
   const int64 max_loop_length = item_.graph.node_size();
   const int64 max_rank = 4;
   const int64 max_loop_iterations =
       max_rank * max_loop_length * std::max<int64>(1, num_loops * num_loops);
-  const int64 num_queues = resources.size();
+  const int64 num_queues = resource_handles.size();
   const int64 max_resource_iterations = num_queues * num_queues * max_rank;
 
   int64 num_resource_iterations = 0;
@@ -1068,22 +991,22 @@ Status GraphProperties::PropagateShapes(
            num_loop_iterations++ < max_loop_iterations) {
       const NodeDef* n = new_shapes->pop();
       bool updated = false;
-      TF_RETURN_IF_ERROR(UpdateShapes(shape_refiner, relax, n, &updated));
+      TF_RETURN_IF_ERROR(
+          UpdateShapes(shape_refiner, resource_handles, n, &updated));
       if (updated) {
-        for (const GraphView::InputPort fanout :
+        for (const GraphView::InputPort& fanout :
              shape_refiner->graph().GetFanouts(*n, false)) {
           new_shapes->push(fanout.node);
         }
+        // Make sure the corresponding queue nodes are (re)processed.
+        if (IsEnqueue(*n)) {
+          auto it = resource_handles.find(n);
+          if (it != resource_handles.end()) {
+            new_shapes->push(it->second);
+          }
+        }
       }
     }
-
-    for (const auto& resource : resources) {
-      // Resources need special handling: since the enqueue nodes are in the
-      // fanout of the queues, we need to manually propagate the shapes from
-      // enqueue node to the corresponding queue.
-      TF_RETURN_IF_ERROR(UpdateResource(resource.first, resource.second,
-                                        shape_refiner, new_shapes));
-    }
   } while (!new_shapes->empty() &&
            num_resource_iterations++ < max_resource_iterations);
 
@@ -1094,54 +1017,48 @@ Status GraphProperties::PropagateShapes(
   return Status::OK();
 }
 
-Status GraphProperties::UpdateResource(
-    const NodeDef* qnode,
-    const std::unordered_set<const NodeDef*>& queue_inputs,
-    SymbolicShapeRefiner* shape_refiner, TopoQueue* new_shapes) {
-  // Proceed only if qnode is a queue or an Enter with queue input.
-  if (!IsQueue(*qnode) && !IsEnterWithQueue(*qnode, shape_refiner->graph())) {
+Status GraphProperties::UpdateEnqueue(
+    const NodeDef* enqueue_node,
+    const std::unordered_map<const NodeDef*, const NodeDef*>& resource_handles,
+    SymbolicShapeRefiner* shape_refiner, bool* new_shapes) {
+  auto ctx = shape_refiner->GetNodeContext(enqueue_node);
+  if (!ctx) {
+    TF_RETURN_IF_ERROR(shape_refiner->AddNode(enqueue_node));
+    ctx = CHECK_NOTNULL(shape_refiner->GetNodeContext(enqueue_node));
+  }
+
+  auto it = resource_handles.find(enqueue_node);
+  if (it == resource_handles.end()) {
+    // The corresponding queue was not found, there isn't much we can do.
     return Status::OK();
   }
+  const NodeDef* qnode = it->second;
   auto qctx = shape_refiner->GetContext(qnode);
   if (!qctx) {
     return Status::OK();
   }
   auto* queue_handle_data = qctx->output_handle_shapes_and_types(0);
 
-  // Merge all inputs into the enqueue node, regardless of which phase we
-  // are in.
-  std::vector<ShapeAndType> queue_shapes_and_types;
-  for (const auto& node : queue_inputs) {
-    auto ctx = shape_refiner->GetNodeContext(node);
-    if (!ctx) {
-      continue;
-    }
-    // TODO(bsteiner): handle EnqueueMany as well.
-    if (node->op().find("Enqueue") != std::string::npos &&
-        node->op().find("EnqueueMany") == std::string::npos) {
-      std::vector<ShapeAndType> shapes_and_types;
-      for (int i = 1; i < ctx->input_types.size(); ++i) {
-        shapes_and_types.push_back(
-            {ctx->inference_context->input(i), ctx->input_types[i]});
-      }
-      if (queue_shapes_and_types.empty()) {
-        queue_shapes_and_types = shapes_and_types;
-      } else {
-        TF_RETURN_IF_ERROR(RelaxEnqueueShapesAndMergeTypes(
-            shape_refiner, qnode, shapes_and_types, &queue_shapes_and_types));
-      }
-    }
+  // TODO(bsteiner): handle EnqueueMany as well.
+  std::vector<ShapeAndType> shapes_and_types;
+  for (int i = 1; i < ctx->input_types.size(); ++i) {
+    GraphView::InputPort inp(enqueue_node, i);
+    GraphView::OutputPort fanin = shape_refiner->graph().GetRegularFanin(inp);
+    InferenceContext* in = shape_refiner->GetContext(fanin.node);
+    ShapeHandle input = in->output(fanin.port_id);
+    ctx->inference_context->SetInput(i, input);
+    shapes_and_types.push_back({input, ctx->input_types[i]});
   }
 
-  if (queue_handle_data == nullptr ||
-      !shape_refiner->EquivalentShapesAndTypes(*queue_handle_data,
-                                               queue_shapes_and_types)) {
-    qctx->set_output_handle_shapes_and_types(0, queue_shapes_and_types);
-
-    for (const GraphView::InputPort fanout :
-         shape_refiner->graph().GetFanouts(*qnode, false)) {
-      new_shapes->push(fanout.node);
-    }
+  if (queue_handle_data == nullptr) {
+    qctx->set_output_handle_shapes_and_types(0, shapes_and_types);
+    *new_shapes = true;
+  } else {
+    TF_RETURN_IF_ERROR(RelaxEnqueueShapesAndMergeTypes(
+        shape_refiner, qnode, *queue_handle_data, &shapes_and_types));
+    *new_shapes |= !shape_refiner->EquivalentShapesAndTypes(*queue_handle_data,
+                                                            shapes_and_types);
+    qctx->set_output_handle_shapes_and_types(0, shapes_and_types);
   }
 
   return Status::OK();
@@ -1159,75 +1076,96 @@ Status GraphProperties::InferStatically(bool assume_valid_feeds) {
     }
   }
 
-  std::unordered_map<const NodeDef*, int> topo_order;
-  TF_RETURN_IF_ERROR(ComputeTopologicalOrder(item_.graph, &topo_order));
-
   GraphView graph_view(&item_.graph);
 
   // List the resources and the nodes using them. Also collect the Merge nodes,
   // fed nodes, and primary inputs.
-  std::unordered_map<const NodeDef*, std::unordered_set<const NodeDef*>>
+  std::unordered_map<const NodeDef*,
+                     std::pair<std::unordered_set<const NodeDef*>,
+                               std::unordered_set<const NodeDef*>>>
       resources;
   std::unordered_set<const NodeDef*> merge_nodes;
   std::unordered_set<const NodeDef*> fed_nodes;
   std::unordered_set<const NodeDef*> primary_inputs;
   int num_loops = 0;
   for (const NodeDef& node : item_.graph.node()) {
+    if (IsQueue(node)) {
+      for (const GraphView::InputPort& fanout :
+           graph_view.GetFanouts(node, false)) {
+        if (IsEnter(*fanout.node)) {
+          const NodeDef& enter = *fanout.node;
+          for (const GraphView::InputPort& fanout :
+               graph_view.GetFanouts(enter, false)) {
+            if (IsEnqueue(*fanout.node)) {
+              resources[&node].first.insert(fanout.node);
+            } else if (IsDequeue(*fanout.node)) {
+              resources[&node].second.insert(fanout.node);
+            }
+          }
+        } else {
+          if (IsEnqueue(*fanout.node)) {
+            resources[&node].first.insert(fanout.node);
+          } else if (IsDequeue(*fanout.node)) {
+            resources[&node].second.insert(fanout.node);
+          }
+        }
+      }
+    }
     if (NumNonControlInputs(node) == 0) {
       primary_inputs.insert(&node);
     } else if (IsMerge(node)) {
       merge_nodes.insert(&node);
     } else if (IsNextIteration(node)) {
       ++num_loops;
-    } else {
-      const OpRegistrationData* op_data;
-      TF_RETURN_IF_ERROR(function_library.LookUp(node.op(), &op_data));
-      DataTypeVector input_types;
-      DataTypeVector output_types;
-      TF_RETURN_IF_ERROR(InOutTypesForNode(node, op_data->op_def, &input_types,
-                                           &output_types));
-      for (int i = 0; i < input_types.size(); ++i) {
-        if (input_types[i] == DataType::DT_RESOURCE) {
-          GraphView::InputPort input(&node, i);
-          const GraphView::OutputPort resource =
-              graph_view.GetRegularFanin(input);
-          resources[resource.node].insert(&node);
-        }
-      }
     }
     if (fed_ports.find(node.name()) != fed_ports.end()) {
       fed_nodes.insert(&node);
     }
   }
 
-  SymbolicShapeRefiner refiner(graph_view, fed_ports);
-
-  // We propagate shapes through the graph in two phases. In the first phase, we
-  // exclusively merge shapes but we do not propagate shapes through the
-  // backedge of loops (i.e. the NextIteration node). Then on the second phase,
-  // we exclusively relax shapes and propagate shapes through loops until
-  // reaching fixed point.
-  for (int relax = 0; relax < 2; relax++) {
-    TopoQueue new_shapes(topo_order);
-    // Seed the propagation of shapes through merge nodes.
-    if (relax) {
-      for (const NodeDef* node : merge_nodes) {
-        new_shapes.push(node);
+  std::unordered_map<const NodeDef*, const NodeDef*> resource_handles;
+  std::vector<std::pair<const NodeDef*, const NodeDef*>> extra_deps;
+  for (const auto& resource : resources) {
+    for (const NodeDef* src : resource.second.first) {
+      resource_handles[src] = resource.first;
+      for (const NodeDef* tgt : resource.second.second) {
+        // Add control edges from enqueue to dequeue nodes to ensure they are
+        // processed in their logical order.
+        extra_deps.emplace_back(src, tgt);
       }
     }
-    // Also seed the propagation of shapes in the fanout of primary inputs.
-    for (const NodeDef* node : primary_inputs) {
-      new_shapes.push(node);
-    }
-    // Also seed the propagation of shapes in the fanout of fed nodes.
-    for (const NodeDef* node : fed_nodes) {
-      new_shapes.push(node);
+  }
+
+  std::unordered_map<const NodeDef*, int> topo_order;
+  Status s = ComputeTopologicalOrder(item_.graph, &topo_order, &extra_deps);
+  if (!s.ok()) {
+    if (extra_deps.empty()) {
+      return s;
+    } else {
+      // There is a loop between queues: we'll just use the graph topological
+      // order. This will make the shape inference less precise but since this
+      // isn't common it's not worth to figure out where to break the loop and
+      // do a proper relaxation.
+      TF_RETURN_IF_ERROR(
+          ComputeTopologicalOrder(item_.graph, &topo_order, nullptr));
     }
-    // Propagate shapes normally.
-    TF_RETURN_IF_ERROR(
-        PropagateShapes(&refiner, relax, &new_shapes, resources, num_loops));
   }
 
+  SymbolicShapeRefiner refiner(graph_view, fed_ports);
+
+  TopoQueue new_shapes(topo_order);
+  // Also seed the propagation of shapes in the fanout of primary inputs.
+  for (const NodeDef* node : primary_inputs) {
+    new_shapes.push(node);
+  }
+  // Also seed the propagation of shapes in the fanout of fed nodes.
+  for (const NodeDef* node : fed_nodes) {
+    new_shapes.push(node);
+  }
+  // Propagate shapes normally.
+  TF_RETURN_IF_ERROR(
+      PropagateShapes(&refiner, &new_shapes, resource_handles, num_loops));
+
   // Track shapes globally across the graph.
   SymbolicShapeManager shape_manager;
   bool found_error = false;
diff --git a/tensorflow/core/grappler/costs/graph_properties.h b/tensorflow/core/grappler/costs/graph_properties.h
index 7d685b5833..ecc10fddb8 100644
--- a/tensorflow/core/grappler/costs/graph_properties.h
+++ b/tensorflow/core/grappler/costs/graph_properties.h
@@ -75,12 +75,6 @@ class GraphProperties {
   void ClearOutputProperties(const string& node_name);
 
  private:
-  // Merges shapes <shapes_and_types>, determined from an EnqueueV2 node, into
-  // <*queue_shapes_and_types>.
-  static Status MergeEnqueueShapesAndTypes(
-      SymbolicShapeRefiner* shape_refiner, const NodeDef* qnode,
-      const std::vector<shape_inference::ShapeAndType>& shapes_and_types,
-      std::vector<shape_inference::ShapeAndType>* queue_shapes_and_types);
   // Relaxes shapes <shapes_and_types>, determined from an EnqueueV2 node, into
   // <*queue_shapes_and_types>.
   static Status RelaxEnqueueShapesAndMergeTypes(
@@ -88,31 +82,33 @@ class GraphProperties {
       const std::vector<shape_inference::ShapeAndType>& shapes_and_types,
       std::vector<shape_inference::ShapeAndType>* queue_shapes_and_types);
 
-  // Update the shapes for qnode. If output shapes of qnode have changed,
-  // enqueue its fanout in 'new_shapes'.
-  static Status UpdateResource(
-      const NodeDef* qnode,
-      const std::unordered_set<const NodeDef*>& queue_inputs,
-      SymbolicShapeRefiner* shape_refiner, TopoQueue* new_shapes);
+  // Update the shapes of the enqueue node, port them over to the corresponding
+  // queue, and schedule the reprocessing of the queue if needed.
+  static Status UpdateEnqueue(
+      const NodeDef* enqueue_node,
+      const std::unordered_map<const NodeDef*, const NodeDef*>&
+          resource_handles,
+      SymbolicShapeRefiner* shape_refiner, bool* new_shapes);
 
   // Update the output shapes of a Merge node, and enqueue its fanout in
   // new_shapes if needed.
   Status UpdateMergeNode(SymbolicShapeRefiner* shape_refiner,
-                         const NodeDef* node, bool relax,
-                         bool* new_shapes) const;
+                         const NodeDef* node, bool* new_shapes) const;
   // Process the Enter node, and enqueue its fanout in new_shapes if needed.
   static Status UpdateEnter(SymbolicShapeRefiner* shape_refiner,
-                            const NodeDef* node, bool relax, bool* new_shapes);
+                            const NodeDef* node, bool* new_shapes);
   // Update the shapes for node 'n'. If output shapes for n have changed,
   // enqueue its fanout in 'new_shapes'.
-  Status UpdateShapes(SymbolicShapeRefiner* shape_refiner, bool relax,
+  Status UpdateShapes(SymbolicShapeRefiner* shape_refiner,
+                      const std::unordered_map<const NodeDef*, const NodeDef*>&
+                          resource_handles,
                       const NodeDef* n, bool* new_shapes) const;
   // Propagate the shapes for the nodes enqueued in new_shapes and their
   // transitive fanout until a fixed point is reached.
   Status PropagateShapes(
-      SymbolicShapeRefiner* shape_refiner, bool relax, TopoQueue* new_shapes,
-      const std::unordered_map<const NodeDef*,
-                               std::unordered_set<const NodeDef*>>& resources,
+      SymbolicShapeRefiner* shape_refiner, TopoQueue* new_shapes,
+      const std::unordered_map<const NodeDef*, const NodeDef*>&
+          resource_handles,
       int num_loops) const;
 
   // Data members
diff --git a/tensorflow/core/grappler/costs/graph_properties_test.cc b/tensorflow/core/grappler/costs/graph_properties_test.cc
index afe334dfa2..a53f6414c3 100644
--- a/tensorflow/core/grappler/costs/graph_properties_test.cc
+++ b/tensorflow/core/grappler/costs/graph_properties_test.cc
@@ -282,20 +282,11 @@ TEST_F(GraphPropertiesTest, Queues) {
   auto dequeue2 =
       ops::QueueDequeue(root.WithOpName("Dequeue2"), q2, {DataType::DT_FLOAT});
 
-  // Create a queue that feeds itself.
-  auto q3 =
-      ops::RandomShuffleQueue(root.WithOpName("Queue3"), {DataType::DT_FLOAT});
-  auto dequeue3 =
-      ops::QueueDequeue(root.WithOpName("Dequeue3"), q3, {DataType::DT_FLOAT});
-  auto merge3 = ops::Merge(root.WithOpName("Merge3"), {dequeue3[0], square2});
-  auto enqueue3 =
-      ops::QueueEnqueue(root.WithOpName("Enqueue3"), q3, {merge3.output});
-
   auto q4 =
       ops::RandomShuffleQueue(root.WithOpName("Queue4"), {DataType::DT_FLOAT});
   auto enqueue4 = ops::QueueEnqueue(root.WithOpName("Enqueue4"), q4, {square2});
   auto enqueue4_2 =
-      ops::QueueEnqueue(root.WithOpName("Enqueue4_2"), q4, {dequeue3[0]});
+      ops::QueueEnqueue(root.WithOpName("Enqueue4_2"), q4, {dequeue2[0]});
   auto dequeue4 =
       ops::QueueDequeue(root.WithOpName("Dequeue4"), q4, {DataType::DT_FLOAT});
 
@@ -327,10 +318,6 @@ TEST_F(GraphPropertiesTest, Queues) {
   ASSERT_EQ(1, props2.size());
   EXPECT_EQ("float: [3,7]", PropToString(props2[0]));
 
-  const auto props3 = properties.GetOutputProperties("Dequeue3");
-  ASSERT_EQ(1, props3.size());
-  EXPECT_EQ("float: [3,7]", PropToString(props3[0]));
-
   // The dequeue3 op shape is unknown. The square2 op shape is known. Verify
   // that we merge the 2 properly to determine the shape of the data coming out
   // of the queue.
diff --git a/tensorflow/core/grappler/op_types.cc b/tensorflow/core/grappler/op_types.cc
index 7a89c26374..839b0bbfc9 100644
--- a/tensorflow/core/grappler/op_types.cc
+++ b/tensorflow/core/grappler/op_types.cc
@@ -250,6 +250,10 @@ bool IsPrint(const NodeDef& node) { return node.op() == "Print"; }
 
 bool IsProd(const NodeDef& node) { return node.op() == "Prod"; }
 
+bool IsQueue(const NodeDef& node) {
+  return str_util::EndsWith(node.op(), "QueueV2");
+}
+
 bool IsRandomShuffle(const NodeDef& node) {
   return node.op() == "RandomShuffle";
 }
diff --git a/tensorflow/core/grappler/op_types.h b/tensorflow/core/grappler/op_types.h
index 976d23e527..bd8d3a44e4 100644
--- a/tensorflow/core/grappler/op_types.h
+++ b/tensorflow/core/grappler/op_types.h
@@ -21,7 +21,6 @@ limitations under the License.
 
 namespace tensorflow {
 namespace grappler {
-
 bool IsAdd(const NodeDef& node);
 bool IsAddN(const NodeDef& node);
 bool IsAll(const NodeDef& node);
@@ -98,6 +97,7 @@ bool IsPolygamma(const NodeDef& node);
 bool IsPrint(const NodeDef& node);
 bool IsProd(const NodeDef& node);
 bool IsPow(const NodeDef& node);
+bool IsQueue(const NodeDef& node);
 bool IsRandomShuffle(const NodeDef& node);
 bool IsReal(const NodeDef& node);
 bool IsRealDiv(const NodeDef& node);
diff --git a/tensorflow/core/grappler/utils.cc b/tensorflow/core/grappler/utils.cc
index 7398d2c896..6db6d71447 100644
--- a/tensorflow/core/grappler/utils.cc
+++ b/tensorflow/core/grappler/utils.cc
@@ -361,8 +361,11 @@ inline void STLSortAndRemoveDuplicates(T* v) {
 }
 }  // namespace
 
-Status SimpleGraphView::Initialize(const GraphDef& graph, bool dedup_inputs,
-                                   bool dedup_outputs) {
+Status SimpleGraphView::Initialize(
+    const GraphDef& graph,
+    const std::vector<std::pair<const NodeDef*, const NodeDef*>>*
+        extra_dependencies,
+    bool dedup_inputs, bool dedup_outputs) {
   graph_ = &graph;
   const int num_nodes = graph.node_size();
   inputs_.clear();
@@ -381,6 +384,23 @@ Status SimpleGraphView::Initialize(const GraphDef& graph, bool dedup_inputs,
     index_to_name_.push_back(node.name());
   }
 
+  if (extra_dependencies) {
+    for (const auto& dep : *extra_dependencies) {
+      auto itr_src = name_to_index_.find(dep.first->name());
+      if (itr_src == name_to_index_.end()) {
+        return errors::InvalidArgument("Non-existent src ", dep.first->name());
+      }
+      auto itr_tgt = name_to_index_.find(dep.second->name());
+      if (itr_tgt == name_to_index_.end()) {
+        return errors::InvalidArgument("Non-existent tgt ", dep.second->name());
+      }
+      const int src_idx = itr_src->second;
+      const int tgt_idx = itr_tgt->second;
+      inputs_[tgt_idx].push_back(src_idx);
+      outputs_[src_idx].push_back(tgt_idx);
+    }
+  }
+
   // Build forward and reverse adjacency lists.
   for (int node_idx = 0; node_idx < num_nodes; ++node_idx) {
     const NodeDef& node = graph.node(node_idx);
diff --git a/tensorflow/core/grappler/utils.h b/tensorflow/core/grappler/utils.h
index 54cb26bafa..15f6b367b0 100644
--- a/tensorflow/core/grappler/utils.h
+++ b/tensorflow/core/grappler/utils.h
@@ -211,11 +211,24 @@ Status SetTensorValue(DataType dtype, int value, Tensor* tensor);
 
 class SimpleGraphView {
  public:
+  // Build a graph view for the specified graphdef.
   Status Initialize(const GraphDef& graph) {
-    return Initialize(graph, true, true);
+    return Initialize(graph, nullptr, true, true);
   }
-  Status Initialize(const GraphDef& graph, bool dedup_inputs,
-                    bool dedup_outputs);
+  // Build a graph view for the specified graphdef augmented with the additional
+  // edges specified in 'extra_dependencies' if any. Note that
+  // extra_dependencies can be null.
+  Status Initialize(
+      const GraphDef& graph,
+      const std::vector<std::pair<const NodeDef*, const NodeDef*>>*
+          extra_dependencies) {
+    return Initialize(graph, extra_dependencies, true, true);
+  }
+  Status Initialize(
+      const GraphDef& graph,
+      const std::vector<std::pair<const NodeDef*, const NodeDef*>>*
+          extra_dependencies,
+      bool dedup_inputs, bool dedup_outputs);
 
   const GraphDef* graph() const { return graph_; }
   inline int num_nodes() const { return index_to_name_.size(); }
diff --git a/tensorflow/core/grappler/utils/topological_sort.cc b/tensorflow/core/grappler/utils/topological_sort.cc
index a8e464d09d..ff89035902 100644
--- a/tensorflow/core/grappler/utils/topological_sort.cc
+++ b/tensorflow/core/grappler/utils/topological_sort.cc
@@ -26,10 +26,12 @@ namespace grappler {
 
 // Kahn's algorithm is implemented.
 // For details, see https://en.wikipedia.org/wiki/Topological_sorting
-Status ComputeTopologicalOrder(const GraphDef& graph,
-                               std::vector<int>* ready_nodes) {
+Status ComputeTopologicalOrder(
+    const GraphDef& graph, std::vector<int>* ready_nodes,
+    const std::vector<std::pair<const NodeDef*, const NodeDef*>>*
+        extra_dependencies) {
   SimpleGraphView graph_view;
-  TF_RETURN_IF_ERROR(graph_view.Initialize(graph));
+  TF_RETURN_IF_ERROR(graph_view.Initialize(graph, extra_dependencies));
 
   ready_nodes->reserve(graph_view.num_nodes());
 
@@ -70,10 +72,12 @@ Status ComputeTopologicalOrder(const GraphDef& graph,
 }
 
 Status ComputeTopologicalOrder(
-    const GraphDef& graph,
-    std::unordered_map<const NodeDef*, int>* topo_order) {
+    const GraphDef& graph, std::unordered_map<const NodeDef*, int>* topo_order,
+    const std::vector<std::pair<const NodeDef*, const NodeDef*>>*
+        extra_dependencies) {
   std::vector<int> ready_nodes;
-  TF_RETURN_IF_ERROR(ComputeTopologicalOrder(graph, &ready_nodes));
+  TF_RETURN_IF_ERROR(
+      ComputeTopologicalOrder(graph, &ready_nodes, extra_dependencies));
   topo_order->reserve(graph.node_size());
   for (int i = 0; i < ready_nodes.size(); ++i) {
     (*topo_order)[&graph.node(ready_nodes[i])] = i;
@@ -83,7 +87,7 @@ Status ComputeTopologicalOrder(
 
 Status TopologicalSort(GraphDef* graph) {
   std::vector<int> ready_nodes;
-  TF_RETURN_IF_ERROR(ComputeTopologicalOrder(*graph, &ready_nodes));
+  TF_RETURN_IF_ERROR(ComputeTopologicalOrder(*graph, &ready_nodes, nullptr));
   PermuteNodesInPlace(graph, &ready_nodes, /*invert_permutation=*/true);
   return Status::OK();
 }
diff --git a/tensorflow/core/grappler/utils/topological_sort.h b/tensorflow/core/grappler/utils/topological_sort.h
index 668c88dc75..bc0299a7b8 100644
--- a/tensorflow/core/grappler/utils/topological_sort.h
+++ b/tensorflow/core/grappler/utils/topological_sort.h
@@ -24,7 +24,9 @@ namespace grappler {
 
 // Compute a topological ordering for the graph nodes.
 Status ComputeTopologicalOrder(
-    const GraphDef& graph, std::unordered_map<const NodeDef*, int>* topo_order);
+    const GraphDef& graph, std::unordered_map<const NodeDef*, int>* topo_order,
+    const std::vector<std::pair<const NodeDef*, const NodeDef*>>*
+        extra_dependencies);
 
 // Sort a graph in topological order.
 Status TopologicalSort(GraphDef* graph);
diff --git a/tensorflow/core/grappler/utils/topological_sort_test.cc b/tensorflow/core/grappler/utils/topological_sort_test.cc
index f5c95009d2..48b7eb50bd 100644
--- a/tensorflow/core/grappler/utils/topological_sort_test.cc
+++ b/tensorflow/core/grappler/utils/topological_sort_test.cc
@@ -53,7 +53,7 @@ TEST_F(TopologicalSortTest, NoLoop) {
   *graph.add_node() = CreateNode("4", {});
 
   std::unordered_map<const NodeDef*, int> topo_order;
-  TF_EXPECT_OK(ComputeTopologicalOrder(graph, &topo_order));
+  TF_EXPECT_OK(ComputeTopologicalOrder(graph, &topo_order, nullptr));
 
   const std::vector<string> order = {"5", "4", "2", "0", "3", "1"};
   for (const auto& topo : topo_order) {
@@ -80,7 +80,7 @@ TEST_F(TopologicalSortTest, WithLoop) {
   *graph.add_node() = CreateNode("1", {});
 
   std::unordered_map<const NodeDef*, int> topo_order;
-  TF_EXPECT_OK(ComputeTopologicalOrder(graph, &topo_order));
+  TF_EXPECT_OK(ComputeTopologicalOrder(graph, &topo_order, nullptr));
 
   const std::vector<string> order = {"1", "2", "3", "4", "5"};
   for (const auto& topo : topo_order) {
@@ -143,6 +143,36 @@ TEST_F(TopologicalSortTest, Idempotent) {
   }
 }
 
+TEST_F(TopologicalSortTest, ExtraDependencies) {
+  GraphDef graph;
+  *graph.add_node() = CreateNode("2", {"5"});
+  *graph.add_node() = CreateNode("0", {"5", "4"});
+  *graph.add_node() = CreateNode("1", {"4", "3"});
+  *graph.add_node() = CreateNode("3", {"2"});
+  *graph.add_node() = CreateNode("5", {});
+  *graph.add_node() = CreateNode("4", {});
+
+  // Add an edge from 4 to 5.
+  std::vector<std::pair<const NodeDef*, const NodeDef*>> extra_dependencies;
+  extra_dependencies.emplace_back(&graph.node(5), &graph.node(4));
+
+  std::unordered_map<const NodeDef*, int> topo_order;
+  TF_EXPECT_OK(
+      ComputeTopologicalOrder(graph, &topo_order, &extra_dependencies));
+
+  const std::vector<string> order = {"4", "5", "2", "0", "3", "1"};
+  for (const auto& topo : topo_order) {
+    const string& node_name = topo.first->name();
+    const int topo_order = topo.second;
+    EXPECT_EQ(node_name, order[topo_order]);
+  }
+
+  // Add an edge from 0 to 4. This will create a loop
+  extra_dependencies.emplace_back(&graph.node(1), &graph.node(5));
+  EXPECT_FALSE(
+      ComputeTopologicalOrder(graph, &topo_order, &extra_dependencies).ok());
+}
+
 }  // namespace
 }  // namespace grappler
 }  // namespace tensorflow
-- 
GitLab


From 59677dc14f5ed28e5d858abc318b1a492f37425f Mon Sep 17 00:00:00 2001
From: Yuefeng Zhou <yuefengz@google.com>
Date: Tue, 1 May 2018 12:24:38 -0700
Subject: [PATCH 099/395] Add device_util.resolve method which merges with
 current device as well.

PiperOrigin-RevId: 194976633
---
 .../distribute/python/cross_tower_ops.py      |  7 +-
 .../distribute/python/mirrored_strategy.py    |  1 -
 tensorflow/python/BUILD                       |  1 +
 tensorflow/python/training/device_util.py     | 27 +++++-
 .../python/training/device_util_test.py       | 89 +++++++++++++++++++
 tensorflow/python/training/distribute.py      |  3 +
 6 files changed, 122 insertions(+), 6 deletions(-)
 create mode 100644 tensorflow/python/training/device_util_test.py

diff --git a/tensorflow/contrib/distribute/python/cross_tower_ops.py b/tensorflow/contrib/distribute/python/cross_tower_ops.py
index cff717db80..c6a1bf6a9f 100644
--- a/tensorflow/contrib/distribute/python/cross_tower_ops.py
+++ b/tensorflow/contrib/distribute/python/cross_tower_ops.py
@@ -53,15 +53,14 @@ def _validate_value_destination_pairs(value_destination_pairs):
   return True
 
 
+# TODO(yuefengz): consider calling this function in the caller of CrossTowerOps.
 def _get_devices_from(destinations):
   if isinstance(destinations, value_lib.DistributedValues):
     return list(destinations.devices)
   elif isinstance(destinations, six.string_types):
-    return [device_util.canonicalize(destinations)]
+    return [device_util.resolve(destinations)]
   else:
-    return [
-        device_util.canonicalize(destination) for destination in destinations
-    ]
+    return [device_util.resolve(destination) for destination in destinations]
 
 
 def _devices_match(left, right):
diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy.py b/tensorflow/contrib/distribute/python/mirrored_strategy.py
index 6efd578a77..2e57b02583 100644
--- a/tensorflow/contrib/distribute/python/mirrored_strategy.py
+++ b/tensorflow/contrib/distribute/python/mirrored_strategy.py
@@ -321,7 +321,6 @@ class MirroredStrategy(distribute_lib.DistributionStrategy):
 
   def _fetch(self, val, destination, fn):
     """Return a copy of `val` or `fn(val)` on `destination`."""
-    assert isinstance(destination, six.string_types)
     if isinstance(val, values.TowerLocalVariable):
       val = self.reduce(val.reduce_method, val, destinations=destination)
       with ops.device(destination):
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 44d9147bb6..087b89b125 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -4032,6 +4032,7 @@ cuda_py_tests(
         "training/basic_loops_test.py",
         "training/coordinator_test.py",
         "training/device_setter_test.py",
+        "training/device_util_test.py",
         "training/ftrl_test.py",
         "training/gradient_descent_test.py",
         "training/learning_rate_decay_test.py",
diff --git a/tensorflow/python/training/device_util.py b/tensorflow/python/training/device_util.py
index f1137e80ab..e31fa02d60 100644
--- a/tensorflow/python/training/device_util.py
+++ b/tensorflow/python/training/device_util.py
@@ -23,17 +23,42 @@ from tensorflow.python.framework import device as tf_device
 from tensorflow.python.framework import ops
 
 
-def canonicalize(d):
+def canonicalize(d, default=None):
+  """Canonicalize device string.
+
+  If d has missing components, the rest would be deduced from the `default`
+  argument or from '/job:localhost/replica:0/task:0/device:CPU:0'. For example:
+    If d = '/cpu:0', default='/job:worker/task:1', it returns
+      '/job:worker/replica:0/task:1/device:CPU:0'.
+    If d = '/cpu:0', default='/job:worker', it returns
+      '/job:worker/replica:0/task:0/device:CPU:0'.
+    If d = '/gpu:0', default=None, it returns
+      '/job:localhost/replica:0/task:0/device:GPU:0'.
+
+  Args:
+    d: a device string.
+    default: a string for default device if d doesn't have all components.
+
+  Returns:
+    a canonicalized device string.
+  """
   d = tf_device.DeviceSpec.from_string(d)
   assert d.device_type is None or d.device_type == d.device_type.upper(), (
       "Device type '%s' must be all-caps." % (d.device_type,))
   # Fill in missing device fields using defaults.
   result = tf_device.DeviceSpec(
       job="localhost", replica=0, task=0, device_type="CPU", device_index=0)
+  if default:
+    result.merge_from(tf_device.DeviceSpec.from_string(default))
   result.merge_from(d)
   return result.to_string()
 
 
+def resolve(d):
+  """Canonicalize `d` with current device as default."""
+  return canonicalize(d, default=current())
+
+
 class _FakeNodeDef(object):
   """A fake NodeDef for _FakeOperation."""
 
diff --git a/tensorflow/python/training/device_util_test.py b/tensorflow/python/training/device_util_test.py
new file mode 100644
index 0000000000..61525e21f5
--- /dev/null
+++ b/tensorflow/python/training/device_util_test.py
@@ -0,0 +1,89 @@
+# Copyright 2015 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 device utilities."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import test
+from tensorflow.python.training import device_util
+
+
+class DeviceUtilTest(test.TestCase):
+
+  def testCurrentDeviceWithGlobalGraph(self):
+    with ops.device("/cpu:0"):
+      self.assertEqual(device_util.current(), "/device:CPU:0")
+
+    with ops.device("/job:worker"):
+      with ops.device("/cpu:0"):
+        self.assertEqual(device_util.current(), "/job:worker/device:CPU:0")
+
+    with ops.device("/cpu:0"):
+      with ops.device("/gpu:0"):
+        self.assertEqual(device_util.current(), "/device:GPU:0")
+
+  def testCurrentDeviceWithNonGlobalGraph(self):
+    with ops.Graph().as_default():
+      with ops.device("/cpu:0"):
+        self.assertEqual(device_util.current(), "/device:CPU:0")
+
+  def testCurrentDeviceWithEager(self):
+    with context.eager_mode():
+      with ops.device("/cpu:0"):
+        self.assertEqual(device_util.current(),
+                         "/job:localhost/replica:0/task:0/device:CPU:0")
+
+  def testCanonicalizeWithoutDefaultDevice(self):
+    self.assertEqual(
+        device_util.canonicalize("/cpu:0"),
+        "/job:localhost/replica:0/task:0/device:CPU:0")
+    self.assertEqual(
+        device_util.canonicalize("/job:worker/cpu:0"),
+        "/job:worker/replica:0/task:0/device:CPU:0")
+    self.assertEqual(
+        device_util.canonicalize("/job:worker/task:1/cpu:0"),
+        "/job:worker/replica:0/task:1/device:CPU:0")
+
+  def testCanonicalizeWithDefaultDevice(self):
+    self.assertEqual(
+        device_util.canonicalize("/job:worker/task:1/cpu:0", default="/gpu:0"),
+        "/job:worker/replica:0/task:1/device:CPU:0")
+    self.assertEqual(
+        device_util.canonicalize("/job:worker/task:1", default="/gpu:0"),
+        "/job:worker/replica:0/task:1/device:GPU:0")
+    self.assertEqual(
+        device_util.canonicalize("/cpu:0", default="/job:worker"),
+        "/job:worker/replica:0/task:0/device:CPU:0")
+
+  def testResolveWithDeviceScope(self):
+    with ops.device("/gpu:0"):
+      self.assertEqual(
+          device_util.resolve("/job:worker/task:1/cpu:0"),
+          "/job:worker/replica:0/task:1/device:CPU:0")
+      self.assertEqual(
+          device_util.resolve("/job:worker/task:1"),
+          "/job:worker/replica:0/task:1/device:GPU:0")
+    with ops.device("/job:worker"):
+      self.assertEqual(
+          device_util.resolve("/cpu:0"),
+          "/job:worker/replica:0/task:0/device:CPU:0")
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensorflow/python/training/distribute.py b/tensorflow/python/training/distribute.py
index 21ec5292ad..6aeecb31dd 100644
--- a/tensorflow/python/training/distribute.py
+++ b/tensorflow/python/training/distribute.py
@@ -19,6 +19,7 @@ from __future__ import division
 from __future__ import print_function
 
 import threading
+import six
 
 from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.framework import ops
@@ -896,6 +897,8 @@ class DistributionStrategy(object):
       A `Tensor` on `destination`.
     """
     _require_cross_tower_context(self)
+    assert isinstance(destination, six.string_types)
+    destination = device_util.resolve(destination)
     return self._fetch(val, destination, fn)
 
   def _fetch(self, val, destination, fn):
-- 
GitLab


From 87ebe118d0c3767d4a3caaef4ba5538f37311ad1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 12:39:52 -0700
Subject: [PATCH 100/395] Implements matrix multiply-accumulate for linear
 no-offset (aka symmetric) quantizer.

PiperOrigin-RevId: 194978865
---
 .../contrib/lite/kernels/internal/BUILD       |   1 +
 .../internal/optimized/neon_tensor_utils.cc   | 125 +++++++
 .../internal/optimized/neon_tensor_utils.h    |   8 +
 .../internal/optimized/tensor_utils_impl.h    |  10 +
 .../reference/portable_tensor_utils.cc        |  24 ++
 .../reference/portable_tensor_utils.h         |  14 +
 .../lite/kernels/internal/tensor_utils.h      |  30 +-
 .../kernels/internal/tensor_utils_test.cc     | 323 ++++++++++++++++++
 8 files changed, 529 insertions(+), 6 deletions(-)

diff --git a/tensorflow/contrib/lite/kernels/internal/BUILD b/tensorflow/contrib/lite/kernels/internal/BUILD
index c5539afb9c..df29172f83 100644
--- a/tensorflow/contrib/lite/kernels/internal/BUILD
+++ b/tensorflow/contrib/lite/kernels/internal/BUILD
@@ -303,6 +303,7 @@ cc_library(
     ],
     hdrs = [
         "common.h",
+        "compatibility.h",
         "optimized/cpu_check.h",
         "optimized/neon_tensor_utils.h",
         "optimized/tensor_utils_impl.h",
diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.cc b/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.cc
index 47dfcbeb01..65f25168e3 100644
--- a/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.cc
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include "tensorflow/contrib/lite/builtin_op_data.h"
 #include "tensorflow/contrib/lite/kernels/activation_functor.h"
 #include "tensorflow/contrib/lite/kernels/internal/common.h"
+#include "tensorflow/contrib/lite/kernels/internal/compatibility.h"
 #include "tensorflow/contrib/lite/kernels/internal/optimized/tensor_utils_impl.h"
 #include "tensorflow/contrib/lite/kernels/internal/round.h"
 
@@ -27,6 +28,22 @@ limitations under the License.
 
 namespace tflite {
 namespace tensor_utils {
+namespace {
+
+// Allocates, at least, size bytes of uninitialized storage whose alignment is
+// specified by alignment. The size parameter must be an integral multiple of
+// alignment.
+// Caller is responsible by freeing the allocated memory by calling free on
+// the passed freeing_buffer pointer.
+void* aligned_alloc(size_t alignment, size_t size, void** freeing_buffer) {
+  *freeing_buffer = malloc(size + alignment);
+  const size_t offset = ((uintptr_t)*freeing_buffer) % alignment;  // NOLINT
+  return offset == 0
+             ? *freeing_buffer
+             : ((char*)*freeing_buffer + (alignment - offset));  // NOLINT
+}
+
+}  // namespace
 
 void NeonMatrixBatchVectorMultiplyAccumulate(const float* matrix, int m_rows,
                                              int m_cols, const float* vector,
@@ -114,6 +131,114 @@ void NeonMatrixBatchVectorMultiplyAccumulate(const float* matrix, int m_rows,
   delete[] vector_cache_float32x4;
 }
 
+void NeonMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors, const float* scaling_factors,
+    int n_batch, float* __restrict__ result, int result_stride) {
+  const int kWeightsPerUint32 = 4;
+  const int kWeightsPerNeonLane = 16;
+  // If the number of rows is not divisible by kWeightsPerUint32, we set a
+  // flag and allocate an aligned memory block. The flag is used to use the
+  // aligned memory block later in the kernel loop.
+  bool unaligned = false;
+  int8* aligned_row = nullptr;
+  void* aligned_row_free = nullptr;
+  if ((m_cols & (kWeightsPerUint32 - 1)) != 0) {
+    unaligned = true;
+    aligned_row = (int8*)aligned_alloc(kWeightsPerUint32, m_cols,  // NOLINT
+                                       &aligned_row_free);
+  }
+  void* aligned_vec_free = nullptr;
+  int8* aligned_vec = (int8*)aligned_alloc(kWeightsPerUint32, m_cols,  // NOLINT
+                                           &aligned_vec_free);
+
+  // If m_cols is not at least kWeightsPerNeonLane, we cannot use the main
+  // vectorized loop, and we need to process sequentially. postamble_start shows
+  // the start index where this should happen.
+  const int postamble_start = m_cols - (m_cols & (kWeightsPerNeonLane - 1));
+
+  int batch, row, col;
+  for (batch = 0; batch < n_batch; ++batch) {
+    const float batch_scaling_factor_inv = 1.0 / scaling_factors[batch];
+    // Copy the vector data to an aligned vector.
+    memcpy(aligned_vec, vectors + batch * m_cols, sizeof(int8) * m_cols);
+    // Compute dot-product for every column.
+    for (row = 0; row < m_rows; ++row, result += result_stride) {
+      // Get the address of the first element of the row.
+      int8* row_ptr = (int8*)matrix + row * m_cols;  // NOLINT
+      if (unaligned) {
+        memcpy(aligned_row, row_ptr, sizeof(int8) * m_cols);
+        row_ptr = aligned_row;
+      }
+
+      // Initialize the dot product sum for the row to 0.
+      int32x4_t dotprod = vmovq_n_s32(0);
+
+      // Prefetch the row to cache.
+      __builtin_prefetch(row_ptr, 0 /* prefetch for read */,
+                         3 /* temporal locality */);
+
+      // For every block of 16 8-bit elements.
+      col = 0;
+      for (; col < postamble_start; col += kWeightsPerNeonLane) {
+        // Load 16 8-bit values from the row and vector, each, to operate on.
+        // Here the assumption is that each buffer is 4-byte aligned.
+        TFLITE_CHECK_EQ((uintptr_t)(&row_ptr[col]) & (kWeightsPerUint32 - 1),
+                        0);
+        const int8x16_t s1_8x16 = vld1q_s8((const int8_t*)(aligned_vec + col));
+        const int8x16_t s2_8x16 = vld1q_s8((const int8_t*)(row_ptr + col));
+        // Multiply the low bits (i.e. the lower 8 8bit numbers in the
+        // registers).
+        int16x8_t prod_16x8 =
+            vmull_s8(vget_low_s8(s1_8x16), vget_low_s8(s2_8x16));
+        // Multiply the high bits (i.e. the lower 8 8bit numbers in the
+        // registers), and accumulate with the result of the low bits product.
+        // The assumption here is that overflow will not happen as we quantize
+        // our values to be in the range [-127, 127]. As such the sum of the 2
+        // products is always strictly smaller than 15-bits (32767 in absolute
+        // value).
+        prod_16x8 =
+            vmlal_s8(prod_16x8, vget_high_s8(s1_8x16), vget_high_s8(s2_8x16));
+
+        dotprod = vpadalq_s16(dotprod, prod_16x8);
+      }  // for col
+
+      int32 postable_sum = 0;
+      // Postamble loop.
+      // TODO(raziel): if (ABSL_PREDICT_FALSE(postamble_start < m_rows))
+      if (postamble_start < m_cols) {
+        col = postamble_start;
+        if ((m_cols - postamble_start) >= (kWeightsPerNeonLane >> 1)) {
+          // Load 8 8-bit values from the row and column each to operate on.
+          // Here the assumption is that each buffer is 4-bytes aligned.
+          TFLITE_CHECK_EQ((uintptr_t)(&row_ptr[col]) & (kWeightsPerUint32 - 1),
+                          0);
+          const int8x8_t s1_8x8 = vld1_s8((const int8_t*)(aligned_vec + col));
+          const int8x8_t s2_8x8 = vld1_s8((const int8_t*)(row_ptr + col));
+          const int16x8_t prod_16x8 = vmull_s8(s1_8x8, s2_8x8);
+          dotprod = vpadalq_s16(dotprod, prod_16x8);
+          col += (kWeightsPerNeonLane >> 1);
+        }
+        for (; col < m_cols; ++col) {
+          postable_sum += row_ptr[col] * aligned_vec[col];
+        }  // for col
+      }
+      // Add the 4 intermediate sum values to get the final dot-prod value for
+      // this row.
+      int64x2_t pairwiseAdded = vpaddlq_s32(dotprod);
+      int32 neon_sum =
+          vgetq_lane_s64(pairwiseAdded, 0) + vgetq_lane_s64(pairwiseAdded, 1);
+
+      *result += ((neon_sum + postable_sum) * batch_scaling_factor_inv);
+    }  // for row
+  }    // for batch
+
+  if (unaligned) {
+    free(aligned_row_free);
+  }
+  free(aligned_vec_free);
+}
+
 void NeonVectorVectorCwiseProduct(const float* vector1, const float* vector2,
                                   int v_size, float* result) {
   // If v_size is not divisible by kWeightsPerNeonLane, we cannot use the main
diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.h b/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.h
index 3b6f4bd583..9e60d0657b 100644
--- a/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.h
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.h
@@ -32,6 +32,14 @@ void MatrixBatchVectorMultiplyAccumulate(const float* matrix, int m_rows,
                    vector, n_batch, result, result_stride);
 }
 
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors, const float* scaling_factors,
+    int n_batch, float* __restrict__ result, int result_stride) {
+  NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
+                   vectors, scaling_factors, n_batch, result, result_stride);
+}
+
 void VectorVectorCwiseProduct(const float* vector1, const float* vector2,
                               int v_size, float* result) {
   NEON_OR_PORTABLE(VectorVectorCwiseProduct, vector1, vector2, v_size, result);
diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/tensor_utils_impl.h b/tensorflow/contrib/lite/kernels/internal/optimized/tensor_utils_impl.h
index 19220470f4..d570dadd86 100644
--- a/tensorflow/contrib/lite/kernels/internal/optimized/tensor_utils_impl.h
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/tensor_utils_impl.h
@@ -40,6 +40,16 @@ void NeonMatrixBatchVectorMultiplyAccumulate(const float* matrix, int m_rows,
                                              int n_batch, float* result,
                                              int result_stride);
 
+// Matrix multiplication for quantized values using symmetric quantization.
+void PortableMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors, const float* scaling_factors,
+    int n_batch, float* __restrict__ result, int result_stride);
+void NeonMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors, const float* scaling_factors,
+    int n_batch, float* __restrict__ result, int result_stride);
+
 // Cwise product of two vectors.
 void PortableVectorVectorCwiseProduct(const float* vector1,
                                       const float* vector2, int v_size,
diff --git a/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.cc b/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.cc
index 5e7586eeda..2607adc0c1 100644
--- a/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.cc
+++ b/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.cc
@@ -69,6 +69,30 @@ void PortableMatrixBatchVectorMultiplyAccumulate(const float* matrix,
   }
 }
 
+void PortableMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors, const float* scaling_factors,
+    int n_batch, float* __restrict__ result, int result_stride) {
+  int batch, row, col;
+  for (batch = 0; batch < n_batch; ++batch, vectors += m_cols) {
+    const float batch_scaling_factor_inv = 1.0 / scaling_factors[batch];
+    // Get the address of the first row.
+    int8_t* row_ptr = (int8_t*)matrix;  // NOLINT
+    for (row = 0; row < m_rows; ++row, result += result_stride) {
+      // Initialize the dot product sum for the row to 0.
+      int32_t dotprod = 0;
+      // Prefetch the row to cache.
+      __builtin_prefetch(row_ptr, 0 /* prefetch for read */,
+                         3 /* temporal locality */);
+      // For every block of 16 8-bit elements (128-bit register) from each row.
+      for (col = 0; col < m_cols; ++col, ++row_ptr) {
+        dotprod += (*row_ptr) * (vectors[col]);
+      }  // for col
+      *result += (dotprod * batch_scaling_factor_inv);
+    }  // for row
+  }    // for batch
+}
+
 void PortableVectorVectorCwiseProduct(const float* vector1,
                                       const float* vector2, int v_size,
                                       float* result) {
diff --git a/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.h b/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.h
index 478cda8e19..1757a9f5e5 100644
--- a/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.h
+++ b/tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.h
@@ -37,6 +37,11 @@ void PortableMatrixBatchVectorMultiplyAccumulate(const float* matrix,
                                                  int n_batch, float* result,
                                                  int result_stride);
 
+void PortableMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors, const float* scaling_factors,
+    int n_batch, float* __restrict__ result, int result_stride);
+
 // Cwise product of two vectors.
 void PortableVectorVectorCwiseProduct(const float* vector1,
                                       const float* vector2, int v_size,
@@ -122,6 +127,15 @@ void MatrixBatchVectorMultiplyAccumulate(const float* matrix, int m_rows,
                                               n_batch, result, result_stride);
 }
 
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vector, const float* scaling_factors,
+    int n_batch, float* __restrict__ result, int result_stride) {
+  PortableMatrixBatchVectorMultiplyAccumulate(matrix, m_rows, m_cols, vector,
+                                              scaling_factors, n_batch, result,
+                                              result_stride);
+}
+
 void VectorVectorCwiseProduct(const float* vector1, const float* vector2,
                               int v_size, float* result) {
   PortableVectorVectorCwiseProduct(vector1, vector2, v_size, result);
diff --git a/tensorflow/contrib/lite/kernels/internal/tensor_utils.h b/tensorflow/contrib/lite/kernels/internal/tensor_utils.h
index 997dc4425d..e1c9ccd84b 100644
--- a/tensorflow/contrib/lite/kernels/internal/tensor_utils.h
+++ b/tensorflow/contrib/lite/kernels/internal/tensor_utils.h
@@ -31,17 +31,35 @@ void SymmetricQuantizeFloats(const float* values, const int size,
                              int8_t* quantized_values, float* min, float* max,
                              float* scaling_factor);
 
-// Multiply a matrix by a batch vector, and store results in a batch-size
-// vector using a stride value provided in result_stride. 'result_stride' shows
-// how the number of elements between consecutive result values. For example
-// result_stride = 1, will cause the output to look like this:
-// [O_1, 0_2, ... O_rows] in memory, but result_stride = 3, will cause it to be
-// arranged like this in memory: [O_1, x, x, 0_2, x, x, ..., O_rows]
+// Multiplies a matrix by a "batched" vector (i.e. a matrix with a batch
+// dimension composed by input vectors independent from each other). The result
+// of the multiplication is accumulated to the passed result buffer.
+// More specifically, for a matrix M of shape [n, i] and a batched-vector
+// of shape [i, batch] it will first compute the product of shape [n, batch].
+// This product will be accumulated to the result buffer, using a stride value
+// provided in result_stride (the number of elements between consecutive result
+// values). For example result_stride = 1, will cause the output to look like
+// this:
+// [O_1, 0_2, ... O_rows]
+// but result_stride = 3, will cause it to be arranged like this in memory:
+// [O_1, x, x, 0_2, x, x, ..., O_rows]
 void MatrixBatchVectorMultiplyAccumulate(const float* matrix, int m_rows,
                                          int m_cols, const float* vector,
                                          int n_batch, float* result,
                                          int result_stride);
 
+// Same as the function above, but for values quantized using symmetric
+// quantization (e.g. by calling SymmetricQuantizeFloats).
+// The passed scaling factors is a buffer of the quantization scaling factors
+// that will be used to dequentize the products into the final result buffer.
+// These scaling factors are the multiplication of the matrix scaling factor
+// by the vector's scaling factor, one per batch (i.e. this allows quantizing
+// each batch in the batch-vector matrix independently).
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors, const float* scaling_factors,
+    int n_batch, float* __restrict__ result, int result_stride);
+
 // Cwise product of two vectors.
 void VectorVectorCwiseProduct(const float* vector1, const float* vector2,
                               int v_size, float* result);
diff --git a/tensorflow/contrib/lite/kernels/internal/tensor_utils_test.cc b/tensorflow/contrib/lite/kernels/internal/tensor_utils_test.cc
index 22b016746f..3d8a2eada0 100644
--- a/tensorflow/contrib/lite/kernels/internal/tensor_utils_test.cc
+++ b/tensorflow/contrib/lite/kernels/internal/tensor_utils_test.cc
@@ -107,6 +107,329 @@ TEST(uKernels, MatrixBatchVectorMultiplyAccumulateTest) {
                                                -1., 3., 7., 3., 23., 3.})));
 }
 
+TEST(uKernels, MatrixBatchVectorMultiplyAccumulateSymmetricQuantizedTest) {
+  // Note we use 29 columns as this exercises all the neon kernel: the
+  // 16-block SIMD code, the 8-block postamble, and the leftover postamble.
+  const int a_rows = 4, a_cols = 29;
+  const int kWeightsPerUint32 = 4;
+  const float a_float_data[] = {
+      /* 1st row */
+      1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.1, 11.11, 12.12, 13.13,
+      14.14, 15.15, 16.16, 17.17, 18.18, 19.19, 20.2, 21.21, 22.22, 23.23,
+      24.24, 25.25, 26.26, 27.27, 28.28, 0,
+      /* 2nd row */
+      -1.1, -2.2, -3.3, -4.4, -5.5, -6.6, -7.7, -8.8, -9.9, -10.1, -11.11,
+      -12.12, -13.13, -14.14, -15.15, -16.16, -17.17, -18.18, -19.19, -20.2,
+      -21.21, -22.22, -23.23, -24.24, -25.25, -26.26, -27.27, -28.28, 0,
+      /* 3rd row */
+      1.1, -2.2, 3.3, -4.4, 5.5, -6.6, 7.7, -8.8, 9.9, -10.1, 11.11, -12.12,
+      13.13, -14.14, 15.15, -16.16, 17.17, -18.18, 19.19, -20.2, 21.21, -22.22,
+      23.23, -24.24, 25.25, -26.26, 27.27, -28.28, 0,
+      /* 4th row */
+      -1.1, 2.2, -3.3, 4.4, -5.5, 6.6, -7.7, 8.8, -9.9, 10.1, -11.11, 12.12,
+      -13.13, 14.14, -15.15, 16.16, -17.17, 18.18, -19.19, 20.2, -21.21, 22.22,
+      -23.23, 24.24, -25.25, 26.26, -27.27, 28.28, 0};
+
+  int8* a_int8_data = reinterpret_cast<int8*>(
+      aligned_malloc(a_rows * a_cols, kWeightsPerUint32));
+  float a_min, a_max;
+  float scaling_factor_a;
+  SymmetricQuantizeFloats(a_float_data, a_rows * a_cols, a_int8_data, &a_min,
+                          &a_max, &scaling_factor_a);
+  const int8 expected_a_int8_data[] = {
+      /* 1st row */
+      5,
+      10,
+      15,
+      20,
+      25,
+      30,
+      35,
+      40,
+      44,
+      45,
+      50,
+      54,
+      59,
+      64,
+      68,
+      73,
+      77,
+      82,
+      86,
+      91,
+      95,
+      100,
+      104,
+      109,
+      113,
+      118,
+      122,
+      127,
+      0,
+      /* 2nd row */
+      -5,
+      -10,
+      -15,
+      -20,
+      -25,
+      -30,
+      -35,
+      -40,
+      -44,
+      -45,
+      -50,
+      -54,
+      -59,
+      -64,
+      -68,
+      -73,
+      -77,
+      -82,
+      -86,
+      -91,
+      -95,
+      -100,
+      -104,
+      -109,
+      -113,
+      -118,
+      -122,
+      -127,
+      0,
+      /* 3rd row */
+      5,
+      -10,
+      15,
+      -20,
+      25,
+      -30,
+      35,
+      -40,
+      44,
+      -45,
+      50,
+      -54,
+      59,
+      -64,
+      68,
+      -73,
+      77,
+      -82,
+      86,
+      -91,
+      95,
+      -100,
+      104,
+      -109,
+      113,
+      -118,
+      122,
+      -127,
+      0,
+      /* 4th row */
+      -5,
+      10,
+      -15,
+      20,
+      -25,
+      30,
+      -35,
+      40,
+      -44,
+      45,
+      -50,
+      54,
+      -59,
+      64,
+      -68,
+      73,
+      -77,
+      82,
+      -86,
+      91,
+      -95,
+      100,
+      -104,
+      109,
+      -113,
+      118,
+      -122,
+      127,
+      0,
+  };
+  for (int i = 0; i < a_rows * a_cols; ++i) {
+    EXPECT_EQ(expected_a_int8_data[i], a_int8_data[i]);
+  }
+
+  const int b_rows = 29, b_cols = 1, batches = 2;
+  const float b_float_data[] = {
+      /* batch 1 */
+      1.0,
+      -1.0,
+      1.0,
+      -1.0,
+      1.0,
+      -1.0,
+      1.0,
+      -1.0,
+      1.0,
+      -1.0,
+      1.0,
+      -1.0,
+      1.0,
+      -1.0,
+      1.0,
+      -1.0,
+      1.0,
+      -1.0,
+      1.0,
+      -1.0,
+      1.0,
+      -1.0,
+      1.0,
+      -1.0,
+      1.0,
+      -1.0,
+      1.0,
+      -1.0,
+      1.0,
+      /* batch 2 */
+      2.5,
+      -2.1,
+      3.0,
+      -1.3,
+      1.3,
+      -1.1,
+      2.0,
+      -1.7,
+      1.9,
+      -1.5,
+      0.5,
+      -0.7,
+      0.8,
+      -0.3,
+      2.8,
+      -2.8,
+      1.1,
+      -2.3,
+      1.9,
+      -1.9,
+      2.1,
+      -0.5,
+      2.4,
+      -0.1,
+      1.0,
+      -2.5,
+      0.7,
+      -1.9,
+      0.2,
+  };
+
+  // Quantized values of B:
+  int8 b_int8_data[b_rows * b_cols * batches];
+  float b_min, b_max;
+  float scaling_factor_b[batches];
+  SymmetricQuantizeFloats(b_float_data, b_rows * b_cols, b_int8_data, &b_min,
+                          &b_max, &scaling_factor_b[0]);
+  SymmetricQuantizeFloats(&b_float_data[b_rows * b_cols], b_rows * b_cols,
+                          &b_int8_data[b_rows * b_cols], &b_min, &b_max,
+                          &scaling_factor_b[1]);
+
+  const int8 expected_b_int8_data[] = {
+      /* batch 1 */
+      127,
+      -127,
+      127,
+      -127,
+      127,
+      -127,
+      127,
+      -127,
+      127,
+      -127,
+      127,
+      -127,
+      127,
+      -127,
+      127,
+      -127,
+      127,
+      -127,
+      127,
+      -127,
+      127,
+      -127,
+      127,
+      -127,
+      127,
+      -127,
+      127,
+      -127,
+      127,
+      /* batch 2 */
+      106,
+      -89,
+      127,
+      -55,
+      55,
+      -47,
+      85,
+      -72,
+      80,
+      -64,
+      21,
+      -30,
+      34,
+      -13,
+      119,
+      -119,
+      47,
+      -97,
+      80,
+      -80,
+      89,
+      -21,
+      102,
+      -4,
+      42,
+      -106,
+      30,
+      -80,
+      8,
+  };
+  for (int i = 0; i < b_rows * b_cols * batches; ++i) {
+    EXPECT_EQ(expected_b_int8_data[i], b_int8_data[i]);
+  }
+
+  // Full float operation results in:
+  // -13.69, 13.69, 414.11, -414.11
+  // -6.325, 6.325, 631.263, -631.263
+  float c_float_data[a_rows * b_cols * batches];
+  for (int i = 0; i < a_rows * b_cols * batches; ++i) {
+    c_float_data[i] = 0.0;
+  }
+
+  // Testing product.
+  const float scaling_factor_c[2] = {
+      scaling_factor_a * scaling_factor_b[0],
+      scaling_factor_a * scaling_factor_b[1],
+  };
+  MatrixBatchVectorMultiplyAccumulate(a_int8_data, a_rows, a_cols, b_int8_data,
+                                      scaling_factor_c, batches, c_float_data,
+                                      /*result_stride=*/1);
+
+  // Assert we obtain the expected recovered float values.
+  const float expected_c_float_data[] = {
+      -14.474, 14.474, 414.402, -414.402, -6.92228, 6.92228, 632.042, -632.042,
+  };
+  for (int i = 0; i < a_rows * b_cols * batches; ++i) {
+    EXPECT_NEAR(expected_c_float_data[i], c_float_data[i], 0.001);
+  }
+
+  aligned_free(a_int8_data);
+}
+
 TEST(uKernels, VectorVectorCwiseProductTest) {
   constexpr int kVectorSize = 10;
   static float input1[kVectorSize] = {0.0,  -0.5, 1.0,  -1.5, 2.0,
-- 
GitLab


From af2d983bcdabc5291ffa919a2c20654e4c0a8c07 Mon Sep 17 00:00:00 2001
From: Sami Kama <skama@nvidia.com>
Date: Tue, 1 May 2018 12:54:04 -0700
Subject: [PATCH 101/395] Review updates

---
 tensorflow/contrib/tensorrt/convert/convert_graph.cc | 2 +-
 tensorflow/contrib/tensorrt/segment/segment.cc       | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.cc b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
index c1979afcf8..8459ad4a61 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_graph.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
@@ -416,7 +416,7 @@ tensorflow::Status ConvertAfterShapes(
   for (auto s : segments) {
     total_num_nodes_in_segments += s.first.size();
   }
-  // Cluster may not be available
+  // We are creating the map here since cluster may not be available in all cases
   std::map<string, tensorflow::Device*> name_to_device_map;
   if (cluster) {
     for (const auto dm : cluster->GetDeviceSet()->devices()) {
diff --git a/tensorflow/contrib/tensorrt/segment/segment.cc b/tensorflow/contrib/tensorrt/segment/segment.cc
index 7e094f552d..4901e30a87 100644
--- a/tensorflow/contrib/tensorrt/segment/segment.cc
+++ b/tensorflow/contrib/tensorrt/segment/segment.cc
@@ -113,7 +113,7 @@ class SimpleGraph {
   const tensorflow::Graph* g_;
   std::vector<SimpleNode*> nodes_;
   std::vector<SimpleEdge*> edges_;
-  // edge_ids_ and node_ids_ contain freed indices.
+  // free_edge_ids_ and free_node_ids_ contain freed indices.
   std::set<int> free_edge_ids_;
   std::set<int> free_node_ids_;
 };
@@ -352,7 +352,7 @@ tensorflow::Status SegmentGraph(
     tensorflow::Graph* tf_graph,
     const std::function<bool(const tensorflow::Node*)>& candidate_fn,
     const SegmentOptions& options, SegmentNodesVector* segments) {
-  // tensorflow::DumpGraph("Pre-Segment", &graph);
+
   auto graph = std::unique_ptr<SimpleGraph>(new SimpleGraph(tf_graph));
   // Use a union-find to collect the nodes that belong to the same
   // segment. A node value of nullptr indicates that the node is not a candidate
-- 
GitLab


From ee236bd4c4251d6a2a87409b4d47470534c975b0 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Tue, 1 May 2018 12:56:29 -0700
Subject: [PATCH 102/395] Add a pointer from Device to its owning DeviceMgr.
 Allow remote function execution on TPU devices.

PiperOrigin-RevId: 194981511
---
 tensorflow/core/common_runtime/device.h               | 11 +++++++++++
 tensorflow/core/common_runtime/device_mgr.cc          |  3 +++
 .../process_function_library_runtime.cc               |  3 ++-
 3 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/common_runtime/device.h b/tensorflow/core/common_runtime/device.h
index 5918cd9bbf..b537666492 100644
--- a/tensorflow/core/common_runtime/device.h
+++ b/tensorflow/core/common_runtime/device.h
@@ -51,6 +51,8 @@ limitations under the License.
 
 namespace tensorflow {
 
+class DeviceMgr;
+
 class Device : public DeviceBase {
  public:
   Device(Env* env, const DeviceAttributes& device_attributes);
@@ -133,6 +135,10 @@ class Device : public DeviceBase {
   // Returns the resource manager associated w/ this device.
   virtual ResourceMgr* resource_manager() { return rmgr_; }
 
+  // Returns the device manager that owns this device, or nullptr if this Device
+  // is not owned by a device manager.
+  DeviceMgr* device_mgr() const { return device_mgr_; }
+
   // Summarizes the status of this Device, for debugging.
   string DebugString() const { return ProtoDebugString(device_attributes_); }
 
@@ -158,6 +164,11 @@ class Device : public DeviceBase {
   }
 
  private:
+  friend class DeviceMgr;
+
+  // Pointer to the device manager that owns this device. Not owned.
+  DeviceMgr* device_mgr_ = nullptr;
+
   const DeviceAttributes device_attributes_;
   DeviceNameUtils::ParsedName parsed_name_;
 
diff --git a/tensorflow/core/common_runtime/device_mgr.cc b/tensorflow/core/common_runtime/device_mgr.cc
index a77601ba79..470abc1431 100644
--- a/tensorflow/core/common_runtime/device_mgr.cc
+++ b/tensorflow/core/common_runtime/device_mgr.cc
@@ -27,6 +27,9 @@ namespace tensorflow {
 DeviceMgr::DeviceMgr(const std::vector<Device*>& devices)
     : name_backing_store_(128) {
   for (Device* d : devices) {
+    CHECK(d->device_mgr_ == nullptr);
+    d->device_mgr_ = this;
+
     devices_.push_back(d);
 
     // Register under the (1) full name and (2) canonical name.
diff --git a/tensorflow/core/common_runtime/process_function_library_runtime.cc b/tensorflow/core/common_runtime/process_function_library_runtime.cc
index e61ed8c479..668ce87749 100644
--- a/tensorflow/core/common_runtime/process_function_library_runtime.cc
+++ b/tensorflow/core/common_runtime/process_function_library_runtime.cc
@@ -144,7 +144,8 @@ Status ProcessFunctionLibraryRuntime::GetDeviceContext(
   }
   Device* device = flr->device();
   string device_type = device->parsed_name().type;
-  if (device_type == "CPU" || device_type == "TPU_SYSTEM") {
+  if (device_type == "CPU" || device_type == "TPU_SYSTEM" ||
+      device_type == "TPU") {
     // "TPU_SYSTEM" indicates that `device` is a CPU.
     return Status::OK();
   }
-- 
GitLab


From 57207f2b9d5bf9edffb72a9fe377492454abd9ec Mon Sep 17 00:00:00 2001
From: Priya Gupta <priyag@google.com>
Date: Tue, 1 May 2018 13:01:41 -0700
Subject: [PATCH 103/395] Add utility to auto shard a dataset pipeline in the
 appropriate place by locating the file readers and sharding their input
 files.

PiperOrigin-RevId: 194982311
---
 .../contrib/data/python/ops/batching.py       |   4 +-
 tensorflow/contrib/distribute/python/BUILD    |  31 ++
 .../contrib/distribute/python/input_ops.py    | 141 ++++++++++
 .../distribute/python/input_ops_test.py       | 265 ++++++++++++++++++
 tensorflow/python/data/ops/readers.py         |  15 +
 5 files changed, 454 insertions(+), 2 deletions(-)
 create mode 100644 tensorflow/contrib/distribute/python/input_ops.py
 create mode 100644 tensorflow/contrib/distribute/python/input_ops_test.py

diff --git a/tensorflow/contrib/data/python/ops/batching.py b/tensorflow/contrib/data/python/ops/batching.py
index 2152bcde84..42ec2b0b01 100644
--- a/tensorflow/contrib/data/python/ops/batching.py
+++ b/tensorflow/contrib/data/python/ops/batching.py
@@ -364,7 +364,7 @@ class _RestructuredDataset(dataset_ops.Dataset):
         with the structure of `dataset`.
     """
     super(_RestructuredDataset, self).__init__()
-    self._dataset = dataset
+    self._input_dataset = dataset
 
     if not allow_unsafe_cast:
       # Validate that the types are compatible.
@@ -408,7 +408,7 @@ class _RestructuredDataset(dataset_ops.Dataset):
       self._output_classes = output_classes
 
   def _as_variant_tensor(self):
-    return self._dataset._as_variant_tensor()  # pylint: disable=protected-access
+    return self._input_dataset._as_variant_tensor()  # pylint: disable=protected-access
 
   @property
   def output_classes(self):
diff --git a/tensorflow/contrib/distribute/python/BUILD b/tensorflow/contrib/distribute/python/BUILD
index aa1a956a2d..cdb3a8d65e 100644
--- a/tensorflow/contrib/distribute/python/BUILD
+++ b/tensorflow/contrib/distribute/python/BUILD
@@ -501,3 +501,34 @@ cuda_py_test(
         "//tensorflow/python/data/ops:iterator_ops",
     ],
 )
+
+py_library(
+    name = "input_ops",
+    srcs = ["input_ops.py"],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/python:framework_ops",
+        "//tensorflow/python/data/util:nest",
+    ],
+)
+
+cuda_py_test(
+    name = "input_ops_test",
+    srcs = ["input_ops_test.py"],
+    additional_deps = [
+        ":input_ops",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/contrib/data/python/ops:batching",
+        "//tensorflow/contrib/data/python/ops:interleave_ops",
+        "//tensorflow/python:errors",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:framework_ops",
+        "//tensorflow/python:framework_test_lib",
+        "//tensorflow/python:io_ops",
+        "//tensorflow/python/data/ops:readers",
+        "//tensorflow/python:util",
+    ],
+    tags = [
+        "no_pip",
+    ],
+)
diff --git a/tensorflow/contrib/distribute/python/input_ops.py b/tensorflow/contrib/distribute/python/input_ops.py
new file mode 100644
index 0000000000..1f24f62947
--- /dev/null
+++ b/tensorflow/contrib/distribute/python/input_ops.py
@@ -0,0 +1,141 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Input-pipeline utilities for Distribution strategies."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.data.ops import readers
+from tensorflow.python.data.util import nest
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import tf_logging
+
+# TODO(priyag): Any other reader datasets to consider here?
+_READER_DATASET_OPS = [
+    "TextLineDataset",
+    "TFRecordDataset",
+    "FixedLengthRecordDataset"
+]
+
+
+# pylint: disable=protected-access
+def auto_shard_dataset(dataset, num_shards, index):
+  """Shard the input pipeline by sharding the underlying list of files.
+
+  Args:
+    dataset: A `tf.data.Dataset` instance, typically the result of a bunch of
+      dataset transformations.
+    num_shards: A `tf.int64` scalar `tf.Tensor`, representing the number of
+        shards operating in parallel. Same usage as in `Dataset.shard`.
+    index: A `tf.int64` scalar `tf.Tensor`, representing the worker index.
+      Same usage as in `Dataset.shard`.
+
+  Returns:
+    A modified `Dataset` obtained by updating the pipeline sharded by the
+    files.
+
+  Raises:
+    NotImplementedError: If we cannot automatically determine a good way to
+      shard the input dataset.
+  """
+
+  # TODO(priyag): Clone datasets instead of updating in place, similar to the
+  # clone method for TFRecordDataset.
+  def _auto_shard_impl(dataset, found_reader_op):
+    """Recursive implementation of auto sharding."""
+
+    if not found_reader_op:
+      # TODO(priyag): Make this check more robust by enforcing some common
+      # property on reader datasets.
+      if (isinstance(dataset, readers.TextLineDataset) or
+          isinstance(dataset, readers.FixedLengthRecordDataset)):
+        filenames_tensor = dataset._filenames
+        num_files = array_ops.size(filenames_tensor)
+        sharded_filenames_tensor = array_ops.gather(
+            filenames_tensor, math_ops.range(index, num_files, num_shards))
+        dataset._filenames = sharded_filenames_tensor
+        return dataset
+      elif isinstance(dataset, readers.TFRecordDataset):
+        # `TFRecordDataset` needs to be handled separately than other readers
+        # because it converts filenames to a dataset first. Also, we clone it
+        # instead of updating in place because it has special logic in the
+        # constructor. Eventually we will change all cases to clone datasets
+        # instead of updating in-place.
+        return dataset._clone(
+            filenames=dataset._filenames.shard(num_shards, index))
+      elif hasattr(dataset, "_map_func"):
+        # TODO(priyag): Make this check more robust by enforcing some common
+        # property on all map/flatmap/interleave datasets.
+        map_func_def = dataset._map_func.definition
+        for node in map_func_def.node_def:
+          if node.op in _READER_DATASET_OPS:
+            found_reader_op = True
+            break
+          elif node.op == "FlatMapDataset":
+            # TODO(priyag): Should this check for other map datasets? Should it
+            # be recursive? It is too specific to implementation of
+            # TFRecordDataset right now.
+            nested_func_name = node.attr["f"].func.name
+            nested_func = ops.get_default_graph()._functions[nested_func_name]
+            for nested_node in nested_func.definition.node_def:
+              if nested_node.op in _READER_DATASET_OPS:
+                found_reader_op = True
+                break
+            if found_reader_op:
+              break
+        if found_reader_op:
+          dataset._input_dataset = _auto_shard_impl(
+              dataset._input_dataset, found_reader_op)
+          return dataset
+
+    # TODO(priyag): Make _input_dataset(s) a common property of all datasets to
+    # make this check more robust.
+    if hasattr(dataset, "_input_dataset"):
+      dataset._input_dataset = _auto_shard_impl(
+          dataset._input_dataset, found_reader_op)
+      if hasattr(dataset, "_dataset_to_concatenate"):
+        # Special case for `ConcatentateDataset`. We want to shard all input
+        # datasets.
+        dataset._dataset_to_concatenate = _auto_shard_impl(
+            dataset._dataset_to_concatenate, found_reader_op)
+      return dataset
+
+    if hasattr(dataset, "_datasets"):
+      # Special case for `ZipDataset`.
+      dataset._datasets = nest.pack_sequence_as(dataset._datasets, [
+          _auto_shard_impl(ds, found_reader_op)
+          for ds in nest.flatten(dataset._datasets)
+      ])
+      return dataset
+
+    if not found_reader_op:
+      tf_logging.warn(
+          "Could not find a standard reader in the input pipeline"
+          "(one of TextLineDataset, TFRecordDataset, FixedLengthRecordDataset)."
+          "Falling back to sharding the dataset anyway. Please verify"
+          "correctness of auto-sharding for your input.")
+
+    # TODO(priyag): What do we want to do if the number of filenames is
+    # uneven in the number of shards? By default, this will just return as
+    # many items it can before throwing OutOfRangeError.
+    # TODO(priyag): This will shard the filenames before any shuffling of the
+    # filename dataset. It might be desirable to shard after shuffling
+    # filenames? If so, how do we achieve that?
+    return dataset.shard(num_shards, index)
+
+  return _auto_shard_impl(dataset=dataset, found_reader_op=False)
diff --git a/tensorflow/contrib/distribute/python/input_ops_test.py b/tensorflow/contrib/distribute/python/input_ops_test.py
new file mode 100644
index 0000000000..16179c3a49
--- /dev/null
+++ b/tensorflow/contrib/distribute/python/input_ops_test.py
@@ -0,0 +1,265 @@
+# Copyright 2018 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 input pipeline modifications for distribution strategies."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from tensorflow.contrib.data.python.ops import batching
+from tensorflow.contrib.data.python.ops import interleave_ops
+from tensorflow.contrib.distribute.python import input_ops
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import readers
+from tensorflow.python.framework import errors
+from tensorflow.python.lib.io import python_io
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+class AutoShardDatasetTest(test.TestCase):
+
+  def setUp(self):
+    super(AutoShardDatasetTest, self).setUp()
+    self._num_files = 10
+    self._num_records = 4
+    self._num_shards = 2
+    self._shard_index = 0
+    self._record_bytes = 10
+
+  def _record(self, r, f):
+    return compat.as_bytes("Record %d of file %d" % (r, f))
+
+  def _text_line(self, r, f):
+    return compat.as_bytes("Text line %d of file %d" % (r, f))
+
+  def _fixed_length_record(self, r, f):
+    return compat.as_bytes(str((r * f) % 10) * self._record_bytes)
+
+  def _createTFRecordFiles(self):
+    filenames = []
+    for i in range(self._num_files):
+      fn = os.path.join(self.get_temp_dir(), "tf_record.%d.txt" % i)
+      filenames.append(fn)
+      writer = python_io.TFRecordWriter(fn)
+      for j in range(self._num_records):
+        record = self._record(j, i)
+        writer.write(record)
+      writer.close()
+    return filenames
+
+  def _createTextFiles(self):
+    filenames = []
+    for i in range(self._num_files):
+      fn = os.path.join(self.get_temp_dir(), "text_line.%d.txt" % i)
+      filenames.append(fn)
+      contents = []
+      for j in range(self._num_records):
+        contents.append(self._text_line(j, i))
+        if j + 1 != self._num_records or i == 0:
+          contents.append(b"\r\n")
+      contents = b"".join(contents)
+
+      with open(fn, "wb") as f:
+        f.write(contents)
+    return filenames
+
+  def _createFixedLengthRecordFiles(self):
+    filenames = []
+    for i in range(self._num_files):
+      fn = os.path.join(self.get_temp_dir(), "fixed_length_record.%d.txt" % i)
+      filenames.append(fn)
+      with open(fn, "wb") as f:
+        for j in range(self._num_records):
+          f.write(self._fixed_length_record(j, i))
+    return filenames
+
+  def _verifySimpleShardingOutput(self, dataset, record_fn):
+    iterator = dataset.make_one_shot_iterator()
+    next_element = iterator.get_next()
+    with self.test_session() as sess:
+      for f in range(self._shard_index, self._num_files, self._num_shards):
+        for r in range(self._num_records):
+          self.assertAllEqual(record_fn(r, f), sess.run(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        sess.run(next_element)
+
+  def testTFRecordDataset(self):
+    dataset = readers.TFRecordDataset(self._createTFRecordFiles())
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    self._verifySimpleShardingOutput(dataset, self._record)
+
+  def testFlatMap(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        self._createTFRecordFiles())
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    self._verifySimpleShardingOutput(dataset, self._record)
+
+  def testInterleave(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        self._createTFRecordFiles())
+    dataset = dataset.interleave(
+        readers.TFRecordDataset, cycle_length=4, block_length=self._num_records)
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    # Since block_length == num records in each file, the output will still
+    # contain records in order of files.
+    self._verifySimpleShardingOutput(dataset, self._record)
+
+  def testParallelInterleave(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        self._createTFRecordFiles())
+    dataset = dataset.apply(interleave_ops.parallel_interleave(
+        readers.TFRecordDataset,
+        cycle_length=4,
+        block_length=self._num_records))
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    # Since block_length == num records in each file, the output will still
+    # contain records in order of files.
+    self._verifySimpleShardingOutput(dataset, self._record)
+
+  def testListfiles(self):
+    filenames = self._createTFRecordFiles()
+    file_pattern = filenames[0].rsplit("/", 1)[0] + "/tf_record.*.txt"
+    dataset = dataset_ops.Dataset.list_files(file_pattern, shuffle=False)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    iterator = dataset.make_one_shot_iterator()
+    next_element = iterator.get_next()
+    with self.test_session() as sess:
+      actual, expected = [], []
+      for f in range(self._shard_index, self._num_files, self._num_shards):
+        for r in range(self._num_records):
+          actual.append(sess.run(next_element))
+          expected.append(self._record(r, f))
+      with self.assertRaises(errors.OutOfRangeError):
+        sess.run(next_element)
+      self.assertAllEqual(expected, actual)
+
+  def testComplexPipeline(self):
+    # Setup a complex input pipeline.
+    batch_size = 2
+    num_epochs = 5
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        self._createTFRecordFiles())
+    dataset = dataset.shuffle(buffer_size=self._num_files)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = dataset.prefetch(buffer_size=batch_size)
+    dataset = dataset.shuffle(2 * self._num_files * self._num_records)
+    dataset = dataset.repeat(num_epochs)
+    dataset = dataset.apply(batching.map_and_batch(
+        lambda x: x, batch_size=batch_size))
+    dataset = dataset.prefetch(buffer_size=None)
+
+    # Auto shard.
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    # Verify output.
+    iterator = dataset.make_one_shot_iterator()
+    next_element = iterator.get_next()
+    with self.test_session() as sess:
+      actual = []
+      num_iterations = (self._num_files * self._num_records * num_epochs) // (
+          self._num_shards * batch_size)
+      for _ in range(num_iterations):
+        actual.extend(sess.run(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        sess.run(next_element)
+
+      expected = []
+      for f in range(0, self._num_files, self._num_shards):
+        for r in range(self._num_records):
+          expected.append(self._record(r, f))
+      expected *= num_epochs
+
+      self.assertAllEqual(sorted(expected), sorted(actual))
+
+  def testZip(self):
+    dataset1 = readers.TFRecordDataset(self._createTFRecordFiles())
+    dataset2 = readers.TextLineDataset(self._createTextFiles())
+    dataset = dataset_ops.Dataset.zip((dataset1, dataset2))
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    record_fn = lambda r, f: (self._record(r, f), self._text_line(r, f))
+    self._verifySimpleShardingOutput(dataset, record_fn)
+
+  def testConcat(self):
+    dataset1 = readers.TFRecordDataset(self._createTFRecordFiles())
+    dataset2 = readers.TextLineDataset(self._createTextFiles())
+    dataset = dataset1.concatenate(dataset2)
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    iterator = dataset.make_one_shot_iterator()
+    next_element = iterator.get_next()
+    with self.test_session() as sess:
+      for f in range(self._shard_index, self._num_files, self._num_shards):
+        for r in range(self._num_records):
+          self.assertAllEqual(self._record(r, f), sess.run(next_element))
+      for f in range(self._shard_index, self._num_files, self._num_shards):
+        for r in range(self._num_records):
+          self.assertAllEqual(self._text_line(r, f), sess.run(next_element))
+      with self.assertRaises(errors.OutOfRangeError):
+        sess.run(next_element)
+
+  def testTextLineReader(self):
+    dataset = readers.TextLineDataset(self._createTextFiles())
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    self._verifySimpleShardingOutput(dataset, self._text_line)
+
+  def testTextLineReaderWithFlatMap(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(self._createTextFiles())
+    dataset = dataset.flat_map(readers.TextLineDataset)
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    self._verifySimpleShardingOutput(dataset, self._text_line)
+
+  def testFixedLengthReader(self):
+    dataset = readers.FixedLengthRecordDataset(
+        self._createFixedLengthRecordFiles(), self._record_bytes)
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    self._verifySimpleShardingOutput(dataset, self._fixed_length_record)
+
+  def testFixedLengthReaderWithFlatMap(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        self._createFixedLengthRecordFiles())
+    dataset = dataset.flat_map(
+        lambda f: readers.FixedLengthRecordDataset(f, self._record_bytes))
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    self._verifySimpleShardingOutput(dataset, self._fixed_length_record)
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensorflow/python/data/ops/readers.py b/tensorflow/python/data/ops/readers.py
index fe033f5546..a73a8b5cdc 100644
--- a/tensorflow/python/data/ops/readers.py
+++ b/tensorflow/python/data/ops/readers.py
@@ -197,6 +197,11 @@ class TFRecordDataset(dataset_ops.Dataset):
       filenames = array_ops.reshape(filenames, [-1], name="flat_filenames")
       filenames = dataset_ops.Dataset.from_tensor_slices(filenames)
 
+    self._filenames = filenames
+    self._compression_type = compression_type
+    self._buffer_size = buffer_size
+    self._num_parallel_reads = num_parallel_reads
+
     def read_one_file(filename):
       return _TFRecordDataset(filename, compression_type, buffer_size)
 
@@ -208,6 +213,16 @@ class TFRecordDataset(dataset_ops.Dataset):
           block_length=1, sloppy=False, buffer_output_elements=None,
           prefetch_input_elements=None)
 
+  def _clone(self,
+             filenames=None,
+             compression_type=None,
+             buffer_size=None,
+             num_parallel_reads=None):
+    return TFRecordDataset(filenames or self._filenames,
+                           compression_type or self._compression_type,
+                           buffer_size or self._buffer_size,
+                           num_parallel_reads or self._num_parallel_reads)
+
   def _as_variant_tensor(self):
     return self._impl._as_variant_tensor()  # pylint: disable=protected-access
 
-- 
GitLab


From 1a50cd4ca8c4fe1c1a9ea14f219fd98be8704a7d Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Tue, 1 May 2018 13:07:03 -0700
Subject: [PATCH 104/395] Open source infeed test

PiperOrigin-RevId: 194983270
---
 .../compiler/xla/service/cpu/tests/BUILD      |  23 ++
 .../xla/service/cpu/tests/cpu_infeed_test.cc  | 294 ++++++++++++++++++
 2 files changed, 317 insertions(+)
 create mode 100644 tensorflow/compiler/xla/service/cpu/tests/cpu_infeed_test.cc

diff --git a/tensorflow/compiler/xla/service/cpu/tests/BUILD b/tensorflow/compiler/xla/service/cpu/tests/BUILD
index 9425b948c1..bfd95c3fe0 100644
--- a/tensorflow/compiler/xla/service/cpu/tests/BUILD
+++ b/tensorflow/compiler/xla/service/cpu/tests/BUILD
@@ -124,3 +124,26 @@ tf_cc_test(
         "//tensorflow/core:test_main",
     ],
 )
+
+tf_cc_test(
+    name = "cpu_infeed_test",
+    srcs = ["cpu_infeed_test.cc"],
+    deps = [
+        "//tensorflow/compiler/xla:literal_util",
+        "//tensorflow/compiler/xla:shape_util",
+        "//tensorflow/compiler/xla:statusor",
+        "//tensorflow/compiler/xla:test_helpers",
+        "//tensorflow/compiler/xla:xla_data_proto",
+        "//tensorflow/compiler/xla/client:global_data",
+        "//tensorflow/compiler/xla/client:local_client",
+        "//tensorflow/compiler/xla/client/lib:arithmetic",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_computation",
+        "//tensorflow/compiler/xla/service:cpu_plugin",
+        "//tensorflow/compiler/xla/tests:client_library_test_base",
+        "//tensorflow/compiler/xla/tests:literal_test_util",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_infeed_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_infeed_test.cc
new file mode 100644
index 0000000000..dd63b998e9
--- /dev/null
+++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_infeed_test.cc
@@ -0,0 +1,294 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include <unistd.h>
+#include <memory>
+
+#include "tensorflow/compiler/xla/client/global_data.h"
+#include "tensorflow/compiler/xla/client/lib/arithmetic.h"
+#include "tensorflow/compiler/xla/client/local_client.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
+#include "tensorflow/compiler/xla/literal_util.h"
+#include "tensorflow/compiler/xla/shape_util.h"
+#include "tensorflow/compiler/xla/statusor.h"
+#include "tensorflow/compiler/xla/test_helpers.h"
+#include "tensorflow/compiler/xla/tests/client_library_test_base.h"
+#include "tensorflow/compiler/xla/tests/literal_test_util.h"
+#include "tensorflow/compiler/xla/xla_data.pb.h"
+#include "tensorflow/core/lib/math/math_util.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace xla {
+namespace {
+
+class InfeedTest : public ClientLibraryTestBase {
+ protected:
+  // Transfers the given literal to the infeed interface of the device, and
+  // check if the returned data from Infeed HLO is same as the literal.
+  void TestInfeedRoundTrip(const Literal& literal) {
+    // TODO(b/31037751) Explicitly reset the Infeed state so that the
+    // test is not affected by the state from the previous tests by
+    // adding ClearInfeed if necessary when it is implemented. For now
+    // don't use ResetDevice since it is not implemented on CPU.
+    ASSERT_IS_OK(client_->TransferToInfeed(literal));
+    XlaBuilder builder(TestName());
+    builder.Infeed(literal.shape());
+    if (ShapeUtil::IsTuple(literal.shape())) {
+      // TODO(b/30609564): Use ComputeAndCompareLiteral instead.
+      ComputeAndCompareTuple(&builder, literal, {});
+    } else {
+      ComputeAndCompareLiteral(&builder, literal, {});
+    }
+  }
+};
+
+TEST_F(InfeedTest, SingleInfeedR0Bool) {
+  TestInfeedRoundTrip(*Literal::CreateR0<bool>(true));
+}
+
+TEST_F(InfeedTest, SingleInfeedR1U32) {
+  TestInfeedRoundTrip(*Literal::CreateR1<uint32>({1, 2, 3}));
+}
+
+TEST_F(InfeedTest, SingleInfeedR2F32) {
+  TestInfeedRoundTrip(*Literal::CreateR2F32Linspace(0.0, 1.0, 128, 64));
+}
+
+TEST_F(InfeedTest, SingleInfeedR3F32) {
+  TestInfeedRoundTrip(
+      *Literal::CreateR3({{{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}},
+                          {{1.1f, 2.1f, 3.1f}, {6.1f, 3.5f, 2.8f}}}));
+}
+
+TEST_F(InfeedTest, SingleInfeedR3F32DifferentLayout) {
+  const Layout r3_dim0minor = LayoutUtil::MakeLayout({0, 1, 2});
+  const Layout r3_dim0major = LayoutUtil::MakeLayout({2, 1, 0});
+
+  TestInfeedRoundTrip(
+      *Literal::CreateR3WithLayout({{{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}},
+                                    {{1.1f, 2.1f, 3.1f}, {6.1f, 3.5f, 2.8f}}},
+                                   r3_dim0minor));
+
+  TestInfeedRoundTrip(
+      *Literal::CreateR3WithLayout({{{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}},
+                                    {{1.1f, 2.1f, 3.1f}, {6.1f, 3.5f, 2.8f}}},
+                                   r3_dim0major));
+}
+
+TEST_F(InfeedTest, SingleInfeedR4S32) {
+  TestInfeedRoundTrip(*Literal::CreateR4(
+      {{{{1, -2}, {-4, 5}, {6, 7}}, {{8, 9}, {10, 11}, {12, 13}}},
+       {{{10, 3}, {7, -2}, {3, 6}}, {{2, 5}, {-11, 5}, {-2, -5}}}}));
+}
+
+TEST_F(InfeedTest, SingleInfeedTuple) {
+  TestInfeedRoundTrip(
+      *Literal::MakeTuple({Literal::CreateR1<uint32>({1, 2, 3}).get(),
+                           Literal::CreateR0<bool>(false).get()}));
+}
+
+TEST_F(InfeedTest, SingleInfeedEmptyTuple) {
+  TestInfeedRoundTrip(*Literal::MakeTuple({}));
+}
+
+// Tests Infeed operation used in a while loop, as in the code below. The
+// computation is launched asynchronously, and then infeed data is transferred.
+//
+// float acc = 0.0f;
+// while (acc < 40.0f) {
+//   acc += reduce_add(Infeed());
+// }
+// return acc;
+// TODO(b/30671675) enable this test once asynchronous execution is
+// implemented for CPU.
+TEST_F(InfeedTest, DISABLED_SingleInfeedInWhile) {
+  XlaBuilder builder(TestName());
+  const auto infeed_shape = ShapeUtil::MakeShape(F32, {3});
+  const auto result_shape = ShapeUtil::MakeShape(F32, {});
+
+  // Create a computation for the condition: repeat until (prev < 40.0f) holds.
+  XlaComputation condition;
+  {
+    XlaBuilder builder("condition");
+    auto prev = builder.Parameter(0, result_shape, "prev");
+    builder.Gt(builder.ConstantR0<float>(40.0f), prev);
+    condition = builder.Build().ConsumeValueOrDie();
+  }
+  // Create a computation for the body: add the reduced value of the Infeed
+  // data to the result variable.
+  XlaComputation body;
+  {
+    XlaBuilder builder("body");
+    auto prev = builder.Parameter(0, result_shape, "prev");
+    auto infeed = builder.Infeed(infeed_shape);
+    auto addend =
+        builder.Reduce(infeed, builder.ConstantR0<float>(0.0f),
+                       CreateScalarAddComputation(F32, &builder), {0});
+    builder.Add(prev, addend);
+    body = builder.Build().ConsumeValueOrDie();
+  }
+  // Create a While node with computations for the condition and the body.
+  auto init = builder.ConstantR0<float>(0.0f);
+  builder.While(condition, body, init);
+
+  // Build and asynchronously launch the computation.
+  auto computation = builder.Build().ConsumeValueOrDie();
+  std::unique_ptr<GlobalData> result;
+  tensorflow::Thread* computation_thread =
+      tensorflow::Env::Default()->StartThread(
+          tensorflow::ThreadOptions{}, "computation_thread", [&] {
+            result = client_->Execute(computation, {}, &execution_options_)
+                         .ValueOrDie();
+          });
+
+  // Send 5 Infeed data of shape F32[3].
+  ASSERT_IS_OK(client_->TransferToInfeed(*Literal::CreateR1<float>({1, 2, 3})));
+  ASSERT_IS_OK(client_->TransferToInfeed(*Literal::CreateR1<float>({4, 5, 6})));
+  ASSERT_IS_OK(client_->TransferToInfeed(*Literal::CreateR1<float>({7, 8, 9})));
+  ASSERT_IS_OK(
+      client_->TransferToInfeed(*Literal::CreateR1<float>({10, 11, 12})));
+  ASSERT_IS_OK(
+      client_->TransferToInfeed(*Literal::CreateR1<float>({13, 14, 15})));
+
+  delete computation_thread;  // Joins the thread.
+  auto result_literal = client_->Transfer(*result).ConsumeValueOrDie();
+
+  // Only the first 3 infeed data should be added.
+  LiteralTestUtil::ExpectR0Near<float>(45.0f, *result_literal, ErrorSpec{1e-7});
+}
+
+// Tests two Infeed operations with a total order. The order is enforced by
+// using the result of the first while loop as the initial value of the second
+// while loop. The shapes of both Infeeds are Tuples, where the first tuple
+// element (R1F32) is for the data to reduce and accumulate, and the second
+// tuple element (PRED) to indicate whether the loop should continue. The
+// computation is launched asynchronously, and then infeed data is transferred.
+//
+// float acc = 0.0f;
+// continue = true;
+// while (!continue) {
+//   (data, continue) = Infeed(shape1);
+//   acc += reduce_add(data)
+// }
+// continue = true;
+// while(!continue) {
+//   (data, continue) = Infeed(shape2);
+//   acc += reduce_add(data)
+// }
+// return acc;
+// TODO(b/30671675) enable this test once asynchronous execution is
+// implemented for CPU.
+TEST_F(InfeedTest, DISABLED_TwoInfeedsInTotalOrder) {
+  XlaBuilder builder(TestName());
+  const auto infeed1_shape = ShapeUtil::MakeTupleShape(
+      {ShapeUtil::MakeShape(F32, {2}), ShapeUtil::MakeShape(PRED, {})});
+  const auto infeed2_shape = ShapeUtil::MakeTupleShape(
+      {ShapeUtil::MakeShape(F32, {3}), ShapeUtil::MakeShape(PRED, {})});
+  const auto result_shape = ShapeUtil::MakeTupleShape(
+      {ShapeUtil::MakeShape(F32, {}), ShapeUtil::MakeShape(PRED, {})});
+
+  // Create a computation for the condition: repeat until the second tuple
+  // element is false.
+  XlaComputation condition;
+  {
+    XlaBuilder builder("condition");
+    auto prev = builder.Parameter(0, result_shape, "prev");
+    builder.GetTupleElement(prev, 1);
+    condition = builder.Build().ConsumeValueOrDie();
+  }
+
+  // A lambda that builds the body computation of a while loop with the given
+  // infeed shape, and returns the computation with the ownership.
+  //
+  // The body adds the reduced value of the Infeed data (first tuple element)
+  // to the previous accumulator, and returns the accumulator and the continue
+  // flag (second tuple element) as a tuple.
+  const auto build_body = [this, &result_shape](const Shape& infeed_shape) {
+    XlaComputation body;
+    XlaBuilder builder("body");
+    auto prev = builder.Parameter(0, result_shape, "prev");
+    auto infeed = builder.Infeed(infeed_shape);
+    auto addend = builder.Reduce(
+        builder.GetTupleElement(infeed, 0), builder.ConstantR0<float>(0.0f),
+        CreateScalarAddComputation(F32, &builder), {0});
+    auto result = builder.Add(builder.GetTupleElement(prev, 0), addend);
+    builder.Tuple({result, builder.GetTupleElement(infeed, 1)});
+    return builder.Build().ConsumeValueOrDie();
+  };
+
+  // Create the first while loop with infeed1_shape.
+  auto init = builder.Tuple(
+      {builder.ConstantR0<float>(0.0f), builder.ConstantR0<bool>(true)});
+  auto while1 = builder.While(condition, build_body(infeed1_shape), init);
+  auto result1 = builder.Tuple(
+      {builder.GetTupleElement(while1, 0), builder.ConstantR0<bool>(true)});
+
+  // Create the second while loop with infeed2_shape. Note that the result from
+  // the first while loop is used as the initial value.
+  auto while2 = builder.While(condition, build_body(infeed2_shape), result1);
+  builder.GetTupleElement(while2, 0);
+
+  // Build the computation.
+  auto computation = builder.Build().ConsumeValueOrDie();
+
+  // Send the first 4 Infeed data of shape Tuple(F32[2], PRED).
+  ASSERT_IS_OK(client_->TransferToInfeed(
+      *Literal::MakeTuple({Literal::CreateR1<float>({1, 2}).get(),
+                           Literal::CreateR0<bool>(true).get()})));
+  ASSERT_IS_OK(client_->TransferToInfeed(
+      *Literal::MakeTuple({Literal::CreateR1<float>({3, 4}).get(),
+                           Literal::CreateR0<bool>(true).get()})));
+  ASSERT_IS_OK(client_->TransferToInfeed(
+      *Literal::MakeTuple({Literal::CreateR1<float>({5, 6}).get(),
+                           Literal::CreateR0<bool>(true).get()})));
+  ASSERT_IS_OK(client_->TransferToInfeed(
+      *Literal::MakeTuple({Literal::CreateR1<float>({7, 8}).get(),
+                           Literal::CreateR0<bool>(false).get()})));
+
+  // Asynchronously launch the execution on the device.
+  std::unique_ptr<GlobalData> result;
+  tensorflow::Thread* computation_thread =
+      tensorflow::Env::Default()->StartThread(
+          tensorflow::ThreadOptions{}, "computation_thread", [&] {
+            result = client_->Execute(computation, {}, &execution_options_)
+                         .ValueOrDie();
+          });
+
+  // Wait for a second to ensure testing that the execution is waiting on the
+  // Infeed data, and send the rest Infeed data of shape Tuple(F32[3], PRED).
+  sleep(1);
+  ASSERT_IS_OK(client_->TransferToInfeed(
+      *Literal::MakeTuple({Literal::CreateR1<float>({1, 2, 3}).get(),
+                           Literal::CreateR0<bool>(true).get()})));
+  ASSERT_IS_OK(client_->TransferToInfeed(
+      *Literal::MakeTuple({Literal::CreateR1<float>({7, 8, 9}).get(),
+                           Literal::CreateR0<bool>(false).get()})));
+  ASSERT_IS_OK(client_->TransferToInfeed(
+      *Literal::MakeTuple({Literal::CreateR1<float>({4, 5, 6}).get(),
+                           Literal::CreateR0<bool>(true).get()})));
+
+  // Wait for the execution to be done, and transfer the result.
+  delete computation_thread;  // Joins the thread.
+  auto result_literal = client_->Transfer(*result).ConsumeValueOrDie();
+
+  // Only the first 6 infeed data should be added.
+  LiteralTestUtil::ExpectR0Near<float>(66.0f, *result_literal, ErrorSpec{1e-7});
+}
+
+}  // namespace
+}  // namespace xla
-- 
GitLab


From 9149558a639efe82baf1b5201feccf2411343a8a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 13:15:53 -0700
Subject: [PATCH 105/395] Collective Ops Part 5

Distributed-mode implementations of DeviceResolverInterface
and ParamResolverInterface.  Extend Worker interface with
new methods in support of these interfaces.

This change is part of a series of changes introducing infrastructure
for collective ops and initial implementations of reduction and broadcast.

PiperOrigin-RevId: 194984585
---
 tensorflow/core/distributed_runtime/BUILD     |  75 ++++
 .../collective_param_resolver_distributed.cc  | 404 ++++++++++++++++++
 .../collective_param_resolver_distributed.h   |  90 ++++
 ...lective_param_resolver_distributed_test.cc | 324 ++++++++++++++
 .../device_resolver_distributed.cc            | 133 ++++++
 .../device_resolver_distributed.h             |  67 +++
 .../device_resolver_distributed_test.cc       | 217 ++++++++++
 .../rpc/grpc_remote_worker.cc                 |  27 ++
 .../rpc/grpc_worker_service.cc                |  47 ++
 .../rpc/grpc_worker_service_impl.cc           |   6 +
 .../rpc/grpc_worker_service_impl.h            |   5 +-
 .../core/distributed_runtime/test_utils.h     | 173 ++++++++
 tensorflow/core/distributed_runtime/worker.cc |  87 ++--
 tensorflow/core/distributed_runtime/worker.h  |  17 +-
 .../core/distributed_runtime/worker_env.h     |   5 +
 .../distributed_runtime/worker_interface.h    |  19 +
 tensorflow/core/protobuf/worker.proto         |  70 +++
 tensorflow/core/protobuf/worker_service.proto |  10 +
 18 files changed, 1744 insertions(+), 32 deletions(-)
 create mode 100644 tensorflow/core/distributed_runtime/collective_param_resolver_distributed.cc
 create mode 100644 tensorflow/core/distributed_runtime/collective_param_resolver_distributed.h
 create mode 100644 tensorflow/core/distributed_runtime/collective_param_resolver_distributed_test.cc
 create mode 100644 tensorflow/core/distributed_runtime/device_resolver_distributed.cc
 create mode 100644 tensorflow/core/distributed_runtime/device_resolver_distributed.h
 create mode 100644 tensorflow/core/distributed_runtime/device_resolver_distributed_test.cc
 create mode 100644 tensorflow/core/distributed_runtime/test_utils.h

diff --git a/tensorflow/core/distributed_runtime/BUILD b/tensorflow/core/distributed_runtime/BUILD
index 343dd5d456..256ce527a4 100644
--- a/tensorflow/core/distributed_runtime/BUILD
+++ b/tensorflow/core/distributed_runtime/BUILD
@@ -452,6 +452,81 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "collective_param_resolver_distributed",
+    srcs = ["collective_param_resolver_distributed.cc"],
+    hdrs = ["collective_param_resolver_distributed.h"],
+    deps = [
+        ":call_options",
+        ":device_resolver_distributed",
+        ":worker_cache",
+        "//tensorflow/core:core_cpu_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:worker_proto_cc",
+    ],
+)
+
+cc_library(
+    name = "test_utils",
+    srcs = [],
+    hdrs = ["test_utils.h"],
+    deps = [
+        ":worker_cache",
+        ":worker_interface",
+    ],
+)
+
+tf_cc_test(
+    name = "collective_param_resolver_distributed_test",
+    size = "small",
+    srcs = ["collective_param_resolver_distributed_test.cc"],
+    deps = [
+        ":collective_param_resolver_distributed",
+        ":device_resolver_distributed",
+        ":test_utils",
+        "//tensorflow/core:core_cpu_lib",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+    ],
+)
+
+cc_library(
+    name = "device_resolver_distributed",
+    srcs = ["device_resolver_distributed.cc"],
+    hdrs = ["device_resolver_distributed.h"],
+    deps = [
+        ":worker_cache",
+        "//tensorflow/core:core_cpu_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:worker_proto_cc",
+    ],
+)
+
+tf_cc_test(
+    name = "device_resolver_distributed_test",
+    size = "small",
+    srcs = ["device_resolver_distributed_test.cc"],
+    deps = [
+        ":device_resolver_distributed",
+        ":test_utils",
+        "//tensorflow/core:core_cpu_lib",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+    ],
+)
+
 # TODO(mrry): Move executor_test.cc to ../common_runtime when once it no longer depends
 # on grpc_testlib.
 tf_cuda_cc_tests(
diff --git a/tensorflow/core/distributed_runtime/collective_param_resolver_distributed.cc b/tensorflow/core/distributed_runtime/collective_param_resolver_distributed.cc
new file mode 100644
index 0000000000..ecf5db8110
--- /dev/null
+++ b/tensorflow/core/distributed_runtime/collective_param_resolver_distributed.cc
@@ -0,0 +1,404 @@
+/* Copyright 2018 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.
+==============================================================================*/
+#include "tensorflow/core/distributed_runtime/collective_param_resolver_distributed.h"
+
+#include "tensorflow/core/distributed_runtime/call_options.h"
+#include "tensorflow/core/distributed_runtime/device_resolver_distributed.h"
+#include "tensorflow/core/distributed_runtime/worker_cache.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+
+// TODO(tucker): When we're ready to enable collectives this const will
+// transition to a settable config member.
+static const char FLAGS_collective_group_leader[] =
+    "/job:worker/replica:0/task:0";
+
+namespace tensorflow {
+namespace {
+// Supports client side cancellation of WorkerInterface calls via
+// registration with a CancellationManager.  Note that ParamResolverInterface
+// calls are done on behalf of an Op execution which needs to abort if the
+// step in which it executes is cancelled.
+class CancellableCall {
+ public:
+  CancellableCall(CancellationManager* cancel_mgr, const string& remote_worker,
+                  WorkerCacheInterface* wc)
+      : cancel_mgr_(cancel_mgr), remote_worker_(remote_worker), wc_(wc) {
+    wi_ = wc_->CreateWorker(remote_worker_);
+  }
+  virtual ~CancellableCall() { wc_->ReleaseWorker(remote_worker_, wi_); }
+
+  virtual void IssueCall(const StatusCallback& done) = 0;
+
+  void Start(const StatusCallback& done) {
+    CancellationToken token = cancel_mgr_->get_cancellation_token();
+    const bool not_yet_cancelled = cancel_mgr_->RegisterCallback(
+        token, [this, token]() { opts_.StartCancel(); });
+    if (not_yet_cancelled) {
+      IssueCall([this, token, done](const Status& s) {
+        cancel_mgr_->DeregisterCallback(token);
+        done(s);
+      });
+    } else {
+      done(errors::Cancelled("RPC Request was cancelled"));
+    }
+  }
+
+ protected:
+  mutable mutex mu_;
+  CancellationManager* cancel_mgr_;  // Not owned
+  const string remote_worker_;
+  WorkerCacheInterface* wc_;  // Not owned
+  WorkerInterface* wi_;       // Owned by wc_, must be released.
+  CallOptions opts_;
+};
+
+class CompleteGroupCall : public CancellableCall {
+ public:
+  CompleteGroupCall(const CollGroupParams& group, const string& device_name,
+                    CancellationManager* cancel_mgr,
+                    const string& remote_worker, WorkerCacheInterface* wc)
+      : CancellableCall(cancel_mgr, remote_worker, wc) {
+    req_.set_group_key(group.group_key);
+    req_.set_group_size(group.group_size);
+    req_.set_device_type(group.device_type.type_string());
+    req_.add_device_name(device_name);
+  }
+  ~CompleteGroupCall() override {}
+
+  void IssueCall(const StatusCallback& done) override {
+    wi_->CompleteGroupAsync(&opts_, &req_, &resp_, done);
+  }
+
+  CompleteGroupRequest req_;
+  CompleteGroupResponse resp_;
+};
+
+class CompleteInstanceCall : public CancellableCall {
+ public:
+  CompleteInstanceCall(const CollGroupParams& group,
+                       const CollInstanceParams& instance,
+                       const string& node_name, const string& device_name,
+                       bool is_source, CancellationManager* cancel_mgr,
+                       const string& remote_worker, WorkerCacheInterface* wc)
+      : CancellableCall(cancel_mgr, remote_worker, wc) {
+    req_.set_name(node_name);
+    req_.set_type(instance.type);
+    req_.set_data_type(instance.data_type);
+    instance.shape.AsProto(req_.mutable_shape());
+    req_.set_group_key(group.group_key);
+    req_.set_group_size(group.group_size);
+    req_.set_instance_key(instance.instance_key);
+    req_.set_device_type(group.device_type.type_string());
+    for (int32 offset : instance.impl_details.subdiv_offsets) {
+      req_.add_subdiv_offset(offset);
+    }
+    req_.set_device(device_name);
+    req_.set_is_source(is_source);
+  }
+
+  ~CompleteInstanceCall() override {}
+
+  void IssueCall(const StatusCallback& done) override {
+    wi_->CompleteInstanceAsync(&opts_, &req_, &resp_, done);
+  }
+
+  CompleteInstanceRequest req_;
+  CompleteInstanceResponse resp_;
+};
+
+}  // namespace
+
+CollectiveParamResolverDistributed::CollectiveParamResolverDistributed(
+    const ConfigProto& config, const DeviceMgr* dev_mgr,
+    DeviceResolverDistributed* dev_resolver, WorkerCacheInterface* worker_cache,
+    const string& task_name)
+    : CollectiveParamResolverLocal(dev_mgr, dev_resolver, task_name),
+      worker_cache_(worker_cache),
+      group_leader_(task_name == FLAGS_collective_group_leader
+                        ? ""
+                        : FLAGS_collective_group_leader) {}
+
+void CollectiveParamResolverDistributed::CompleteParamsAsync(
+    const string& device, CollectiveParams* cp, CancellationManager* cancel_mgr,
+    const StatusCallback& done) {
+  CompleteGroupDistributed(device, cp, cancel_mgr,
+                           [this, device, cp, cancel_mgr, done](
+                               const Status& s, const GroupRec* gr) {
+                             if (s.ok()) {
+                               CompleteInstanceDistributed(device, gr, cp,
+                                                           cancel_mgr, done);
+                             } else {
+                               done(s);
+                             }
+                           });
+}
+
+void CollectiveParamResolverDistributed::CompleteGroupAsync(
+    const CompleteGroupRequest* request, CompleteGroupResponse* response,
+    CancellationManager* cancel_mgr, const StatusCallback& done) {
+  CollectiveParams cp;
+  cp.group.group_key = request->group_key();
+  cp.group.group_size = request->group_size();
+  cp.group.device_type = DeviceType(request->device_type());
+  for (const string& dn : request->device_name()) {
+    cp.instance.device_names.push_back(dn);
+  }
+  CompleteGroupDistributed(
+      cp.instance.device_names[0], &cp, cancel_mgr,
+      [this, response, done](const Status& s, const GroupRec* gr) {
+        if (s.ok()) {
+          mutex_lock l(gr->mu);
+          response->set_group_key(gr->group.group_key);
+          response->set_group_size(gr->group.group_size);
+          response->set_device_type(gr->group.device_type.type_string());
+          response->set_num_tasks(gr->task_set.size());
+          for (const string& dn : gr->device_list) {
+            response->add_device_name(dn);
+          }
+          for (const string& tn : gr->task_list) {
+            response->add_task_name(tn);
+          }
+        } else {
+          LOG(ERROR) << "Bad status from CompleteGroupDistributed: " << s;
+        }
+        done(s);
+      });
+}
+
+void CollectiveParamResolverDistributed::CompleteInstanceAsync(
+    const CompleteInstanceRequest* request, CompleteInstanceResponse* response,
+    CancellationManager* cancel_mgr, const StatusCallback& done) {
+  CollectiveParams* cp = new CollectiveParams;
+  cp->name = request->name();
+  cp->group.group_key = request->group_key();
+  cp->group.group_size = request->group_size();
+  cp->group.device_type = DeviceType(request->device_type());
+  cp->instance.type = CollectiveType(request->type());
+  cp->instance.instance_key = request->instance_key();
+  cp->instance.data_type = request->data_type();
+  cp->instance.shape = TensorShape(request->shape());
+  for (int32 offset : request->subdiv_offset()) {
+    cp->instance.impl_details.subdiv_offsets.push_back(offset);
+  }
+  VLOG(1) << "New cp " << cp << " for device " << request->device() << " : "
+          << cp->ToString();
+  StatusCallback done_and_cleanup = [this, cp, done](const Status& s) {
+    done(s);
+    delete cp;
+  };
+  // Start by completing the group.
+  CompleteGroupDistributed(
+      request->device(), cp, cancel_mgr,
+      [this, cp, request, response, cancel_mgr, done_and_cleanup](
+          const Status& cg_status, const GroupRec* gr) {
+        if (cg_status.ok()) {
+          // Then complete the instance.
+          CompleteInstanceDistributed(
+              request->device(), gr, cp, cancel_mgr,
+              [this, gr, cp, response,
+               done_and_cleanup](const Status& ci_status) {
+                if (ci_status.ok()) {
+                  // Now source_rank should be known, so
+                  // retrieve it.
+                  FindInstanceRec(
+                      gr, cp,
+                      [this, gr, cp, response, done_and_cleanup](
+                          const Status& fi_status, InstanceRec* ir) {
+                        if (fi_status.ok()) {
+                          mutex_lock l(ir->out_mu);
+                          response->set_instance_key(cp->instance.instance_key);
+                          response->set_source_rank(ir->source_rank);
+                          done_and_cleanup(fi_status);
+                        } else {
+                          done_and_cleanup(fi_status);
+                        }
+                      });
+                } else {
+                  done_and_cleanup(ci_status);
+                }
+              });
+        } else {
+          done_and_cleanup(cg_status);
+        }
+      });
+}
+
+bool CollectiveParamResolverDistributed::GroupIsCached(int32 group_key) {
+  mutex_lock l(group_mu_);
+  const auto& it = group_table_.find(group_key);
+  return it != group_table_.end();
+}
+
+Status CollectiveParamResolverDistributed::UpdateGroupCache(
+    const CompleteGroupResponse& resp) {
+  // Build a new record from resp.
+  std::unique_ptr<GroupRec> gr(new GroupRec);
+  mutex_lock grl(gr->mu);
+  gr->group.device_type = DeviceType(resp.device_type());
+  gr->group.group_key = resp.group_key();
+  gr->group.group_size = resp.group_size();
+  gr->group.num_tasks = resp.num_tasks();
+  if (resp.device_name_size() != gr->group.group_size) {
+    return errors::Internal(
+        "CompleteGroupResponse group_size doesn't match device_name list");
+  }
+  for (const string& dn : resp.device_name()) {
+    gr->device_set.insert(dn);
+    gr->device_list.push_back(dn);
+  }
+  if (resp.task_name_size() != gr->group.group_size) {
+    return errors::Internal(
+        "CompleteGroupResponse group_size doesn't match task_name list");
+  }
+  for (const string& tn : resp.task_name()) {
+    gr->task_list.push_back(tn);
+    gr->task_set.insert(tn);
+  }
+  CHECK_EQ(gr->task_set.size(), gr->group.num_tasks);
+  {
+    // Group membership should never change. Once a record is in group_table_
+    // it never gets removed.
+    mutex_lock l(group_mu_);
+    auto it = group_table_.find(gr->group.group_key);
+    if (it == group_table_.end()) {
+      group_table_[gr->group.group_key] = std::move(gr);
+    }
+  }
+  return Status::OK();
+}
+
+void CollectiveParamResolverDistributed::CompleteGroupDistributed(
+    const string& device, CollectiveParams* cp, CancellationManager* cancel_mgr,
+    const GroupRecCallback& done) {
+  VLOG(1) << "CompleteGroupDistributed group_key=" << cp->group.group_key
+          << " dev: " << device << " is_leader=" << (group_leader_.empty());
+  VLOG(0) << "cp: " << cp->ToString();
+  if (group_leader_.empty()) {
+    // This is the group leader, so resolution is local.
+    return CompleteGroupLocal(device, cp, done);
+  } else if (!GroupIsCached(cp->group.group_key)) {
+    // Need to update Group cache from the leader.
+    CompleteGroupCall* call = new CompleteGroupCall(
+        cp->group, device, cancel_mgr, group_leader_, worker_cache_);
+    call->Start([this, device, cp, call, done](const Status& s) {
+      if (s.ok()) {
+        Status status = UpdateGroupCache(call->resp_);
+        if (status.ok()) {
+          CompleteGroupLocal(device, cp, done);
+        } else {
+          done(status, nullptr);
+        }
+      } else {
+        done(s, nullptr);
+      }
+      delete call;
+    });
+    return;
+  } else {
+    return CompleteGroupLocal(device, cp, done);
+  }
+}
+
+bool CollectiveParamResolverDistributed::InstanceIsCached(int32 instance_key) {
+  mutex_lock l(instance_mu_);
+  const auto& it = instance_table_.find(instance_key);
+  return it != instance_table_.end();
+}
+
+void CollectiveParamResolverDistributed::UpdateInstanceCache(
+    const GroupRec* gr, CollectiveParams* cp,
+    const CompleteInstanceResponse& resp, const StatusCallback& done) {
+  Notification note;
+  InstanceRec* ir = nullptr;
+  int32 source_rank = resp.source_rank();
+
+  auto continue_with_ir = [this, cp, &ir, source_rank, done](const Status& s) {
+    if (!s.ok()) {
+      done(s);
+      return;
+    }
+    Status status;
+    do {
+      mutex_lock l(ir->out_mu);
+      if (ir->source_rank != source_rank) {
+        if (ir->source_rank >= 0) {
+          ir->status = errors::Internal(
+              "UpdateInstanceCache: CompleteInstanceResponse for instance ",
+              cp->instance.instance_key, " gives source_rank=", source_rank,
+              " but cache already holds value=", ir->source_rank);
+          status = ir->status;
+          break;
+        }
+        ir->source_rank = source_rank;
+      }
+      if (ir->known_count < cp->group.group_size) {
+        ir->known_count = cp->group.group_size;
+        if (ir->known.size() != cp->group.group_size) {
+          ir->status = errors::Internal(
+              "UpdateInstanceCache:: CompleteInstanceResponse for instance ",
+              cp->instance.instance_key, " has known.size()=", ir->known.size(),
+              " < group_size=", cp->group.group_size);
+          status = ir->status;
+          break;
+        }
+        for (int i = 0; i < ir->known.size(); ++i) {
+          ir->known[i] = true;
+        }
+      }
+      status = ir->status;
+    } while (false);
+    // Callback outside of lock.
+    done(status);
+  };
+
+  FindInstanceRec(
+      gr, cp, [this, &ir, continue_with_ir](const Status s, InstanceRec* irec) {
+        ir = irec;
+        continue_with_ir(s);
+      });
+}
+
+void CollectiveParamResolverDistributed::CompleteInstanceDistributed(
+    const string& device, const GroupRec* gr, CollectiveParams* cp,
+    CancellationManager* cancel_mgr, const StatusCallback& done) {
+  if (group_leader_.empty()) {
+    // This is the group leader so resolution is local.
+    return CompleteInstanceLocal(device, gr, cp, cp->is_source, done);
+  } else if (InstanceIsCached(cp->instance.instance_key)) {
+    return CompleteInstanceLocal(device, gr, cp, cp->is_source, done);
+  } else {
+    CompleteInstanceCall* call = new CompleteInstanceCall(
+        cp->group, cp->instance, cp->name, device, cp->is_source, cancel_mgr,
+        group_leader_, worker_cache_);
+    call->Start([this, device, gr, cp, call, done](const Status& s) {
+      if (s.ok()) {
+        UpdateInstanceCache(
+            gr, cp, call->resp_, [this, device, gr, cp, done](const Status& s) {
+              if (!s.ok()) {
+                done(s);
+              } else {
+                CompleteInstanceLocal(device, gr, cp, cp->is_source, done);
+              }
+            });
+      } else {
+        done(s);
+      }
+      delete call;
+    });
+    return;
+  }
+}
+
+}  // namespace tensorflow
diff --git a/tensorflow/core/distributed_runtime/collective_param_resolver_distributed.h b/tensorflow/core/distributed_runtime/collective_param_resolver_distributed.h
new file mode 100644
index 0000000000..a35131d835
--- /dev/null
+++ b/tensorflow/core/distributed_runtime/collective_param_resolver_distributed.h
@@ -0,0 +1,90 @@
+/* Copyright 2018 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.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_COLLECTIVE_PARAM_RESOLVER_DISTRIBUTED_H_
+#define TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_COLLECTIVE_PARAM_RESOLVER_DISTRIBUTED_H_
+
+#include "tensorflow/core/common_runtime/collective_param_resolver_local.h"
+
+namespace tensorflow {
+class ConfigProto;
+class WorkerCacheInterface;
+class DeviceResolverDistributed;
+class DeviceMgr;
+
+class CollectiveParamResolverDistributed : public CollectiveParamResolverLocal {
+ public:
+  CollectiveParamResolverDistributed(const ConfigProto& config,
+                                     const DeviceMgr* dev_mgr,
+                                     DeviceResolverDistributed* dev_resolver,
+                                     WorkerCacheInterface* worker_cache,
+                                     const string& task_name);
+
+  void CompleteParamsAsync(const string& device, CollectiveParams* cp,
+                           CancellationManager* cancel_mgr,
+                           const StatusCallback& done) override;
+
+  void CompleteGroupAsync(const CompleteGroupRequest* request,
+                          CompleteGroupResponse* response,
+                          CancellationManager* cancel_mgr,
+                          const StatusCallback& done) override;
+
+  void CompleteInstanceAsync(const CompleteInstanceRequest* request,
+                             CompleteInstanceResponse* response,
+                             CancellationManager* cancel_mgr,
+                             const StatusCallback& done) override;
+
+ protected:
+  // Returns true iff there's an entry for this group_key in the
+  // local group_table_.
+  bool GroupIsCached(int32 group_key) LOCKS_EXCLUDED(group_mu_);
+
+  // Updates group_table_ with contents of resp.
+  Status UpdateGroupCache(const CompleteGroupResponse& resp)
+      LOCKS_EXCLUDED(group_mu_);
+
+  // Finds the GroupRec that corresponds to cp->group_key and also
+  // populates cp->group from that GroupRec.
+  //
+  // Semantics are like those of CompleteGroupLocal but will make a
+  // remote call to the group leader if necessary.
+  void CompleteGroupDistributed(const string& device, CollectiveParams* cp,
+                                CancellationManager* cancel_mgr,
+                                const GroupRecCallback& done);
+
+  // Returns true iff there's an entry for this instance_key in the
+  // local instance_table_.
+  bool InstanceIsCached(int32 instance_key) LOCKS_EXCLUDED(instance_mu_);
+
+  // Updates instance_table_ with contents of resp.
+  void UpdateInstanceCache(const GroupRec* gr, CollectiveParams* cp,
+                           const CompleteInstanceResponse& resp,
+                           const StatusCallback& done)
+      LOCKS_EXCLUDED(instance_mu_, gr->mu, group_mu_);
+
+  // Finish populating *cp.  Semantics are like those of
+  // CompleteInstanceLocal but will make a remote call to the group
+  // leader if necessary.
+  void CompleteInstanceDistributed(const string& device, const GroupRec* gr,
+                                   CollectiveParams* cp,
+                                   CancellationManager* cancel_mgr,
+                                   const StatusCallback& done)
+      LOCKS_EXCLUDED(instance_mu_, gr->mu, group_mu_);
+
+  WorkerCacheInterface* worker_cache_;  // Not owned
+  const string group_leader_;
+};
+
+}  // namespace tensorflow
+#endif  // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_COLLECTIVE_PARAM_RESOLVER_DISTRIBUTED_H_
diff --git a/tensorflow/core/distributed_runtime/collective_param_resolver_distributed_test.cc b/tensorflow/core/distributed_runtime/collective_param_resolver_distributed_test.cc
new file mode 100644
index 0000000000..95a010286d
--- /dev/null
+++ b/tensorflow/core/distributed_runtime/collective_param_resolver_distributed_test.cc
@@ -0,0 +1,324 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/core/distributed_runtime/collective_param_resolver_distributed.h"
+
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/distributed_runtime/device_resolver_distributed.h"
+#include "tensorflow/core/distributed_runtime/test_utils.h"
+#include "tensorflow/core/framework/cancellation.h"
+#include "tensorflow/core/lib/core/notification.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/util/device_name_utils.h"
+
+namespace tensorflow {
+namespace {
+
+static Device* NewDevice(const string& type, const string& name) {
+  class FakeDevice : public Device {
+   public:
+    explicit FakeDevice(const DeviceAttributes& attr) : Device(nullptr, attr) {}
+    Status Sync() override { return Status::OK(); }
+    Allocator* GetAllocator(AllocatorAttributes) override { return nullptr; }
+  };
+  DeviceAttributes attr;
+  attr.set_name(name);
+  attr.set_device_type(type);
+  attr.mutable_locality()->set_numa_node(3);  // a non-default value
+  return new FakeDevice(attr);
+}
+
+class FakeWorker : public TestWorkerInterface {
+ public:
+  FakeWorker(const string& name, DeviceMgr* dev_mgr,
+             CollectiveParamResolverDistributed* cpres)
+      : name_(name), device_mgr_(dev_mgr), param_resolver_(cpres) {}
+
+  void GetStatusAsync(const GetStatusRequest* request,
+                      GetStatusResponse* response,
+                      StatusCallback done) override {
+    std::vector<DeviceAttributes> dev_attr;
+    device_mgr_->ListDeviceAttributes(&dev_attr);
+    for (const auto& da : dev_attr) {
+      *response->add_device_attributes() = da;
+    }
+    done(Status::OK());
+  }
+
+  void CompleteGroupAsync(CallOptions* opts,
+                          const CompleteGroupRequest* request,
+                          CompleteGroupResponse* response,
+                          StatusCallback done) override {
+    param_resolver_->CompleteGroupAsync(request, response, &cm_, done);
+  }
+
+  void CompleteInstanceAsync(CallOptions* ops,
+                             const CompleteInstanceRequest* request,
+                             CompleteInstanceResponse* response,
+                             StatusCallback done) override {
+    param_resolver_->CompleteInstanceAsync(request, response, &cm_, done);
+  }
+
+ private:
+  string name_;
+  DeviceMgr* device_mgr_;
+  CancellationManager cm_;
+  CollectiveParamResolverDistributed* param_resolver_;
+};
+
+class FakeCache : public TestWorkerCache {
+ public:
+  // Override the Locality methods to actually pass through to the
+  // worker.
+  bool GetDeviceLocalityNonBlocking(const string& device,
+                                    DeviceLocality* locality) override {
+    return false;
+  }
+
+  void GetDeviceLocalityAsync(const string& device, DeviceLocality* locality,
+                              StatusCallback done) override {
+    string task_name;
+    string dev_part;
+    if (!DeviceNameUtils::SplitDeviceName(device, &task_name, &dev_part)) {
+      done(errors::Internal("failed to parse device name"));
+      return;
+    }
+    auto it = workers_.find(task_name);
+    if (it == workers_.end()) {
+      done(errors::Internal("failed to find worker ", task_name));
+      return;
+    }
+    WorkerInterface* wi = it->second;
+    GetStatusRequest req;
+    GetStatusResponse resp;
+    Notification note;
+    Status status = wi->GetStatus(&req, &resp);
+    if (!status.ok()) {
+      done(status);
+      return;
+    }
+    for (const auto& it : resp.device_attributes()) {
+      if (it.name() == device) {
+        *locality = it.locality();
+        done(Status::OK());
+        return;
+      }
+    }
+    done(errors::Internal("device not found: ", device));
+  }
+};
+
+class DeviceResDistTest : public ::testing::Test {
+ protected:
+  DeviceResDistTest() {}
+
+  ~DeviceResDistTest() override {
+    for (DeviceMgr* dm : device_mgrs_) {
+      delete dm;
+    }
+    for (auto it : dev_resolvers_) {
+      delete it.second;
+    }
+    for (auto it : cp_resolvers_) {
+      delete it.second;
+    }
+    for (FakeWorker* w : workers_) {
+      delete w;
+    }
+  }
+
+  void DefineWorkers(int num_workers, int num_devices,
+                     const string& device_type) {
+    ConfigProto config;
+    for (int w = 0; w < num_workers; ++w) {
+      string name = strings::StrCat("/job:worker/replica:0/task:", w);
+      // TODO(tucker): When config option becomes available, set here.
+      // if (w == 0) {
+      //   config.set_collective_group_leader(name);
+      // }
+      DefineWorker(config, name, device_type, num_devices);
+    }
+  }
+
+  void DefineWorker(const ConfigProto& config, const string& worker_name,
+                    const string& device_type, int num_devices) {
+    std::vector<Device*> devices;
+    for (int i = 0; i < num_devices; ++i) {
+      devices.push_back(NewDevice(
+          device_type,
+          strings::StrCat(worker_name, "/device:", device_type, ":", i)));
+    }
+    DeviceMgr* dev_mgr = new DeviceMgr(devices);
+    device_mgrs_.push_back(dev_mgr);
+    std::vector<string>* dv = &dev_by_task_[worker_name];
+    for (auto d : devices) {
+      dv->push_back(d->name());
+    }
+    DeviceResolverDistributed* dev_res =
+        new DeviceResolverDistributed(dev_mgr, &wc_, worker_name);
+    dev_resolvers_[worker_name] = dev_res;
+    CollectiveParamResolverDistributed* cp_res =
+        new CollectiveParamResolverDistributed(config, dev_mgr, dev_res, &wc_,
+                                               worker_name);
+    cp_resolvers_[worker_name] = cp_res;
+    FakeWorker* fw = new FakeWorker(worker_name, dev_mgr, cp_res);
+    workers_.push_back(fw);
+    wc_.AddWorker(worker_name, fw);
+  }
+
+  void DefineCollectiveParams(int num_workers, int num_devices) {
+    const int kGroupKey = 5;
+    const int kInstanceKey = 3;
+    for (int wi = 0; wi < num_workers; ++wi) {
+      string task_name = strings::StrCat("/job:worker/replica:0/task:", wi);
+      for (int di = 0; di < num_devices; ++di) {
+        string device_name = strings::StrCat(task_name, "/device:CPU:", di);
+        cp_.push_back(CollectiveParams());
+        CollectiveParams& cp = cp_.back();
+        cp.group.group_key = kGroupKey;
+        cp.group.group_size = num_workers * num_devices;
+        cp.group.device_type = DEVICE_CPU;
+        cp.group.num_tasks = num_workers;
+        cp.instance.instance_key = kInstanceKey;
+        cp.instance.type = REDUCTION_COLLECTIVE;
+        cp.instance.data_type = DT_FLOAT;
+        cp.instance.shape = TensorShape({64});
+        cp.instance.impl_details.subdiv_offsets.push_back(0);
+      }
+    }
+  }
+
+  void IssueRequests(int num_workers, int num_devices) {
+    const int device_count = num_workers * num_devices;
+    {
+      mutex_lock l(mu_);
+      num_done_ = 0;
+    }
+    cp_.resize(device_count);
+    status_.resize(device_count);
+    int idx = 0;
+    for (int wi = 0; wi < num_workers; ++wi) {
+      for (int di = 0; di < num_devices; ++di) {
+        IssueRequest(num_workers, num_devices, idx);
+        ++idx;
+      }
+    }
+  }
+
+  void IssueRequest(int num_workers, int num_devices, int idx) {
+    int device_count = num_workers * num_devices;
+    int wi = idx / num_devices;
+    int di = idx % num_devices;
+    string task_name = strings::StrCat("/job:worker/replica:0/task:", wi);
+    string device_name = strings::StrCat(task_name, "/device:CPU:", di);
+    while (idx >= cp_.size()) {
+      status_.resize(idx + 1);
+      cp_.resize(idx + 1);
+    }
+    CollectiveParams* cp = &cp_[idx];
+    CollectiveParamResolverDistributed* cp_res = cp_resolvers_[task_name];
+    CHECK(cp_res);
+    cp_res->CompleteParamsAsync(device_name, cp, &cm_,
+                                [this, idx, device_count](const Status& s) {
+                                  status_[idx] = s;
+                                  {
+                                    mutex_lock l(mu_);
+                                    ++num_done_;
+                                    if (num_done_ == device_count) {
+                                      done_.notify_all();
+                                    }
+                                  }
+                                });
+  }
+
+  void ValidateCollectiveParams(int num_workers, int num_devices) {
+    int device_count = num_workers * num_devices;
+    {
+      mutex_lock l(mu_);
+      if (num_done_ < device_count) {
+        done_.wait(l);
+      }
+    }
+    // Verify that all cp_ values get the same set of task and device
+    // names, with unique default_rank in the expected order.
+    const int dev_count = num_workers * num_devices;
+    for (int wi = 0; wi < num_workers; ++wi) {
+      string task_name = strings::StrCat("/job:worker/replica:0/task:", wi);
+      for (int di = 0; di < num_devices; ++di) {
+        string device_name = strings::StrCat(task_name, "/device:CPU:", di);
+        int idx = wi * num_devices + di;
+        TF_ASSERT_OK(status_[idx]);
+        EXPECT_EQ(cp_[idx].default_rank, idx);
+        EXPECT_EQ(cp_[idx].instance.device_names.size(), dev_count);
+        EXPECT_EQ(cp_[idx].instance.device_names[idx], device_name);
+        EXPECT_EQ(cp_[idx].instance.task_names[idx], task_name);
+        if (idx > 0) {
+          for (int i = 0; i < dev_count; ++i) {
+            EXPECT_EQ(cp_[0].instance.device_names[i],
+                      cp_[idx].instance.device_names[i]);
+            EXPECT_EQ(cp_[0].instance.task_names[i],
+                      cp_[idx].instance.task_names[i]);
+          }
+        }
+      }
+    }
+  }
+
+  FakeCache wc_;
+  CancellationManager cm_;
+  std::vector<DeviceMgr*> device_mgrs_;
+  std::unordered_map<string, DeviceResolverDistributed*> dev_resolvers_;
+  std::unordered_map<string, CollectiveParamResolverDistributed*> cp_resolvers_;
+  std::unordered_map<string, std::vector<string>> dev_by_task_;
+  std::vector<FakeWorker*> workers_;
+  std::vector<CollectiveParams> cp_;
+  std::vector<Status> status_;
+  mutex mu_;
+  int num_done_ GUARDED_BY(mu_);
+  condition_variable done_;
+};
+
+TEST_F(DeviceResDistTest, Workers1Devices1) {
+  const int num_workers = 1;
+  const int num_devices = 1;
+  DefineWorkers(num_workers, num_devices, "CPU");
+  DefineCollectiveParams(num_workers, num_devices);
+  IssueRequests(num_workers, num_devices);
+  ValidateCollectiveParams(num_workers, num_devices);
+}
+
+TEST_F(DeviceResDistTest, Workers2Devices2) {
+  const int num_workers = 2;
+  const int num_devices = 2;
+  DefineWorkers(num_workers, num_devices, "CPU");
+  DefineCollectiveParams(num_workers, num_devices);
+  IssueRequests(num_workers, num_devices);
+  ValidateCollectiveParams(num_workers, num_devices);
+}
+
+TEST_F(DeviceResDistTest, Workers4Devices3) {
+  const int num_workers = 4;
+  const int num_devices = 3;
+  DefineWorkers(num_workers, num_devices, "CPU");
+  DefineCollectiveParams(num_workers, num_devices);
+  IssueRequests(num_workers, num_devices);
+  ValidateCollectiveParams(num_workers, num_devices);
+}
+
+}  // namespace
+}  // namespace tensorflow
diff --git a/tensorflow/core/distributed_runtime/device_resolver_distributed.cc b/tensorflow/core/distributed_runtime/device_resolver_distributed.cc
new file mode 100644
index 0000000000..038974cb39
--- /dev/null
+++ b/tensorflow/core/distributed_runtime/device_resolver_distributed.cc
@@ -0,0 +1,133 @@
+/* Copyright 2018 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.
+==============================================================================*/
+#include "tensorflow/core/distributed_runtime/device_resolver_distributed.h"
+
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/distributed_runtime/worker_cache.h"
+
+namespace tensorflow {
+DeviceResolverDistributed::DeviceResolverDistributed(
+    const DeviceMgr* dev_mgr, WorkerCacheInterface* worker_cache,
+    const string& task_name)
+    : dev_mgr_(dev_mgr), worker_cache_(worker_cache), task_name_(task_name) {}
+
+void DeviceResolverDistributed::GetLocalityAsync(const string& device,
+                                                 const string& task,
+                                                 DeviceLocality* locality,
+                                                 const StatusCallback& done) {
+  if (task.empty() || task == task_name_) {
+    // Device is local to this task.
+    Device* dev;
+    Status s = dev_mgr_->LookupDevice(device, &dev);
+    if (s.ok()) {
+      *locality = dev->attributes().locality();
+    }
+    done(s);
+    return;
+  } else {
+    // Lookup of a remote device: first try the local cache.
+    bool found = false;
+    {
+      mutex_lock l(mu_);
+      auto it = attr_table_.find(device);
+      if (it != attr_table_.end()) {
+        *locality = it->second.locality();
+        found = true;
+      }
+    }
+    if (found) {
+      done(Status::OK());
+      return;
+    }
+  }
+  // Device is remote and no cache entry was found.  Refresh the cache
+  // then retry the lookup.
+  RefreshRemoteAttributes(
+      device, task, [this, device, task, locality, done](const Status& s) {
+        if (!s.ok()) {
+          done(s);
+        } else {
+          GetLocalityAsync(device, task, locality, done);
+        }
+      });
+}
+
+void DeviceResolverDistributed::GetDeviceLocalitiesAsync(
+    const CollInstanceParams& inst_params,
+    std::vector<DeviceLocality>* localities, const StatusCallback& done) {
+  localities->clear();
+  GetDeviceLocalitiesRecursive(inst_params, localities, done);
+}
+
+void DeviceResolverDistributed::GetDeviceLocalitiesRecursive(
+    const CollInstanceParams& inst_params,
+    std::vector<DeviceLocality>* localities, const StatusCallback& done) {
+  size_t i = localities->size();
+  if (i < inst_params.device_names.size()) {
+    localities->push_back(DeviceLocality());
+    GetLocalityAsync(inst_params.device_names[i], inst_params.task_names[i],
+                     &localities->back(),
+                     [this, &inst_params, localities, done](const Status& s) {
+                       if (!s.ok()) {
+                         done(s);
+                         return;
+                       } else {
+                         GetDeviceLocalitiesRecursive(inst_params, localities,
+                                                      done);
+                       }
+                     });
+  } else {
+    done(Status::OK());
+  }
+}
+
+void DeviceResolverDistributed::RefreshRemoteAttributes(
+    const string& device, const string& task, const StatusCallback& done) {
+  GetStatusRequest* req = new GetStatusRequest;
+  GetStatusResponse* resp = new GetStatusResponse;
+  WorkerInterface* worker = worker_cache_->CreateWorker(task);
+  CHECK(worker) << "Failed to get worker for " << task;
+  worker->GetStatusAsync(
+      req, resp, [this, device, task, req, resp, worker, done](Status s) {
+        if (s.ok()) {
+          mutex_lock l(mu_);
+          for (const DeviceAttributes& da : resp->device_attributes()) {
+            attr_table_[da.name()] = da;
+          }
+        }
+        done(s);
+        delete req;
+        delete resp;
+        worker_cache_->ReleaseWorker(task, worker);
+      });
+}
+
+void DeviceResolverDistributed::ClearTask(const string& task) {
+  mutex_lock l(mu_);
+  // First find all the keys belonging to the task.
+  std::unordered_set<string> task_keys;
+  for (const auto& it : attr_table_) {
+    const string& device_name = it.first;
+    if (DeviceNameUtils::IsSameAddressSpace(task, device_name)) {
+      task_keys.insert(device_name);
+    }
+  }
+  // Then delete them.
+  for (const string& key : task_keys) {
+    attr_table_.erase(key);
+  }
+}
+
+}  // namespace tensorflow
diff --git a/tensorflow/core/distributed_runtime/device_resolver_distributed.h b/tensorflow/core/distributed_runtime/device_resolver_distributed.h
new file mode 100644
index 0000000000..ac68ec6873
--- /dev/null
+++ b/tensorflow/core/distributed_runtime/device_resolver_distributed.h
@@ -0,0 +1,67 @@
+/* Copyright 2018 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.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_DEVICE_RESOLVER_DISTRIBUTED_H_
+#define TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_DEVICE_RESOLVER_DISTRIBUTED_H_
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/core/framework/collective.h"
+#include "tensorflow/core/framework/device_attributes.pb.h"
+#include "tensorflow/core/lib/gtl/flatmap.h"
+
+namespace tensorflow {
+class DeviceMgr;
+class WorkerCacheInterface;
+
+class DeviceResolverDistributed : public DeviceResolverInterface {
+ public:
+  DeviceResolverDistributed(const DeviceMgr* dev_mgr,
+                            WorkerCacheInterface* worker_cache,
+                            const string& task_name);
+
+  virtual ~DeviceResolverDistributed() {}
+
+  void GetDeviceLocalitiesAsync(const CollInstanceParams& inst_params,
+                                std::vector<DeviceLocality>* localities,
+                                const StatusCallback& done) override;
+
+  void GetLocalityAsync(const string& device, const string& task,
+                        DeviceLocality* locality,
+                        const StatusCallback& done) override;
+
+  void ClearTask(const string& task) override;
+
+ protected:
+  // Loads attr_table_ with device attributes retrieved from remote task.
+  void RefreshRemoteAttributes(const string& device, const string& task,
+                               const StatusCallback& done) LOCKS_EXCLUDED(mu_);
+
+  // Subroutine used by GetDeviceLocalitiesAsync.  Recursively extends
+  // *localities with DeviceLocality of the corresponding device named
+  // by inst_params.instance.device_names.
+  void GetDeviceLocalitiesRecursive(const CollInstanceParams& inst_params,
+                                    std::vector<DeviceLocality>* localities,
+                                    const StatusCallback& done);
+
+  const DeviceMgr* dev_mgr_;            // Not owned
+  WorkerCacheInterface* worker_cache_;  // Not owned
+  const string task_name_;
+  mutex mu_;
+  gtl::FlatMap<string, DeviceAttributes> attr_table_ GUARDED_BY(mu_);
+};
+
+}  // namespace tensorflow
+#endif  // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_DEVICE_RESOLVER_DISTRIBUTED_H_
diff --git a/tensorflow/core/distributed_runtime/device_resolver_distributed_test.cc b/tensorflow/core/distributed_runtime/device_resolver_distributed_test.cc
new file mode 100644
index 0000000000..ae44b98bd5
--- /dev/null
+++ b/tensorflow/core/distributed_runtime/device_resolver_distributed_test.cc
@@ -0,0 +1,217 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/core/distributed_runtime/device_resolver_distributed.h"
+
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/distributed_runtime/test_utils.h"
+#include "tensorflow/core/lib/core/notification.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/util/device_name_utils.h"
+
+namespace tensorflow {
+namespace {
+
+// Subclass of DeviceResolverDistributed which behaves identically but
+// allows access to the attr_table_.
+class TestableDeviceResolverDistributed : public DeviceResolverDistributed {
+ public:
+  TestableDeviceResolverDistributed(const DeviceMgr* dev_mgr,
+                                    WorkerCacheInterface* worker_cache,
+                                    const string& task)
+      : DeviceResolverDistributed(dev_mgr, worker_cache, task) {}
+
+  gtl::FlatMap<string, DeviceAttributes>& attr_table() { return attr_table_; }
+};
+
+// Create a fake 'Device' whose only interesting attribute is a non-default
+// DeviceLocality.
+static Device* NewDevice(const string& type, const string& name,
+                         int numa_node) {
+  class FakeDevice : public Device {
+   public:
+    explicit FakeDevice(const DeviceAttributes& attr) : Device(nullptr, attr) {}
+    Status Sync() override { return Status::OK(); }
+    Allocator* GetAllocator(AllocatorAttributes) override { return nullptr; }
+  };
+  DeviceAttributes attr;
+  attr.set_name(name);
+  attr.set_device_type(type);
+  attr.mutable_locality()->set_numa_node(numa_node);
+  return new FakeDevice(attr);
+}
+
+// Create a fake WorkerInterface that responds to requests without RPCs,
+// in this case returning the DeviceAttributes of a fake remote worker.
+class FakeWorker : public TestWorkerInterface {
+ public:
+  FakeWorker(const string& name, DeviceMgr* dev_mgr,
+             DeviceResolverDistributed* dres)
+      : name_(name), device_mgr_(dev_mgr), device_resolver_(dres) {}
+
+  void GetStatusAsync(const GetStatusRequest* request,
+                      GetStatusResponse* response,
+                      StatusCallback done) override {
+    std::vector<DeviceAttributes> dev_attr;
+    device_mgr_->ListDeviceAttributes(&dev_attr);
+    for (const auto& da : dev_attr) {
+      *response->add_device_attributes() = da;
+    }
+    done(Status::OK());
+  }
+
+ private:
+  string name_;
+  DeviceMgr* device_mgr_;
+  DeviceResolverDistributed* device_resolver_;
+};
+
+// An implementation of WorkerCacheInterface that routes all requests
+// to local FakeWorkers, implementing only the methods needed for tests.
+class FakeCache : public TestWorkerCache {
+ public:
+  // Override the Locality methods to actually pass through to the
+  // worker.
+  bool GetDeviceLocalityNonBlocking(const string& device,
+                                    DeviceLocality* locality) override {
+    return false;
+  }
+
+  void GetDeviceLocalityAsync(const string& device, DeviceLocality* locality,
+                              StatusCallback done) override {
+    string task_name;
+    string dev_part;
+    if (!DeviceNameUtils::SplitDeviceName(device, &task_name, &dev_part)) {
+      done(errors::Internal("failed to parse device name"));
+      return;
+    }
+    auto it = workers_.find(task_name);
+    if (it == workers_.end()) {
+      done(errors::Internal("failed to find worker ", task_name));
+      return;
+    }
+    WorkerInterface* wi = it->second;
+    GetStatusRequest req;
+    GetStatusResponse resp;
+    Notification note;
+    Status status = wi->GetStatus(&req, &resp);
+    if (!status.ok()) {
+      done(status);
+      return;
+    }
+    for (const auto& it : resp.device_attributes()) {
+      if (it.name() == device) {
+        *locality = it.locality();
+        done(Status::OK());
+        return;
+      }
+    }
+    done(errors::Internal("device not found: ", device));
+  }
+};
+
+class DeviceResDistTest : public ::testing::Test {
+ protected:
+  DeviceResDistTest() {}
+
+  ~DeviceResDistTest() override {
+    for (DeviceMgr* dm : device_mgrs_) {
+      delete dm;
+    }
+    for (auto it : resolvers_) {
+      delete it.second;
+    }
+    for (FakeWorker* w : workers_) {
+      delete w;
+    }
+  }
+
+  void DefineWorkers(int num_workers, int num_devices,
+                     const string& device_type) {
+    for (int w = 0; w < num_workers; ++w) {
+      string name = strings::StrCat("/job:worker/replica:0/task:", w);
+      DefineWorker(name, device_type, num_devices);
+    }
+  }
+
+  void DefineWorker(const string& worker_name, const string& device_type,
+                    int num_devices) {
+    std::vector<Device*> devices;
+    for (int i = 0; i < num_devices; ++i) {
+      devices.push_back(NewDevice(
+          device_type,
+          strings::StrCat(worker_name, "/device:", device_type, ":", i), i));
+    }
+    DeviceMgr* dev_mgr = new DeviceMgr(devices);
+    TestableDeviceResolverDistributed* dev_res =
+        new TestableDeviceResolverDistributed(dev_mgr, &wc_, worker_name);
+    resolvers_[worker_name] = dev_res;
+    device_mgrs_.push_back(dev_mgr);
+    std::vector<string>* dv = &dev_by_task_[worker_name];
+    for (auto d : devices) {
+      dv->push_back(d->name());
+    }
+    FakeWorker* fw = new FakeWorker(worker_name, dev_mgr, dev_res);
+    workers_.push_back(fw);
+    wc_.AddWorker(worker_name, fw);
+  }
+
+  FakeCache wc_;
+  std::vector<DeviceMgr*> device_mgrs_;
+  std::unordered_map<string, TestableDeviceResolverDistributed*> resolvers_;
+  std::unordered_map<string, std::vector<string>> dev_by_task_;
+  std::vector<FakeWorker*> workers_;
+};
+
+TEST_F(DeviceResDistTest, Workers3Devices4) {
+  DefineWorkers(3, 4, "CPU");
+  // Check that every device is available from every task.
+  for (auto it : resolvers_) {
+    DeviceResolverDistributed* dres = it.second;
+    for (auto it2 : dev_by_task_) {
+      const string& task_name = it2.first;
+      for (const auto& dev_name : it2.second) {
+        DeviceNameUtils::ParsedName parsed;
+        ASSERT_TRUE(DeviceNameUtils::ParseFullName(dev_name, &parsed));
+        Notification note;
+        Status status;
+        DeviceLocality locality;
+        dres->GetLocalityAsync(dev_name, task_name, &locality,
+                               [this, &note, &status](const Status& s) {
+                                 status = s;
+                                 note.Notify();
+                               });
+        note.WaitForNotification();
+        TF_EXPECT_OK(status);
+        EXPECT_EQ(parsed.id, locality.numa_node());
+      }
+    }
+  }
+  // Clear just task 0 from all.
+  const string w0_name = "/job:worker/replica:0/task:0";
+  for (auto it : resolvers_) {
+    if (it.first == w0_name) continue;
+    TestableDeviceResolverDistributed* dres = it.second;
+    EXPECT_EQ(8, it.second->attr_table().size());
+    dres->ClearTask("/job:worker/replica:0/task:0");
+    EXPECT_EQ(4, it.second->attr_table().size());
+  }
+}
+
+}  // namespace
+}  // namespace tensorflow
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_remote_worker.cc b/tensorflow/core/distributed_runtime/rpc/grpc_remote_worker.cc
index 895bbd97b7..5b7b74ce63 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_remote_worker.cc
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_remote_worker.cc
@@ -56,6 +56,9 @@ class GrpcRemoteWorker : public WorkerInterface {
         recvtensor_(Method(GrpcWorkerMethod::kRecvTensor)),
         logging_(Method(GrpcWorkerMethod::kLogging)),
         tracing_(Method(GrpcWorkerMethod::kTracing)),
+        completegroup_(Method(GrpcWorkerMethod::kCompleteGroup)),
+        instancesource_(Method(GrpcWorkerMethod::kCompleteInstance)),
+        getstepsequence_(Method(GrpcWorkerMethod::kGetStepSequence)),
         logger_(logger) {}
 
   ~GrpcRemoteWorker() override {}
@@ -115,6 +118,27 @@ class GrpcRemoteWorker : public WorkerInterface {
     IssueRequest(request, response, cleanupall_, std::move(done));
   }
 
+  void CompleteGroupAsync(CallOptions* call_opts,
+                          const CompleteGroupRequest* request,
+                          CompleteGroupResponse* response,
+                          StatusCallback done) override {
+    IssueRequest(request, response, completegroup_, std::move(done), call_opts);
+  }
+
+  void CompleteInstanceAsync(CallOptions* call_opts,
+                             const CompleteInstanceRequest* request,
+                             CompleteInstanceResponse* response,
+                             StatusCallback done) override {
+    IssueRequest(request, response, instancesource_, std::move(done),
+                 call_opts);
+  }
+
+  void GetStepSequenceAsync(const GetStepSequenceRequest* request,
+                            GetStepSequenceResponse* response,
+                            StatusCallback done) override {
+    IssueRequest(request, response, getstepsequence_, std::move(done));
+  }
+
   void RecvTensorAsync(CallOptions* call_opts, const RecvTensorRequest* request,
                        TensorResponse* response, StatusCallback done) override {
     VLOG(1) << "RecvTensorAsync req: " << request->DebugString();
@@ -217,6 +241,9 @@ class GrpcRemoteWorker : public WorkerInterface {
   const ::grpc::string recvtensor_;
   const ::grpc::string logging_;
   const ::grpc::string tracing_;
+  const ::grpc::string completegroup_;
+  const ::grpc::string instancesource_;
+  const ::grpc::string getstepsequence_;
 
   // Support for logging.
   WorkerCacheLogger* logger_;
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_worker_service.cc b/tensorflow/core/distributed_runtime/rpc/grpc_worker_service.cc
index b20e744a97..bbf7391377 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_worker_service.cc
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_worker_service.cc
@@ -172,6 +172,12 @@ class GrpcWorkerService : public AsyncServiceInterface {
       ENQUEUE_REQUEST(Logging, false);
       ENQUEUE_REQUEST(Tracing, false);
 
+      for (int i = 0; i < 10; ++i) {
+        ENQUEUE_REQUEST(CompleteGroup, false);
+        ENQUEUE_REQUEST(CompleteInstance, false);
+        ENQUEUE_REQUEST(GetStepSequence, false);
+      }
+
       void* tag;
       bool ok;
 
@@ -318,6 +324,47 @@ class GrpcWorkerService : public AsyncServiceInterface {
       });
       ENQUEUE_REQUEST(Tracing, false);
     }
+
+    void CompleteGroupHandler(
+        WorkerCall<CompleteGroupRequest, CompleteGroupResponse>* call) {
+      Schedule([this, call]() {
+        CallOptions* call_opts = new CallOptions;
+        call->SetCancelCallback([call_opts]() { call_opts->StartCancel(); });
+        worker_->CompleteGroupAsync(call_opts, &call->request, &call->response,
+                                    [call, call_opts](const Status& s) {
+                                      call->ClearCancelCallback();
+                                      delete call_opts;
+                                      call->SendResponse(ToGrpcStatus(s));
+                                    });
+      });
+      ENQUEUE_REQUEST(CompleteGroup, false);
+    }
+
+    void CompleteInstanceHandler(
+        WorkerCall<CompleteInstanceRequest, CompleteInstanceResponse>* call) {
+      Schedule([this, call]() {
+        CallOptions* call_opts = new CallOptions;
+        call->SetCancelCallback([call_opts]() { call_opts->StartCancel(); });
+        worker_->CompleteInstanceAsync(call_opts, &call->request,
+                                       &call->response,
+                                       [call, call_opts](const Status& s) {
+                                         call->ClearCancelCallback();
+                                         delete call_opts;
+                                         call->SendResponse(ToGrpcStatus(s));
+                                       });
+      });
+      ENQUEUE_REQUEST(CompleteInstance, false);
+    }
+
+    void GetStepSequenceHandler(
+        WorkerCall<GetStepSequenceRequest, GetStepSequenceResponse>* call) {
+      Schedule([this, call]() {
+        worker_->GetStepSequenceAsync(
+            &call->request, &call->response,
+            [call](const Status& s) { call->SendResponse(ToGrpcStatus(s)); });
+      });
+      ENQUEUE_REQUEST(GetStepSequence, false);
+    }
 #undef ENQUEUE_REQUEST
 
     void EnqueueRecvTensorRequestRaw() {
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_worker_service_impl.cc b/tensorflow/core/distributed_runtime/rpc/grpc_worker_service_impl.cc
index 05a9db10d3..a91cc0692a 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_worker_service_impl.cc
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_worker_service_impl.cc
@@ -50,6 +50,12 @@ const char* GrpcWorkerMethodName(GrpcWorkerMethod id) {
       return "/tensorflow.WorkerService/Logging";
     case GrpcWorkerMethod::kTracing:
       return "/tensorflow.WorkerService/Tracing";
+    case GrpcWorkerMethod::kCompleteGroup:
+      return "/tensorflow.WorkerService/CompleteGroup";
+    case GrpcWorkerMethod::kCompleteInstance:
+      return "/tensorflow.WorkerService/CompleteInstance";
+    case GrpcWorkerMethod::kGetStepSequence:
+      return "/tensorflow.WorkerService/GetStepSequence";
   }
   // Shouldn't be reached.
   LOG(FATAL) << "Invalid id: this line shouldn't be reached.";
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_worker_service_impl.h b/tensorflow/core/distributed_runtime/rpc/grpc_worker_service_impl.h
index a54ea93796..c5104c6a50 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_worker_service_impl.h
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_worker_service_impl.h
@@ -83,9 +83,12 @@ enum class GrpcWorkerMethod {
   kRecvTensor,
   kLogging,
   kTracing,
+  kCompleteGroup,
+  kCompleteInstance,
+  kGetStepSequence,
 };
 static const int kGrpcNumWorkerMethods =
-    static_cast<int>(GrpcWorkerMethod::kTracing) + 1;
+    static_cast<int>(GrpcWorkerMethod::kGetStepSequence) + 1;
 
 const char* GrpcWorkerMethodName(GrpcWorkerMethod id);
 
diff --git a/tensorflow/core/distributed_runtime/test_utils.h b/tensorflow/core/distributed_runtime/test_utils.h
new file mode 100644
index 0000000000..0ed078241f
--- /dev/null
+++ b/tensorflow/core/distributed_runtime/test_utils.h
@@ -0,0 +1,173 @@
+/* Copyright 2018 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.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_TEST_UTILS_H_
+#define TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_TEST_UTILS_H_
+
+#include <unordered_map>
+#include "tensorflow/core/distributed_runtime/worker_cache.h"
+#include "tensorflow/core/distributed_runtime/worker_interface.h"
+
+namespace tensorflow {
+
+// Some utilities for testing distributed-mode components in a single process
+// without RPCs.
+
+// Implements the worker interface with methods that just respond with
+// "unimplemented" status.  Override just the methods needed for
+// testing.
+class TestWorkerInterface : public WorkerInterface {
+ public:
+  void GetStatusAsync(const GetStatusRequest* request,
+                      GetStatusResponse* response,
+                      StatusCallback done) override {
+    done(errors::Unimplemented("GetStatusAsync"));
+  }
+
+  void CreateWorkerSessionAsync(const CreateWorkerSessionRequest* request,
+                                CreateWorkerSessionResponse* response,
+                                StatusCallback done) override {
+    done(errors::Unimplemented("CreateWorkerSessionAsync"));
+  }
+
+  void DeleteWorkerSessionAsync(CallOptions* opts,
+                                const DeleteWorkerSessionRequest* request,
+                                DeleteWorkerSessionResponse* response,
+                                StatusCallback done) override {
+    done(errors::Unimplemented("DeleteWorkerSessionAsync"));
+  }
+
+  void RegisterGraphAsync(const RegisterGraphRequest* request,
+                          RegisterGraphResponse* response,
+                          StatusCallback done) override {
+    done(errors::Unimplemented("RegisterGraphAsync"));
+  }
+
+  void DeregisterGraphAsync(const DeregisterGraphRequest* request,
+                            DeregisterGraphResponse* response,
+                            StatusCallback done) override {
+    done(errors::Unimplemented("DeregisterGraphAsync"));
+  }
+
+  void RunGraphAsync(CallOptions* opts, RunGraphRequestWrapper* request,
+                     MutableRunGraphResponseWrapper* repsonse,
+                     StatusCallback done) override {
+    done(errors::Unimplemented("RunGraphAsync"));
+  }
+
+  void CleanupGraphAsync(const CleanupGraphRequest* request,
+                         CleanupGraphResponse* response,
+                         StatusCallback done) override {
+    done(errors::Unimplemented("RunGraphAsync"));
+  }
+
+  void CleanupAllAsync(const CleanupAllRequest* request,
+                       CleanupAllResponse* response,
+                       StatusCallback done) override {
+    done(errors::Unimplemented("RunGraphAsync"));
+  }
+
+  void RecvTensorAsync(CallOptions* opts, const RecvTensorRequest* request,
+                       TensorResponse* response, StatusCallback done) override {
+    done(errors::Unimplemented("RunGraphAsync"));
+  }
+
+  void LoggingAsync(const LoggingRequest* request, LoggingResponse* response,
+                    StatusCallback done) override {
+    done(errors::Unimplemented("RunGraphAsync"));
+  }
+
+  void TracingAsync(const TracingRequest* request, TracingResponse* response,
+                    StatusCallback done) override {
+    done(errors::Unimplemented("RunGraphAsync"));
+  }
+
+  void CompleteGroupAsync(CallOptions* opts,
+                          const CompleteGroupRequest* request,
+                          CompleteGroupResponse* response,
+                          StatusCallback done) override {
+    done(errors::Unimplemented("RunGraphAsync"));
+  }
+
+  void CompleteInstanceAsync(CallOptions* ops,
+                             const CompleteInstanceRequest* request,
+                             CompleteInstanceResponse* response,
+                             StatusCallback done) override {
+    done(errors::Unimplemented("RunGraphAsync"));
+  }
+
+  void GetStepSequenceAsync(const GetStepSequenceRequest* request,
+                            GetStepSequenceResponse* response,
+                            StatusCallback done) override {
+    done(errors::Unimplemented("RunGraphAsync"));
+  }
+};
+
+class TestWorkerCache : public WorkerCacheInterface {
+ public:
+  virtual ~TestWorkerCache() {}
+
+  void AddWorker(const string& target, WorkerInterface* wi) {
+    workers_[target] = wi;
+  }
+
+  void AddDevice(const string& device_name, const DeviceLocality& dev_loc) {
+    localities_[device_name] = dev_loc;
+  }
+
+  void ListWorkers(std::vector<string>* workers) const override {
+    workers->clear();
+    for (auto it : workers_) {
+      workers->push_back(it.first);
+    }
+  }
+
+  WorkerInterface* CreateWorker(const string& target) override {
+    auto it = workers_.find(target);
+    if (it != workers_.end()) {
+      return it->second;
+    }
+    return nullptr;
+  }
+
+  void ReleaseWorker(const string& target, WorkerInterface* worker) override {}
+
+  bool GetDeviceLocalityNonBlocking(const string& device,
+                                    DeviceLocality* locality) override {
+    auto it = localities_.find(device);
+    if (it != localities_.end()) {
+      *locality = it->second;
+      return true;
+    }
+    return false;
+  }
+
+  void GetDeviceLocalityAsync(const string& device, DeviceLocality* locality,
+                              StatusCallback done) override {
+    auto it = localities_.find(device);
+    if (it != localities_.end()) {
+      *locality = it->second;
+      done(Status::OK());
+      return;
+    }
+    done(errors::Internal("Device not found: ", device));
+  }
+
+ protected:
+  std::unordered_map<string, WorkerInterface*> workers_;
+  std::unordered_map<string, DeviceLocality> localities_;
+};
+
+}  // namespace tensorflow
+#endif  // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_TEST_UTILS_H_
diff --git a/tensorflow/core/distributed_runtime/worker.cc b/tensorflow/core/distributed_runtime/worker.cc
index e9073ef9f6..d682ac8f34 100644
--- a/tensorflow/core/distributed_runtime/worker.cc
+++ b/tensorflow/core/distributed_runtime/worker.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "tensorflow/core/distributed_runtime/worker.h"
 
+#include "tensorflow/core/common_runtime/collective_executor_mgr.h"
 #include "tensorflow/core/common_runtime/device_mgr.h"
 #include "tensorflow/core/common_runtime/process_util.h"
 #include "tensorflow/core/common_runtime/step_stats_collector.h"
@@ -25,8 +26,7 @@ limitations under the License.
 
 namespace tensorflow {
 
-Worker::Worker(WorkerEnv* env)
-    : env_(env), cancellation_manager_(new CancellationManager) {}
+Worker::Worker(WorkerEnv* env) : env_(env) {}
 
 void Worker::GetStatusAsync(const GetStatusRequest* request,
                             GetStatusResponse* response, StatusCallback done) {
@@ -185,19 +185,16 @@ void Worker::DoRunGraph(CallOptions* opts, RunGraphRequestWrapper* request,
     AbortStep(step_id);
   });
   CancellationToken token;
-  {
-    mutex_lock l(mu_);
-    token = cancellation_manager_->get_cancellation_token();
-    bool already_cancelled = !cancellation_manager_->RegisterCallback(
-        token, [cm]() { cm->StartCancel(); });
-    if (already_cancelled) {
-      opts->ClearCancelCallback();
-      delete cm;
-      delete collector;
-      delete out;
-      done(errors::Aborted("Call was aborted"));
-      return;
-    }
+  token = cancellation_manager_.get_cancellation_token();
+  bool already_cancelled = !cancellation_manager_.RegisterCallback(
+      token, [cm]() { cm->StartCancel(); });
+  if (already_cancelled) {
+    opts->ClearCancelCallback();
+    delete cm;
+    delete collector;
+    delete out;
+    done(errors::Aborted("Call was aborted"));
+    return;
   }
   session->graph_mgr->ExecuteAsync(
       request->graph_handle(), step_id, session.get(), request->exec_opts(),
@@ -208,10 +205,7 @@ void Worker::DoRunGraph(CallOptions* opts, RunGraphRequestWrapper* request,
           s = session->graph_mgr->RecvOutputs(step_id, out);
         }
         opts->ClearCancelCallback();
-        {
-          mutex_lock l(mu_);
-          cancellation_manager_->DeregisterCallback(token);
-        }
+        cancellation_manager_.DeregisterCallback(token);
         delete cm;
 
         if (s.ok()) {
@@ -276,20 +270,14 @@ void Worker::DoPartialRunGraph(CallOptions* opts,
   // executors.
   if (is_new_partial_run) {
     CancellationToken token;
-    {
-      mutex_lock l(mu_);
-      token = cancellation_manager_->get_cancellation_token();
-      cancellation_manager_->RegisterCallback(token,
-                                              [cm]() { cm->StartCancel(); });
-    }
+    token = cancellation_manager_.get_cancellation_token();
+    cancellation_manager_.RegisterCallback(token,
+                                           [cm]() { cm->StartCancel(); });
     session->graph_mgr->ExecuteAsync(
         graph_handle, step_id, session.get(), request->exec_opts(),
         nullptr /* collector */, nullptr /* response */, cm, in,
         [this, token, step_id, session](Status s) {
-          {
-            mutex_lock l(mu_);
-            cancellation_manager_->DeregisterCallback(token);
-          }
+          cancellation_manager_.DeregisterCallback(token);
           partial_run_mgr_.ExecutorDone(step_id, s);
         });
   } else {
@@ -324,6 +312,9 @@ void Worker::CleanupGraphAsync(const CleanupGraphRequest* request,
                                StatusCallback done) {
   const int64 step_id = request->step_id();
   env_->rendezvous_mgr->Cleanup(step_id);
+  if (env_->collective_executor_mgr) {
+    env_->collective_executor_mgr->Cleanup(step_id);
+  }
   done(Status::OK());
 }
 
@@ -346,6 +337,44 @@ void Worker::TracingAsync(const TracingRequest* request,
   done(errors::Unimplemented("Tracing"));
 }
 
+void Worker::CompleteGroupAsync(CallOptions* opts,
+                                const CompleteGroupRequest* request,
+                                CompleteGroupResponse* response,
+                                StatusCallback done) {
+  if (env_->collective_executor_mgr) {
+    env_->collective_executor_mgr->GetParamResolver()->CompleteGroupAsync(
+        request, response, &cancellation_manager_, done);
+  } else {
+    done(
+        errors::Internal("Runtime not initialized with CollectiveExecutorMgr"));
+  }
+}
+
+void Worker::CompleteInstanceAsync(CallOptions* opts,
+                                   const CompleteInstanceRequest* request,
+                                   CompleteInstanceResponse* response,
+                                   StatusCallback done) {
+  if (env_->collective_executor_mgr) {
+    env_->collective_executor_mgr->GetParamResolver()->CompleteInstanceAsync(
+        request, response, &cancellation_manager_, done);
+  } else {
+    done(
+        errors::Internal("Runtime not initialized with CollectiveExecutorMgr"));
+  }
+}
+
+void Worker::GetStepSequenceAsync(const GetStepSequenceRequest* request,
+                                  GetStepSequenceResponse* response,
+                                  StatusCallback done) {
+  if (env_->collective_executor_mgr) {
+    env_->collective_executor_mgr->GetStepSequenceAsync(request, response,
+                                                        done);
+  } else {
+    done(
+        errors::Internal("Runtime not initialized with CollectiveExecutorMgr"));
+  }
+}
+
 // Helper for RecvTensor. Validates "key" and returns the source
 // device in "*src_dev".
 Status Worker::PrepareRecvTensor(const Rendezvous::ParsedKey& parsed,
diff --git a/tensorflow/core/distributed_runtime/worker.h b/tensorflow/core/distributed_runtime/worker.h
index 19aeeb752c..b5a9ada502 100644
--- a/tensorflow/core/distributed_runtime/worker.h
+++ b/tensorflow/core/distributed_runtime/worker.h
@@ -90,6 +90,20 @@ class Worker : public WorkerInterface {
   void TracingAsync(const TracingRequest* request, TracingResponse* response,
                     StatusCallback done) override;
 
+  void CompleteGroupAsync(CallOptions* opts,
+                          const CompleteGroupRequest* request,
+                          CompleteGroupResponse* response,
+                          StatusCallback done) override;
+
+  void CompleteInstanceAsync(CallOptions* opts,
+                             const CompleteInstanceRequest* request,
+                             CompleteInstanceResponse* response,
+                             StatusCallback done) override;
+
+  void GetStepSequenceAsync(const GetStepSequenceRequest* request,
+                            GetStepSequenceResponse* response,
+                            StatusCallback done) override;
+
  protected:
   WorkerEnv* const env_;  // Not owned.
 
@@ -101,8 +115,7 @@ class Worker : public WorkerInterface {
  private:
   PartialRunMgr partial_run_mgr_;
 
-  mutex mu_;
-  CancellationManager* cancellation_manager_ GUARDED_BY(mu_);
+  CancellationManager cancellation_manager_;
 
   Status PrepareRunGraph(RunGraphRequestWrapper* req,
                          GraphMgr::NamedTensors* in,
diff --git a/tensorflow/core/distributed_runtime/worker_env.h b/tensorflow/core/distributed_runtime/worker_env.h
index 793d58c8a1..93d933bfa6 100644
--- a/tensorflow/core/distributed_runtime/worker_env.h
+++ b/tensorflow/core/distributed_runtime/worker_env.h
@@ -25,6 +25,7 @@ namespace thread {
 class ThreadPool;
 }  // namespace thread
 
+class CollectiveExecutorMgrInterface;
 class Device;
 class DeviceMgr;
 class Env;
@@ -57,6 +58,10 @@ struct WorkerEnv {
   // A set of rendezvous keyed by step ids.
   RendezvousMgrInterface* rendezvous_mgr = nullptr;
 
+  // Generates per-step CollectiveExecutors and has access to utilities
+  // supporting collective operations.
+  CollectiveExecutorMgrInterface* collective_executor_mgr = nullptr;
+
   // A pool of threads for scheduling compute work.
   thread::ThreadPool* compute_pool = nullptr;
 };
diff --git a/tensorflow/core/distributed_runtime/worker_interface.h b/tensorflow/core/distributed_runtime/worker_interface.h
index a1597ee798..bad31d27b2 100644
--- a/tensorflow/core/distributed_runtime/worker_interface.h
+++ b/tensorflow/core/distributed_runtime/worker_interface.h
@@ -112,6 +112,20 @@ class WorkerInterface {
   virtual void TracingAsync(const TracingRequest* request,
                             TracingResponse* response, StatusCallback done) = 0;
 
+  virtual void CompleteGroupAsync(CallOptions* opts,
+                                  const CompleteGroupRequest* request,
+                                  CompleteGroupResponse* response,
+                                  StatusCallback done) = 0;
+
+  virtual void CompleteInstanceAsync(CallOptions* ops,
+                                     const CompleteInstanceRequest* request,
+                                     CompleteInstanceResponse* response,
+                                     StatusCallback done) = 0;
+
+  virtual void GetStepSequenceAsync(const GetStepSequenceRequest* request,
+                                    GetStepSequenceResponse* response,
+                                    StatusCallback done) = 0;
+
   Status GetStatus(const GetStatusRequest* request,
                    GetStatusResponse* response) {
     return CallAndWait(&ME::GetStatusAsync, request, response);
@@ -156,6 +170,11 @@ class WorkerInterface {
     return CallAndWait(&ME::TracingAsync, request, response);
   }
 
+  Status GetStepSequence(const GetStepSequenceRequest* request,
+                         GetStepSequenceResponse* response) {
+    return CallAndWait(&ME::GetStepSequenceAsync, request, response);
+  }
+
  protected:
   // Instances of WorkerInterface must be deleted by a call to
   // WorkerCacheInterface::ReleaseWorker().
diff --git a/tensorflow/core/protobuf/worker.proto b/tensorflow/core/protobuf/worker.proto
index 1819a35248..602f6a1ef1 100644
--- a/tensorflow/core/protobuf/worker.proto
+++ b/tensorflow/core/protobuf/worker.proto
@@ -27,6 +27,8 @@ import "tensorflow/core/framework/step_stats.proto";
 import "tensorflow/core/framework/device_attributes.proto";
 import "tensorflow/core/framework/graph.proto";
 import "tensorflow/core/framework/tensor.proto";
+import "tensorflow/core/framework/tensor_shape.proto";
+import "tensorflow/core/framework/types.proto";
 import "tensorflow/core/lib/core/error_codes.proto";
 import "tensorflow/core/protobuf/config.proto";
 import "tensorflow/core/protobuf/debug.proto";
@@ -413,3 +415,71 @@ message TracingRequest {
 
 message TracingResponse {
 }
+
+////////////////////////////////////////////////////////////////////////////////
+//
+// Collective Op dynamic group resolution messages.
+//
+////////////////////////////////////////////////////////////////////////////////
+
+// Supplies one or more device names as members of the group identified by
+// group_key.  Service will respond when all group_size devices become known.
+// All devices in group must have same type.
+message CompleteGroupRequest {
+  int32 group_key = 1;
+  int32 group_size = 2;
+  string device_type = 3;
+  repeated string device_name = 4;
+}
+
+// Gives the complete membership of the group identified by group_key.
+message CompleteGroupResponse {
+  int32 group_key = 1;
+  int32 group_size = 2;
+  string device_type = 3;
+  int32 num_tasks = 4;  // number of distinct tasks hosting the devices
+  repeated string device_name = 5;
+  repeated string task_name = 6;  // task name prefixes of device_names
+}
+
+// Supplies data about one collective op belonging to the instance identified
+// by instance_key.  Service will respond when all group_size ops have
+// become known.  Most of the data being sent is for correctness checking,
+// to ensure that all ops in the instance share common attributes.
+message CompleteInstanceRequest {
+  string name = 1;
+  int32 type = 2;
+  DataType data_type = 3;
+  TensorShapeProto shape = 4;
+  int32 group_key = 5;
+  int32 group_size = 6;
+  int32 instance_key = 7;
+  string device_type = 8;
+  repeated int32 subdiv_offset = 9;
+  string device = 10;
+  bool is_source = 11;
+}
+
+// Confirms that every op in the instance has consistently declared itself.
+// Also gives the source_rank in case of broadcast.
+message CompleteInstanceResponse {
+  int32 instance_key = 1;
+  int32 source_rank = 2;
+}
+
+// Request for next agreed-upon step_id for the specified graph_keys.
+// This is used to enable multiple graphs containing nodes from
+// a common collective instance to coordinate using the same step_ids.
+message GetStepSequenceRequest {
+  repeated int64 graph_key = 1;
+}
+
+message StepSequence {
+  int64 graph_key = 1;
+  int64 next_step_id = 2;
+}
+
+// Next valid step_ids for one or more graph_keys.
+message GetStepSequenceResponse {
+  repeated StepSequence step_sequence = 1;
+}
diff --git a/tensorflow/core/protobuf/worker_service.proto b/tensorflow/core/protobuf/worker_service.proto
index e1bfb04d7c..01c76c01a9 100644
--- a/tensorflow/core/protobuf/worker_service.proto
+++ b/tensorflow/core/protobuf/worker_service.proto
@@ -72,4 +72,14 @@ service WorkerService {
 
   // See worker.proto for details.
   rpc Tracing(TracingRequest) returns (TracingResponse);
+
+  // See worker.proto for details.
+  rpc GetStepSequence(GetStepSequenceRequest) returns (GetStepSequenceResponse);
+
+  // See worker.proto for details.
+  rpc CompleteGroup(CompleteGroupRequest) returns (CompleteGroupResponse);
+
+  // See worker.proto for details.
+  rpc CompleteInstance(CompleteInstanceRequest)
+      returns (CompleteInstanceResponse);
 }
-- 
GitLab


From 3b7f22f9180935919bab478adb45037b1f0d38c2 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 13:34:39 -0700
Subject: [PATCH 106/395] Relax the stringent memory allocator constraints in
 AssignOp if a Grappler graph analysis determines it to be safe. This will
 allow Assign to reuse the input buffer to initialize the variable in many
 cases.

PiperOrigin-RevId: 194988134
---
 tensorflow/core/grappler/op_types.cc          |   4 +
 tensorflow/core/grappler/op_types.h           |   1 +
 .../grappler/optimizers/memory_optimizer.cc   |  76 ++++++++++
 .../optimizers/memory_optimizer_test.cc       | 134 ++++++++++++++++++
 tensorflow/core/grappler/utils.cc             |   6 +-
 tensorflow/core/grappler/utils.h              |   1 +
 tensorflow/core/kernels/assign_op.h           |  73 +++++-----
 .../core/kernels/resource_variable_ops.cc     |  18 ++-
 8 files changed, 274 insertions(+), 39 deletions(-)

diff --git a/tensorflow/core/grappler/op_types.cc b/tensorflow/core/grappler/op_types.cc
index 839b0bbfc9..bf6d4c0921 100644
--- a/tensorflow/core/grappler/op_types.cc
+++ b/tensorflow/core/grappler/op_types.cc
@@ -54,6 +54,10 @@ bool IsApproximateEqual(const NodeDef& node) {
 
 bool IsAvgPoolGrad(const NodeDef& node) { return node.op() == "AvgPoolGrad"; }
 
+bool IsAssign(const NodeDef& node) {
+  return node.op() == "Assign" || node.op() == "AssignVariableOp";
+}
+
 bool IsAssert(const NodeDef& node) { return node.op() == "Assert"; }
 
 bool IsAtan2(const NodeDef& node) { return node.op() == "Atan2"; }
diff --git a/tensorflow/core/grappler/op_types.h b/tensorflow/core/grappler/op_types.h
index bd8d3a44e4..3dddf3f1ea 100644
--- a/tensorflow/core/grappler/op_types.h
+++ b/tensorflow/core/grappler/op_types.h
@@ -30,6 +30,7 @@ bool IsAnyDiv(const NodeDef& node);
 bool IsApproximateEqual(const NodeDef& node);
 bool IsAvgPoolGrad(const NodeDef& node);
 bool IsAssert(const NodeDef& node);
+bool IsAssign(const NodeDef& node);
 bool IsAtan2(const NodeDef& node);
 bool IsBetainc(const NodeDef& node);
 bool IsBiasAdd(const NodeDef& node);
diff --git a/tensorflow/core/grappler/optimizers/memory_optimizer.cc b/tensorflow/core/grappler/optimizers/memory_optimizer.cc
index c1fee0e993..7c6468bfcb 100644
--- a/tensorflow/core/grappler/optimizers/memory_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/memory_optimizer.cc
@@ -1219,6 +1219,80 @@ bool SwappingPass(RewriterConfig::MemOptType optimization_level,
   return updated_graph;
 }
 
+// TODO(rmlarsen): Add distributed TF test.
+Status RelaxAllocatorConstraints(GraphDef* optimized_graph) {
+  std::unordered_set<string> devices;
+  std::vector<int> assign_nodes;
+  bool found_send = false;
+  for (int i = 0; i < optimized_graph->node_size(); ++i) {
+    const NodeDef& node = optimized_graph->node(i);
+    devices.insert(node.device());
+    if (IsAssign(node)) {
+      assign_nodes.push_back(i);
+    }
+    if (IsSend(node)) {
+      found_send = true;
+      break;
+    }
+  }
+  if (!found_send && devices.size() == 1) {
+    for (int assign_idx : assign_nodes) {
+      // Set an attribute telling AssignOp to ignore allocator constraints.
+      NodeDef* assign_node = optimized_graph->mutable_node(assign_idx);
+      (*assign_node->mutable_attr())["_grappler_relax_allocator_constraints"]
+          .set_b(true);
+    }
+    return Status::OK();
+  }
+
+  std::unordered_set<int> optimized_nodes;
+  SimpleGraphView graph_view;
+  TF_RETURN_IF_ERROR(graph_view.Initialize(*optimized_graph));
+  for (int i : assign_nodes) {
+    if (optimized_nodes.find(i) == optimized_nodes.end()) {
+      const NodeDef& node = optimized_graph->node(i);
+      optimized_nodes.insert(i);
+      std::vector<int> assign_nodes_in_fanout;
+      assign_nodes_in_fanout.push_back(i);
+      std::set<int> transitive_fanout;
+      graph_view.DepthFirstSearch(std::unordered_set<string>{}, i,
+                                  &transitive_fanout);
+      const string& assign_device = node.device();
+      bool relax_constraint = true;
+      // If all nodes in the transitive fanout are on the same device as the
+      // assign node, there is no need to allocate the output in pinned memory.
+      for (int fanout : transitive_fanout) {
+        const NodeDef& fanout_node = optimized_graph->node(fanout);
+        if (relax_constraint &&
+            (fanout_node.device() != assign_device || IsSend(fanout_node))) {
+          relax_constraint = false;
+        }
+        if (optimized_nodes.find(fanout) == optimized_nodes.end() &&
+            IsAssign(fanout_node)) {
+          assign_nodes_in_fanout.push_back(fanout);
+        }
+      }
+
+      for (int assign_idx : assign_nodes_in_fanout) {
+        if (relax_constraint) {
+          // If all devices match in fanout of node(i) then, by transitivity,
+          // they must also match in the fanout of other assign nodes
+          // node(assign_idx) in the fanout, so we can process them here,
+          // and save computing their transitive fanout later.
+          optimized_nodes.insert(assign_idx);
+
+          // Set an attribute telling AssignOp to ignore allocator constraints.
+          NodeDef* assign_node = optimized_graph->mutable_node(assign_idx);
+          (*assign_node
+                ->mutable_attr())["_grappler_relax_allocator_constraints"]
+              .set_b(true);
+        }
+      }
+    }
+  }
+  return Status::OK();
+}
+
 Status MemoryOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
                                  GraphDef* optimized_graph) {
   *optimized_graph = item.graph;
@@ -1251,6 +1325,8 @@ Status MemoryOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
     }
   }
 
+  TF_RETURN_IF_ERROR(RelaxAllocatorConstraints(&optimized_item.graph));
+
   optimized_graph->Swap(&optimized_item.graph);
   return Status::OK();
 }
diff --git a/tensorflow/core/grappler/optimizers/memory_optimizer_test.cc b/tensorflow/core/grappler/optimizers/memory_optimizer_test.cc
index a1f80802dd..a3f0e07861 100644
--- a/tensorflow/core/grappler/optimizers/memory_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/memory_optimizer_test.cc
@@ -440,6 +440,140 @@ TEST_F(MemoryOptimizerTest, AccumulationRewrites) {
   }
 }
 
+class RelaxAllocatorConstraintsTest : public GrapplerTest {};
+
+TEST_F(RelaxAllocatorConstraintsTest, SameDevice) {
+  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+  Output constant = ops::Const(s.WithOpName("constant").WithDevice("/cpu:0"),
+                               -3.14f, {128, 128});
+  Output variable = ops::Variable(s.WithOpName("variable").WithDevice("/cpu:0"),
+                                  {128, 128}, DT_FLOAT);
+  Output assign = ops::Assign(s.WithOpName("assign").WithDevice("/cpu:0"),
+                              variable, constant);
+  Output exp = ops::Exp(s.WithOpName("exp").WithDevice("/cpu:0"), assign);
+
+  GrapplerItem item;
+  TF_CHECK_OK(s.ToGraphDef(&item.graph));
+
+  MemoryOptimizer optimizer(RewriterConfig::MANUAL);
+  GraphDef output;
+  TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
+
+  auto node = output.node(2);
+  EXPECT_EQ("assign", node.name());
+  EXPECT_EQ(1, node.attr().count("_grappler_relax_allocator_constraints"));
+  EXPECT_EQ(true, node.attr().at("_grappler_relax_allocator_constraints").b());
+
+  item.fetch = {"exp"};
+  item.init_ops = {"variable"};
+  auto tensors_expected = EvaluateFetchNodes(item);
+  GrapplerItem optimized(item, std::move(output));
+  auto tensors = EvaluateFetchNodes(optimized);
+  test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
+}
+
+TEST_F(RelaxAllocatorConstraintsTest, DifferentDevice) {
+  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+  Output constant = ops::Const(s.WithOpName("constant").WithDevice("/cpu:0"),
+                               -3.14f, {128, 128});
+  Output variable = ops::Variable(s.WithOpName("variable").WithDevice("/cpu:0"),
+                                  {128, 128}, DT_FLOAT);
+  Output assign = ops::Assign(s.WithOpName("assign").WithDevice("/cpu:0"),
+                              variable, constant);
+  // exp runs on a different device, so we cannot relax the allocation
+  // constraints on assign.
+  Output exp = ops::Exp(s.WithOpName("exp").WithDevice("/gpu:0"), assign);
+
+  GrapplerItem item;
+  TF_CHECK_OK(s.ToGraphDef(&item.graph));
+
+  MemoryOptimizer optimizer(RewriterConfig::MANUAL);
+  GraphDef output;
+  TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
+
+  auto node = output.node(2);
+  EXPECT_EQ("assign", node.name());
+  EXPECT_EQ(0, node.attr().count("_grappler_relax_allocator_constraints"));
+#if GOOGLE_CUDA
+  item.fetch = {"exp"};
+  item.init_ops = {"variable"};
+  auto tensors_expected = EvaluateFetchNodes(item);
+  GrapplerItem optimized(item, std::move(output));
+  auto tensors = EvaluateFetchNodes(optimized);
+  test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
+#endif
+}
+
+TEST_F(RelaxAllocatorConstraintsTest, SendNode) {
+  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+  Output constant = ops::Const(s.WithOpName("constant").WithDevice("/cpu:0"),
+                               -3.14f, {128, 128});
+  Output variable = ops::Variable(s.WithOpName("variable").WithDevice("/cpu:0"),
+                                  {128, 128}, DT_FLOAT);
+  Output assign = ops::Assign(s.WithOpName("assign").WithDevice("/cpu:0"),
+                              variable, constant);
+
+  GrapplerItem item;
+  TF_CHECK_OK(s.ToGraphDef(&item.graph));
+  NodeDef* send = item.graph.add_node();
+  // Add a send node to the graph in the fanout of "assign".
+  send->set_name("send");
+  send->set_op("_Send");
+  send->add_input("assign");
+
+  MemoryOptimizer optimizer(RewriterConfig::MANUAL);
+  GraphDef output;
+  TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
+
+  auto node = output.node(2);
+  EXPECT_EQ("assign", node.name());
+  EXPECT_EQ(0, node.attr().count("_grappler_relax_allocator_constraints"));
+}
+
+TEST_F(RelaxAllocatorConstraintsTest, AssignNodeInFanout) {
+  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+  Output constant0 = ops::Const(s.WithOpName("constant0").WithDevice("/cpu:0"),
+                                -42.0f, {128, 128});
+  Output variable0 = ops::Variable(
+      s.WithOpName("variable0").WithDevice("/cpu:0"), {128, 128}, DT_FLOAT);
+  Output assign0 = ops::Assign(s.WithOpName("assign0").WithDevice("/cpu:0"),
+                               variable0, constant0);
+  // The rest of the graph is on a second device, so we can relax the
+  // constraint for assign1, but not for assign0.
+  Output exp1 = ops::Exp(s.WithOpName("exp1").WithDevice("/gpu:0"), assign0);
+  Output variable1 = ops::Variable(
+      s.WithOpName("variable1").WithDevice("/gpu:0"), {128, 128}, DT_FLOAT);
+  Output assign1 = ops::Assign(s.WithOpName("assign1").WithDevice("/gpu:0"),
+                               variable1, exp1);
+
+  GrapplerItem item;
+  TF_CHECK_OK(s.ToGraphDef(&item.graph));
+
+  MemoryOptimizer optimizer(RewriterConfig::MANUAL);
+  GraphDef output;
+  TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
+
+  auto node = output.node(3);
+  EXPECT_EQ("assign0", node.name());
+  EXPECT_EQ(0, node.attr().count("_grappler_relax_allocator_constraints"));
+
+  node = output.node(5);
+  EXPECT_EQ("assign1", node.name());
+  EXPECT_EQ(1, node.attr().count("_grappler_relax_allocator_constraints"));
+  EXPECT_EQ(true, node.attr().at("_grappler_relax_allocator_constraints").b());
+
+#if GOOGLE_CUDA
+  item.fetch = {"assign0", "assign1"};
+  item.init_ops = {"exp1", "variable1"};
+  auto tensors_expected = EvaluateFetchNodes(item);
+  GrapplerItem optimized(item, std::move(output));
+  auto tensors = EvaluateFetchNodes(optimized);
+  for (int i = 0; i < tensors_expected.size(); ++i) {
+    test::ExpectTensorEqual<float>(tensors_expected[i], tensors[i]);
+  }
+#endif
+}
+
 }  // namespace
 }  // namespace grappler
 }  // namespace tensorflow
diff --git a/tensorflow/core/grappler/utils.cc b/tensorflow/core/grappler/utils.cc
index 6db6d71447..c8e63f95e1 100644
--- a/tensorflow/core/grappler/utils.cc
+++ b/tensorflow/core/grappler/utils.cc
@@ -435,7 +435,8 @@ void SimpleGraphView::DepthFirstSearch(
     std::set<int>* nodes_found) const {
   nodes_found->clear();
   const string& op_type = graph_->node(root_node).op();
-  if (op_types_to_traverse.find(op_type) == op_types_to_traverse.end()) {
+  if (!op_types_to_traverse.empty() &&
+      op_types_to_traverse.find(op_type) == op_types_to_traverse.end()) {
     return;
   }
   std::vector<int> stack;
@@ -446,7 +447,8 @@ void SimpleGraphView::DepthFirstSearch(
     stack.pop_back();
     nodes_found->insert(node_idx);
     const string& op_type = graph_->node(node_idx).op();
-    if (op_types_to_traverse.find(op_type) != op_types_to_traverse.end()) {
+    if (op_types_to_traverse.empty() ||
+        op_types_to_traverse.find(op_type) != op_types_to_traverse.end()) {
       for (auto output_idx : this->outputs(node_idx)) {
         if (nodes_found->find(output_idx) == nodes_found->end()) {
           stack.push_back(output_idx);
diff --git a/tensorflow/core/grappler/utils.h b/tensorflow/core/grappler/utils.h
index 15f6b367b0..9776e99f20 100644
--- a/tensorflow/core/grappler/utils.h
+++ b/tensorflow/core/grappler/utils.h
@@ -251,6 +251,7 @@ class SimpleGraphView {
   // visited in nodes_found. If a node has an op in `op_types_to_traverse`, the
   // walk continues to its children. It is assumed that *graph_ was not modified
   // after the call to Initialize().
+  // If `op_types_to_traverse` is empty the DFS will traverse any node type.
   void DepthFirstSearch(const std::unordered_set<string>& op_types_to_traverse,
                         int node_idx, std::set<int>* nodes_found) const;
 
diff --git a/tensorflow/core/kernels/assign_op.h b/tensorflow/core/kernels/assign_op.h
index 19b38f9e68..a450b1d1ee 100644
--- a/tensorflow/core/kernels/assign_op.h
+++ b/tensorflow/core/kernels/assign_op.h
@@ -36,6 +36,12 @@ class AssignOp : public OpKernel {
                    context->GetAttr("validate_shape", &validate_shape_));
     OP_REQUIRES(context, IsRefType(context->input_type(0)),
                 errors::InvalidArgument("lhs input needs to be a ref type"));
+    if (!context
+             ->GetAttr("_grappler_relax_allocator_constraints",
+                       &relax_constraints_)
+             .ok()) {
+      relax_constraints_ = false;
+    }
   }
 
   void Compute(OpKernelContext* context) override {
@@ -44,48 +50,37 @@ class AssignOp : public OpKernel {
     // We always return the input ref.
     context->forward_ref_input_to_ref_output(0, 0);
 
-    // We can't always know how this value will be used downstream,
-    // so make conservative assumptions in specifying constraints on
-    // the memory allocation attributes.
-    // TODO(rmlarsen): These conservative constraints make buffer
-    // forwarding unlikely to happen very often. Try to use graph analysis
-    // (possibly the InferAllocAttr pass in the executer) to improve the
-    // situation.
+    // We can't always know how this value will be used downstream, so make
+    // conservative assumptions in specifying constraints on the memory
+    // allocation attributes, unless the Grappler graph analysis determined that
+    // it was safe not to.
     AllocatorAttributes attr;
-    attr.set_gpu_compatible(true);
-    attr.set_nic_compatible(true);
+    if (!relax_constraints_) {
+      attr.set_gpu_compatible(true);
+      attr.set_nic_compatible(true);
+    }
 
     {
       mutex_lock l(*context->input_ref_mutex(0));
       const Tensor& old_lhs = context->mutable_input(0, /* lock_held */ true);
       const bool same_shape = old_lhs.shape().IsSameSize(rhs.shape());
       if (validate_shape_) {
-        OP_REQUIRES(
-            context, same_shape,
-            errors::InvalidArgument(
-                "Assign requires shapes of both tensors to match. lhs shape= ",
-                old_lhs.shape().DebugString(),
-                " rhs shape= ", rhs.shape().DebugString()));
+        OP_REQUIRES(context, same_shape,
+                    errors::InvalidArgument(
+                        "Assign requires shapes of both tensors to match. "
+                        "lhs shape= ",
+                        old_lhs.shape().DebugString(),
+                        " rhs shape= ", rhs.shape().DebugString()));
       }
 
       // In the code below we try to minimize the amount of memory allocation
       // and copying by trying the following two shortcuts:
-      // 1. If we can reuse the rhs buffer we avoid both a memory allocation
-      //   and copying.
-      // 2. If the lhs is initialized and has the same number of elements as the
-      //    rhs we can avoid a memory allocation.
-
-      // 1. Try to reuse the rhs.
-      std::unique_ptr<Tensor> input_alias = context->forward_input(
-          1, OpKernelContext::Params::kNoReservation /*output_index*/,
-          rhs.dtype(), rhs.shape(), DEVICE_MEMORY, attr);
-      if (input_alias != nullptr) {
-        // Transfer ownership to the ref.
-        context->replace_ref_input(0, *input_alias, /* lock_held */ true);
-        return;
-      }
+      // 1. If the lhs is initialized and has the same number of elements as
+      //    the rhs we can avoid a memory allocation.
+      // 2. If we can reuse the rhs buffer we avoid both a memory allocation
+      //    and copying.
 
-      // 2. Try to copy into an existing buffer.
+      // 1. Try to copy into an existing buffer.
       if (old_lhs.IsInitialized() &&
           old_lhs.shape().num_elements() == rhs.shape().num_elements()) {
         // The existing lhs tensor has already been initialized and the right
@@ -95,15 +90,26 @@ class AssignOp : public OpKernel {
           reshaped_old_lhs = old_lhs;
         } else {
           CHECK(reshaped_old_lhs.CopyFrom(old_lhs, rhs.shape()));
-          context->replace_ref_input(0, reshaped_old_lhs, /* lock_held */ true);
+          context->replace_ref_input(0, reshaped_old_lhs,
+                                     /* lock_held */ true);
         }
         if (use_exclusive_lock_) {
           Copy(context, &reshaped_old_lhs, rhs);
           return;
         }
       } else {
-        // Create a new persistent tensor whose shape matches the right hand
-        // side, hand off to lhs and copy the rhs into it.
+        // 2. Try to reuse the rhs.
+        std::unique_ptr<Tensor> input_alias = context->forward_input(
+            1, OpKernelContext::Params::kNoReservation /*output_index*/,
+            rhs.dtype(), rhs.shape(), DEVICE_MEMORY, attr);
+        if (input_alias != nullptr) {
+          // Update the ref to point to the new buffer.
+          context->replace_ref_input(0, *input_alias, /* lock_held */ true);
+          return;
+        }
+
+        // Otherwise, create a new persistent tensor whose shape matches the
+        // right hand side, hand off to lhs and copy the rhs into it.
         PersistentTensor copy;
         Tensor* copyTensor = nullptr;
         OP_REQUIRES_OK(
@@ -132,6 +138,7 @@ class AssignOp : public OpKernel {
 
   bool use_exclusive_lock_;
   bool validate_shape_;
+  bool relax_constraints_;
 };
 
 }  // end namespace tensorflow
diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc
index 916869fb56..a8bcc7f7dc 100644
--- a/tensorflow/core/kernels/resource_variable_ops.cc
+++ b/tensorflow/core/kernels/resource_variable_ops.cc
@@ -211,6 +211,11 @@ class AssignVariableOp : public OpKernel {
  public:
   explicit AssignVariableOp(OpKernelConstruction* c) : OpKernel(c) {
     OP_REQUIRES_OK(c, c->GetAttr("dtype", &dtype_));
+    if (!c->GetAttr("_grappler_relax_allocator_constraints",
+                    &relax_constraints_)
+             .ok()) {
+      relax_constraints_ = false;
+    }
   }
 
   void Compute(OpKernelContext* context) override {
@@ -228,8 +233,10 @@ class AssignVariableOp : public OpKernel {
               PersistentTensor unused;
               Tensor* tmp;
               AllocatorAttributes attr;
-              attr.set_gpu_compatible(true);
-              attr.set_nic_compatible(true);
+              if (!relax_constraints_) {
+                attr.set_gpu_compatible(true);
+                attr.set_nic_compatible(true);
+              }
               TF_RETURN_IF_ERROR(context->allocate_persistent(
                   dtype_, context->input(1).shape(), &unused, &tmp, attr));
               *(*ptr)->tensor() = *tmp;
@@ -245,8 +252,10 @@ class AssignVariableOp : public OpKernel {
 
     const Tensor& value = context->input(1);
     AllocatorAttributes attr;
-    attr.set_gpu_compatible(true);
-    attr.set_nic_compatible(true);
+    if (!relax_constraints_) {
+      attr.set_gpu_compatible(true);
+      attr.set_nic_compatible(true);
+    }
 
     // Copying is unnecessary if we are the last user of the value
     // tensor, we can just adopt the input tensor's buffer instead.
@@ -277,6 +286,7 @@ class AssignVariableOp : public OpKernel {
 
  private:
   DataType dtype_;
+  bool relax_constraints_;
 };
 
 template <typename Device>
-- 
GitLab


From 594f970f81089c91f713bbdda48d44ef99f80c9e Mon Sep 17 00:00:00 2001
From: Shashi Shekhar <shashishekhar@google.com>
Date: Tue, 1 May 2018 13:44:58 -0700
Subject: [PATCH 107/395] Update schema.

PiperOrigin-RevId: 194989704
---
 .../contrib/lite/schema/schema_generated.h    | 27 +++++++++++++++----
 1 file changed, 22 insertions(+), 5 deletions(-)

diff --git a/tensorflow/contrib/lite/schema/schema_generated.h b/tensorflow/contrib/lite/schema/schema_generated.h
index 25ed9abd9f..57af973460 100755
--- a/tensorflow/contrib/lite/schema/schema_generated.h
+++ b/tensorflow/contrib/lite/schema/schema_generated.h
@@ -4711,6 +4711,7 @@ struct ModelT : public flatbuffers::NativeTable {
   std::vector<std::unique_ptr<SubGraphT>> subgraphs;
   std::string description;
   std::vector<std::unique_ptr<BufferT>> buffers;
+  std::vector<int32_t> metadata_buffer;
   ModelT()
       : version(0) {
   }
@@ -4723,7 +4724,8 @@ struct Model FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
     VT_OPERATOR_CODES = 6,
     VT_SUBGRAPHS = 8,
     VT_DESCRIPTION = 10,
-    VT_BUFFERS = 12
+    VT_BUFFERS = 12,
+    VT_METADATA_BUFFER = 14
   };
   uint32_t version() const {
     return GetField<uint32_t>(VT_VERSION, 0);
@@ -4740,6 +4742,9 @@ struct Model FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
   const flatbuffers::Vector<flatbuffers::Offset<Buffer>> *buffers() const {
     return GetPointer<const flatbuffers::Vector<flatbuffers::Offset<Buffer>> *>(VT_BUFFERS);
   }
+  const flatbuffers::Vector<int32_t> *metadata_buffer() const {
+    return GetPointer<const flatbuffers::Vector<int32_t> *>(VT_METADATA_BUFFER);
+  }
   bool Verify(flatbuffers::Verifier &verifier) const {
     return VerifyTableStart(verifier) &&
            VerifyField<uint32_t>(verifier, VT_VERSION) &&
@@ -4754,6 +4759,8 @@ struct Model FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
            VerifyOffset(verifier, VT_BUFFERS) &&
            verifier.Verify(buffers()) &&
            verifier.VerifyVectorOfTables(buffers()) &&
+           VerifyOffset(verifier, VT_METADATA_BUFFER) &&
+           verifier.Verify(metadata_buffer()) &&
            verifier.EndTable();
   }
   ModelT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const;
@@ -4779,6 +4786,9 @@ struct ModelBuilder {
   void add_buffers(flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Buffer>>> buffers) {
     fbb_.AddOffset(Model::VT_BUFFERS, buffers);
   }
+  void add_metadata_buffer(flatbuffers::Offset<flatbuffers::Vector<int32_t>> metadata_buffer) {
+    fbb_.AddOffset(Model::VT_METADATA_BUFFER, metadata_buffer);
+  }
   explicit ModelBuilder(flatbuffers::FlatBufferBuilder &_fbb)
         : fbb_(_fbb) {
     start_ = fbb_.StartTable();
@@ -4797,8 +4807,10 @@ inline flatbuffers::Offset<Model> CreateModel(
     flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<OperatorCode>>> operator_codes = 0,
     flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<SubGraph>>> subgraphs = 0,
     flatbuffers::Offset<flatbuffers::String> description = 0,
-    flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Buffer>>> buffers = 0) {
+    flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Buffer>>> buffers = 0,
+    flatbuffers::Offset<flatbuffers::Vector<int32_t>> metadata_buffer = 0) {
   ModelBuilder builder_(_fbb);
+  builder_.add_metadata_buffer(metadata_buffer);
   builder_.add_buffers(buffers);
   builder_.add_description(description);
   builder_.add_subgraphs(subgraphs);
@@ -4813,14 +4825,16 @@ inline flatbuffers::Offset<Model> CreateModelDirect(
     const std::vector<flatbuffers::Offset<OperatorCode>> *operator_codes = nullptr,
     const std::vector<flatbuffers::Offset<SubGraph>> *subgraphs = nullptr,
     const char *description = nullptr,
-    const std::vector<flatbuffers::Offset<Buffer>> *buffers = nullptr) {
+    const std::vector<flatbuffers::Offset<Buffer>> *buffers = nullptr,
+    const std::vector<int32_t> *metadata_buffer = nullptr) {
   return tflite::CreateModel(
       _fbb,
       version,
       operator_codes ? _fbb.CreateVector<flatbuffers::Offset<OperatorCode>>(*operator_codes) : 0,
       subgraphs ? _fbb.CreateVector<flatbuffers::Offset<SubGraph>>(*subgraphs) : 0,
       description ? _fbb.CreateString(description) : 0,
-      buffers ? _fbb.CreateVector<flatbuffers::Offset<Buffer>>(*buffers) : 0);
+      buffers ? _fbb.CreateVector<flatbuffers::Offset<Buffer>>(*buffers) : 0,
+      metadata_buffer ? _fbb.CreateVector<int32_t>(*metadata_buffer) : 0);
 }
 
 flatbuffers::Offset<Model> CreateModel(flatbuffers::FlatBufferBuilder &_fbb, const ModelT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
@@ -6207,6 +6221,7 @@ inline void Model::UnPackTo(ModelT *_o, const flatbuffers::resolver_function_t *
   { auto _e = subgraphs(); if (_e) { _o->subgraphs.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->subgraphs[_i] = std::unique_ptr<SubGraphT>(_e->Get(_i)->UnPack(_resolver)); } } };
   { auto _e = description(); if (_e) _o->description = _e->str(); };
   { auto _e = buffers(); if (_e) { _o->buffers.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->buffers[_i] = std::unique_ptr<BufferT>(_e->Get(_i)->UnPack(_resolver)); } } };
+  { auto _e = metadata_buffer(); if (_e) { _o->metadata_buffer.resize(_e->size()); for (flatbuffers::uoffset_t _i = 0; _i < _e->size(); _i++) { _o->metadata_buffer[_i] = _e->Get(_i); } } };
 }
 
 inline flatbuffers::Offset<Model> Model::Pack(flatbuffers::FlatBufferBuilder &_fbb, const ModelT* _o, const flatbuffers::rehasher_function_t *_rehasher) {
@@ -6222,13 +6237,15 @@ inline flatbuffers::Offset<Model> CreateModel(flatbuffers::FlatBufferBuilder &_f
   auto _subgraphs = _o->subgraphs.size() ? _fbb.CreateVector<flatbuffers::Offset<SubGraph>> (_o->subgraphs.size(), [](size_t i, _VectorArgs *__va) { return CreateSubGraph(*__va->__fbb, __va->__o->subgraphs[i].get(), __va->__rehasher); }, &_va ) : 0;
   auto _description = _o->description.empty() ? 0 : _fbb.CreateString(_o->description);
   auto _buffers = _o->buffers.size() ? _fbb.CreateVector<flatbuffers::Offset<Buffer>> (_o->buffers.size(), [](size_t i, _VectorArgs *__va) { return CreateBuffer(*__va->__fbb, __va->__o->buffers[i].get(), __va->__rehasher); }, &_va ) : 0;
+  auto _metadata_buffer = _o->metadata_buffer.size() ? _fbb.CreateVector(_o->metadata_buffer) : 0;
   return tflite::CreateModel(
       _fbb,
       _version,
       _operator_codes,
       _subgraphs,
       _description,
-      _buffers);
+      _buffers,
+      _metadata_buffer);
 }
 
 inline bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *obj, BuiltinOptions type) {
-- 
GitLab


From 75c1896fb26a91ae8d895e24bfc128084cba4e9e Mon Sep 17 00:00:00 2001
From: Mark Daoust <markdaoust@google.com>
Date: Tue, 1 May 2018 13:54:34 -0700
Subject: [PATCH 108/395] Update community/swift

PiperOrigin-RevId: 194991305
---
 tensorflow/docs_src/community/swift.md | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/docs_src/community/swift.md b/tensorflow/docs_src/community/swift.md
index a7da189a5c..e5a0f02a8c 100644
--- a/tensorflow/docs_src/community/swift.md
+++ b/tensorflow/docs_src/community/swift.md
@@ -8,7 +8,7 @@ Welcome to the Swift for TensorFlow development community!
 
 Swift for TensorFlow is a new way to develop machine learning models. It
 gives you the power of
-[TensorFlow](https://www.tensorflow.org/programmers_guide/eager) directly
+[TensorFlow](programmers_guide/eager) directly
 integrated into the [Swift programming language](https://swift.org/about).
 With Swift, you can write the following imperative code, and Swift
 automatically turns it into **a single TensorFlow Graph** and runs it
@@ -28,15 +28,15 @@ print(x)
 ```
 
 Swift combines the flexibility of
-[Eager Execution](https://www.tensorflow.org/programmers_guide/eager) with the
-high performance of [Graphs and Sessions](https://www.tensorflow.org/programmers_guide/graphs).
+[Eager Execution](programmers_guide/eager) with the
+high performance of [Graphs and Sessions](programmers_guide/graphs).
 Behind the scenes, Swift analyzes your Tensor code and automatically builds
 graphs for you. Swift also catches type errors and shape mismatches before
 running your code, and has [Automatic Differentiation](https://en.wikipedia.org/wiki/Automatic_differentiation)
 built right in. We believe that machine learning tools are so important that
 they deserve **a first-class language and a compiler**.
 
-**Note:** Swift for TensorFlow is an early stage research project. It has been
+Note: Swift for TensorFlow is an early stage research project. It has been
 released to enable open source development and is not yet ready for general use
 by machine learning developers.
 
-- 
GitLab


From 7cbbd3525b4232f2dc8cd117852c26ec472aa9b2 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 14:04:59 -0700
Subject: [PATCH 109/395] Enable checkpointless eval and predict for
 tf.estimator.

PiperOrigin-RevId: 194993191
---
 tensorflow/python/estimator/estimator.py      | 17 ++++++----
 tensorflow/python/estimator/estimator_test.py | 32 +++++++++++++------
 2 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/tensorflow/python/estimator/estimator.py b/tensorflow/python/estimator/estimator.py
index 2363845110..63099b44bb 100644
--- a/tensorflow/python/estimator/estimator.py
+++ b/tensorflow/python/estimator/estimator.py
@@ -400,7 +400,9 @@ class Estimator(object):
       hooks: List of `SessionRunHook` subclass instances. Used for callbacks
         inside the evaluation call.
       checkpoint_path: Path of a specific checkpoint to evaluate. If `None`, the
-        latest checkpoint in `model_dir` is used.
+        latest checkpoint in `model_dir` is used.  If there are no checkpoints
+        in `model_dir`, evaluation is run with newly initialized `Variables`
+        instead of restored from checkpoint.
       name: Name of the evaluation if user needs to run multiple evaluations on
         different data sets, such as on training data vs test data. Metrics for
         different evaluations are saved in separate folders, and appear
@@ -464,7 +466,9 @@ class Estimator(object):
       hooks: List of `SessionRunHook` subclass instances. Used for callbacks
         inside the prediction call.
       checkpoint_path: Path of a specific checkpoint to predict. If `None`, the
-        latest checkpoint in `model_dir` is used.
+        latest checkpoint in `model_dir` is used.  If there are no checkpoints
+        in `model_dir`, prediction is run with newly initialized `Variables`
+        instead of restored from checkpoint.
       yield_single_examples: If False, yield the whole batch as returned by the
         `model_fn` instead of decomposing the batch into individual elements.
         This is useful if `model_fn` returns some tensors whose first dimension
@@ -487,9 +491,8 @@ class Estimator(object):
       if not checkpoint_path:
         checkpoint_path = saver.latest_checkpoint(self._model_dir)
       if not checkpoint_path:
-        raise ValueError(
-            'Could not find trained model in model_dir: {}.'.format(
-                self._model_dir))
+        logging.info('Could not find trained model in model_dir: {}, running '
+                     'initialization to predict.'.format(self._model_dir))
 
       with ops.Graph().as_default() as g:
         random_seed.set_random_seed(self._config.tf_random_seed)
@@ -1068,8 +1071,8 @@ class Estimator(object):
     if not checkpoint_path:
       latest_path = saver.latest_checkpoint(self._model_dir)
       if not latest_path:
-        raise ValueError('Could not find trained model in model_dir: {}.'.
-                         format(self._model_dir))
+        logging.info('Could not find trained model in model_dir: {}, running '
+                     'initialization to evaluate.'.format(self._model_dir))
       checkpoint_path = latest_path
 
     # Setup output directory.
diff --git a/tensorflow/python/estimator/estimator_test.py b/tensorflow/python/estimator/estimator_test.py
index 0fea86124c..74114fab3b 100644
--- a/tensorflow/python/estimator/estimator_test.py
+++ b/tensorflow/python/estimator/estimator_test.py
@@ -1067,11 +1067,19 @@ class EstimatorEvaluateTest(test.TestCase):
         ValueError, 'model_fn should return an EstimatorSpec'):
       est.evaluate(dummy_input_fn, steps=1)
 
-  def test_no_trained_model(self):
-    est = estimator.Estimator(model_fn=_model_fn_with_eval_metric_ops)
-    with self.assertRaisesRegexp(
-        ValueError, 'Could not find trained model in model_dir'):
-      est.evaluate(dummy_input_fn, steps=1)
+  def test_no_checkpoint_uses_init(self):
+    def _model_fn(features, labels, mode, params):
+      del features, labels, params
+      return model_fn_lib.EstimatorSpec(
+          mode,
+          loss=constant_op.constant(1.),
+          eval_metric_ops={'metric': metrics_lib.mean(
+              variables.Variable(2.) + 1)})
+    est = estimator.Estimator(model_fn=_model_fn)
+    metrics = est.evaluate(dummy_input_fn, steps=1)
+    # Metric value here is set to 1 + the value of the Variable that is newly
+    # initialized (since there is no checkpoint).
+    self.assertEqual(3., metrics['metric'])
 
   def test_scores(self):
     est = estimator.Estimator(
@@ -1331,11 +1339,15 @@ class EstimatorPredictTest(test.TestCase):
     next(est.predict(_input_fn))
     self.assertEqual(1, input_fn_call_count[0])
 
-  def test_no_trained_model_in_model_dir(self):
-    est = estimator.Estimator(model_fn=model_fn_global_step_incrementer)
-    with self.assertRaisesRegexp(ValueError,
-                                 'Could not find trained model in model_dir'):
-      next(est.predict(dummy_input_fn))
+  def test_no_checkpoint_uses_init(self):
+    def _model_fn(features, labels, mode, params, config):
+      del features, labels, params, config
+      x = variables.Variable([[3.]], name='x')
+      return model_fn_lib.EstimatorSpec(mode, predictions=math_ops.add(x, 1.))
+    est = estimator.Estimator(model_fn=_model_fn)
+    # Expected prediction value is 1 + the value of the Variable that is newly
+    # initialized (since there is no checkpoint).
+    self.assertEqual(4., next(est.predict(dummy_input_fn)))
 
   def test_no_trained_model_invalid_checkpoint_path(self):
     est = estimator.Estimator(model_fn=model_fn_global_step_incrementer)
-- 
GitLab


From 46bf1e8934b3bc8edeff3f218a50b0ee5806e96b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 14:27:33 -0700
Subject: [PATCH 110/395] Make tower-local variables non-trainable even with
 the default DistributionStrategy.

PiperOrigin-RevId: 194996819
---
 tensorflow/python/training/distribute.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/training/distribute.py b/tensorflow/python/training/distribute.py
index 6aeecb31dd..c16b05102e 100644
--- a/tensorflow/python/training/distribute.py
+++ b/tensorflow/python/training/distribute.py
@@ -1127,8 +1127,7 @@ class _DefaultDistributionStrategy(DistributionStrategy):
 
     def creator(next_creator, *args, **kwargs):
       _require_distribution_strategy_scope(self)
-      if kwargs.pop("tower_local_reduce_method", None) is not None:
-        kwargs["trainable"] = False
+      kwargs.pop("tower_local_reduce_method", None)
       return next_creator(*args, **kwargs)
 
     return _CurrentDistributionContext(
@@ -1138,7 +1137,7 @@ class _DefaultDistributionStrategy(DistributionStrategy):
     """Does not set to resource variables."""
     def create_tower_local_variable(next_creator, *args, **kwargs):
       _require_distribution_strategy_scope(self)
-      kwargs["tower_local_reduce_method"] = reduce_method
+      kwargs["trainable"] = False
       return next_creator(*args, **kwargs)
 
     _require_distribution_strategy_scope(self)
-- 
GitLab


From 325d0ef21a48bea1cc618a2bd24a9776de417ce5 Mon Sep 17 00:00:00 2001
From: Patrick Nguyen <drpng@google.com>
Date: Tue, 1 May 2018 14:28:36 -0700
Subject: [PATCH 111/395] Merge changes from github.

PiperOrigin-RevId: 194997009
---
 .gitignore                                    |   1 +
 tensorflow/c/c_api_test.cc                    |   2 +-
 tensorflow/cc/gradients/array_grad.cc         |  36 ++
 tensorflow/cc/gradients/array_grad_test.cc    |  24 +
 tensorflow/contrib/autograph/README.md        |   9 +-
 tensorflow/contrib/cmake/CMakeLists.txt       |   6 +-
 .../contrib/cmake/tf_core_kernels.cmake       |  10 +
 .../contrib/cmake/tf_stream_executor.cmake    |   2 +
 .../crf/python/kernel_tests/crf_test.py       |  24 +-
 tensorflow/contrib/crf/python/ops/crf.py      |  23 +-
 .../kernel_tests/bijectors/ordered_test.py    | 109 ++++
 .../python/ops/bijectors/__init__.py          |   2 +
 .../ops/bijectors/cholesky_outer_product.py   |   2 +-
 .../python/ops/bijectors/invert.py            |   4 +-
 .../ops/bijectors/masked_autoregressive.py    |   4 +-
 .../python/ops/bijectors/ordered.py           | 125 ++++
 .../python/ops/bijectors/permute.py           |   4 +-
 .../python/ops/bijectors/real_nvp.py          |   4 +-
 .../python/ops/bijectors/reshape.py           |   4 +-
 .../python/ops/bijectors/weibull.py           |   2 +-
 .../contrib/distributions/python/ops/shape.py |   2 +-
 .../factorization/python/ops/gmm_ops.py       |   2 +-
 .../contrib/layers/python/layers/layers.py    | 142 ++++-
 .../layers/python/layers/layers_test.py       |  15 +-
 .../layers/python/layers/target_column.py     |   4 +-
 .../learn/python/learn/estimators/head.py     |   4 +-
 .../learn/python/learn/ops/losses_ops.py      |   2 +-
 .../lite/examples/label_image/label_image.cc  |  45 +-
 .../lite/examples/label_image/label_image.h   |   1 +
 .../res/layout/fragment_camera2_basic.xml     |  28 +
 tensorflow/contrib/lite/kernels/topk_v2.cc    |   4 +-
 .../contrib/lite/kernels/topk_v2_test.cc      |   2 +-
 .../contrib/lite/nnapi/NeuralNetworksShim.h   |   2 +-
 .../contrib/lite/profiling/profile_buffer.h   |   8 +-
 .../propagate_array_data_types.cc             |   9 +
 .../propagate_fixed_sizes.cc                  |   4 +-
 .../contrib/lite/toco/import_tensorflow.cc    |   2 +-
 tensorflow/contrib/lite/toco/tooling_util.cc  |   9 +-
 tensorflow/contrib/mpi/mpi_utils.h            |   1 +
 .../python/training/lazy_adam_optimizer.py    |   6 +-
 tensorflow/contrib/optimizer_v2/adam.py       |  20 +-
 .../rnn/python/kernel_tests/core_rnn_test.py  |  15 +
 .../tensor_forest/client/eval_metrics.py      |   4 +-
 .../hybrid/python/layers/fully_connected.py   |   2 +-
 .../tensor_forest/python/tensor_forest.py     |   2 +-
 .../contrib/tensorrt/convert/convert_graph.cc |  13 +-
 .../timeseries/state_management_test.py       |   2 +-
 .../state_space_models/kalman_filter.py       |   6 +-
 tensorflow/core/BUILD                         |   4 +-
 .../api_def/base_api/api_def_ApplyAdam.pbtxt  |   8 +-
 .../core/api_def/base_api/api_def_Pad.pbtxt   |   1 +
 .../api_def/base_api/api_def_QuantizeV2.pbtxt |   6 +
 .../base_api/api_def_ResourceApplyAdam.pbtxt  |   8 +-
 .../api_def/base_api/api_def_ScatterNd.pbtxt  |  10 +-
 .../common_runtime/graph_execution_state.cc   |   2 +-
 tensorflow/core/graph/mkl_layout_pass.cc      |   8 +-
 tensorflow/core/grappler/clusters/BUILD       |   2 +
 tensorflow/core/grappler/clusters/cluster.h   |   6 +
 .../core/grappler/clusters/virtual_cluster.cc |   8 +
 .../core/grappler/clusters/virtual_cluster.h  |   4 +
 .../core/grappler/costs/virtual_scheduler.h   |   4 +-
 .../custom_graph_optimizer_registry.h         |   2 +-
 .../grappler/optimizers/meta_optimizer.cc     |  17 +-
 tensorflow/core/kernels/batch_util.cc         |   2 +
 tensorflow/core/kernels/cwise_op_floor_div.cc |   4 +-
 tensorflow/core/kernels/mkl_conv_ops.cc       | 414 +++++++++----
 tensorflow/core/kernels/scatter_nd_op.cc      |   1 +
 .../core/kernels/scatter_nd_op_cpu_impl.h     |   1 +
 .../core/kernels/segment_reduction_ops.h      |  29 -
 .../core/platform/default/gpu/cupti_wrapper.h |   2 +-
 tensorflow/core/public/version.h              |   2 +-
 tensorflow/core/util/mkl_util.h               |  87 ++-
 tensorflow/docs_src/community/roadmap.md      |  74 ++-
 .../docs_src/get_started/checkpoints.md       |   6 +-
 .../docs_src/get_started/feature_columns.md   |   2 +-
 tensorflow/docs_src/get_started/index.md      |   2 +-
 tensorflow/docs_src/install/install_c.md      |   2 +-
 tensorflow/docs_src/install/install_go.md     |   2 +-
 tensorflow/docs_src/install/install_java.md   |  22 +-
 tensorflow/docs_src/install/install_linux.md  | 571 +++++++++---------
 tensorflow/docs_src/install/install_mac.md    |  10 +-
 .../docs_src/install/install_sources.md       |   4 +-
 .../docs_src/performance/xla/tfcompile.md     |   6 +-
 .../examples/tutorials/estimators/__init__.py |   0
 .../examples/tutorials/input_fn/__init__.py   |   0
 .../examples/tutorials/layers/__init__.py     |   0
 .../examples/tutorials/monitors/__init__.py   |   0
 .../tutorials/monitors/iris_monitors.py       |   6 +-
 tensorflow/go/README.md                       |   2 +-
 tensorflow/go/op/wrappers.go                  |   2 +-
 tensorflow/python/estimator/estimator.py      |   4 +-
 .../python/keras/_impl/keras/estimator.py     |  22 +-
 .../keras/_impl/keras/estimator_test.py       |  24 +-
 .../python/kernel_tests/division_past_test.py |   3 +-
 .../kernel_tests/reduce_join_op_test.py       |   2 +-
 .../python/kernel_tests/reduction_ops_test.py |  35 +-
 .../kernel_tests/scatter_nd_ops_test.py       |  40 ++
 tensorflow/python/ops/array_grad.py           |   2 +-
 tensorflow/python/ops/array_ops.py            |  10 +-
 tensorflow/python/ops/image_ops_impl.py       |  10 +-
 tensorflow/python/ops/math_ops.py             |  15 +-
 tensorflow/python/ops/nn_impl.py              |   4 +-
 tensorflow/python/ops/rnn_cell_impl.py        |   6 +-
 tensorflow/python/training/adam.py            |  20 +-
 tensorflow/python/training/input_test.py      |  22 +
 .../python/training/monitored_session.py      |  14 +-
 .../golden/tensorflow.train.-scaffold.pbtxt   |   4 +
 .../tools/ci_build/ci_parameterized_build.sh  |   2 +-
 .../install/install_python3.5_pip_packages.sh |   3 +
 .../install/install_python3.6_pip_packages.sh |   4 +
 .../windows/cpu/bazel/run_cc_test_windows.sh  |   2 +-
 .../windows/gpu/bazel/run_cc_test_windows.sh  |   2 +-
 tensorflow/tools/docker/Dockerfile            |   1 +
 tensorflow/tools/docker/Dockerfile.devel      |   4 +
 tensorflow/tools/docker/Dockerfile.devel-gpu  |   4 +
 tensorflow/tools/docker/Dockerfile.gpu        |   1 +
 tensorflow/tools/docker/README.md             |  12 +-
 .../graph_transforms/fold_old_batch_norms.cc  |   3 +
 tensorflow/tools/pip_package/setup.py         |   4 +-
 tensorflow/workspace.bzl                      |  32 +-
 .../BackwardSpatialConvolutions.h             |   4 +-
 121 files changed, 1748 insertions(+), 663 deletions(-)
 create mode 100644 tensorflow/contrib/distributions/python/kernel_tests/bijectors/ordered_test.py
 create mode 100644 tensorflow/contrib/distributions/python/ops/bijectors/ordered.py
 create mode 100644 tensorflow/examples/tutorials/estimators/__init__.py
 create mode 100644 tensorflow/examples/tutorials/input_fn/__init__.py
 create mode 100644 tensorflow/examples/tutorials/layers/__init__.py
 create mode 100644 tensorflow/examples/tutorials/monitors/__init__.py

diff --git a/.gitignore b/.gitignore
index be75938ec4..828bbe9bd3 100644
--- a/.gitignore
+++ b/.gitignore
@@ -27,6 +27,7 @@ Podfile.lock
 /tensorflow/contrib/lite/examples/ios/simple/data/*.txt
 /tensorflow/contrib/lite/examples/ios/simple/data/*.tflite
 xcuserdata/**
+/api_init_files_list.txt
 
 # Android
 .gradle
diff --git a/tensorflow/c/c_api_test.cc b/tensorflow/c/c_api_test.cc
index ca80db23ed..9b86425aa5 100644
--- a/tensorflow/c/c_api_test.cc
+++ b/tensorflow/c/c_api_test.cc
@@ -1700,7 +1700,7 @@ TEST_F(CApiGradientsTest, OpWithNoGradientRegistered_NoGradInputs) {
   TestGradientsError(false);
 }
 
-// REGISTER_OP for CApiTestAttributesTest test cases.
+// REGISTER_OP for CApiAttributesTest test cases.
 // Registers two ops, each with a single attribute called 'v'.
 // The attribute in one op will have a type 'type', the other
 // will have list(type).
diff --git a/tensorflow/cc/gradients/array_grad.cc b/tensorflow/cc/gradients/array_grad.cc
index 6545e4ee3e..ff348fadb2 100644
--- a/tensorflow/cc/gradients/array_grad.cc
+++ b/tensorflow/cc/gradients/array_grad.cc
@@ -385,6 +385,42 @@ Status MirrorPadGradGrad(const Scope& scope, const Operation& op,
 }
 REGISTER_GRADIENT_OP("MirrorPadGrad", MirrorPadGradGrad);
 
+Status StridedSliceGradHelper(const Scope& scope, const Operation& op,
+                              const std::vector<Output>& grad_inputs,
+                              std::vector<Output>* grad_outputs) {
+  Input x = Shape(scope, op.input(0));
+  Input begin = op.input(1);
+  Input end = op.input(2);
+  Input strides = op.input(3);
+  int64 begin_mask;
+  int64 end_mask;
+  int64 ellipsis_mask;
+  int64 new_axis_mask;
+  int64 shrink_axis_mask;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "begin_mask", &begin_mask));
+  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "end_mask", &end_mask));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "ellipsis_mask", &ellipsis_mask));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "new_axis_mask", &new_axis_mask));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(op.node()->attrs(), "shrink_axis_mask", &shrink_axis_mask));
+  grad_outputs->push_back(
+      StridedSliceGrad(scope, x, begin, end, strides, grad_inputs[0],
+                       StridedSliceGrad::BeginMask(begin_mask)
+                           .EndMask(end_mask)
+                           .EllipsisMask(ellipsis_mask)
+                           .NewAxisMask(new_axis_mask)
+                           .ShrinkAxisMask(shrink_axis_mask)));
+  // No gradients returned for begin, end and strides
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  grad_outputs->push_back(NoGradient());
+  return scope.status();
+}
+REGISTER_GRADIENT_OP("StridedSlice", StridedSliceGradHelper);
+
 }  // anonymous namespace
 }  // namespace ops
 }  // namespace tensorflow
diff --git a/tensorflow/cc/gradients/array_grad_test.cc b/tensorflow/cc/gradients/array_grad_test.cc
index 4a215fcc92..de3bd0fc9e 100644
--- a/tensorflow/cc/gradients/array_grad_test.cc
+++ b/tensorflow/cc/gradients/array_grad_test.cc
@@ -354,5 +354,29 @@ TEST_F(ArrayGradTest, MirrorPadGradGrad_Symmetric) {
   RunTest(x, x_shape, y, y_shape);
 }
 
+TEST_F(ArrayGradTest, StridedSliceGrad) {
+  TensorShape x_shape({6, 4, 4});
+  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+
+  // y = x[2:6:2, 1:3, 1:3]
+  auto y = StridedSlice(scope_, x, {2, 1, 1}, {6, 3, 3}, {2, 1, 1});
+  // y.shape = [2, 2, 2];
+  RunTest(x, x_shape, y, {2, 2, 2});
+
+  // y = x[2:6:2, 1:3, 1:3]
+  // begin_mask = 1<<1 (ignore begin_index = 1)
+  // end_mask = 1<<2 (ignore end_index = 2)
+  y = StridedSlice(scope_, x, {2, 1, 1}, {6, 3, 3}, {2, 1, 1},
+                   StridedSlice::BeginMask(1 << 1).EndMask(1 << 2));
+  // y.shape = [2, 3, 3];
+  RunTest(x, x_shape, y, {2, 3, 3});
+
+  // y = [tf.newaxis, 2:6:2, 1:3, 1:3]
+  y = StridedSlice(scope_, x, {0, 2, 1, 1}, {0, 6, 3, 3}, {1, 2, 1, 1},
+                   StridedSlice::NewAxisMask(1 << 0));
+  // y.shape = [1, 2, 2, 2];
+  RunTest(x, x_shape, y, {1, 2, 2, 2});
+}
+
 }  // namespace
 }  // namespace tensorflow
diff --git a/tensorflow/contrib/autograph/README.md b/tensorflow/contrib/autograph/README.md
index 0fcbf5dd59..0ba99c396f 100644
--- a/tensorflow/contrib/autograph/README.md
+++ b/tensorflow/contrib/autograph/README.md
@@ -56,8 +56,6 @@ Use AutoGraph in one of the following ways, described below:
  1. Annotations (simpler)
  2. Functional API (more flexible)
 
-NOTE: You can find more examples in this [interactive notebook](https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/contrib/autograph/examples/notebooks/dev_summit_2018_demo.ipynb).
-
 To get started, install the latest nightly TensorFlow build:
 
 ```shell
@@ -70,6 +68,13 @@ Then import the `autograph` module from `tf.contrib`:
 from tensorflow.contrib import autograph as ag
 ```
 
+### Interactive demo notebooks
+
+For more extensive examples, check out these interactive notebooks:
+
+ * [RNN trained using Keras and Estimators](https://colab.sandbox.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/contrib/autograph/examples/notebooks/rnn_keras_estimator.ipynb)
+ * [Demo from the TF Dev Summit 2018](https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/contrib/autograph/examples/notebooks/dev_summit_2018_demo.ipynb)
+
 ## Using with annotations
 
 Annotating a function or class with `@convert` converts it in place:
diff --git a/tensorflow/contrib/cmake/CMakeLists.txt b/tensorflow/contrib/cmake/CMakeLists.txt
index 5f38a8e5c7..44e39f7f7b 100644
--- a/tensorflow/contrib/cmake/CMakeLists.txt
+++ b/tensorflow/contrib/cmake/CMakeLists.txt
@@ -84,7 +84,7 @@ if (NOT WIN32)
 
   option(systemlib_ALL "Turn on every possible systemlib_* options" OFF)
   if (systemlib_ALL)
-    set (systmelib_ZLIB ON)
+    set (systemlib_ZLIB ON)
   endif (systemlib_ALL)
 endif()
 
@@ -471,6 +471,10 @@ if (tensorflow_ENABLE_GPU)
   include_directories(${tensorflow_source_dir}/third_party/gpus)
   # add cuda libraries to tensorflow_EXTERNAL_LIBRARIES
   list(APPEND tensorflow_EXTERNAL_LIBRARIES ${CUDA_LIBRARIES})
+  if(NOT WIN32)
+    # add gomp to tensorflow_EXTERNAL_LIBRARIES, needed by libcusolver.so
+    list(APPEND tensorflow_EXTERNAL_LIBRARIES gomp)
+  endif()
 
   # NOTE(mrry): Update these flags when the version of CUDA or cuDNN used
   # in the default build is upgraded.
diff --git a/tensorflow/contrib/cmake/tf_core_kernels.cmake b/tensorflow/contrib/cmake/tf_core_kernels.cmake
index 376496b33f..f38c9e0513 100644
--- a/tensorflow/contrib/cmake/tf_core_kernels.cmake
+++ b/tensorflow/contrib/cmake/tf_core_kernels.cmake
@@ -177,6 +177,16 @@ if(WIN32)
       "${tensorflow_source_dir}/tensorflow/contrib/nccl/ops/nccl_ops.cc"
   )
   list(REMOVE_ITEM tf_core_kernels_srcs ${tf_core_kernels_windows_exclude_srcs})
+else(WIN32)
+  if(tensorflow_ENABLE_GPU)
+    file(GLOB_RECURSE tf_core_kernels_gpu_exclude_srcs
+        # temporarily disable nccl as it needs to be ported with gpu
+        "${tensorflow_source_dir}/tensorflow/contrib/nccl/kernels/nccl_manager.cc"
+        "${tensorflow_source_dir}/tensorflow/contrib/nccl/kernels/nccl_ops.cc"
+        "${tensorflow_source_dir}/tensorflow/contrib/nccl/ops/nccl_ops.cc"
+    )
+    list(REMOVE_ITEM tf_core_kernels_srcs ${tf_core_kernels_gpu_exclude_srcs})
+  endif(tensorflow_ENABLE_GPU)
 endif(WIN32)
 
 file(GLOB_RECURSE tf_core_gpu_kernels_srcs
diff --git a/tensorflow/contrib/cmake/tf_stream_executor.cmake b/tensorflow/contrib/cmake/tf_stream_executor.cmake
index af48ef1fd4..9a37b68119 100644
--- a/tensorflow/contrib/cmake/tf_stream_executor.cmake
+++ b/tensorflow/contrib/cmake/tf_stream_executor.cmake
@@ -64,6 +64,8 @@ file(GLOB tf_stream_executor_srcs
 if (tensorflow_ENABLE_GPU)
     file(GLOB tf_stream_executor_gpu_srcs
         "${tensorflow_source_dir}/tensorflow/stream_executor/cuda/*.cc"
+        "${tensorflow_source_dir}/tensorflow/compiler/xla/statusor.h"
+        "${tensorflow_source_dir}/tensorflow/compiler/xla/statusor.cc"
     )
     if (NOT tensorflow_BUILD_CC_TESTS)
         file(GLOB tf_stream_executor_gpu_tests
diff --git a/tensorflow/contrib/crf/python/kernel_tests/crf_test.py b/tensorflow/contrib/crf/python/kernel_tests/crf_test.py
index a5e065b93a..74f2ec22ff 100644
--- a/tensorflow/contrib/crf/python/kernel_tests/crf_test.py
+++ b/tensorflow/contrib/crf/python/kernel_tests/crf_test.py
@@ -152,6 +152,22 @@ class CrfTest(test.TestCase):
 
         self.assertAllClose(tf_log_norm, tf_brute_force_log_norm)
 
+  def testCrfLogNormZeroSeqLength(self):
+    """
+    Test `crf_log_norm` when `sequence_lengths` contains one or more zeros.
+    """
+    with self.test_session() as sess:
+      inputs = constant_op.constant(np.ones([2, 10, 5],
+                                            dtype=np.float32))
+      transition_params = constant_op.constant(np.ones([5, 5],
+                                                       dtype=np.float32))
+      sequence_lengths = constant_op.constant(np.zeros([2],
+                                                       dtype=np.int32))
+      expected_log_norm = np.zeros([2], dtype=np.float32)
+      log_norm = crf.crf_log_norm(inputs, sequence_lengths, transition_params)
+      tf_log_norm = sess.run(log_norm)
+      self.assertAllClose(tf_log_norm, expected_log_norm)
+
   def testCrfLogLikelihood(self):
     inputs = np.array(
         [[4, 5, -3], [3, -1, 3], [-1, 2, 1], [0, 0, 0]], dtype=np.float32)
@@ -292,10 +308,10 @@ class CrfTest(test.TestCase):
                                                        dtype=np.float32))
       sequence_lengths = constant_op.constant(np.zeros([2],
                                                        dtype=np.int32))
-      values = crf.crf_decode(inputs, transition_params, sequence_lengths)
-      tags, scores = sess.run(values)
-      self.assertEqual(len(tags.shape), 2)
-      self.assertEqual(len(scores.shape), 1)
+      tags, scores = crf.crf_decode(inputs, transition_params, sequence_lengths)
+      tf_tags, tf_scores = sess.run([tags, scores])
+      self.assertEqual(len(tf_tags.shape), 2)
+      self.assertEqual(len(tf_scores.shape), 1)
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/contrib/crf/python/ops/crf.py b/tensorflow/contrib/crf/python/ops/crf.py
index e37c029ceb..d2beff849e 100644
--- a/tensorflow/contrib/crf/python/ops/crf.py
+++ b/tensorflow/contrib/crf/python/ops/crf.py
@@ -90,9 +90,13 @@ def crf_sequence_score(inputs, tag_indices, sequence_lengths,
     batch_size = array_ops.shape(inputs, out_type=tag_indices.dtype)[0]
     example_inds = array_ops.reshape(
         math_ops.range(batch_size, dtype=tag_indices.dtype), [-1, 1])
-    return array_ops.gather_nd(
+    sequence_scores = array_ops.gather_nd(
         array_ops.squeeze(inputs, [1]),
         array_ops.concat([example_inds, tag_indices], axis=1))
+    sequence_scores = array_ops.where(math_ops.less_equal(sequence_lengths, 0),
+                                      array_ops.zeros_like(sequence_scores),
+                                      sequence_scores)
+    return sequence_scores
 
   def _multi_seq_fn():
     # Compute the scores of the given tag sequence.
@@ -128,7 +132,12 @@ def crf_log_norm(inputs, sequence_lengths, transition_params):
   # If max_seq_len is 1, we skip the algorithm and simply reduce_logsumexp over
   # the "initial state" (the unary potentials).
   def _single_seq_fn():
-    return math_ops.reduce_logsumexp(first_input, [1])
+    log_norm = math_ops.reduce_logsumexp(first_input, [1])
+    # Mask `log_norm` of the sequences with length <= zero.
+    log_norm = array_ops.where(math_ops.less_equal(sequence_lengths, 0),
+                               array_ops.zeros_like(log_norm),
+                               log_norm)
+    return log_norm
 
   def _multi_seq_fn():
     """Forward computation of alpha values."""
@@ -137,13 +146,19 @@ def crf_log_norm(inputs, sequence_lengths, transition_params):
     # Compute the alpha values in the forward algorithm in order to get the
     # partition function.
     forward_cell = CrfForwardRnnCell(transition_params)
+    # Sequence length is not allowed to be less than zero.
+    sequence_lengths_less_one = math_ops.maximum(0, sequence_lengths - 1)
     _, alphas = rnn.dynamic_rnn(
         cell=forward_cell,
         inputs=rest_of_input,
-        sequence_length=sequence_lengths - 1,
+        sequence_length=sequence_lengths_less_one,
         initial_state=first_input,
         dtype=dtypes.float32)
     log_norm = math_ops.reduce_logsumexp(alphas, [1])
+    # Mask `log_norm` of the sequences with length <= zero.
+    log_norm = array_ops.where(math_ops.less_equal(sequence_lengths, 0),
+                               array_ops.zeros_like(log_norm),
+                               log_norm)
     return log_norm
 
   max_seq_len = array_ops.shape(inputs)[1]
@@ -479,7 +494,7 @@ def crf_decode(potentials, transition_params, sequence_length):
     initial_state = array_ops.slice(potentials, [0, 0, 0], [-1, 1, -1])
     initial_state = array_ops.squeeze(initial_state, axis=[1])  # [B, O]
     inputs = array_ops.slice(potentials, [0, 1, 0], [-1, -1, -1])  # [B, T-1, O]
-    # sequence length is not allowed to be less than zero
+    # Sequence length is not allowed to be less than zero.
     sequence_length_less_one = math_ops.maximum(0, sequence_length - 1)
     backpointers, last_score = rnn.dynamic_rnn(  # [B, T - 1, O], [B, O]
         crf_fwd_cell,
diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/ordered_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/ordered_test.py
new file mode 100644
index 0000000000..a5f5219588
--- /dev/null
+++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/ordered_test.py
@@ -0,0 +1,109 @@
+# Copyright 2018 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 Bijector."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.contrib.distributions.python.ops.bijectors.ordered import Ordered
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.distributions.bijector_test_util import assert_bijective_and_finite
+from tensorflow.python.platform import test
+
+
+
+class OrderedBijectorTest(test.TestCase):
+  """Tests correctness of the ordered transformation."""
+
+  def setUp(self):
+    self._rng = np.random.RandomState(42)
+
+  @test_util.run_in_graph_and_eager_modes()
+  def testBijectorVector(self):
+    with self.test_session():
+      ordered = Ordered()
+      self.assertEqual("ordered", ordered.name)
+      x = np.asarray([[2., 3, 4], [4., 8, 13]])
+      y = [[2., 0, 0], [4., np.log(4.), np.log(5.)]]
+      self.assertAllClose(y, self.evaluate(ordered.forward(x)))
+      self.assertAllClose(x, self.evaluate(ordered.inverse(y)))
+      self.assertAllClose(
+          np.sum(np.asarray(y)[..., 1:], axis=-1),
+          self.evaluate(ordered.inverse_log_det_jacobian(y, event_ndims=1)),
+          atol=0.,
+          rtol=1e-7)
+      self.assertAllClose(
+          self.evaluate(-ordered.inverse_log_det_jacobian(y, event_ndims=1)),
+          self.evaluate(ordered.forward_log_det_jacobian(x, event_ndims=1)),
+          atol=0.,
+          rtol=1e-7)
+
+  def testBijectorUnknownShape(self):
+    with self.test_session():
+      ordered = Ordered()
+      self.assertEqual("ordered", ordered.name)
+      x = array_ops.placeholder(shape=[2, None], dtype=dtypes.float32)
+      real_x = np.asarray([[2., 3, 4], [4., 8, 13]])
+      y = array_ops.placeholder(shape=[2, None], dtype=dtypes.float32)
+      real_y = [[2., 0, 0], [4., np.log(4.), np.log(5.)]]
+      self.assertAllClose(real_y, ordered.forward(x).eval(
+          feed_dict={x: real_x}))
+      self.assertAllClose(real_x, ordered.inverse(y).eval(
+          feed_dict={y: real_y}))
+      self.assertAllClose(
+          np.sum(np.asarray(real_y)[..., 1:], axis=-1),
+          ordered.inverse_log_det_jacobian(y, event_ndims=1).eval(
+              feed_dict={y: real_y}),
+          atol=0.,
+          rtol=1e-7)
+      self.assertAllClose(
+          -ordered.inverse_log_det_jacobian(y, event_ndims=1).eval(
+              feed_dict={y: real_y}),
+          ordered.forward_log_det_jacobian(x, event_ndims=1).eval(
+              feed_dict={x: real_x}),
+          atol=0.,
+          rtol=1e-7)
+
+  @test_util.run_in_graph_and_eager_modes()
+  def testShapeGetters(self):
+    with self.test_session():
+      x = tensor_shape.TensorShape([4])
+      y = tensor_shape.TensorShape([4])
+      bijector = Ordered(validate_args=True)
+      self.assertAllEqual(y, bijector.forward_event_shape(x))
+      self.assertAllEqual(y.as_list(),
+                          self.evaluate(bijector.forward_event_shape_tensor(
+                              x.as_list())))
+      self.assertAllEqual(x, bijector.inverse_event_shape(y))
+      self.assertAllEqual(x.as_list(),
+                          self.evaluate(bijector.inverse_event_shape_tensor(
+                              y.as_list())))
+
+  def testBijectiveAndFinite(self):
+    with self.test_session():
+      ordered = Ordered()
+      x = np.sort(self._rng.randn(3, 10), axis=-1).astype(np.float32)
+      y = (self._rng.randn(3, 10)).astype(np.float32)
+      assert_bijective_and_finite(ordered, x, y, event_ndims=1)
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py b/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py
index babce80396..51478dbeff 100644
--- a/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py
+++ b/tensorflow/contrib/distributions/python/ops/bijectors/__init__.py
@@ -30,6 +30,7 @@
 @@Invert
 @@Kumaraswamy
 @@MaskedAutoregressiveFlow
+@@Ordered
 @@Permute
 @@PowerTransform
 @@RealNVP
@@ -67,6 +68,7 @@ from tensorflow.contrib.distributions.python.ops.bijectors.inline import *
 from tensorflow.contrib.distributions.python.ops.bijectors.invert import *
 from tensorflow.contrib.distributions.python.ops.bijectors.kumaraswamy import *
 from tensorflow.contrib.distributions.python.ops.bijectors.masked_autoregressive import *
+from tensorflow.contrib.distributions.python.ops.bijectors.ordered import *
 from tensorflow.contrib.distributions.python.ops.bijectors.permute import *
 from tensorflow.contrib.distributions.python.ops.bijectors.power_transform import *
 from tensorflow.contrib.distributions.python.ops.bijectors.real_nvp import *
diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py b/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py
index caae2adcfa..ecdb8967f4 100644
--- a/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py
+++ b/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py
@@ -170,7 +170,7 @@ class CholeskyOuterProduct(bijector.Bijector):
     sum_weighted_log_diag = array_ops.squeeze(
         math_ops.matmul(math_ops.log(diag),
                         exponents[..., array_ops.newaxis]),
-        squeeze_dims=-1)
+        axis=-1)
     fldj = p_float * np.log(2.) + sum_weighted_log_diag
 
     return fldj
diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/invert.py b/tensorflow/contrib/distributions/python/ops/bijectors/invert.py
index 1904239a0e..84a3289ba2 100644
--- a/tensorflow/contrib/distributions/python/ops/bijectors/invert.py
+++ b/tensorflow/contrib/distributions/python/ops/bijectors/invert.py
@@ -18,14 +18,14 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensorflow.python.ops.distributions import bijector as bijector_lib
+from tensorflow.python.ops.distributions import bijector
 
 __all__ = [
     "Invert",
 ]
 
 
-class Invert(bijector_lib.Bijector):
+class Invert(bijector.Bijector):
   """Bijector which inverts another Bijector.
 
   Example Use: [ExpGammaDistribution (see Background & Context)](
diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py b/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py
index ef56cf6ddd..83667b0e80 100644
--- a/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py
+++ b/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py
@@ -32,7 +32,7 @@ from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import nn_ops
 from tensorflow.python.ops import template as template_ops
 from tensorflow.python.ops import variable_scope as variable_scope_lib
-from tensorflow.python.ops.distributions import bijector as bijector_lib
+from tensorflow.python.ops.distributions import bijector
 
 
 __all__ = [
@@ -42,7 +42,7 @@ __all__ = [
 ]
 
 
-class MaskedAutoregressiveFlow(bijector_lib.Bijector):
+class MaskedAutoregressiveFlow(bijector.Bijector):
   """Affine MaskedAutoregressiveFlow bijector for vector-valued events.
 
   The affine autoregressive flow [(Papamakarios et al., 2016)][3] provides a
diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/ordered.py b/tensorflow/contrib/distributions/python/ops/bijectors/ordered.py
new file mode 100644
index 0000000000..3f03592f31
--- /dev/null
+++ b/tensorflow/contrib/distributions/python/ops/bijectors/ordered.py
@@ -0,0 +1,125 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Ordered bijector."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.distributions import bijector
+
+
+__all__ = [
+    "Ordered",
+]
+
+
+class Ordered(bijector.Bijector):
+  """Bijector which maps a tensor x_k that has increasing elements in the last
+  dimension to an unconstrained tensor y_k.
+
+  Both the domain and the codomain of the mapping is [-inf, inf], however,
+  the input of the forward mapping must be strictly increasing.
+  The inverse of the bijector applied to a normal random vector `y ~ N(0, 1)`
+  gives back a sorted random vector with the same distribution `x ~ N(0, 1)`
+  where `x = sort(y)`
+
+  On the last dimension of the tensor, Ordered bijector performs:
+  `y[0] = x[0]`
+  `y[1:] = math_ops.log(x[1:] - x[:-1])`
+
+  #### Example Use:
+
+  ```python
+  bijector.Ordered().forward([2, 3, 4])
+  # Result: [2., 0., 0.]
+
+  bijector.Ordered().inverse([0.06428002, -1.07774478, -0.71530371])
+  # Result: [0.06428002, 0.40464228, 0.8936858]
+  ```
+  """
+
+  def __init__(self, validate_args=False, name="ordered"):
+    super(Ordered, self).__init__(
+        forward_min_event_ndims=1,
+        validate_args=validate_args,
+        name=name)
+
+  def _forward_event_shape(self, input_shape):
+    if input_shape.ndims is None or input_shape[-1] is None:
+      return input_shape
+    return tensor_shape.TensorShape([input_shape[-1]])
+
+  def _forward_event_shape_tensor(self, input_shape):
+    return (input_shape[-1])[..., array_ops.newaxis]
+
+  def _inverse_event_shape(self, output_shape):
+    if output_shape.ndims is None or output_shape[-1] is None:
+      return output_shape
+    if output_shape[-1] <= 1:
+      raise ValueError("output_shape[-1] = %d <= 1" % output_shape[-1])
+    return tensor_shape.TensorShape([output_shape[-1]])
+
+  def _inverse_event_shape_tensor(self, output_shape):
+    if self.validate_args:
+      is_greater_one = check_ops.assert_greater(
+          output_shape[-1], 1, message="Need last dimension greater than 1.")
+      output_shape = control_flow_ops.with_dependencies(
+          [is_greater_one], output_shape)
+    return (output_shape[-1])[..., array_ops.newaxis]
+
+  def _forward(self, x):
+    x = self._maybe_assert_valid_x(x)
+    y0 = x[..., 0, array_ops.newaxis]
+    yk = math_ops.log(x[..., 1:] - x[..., :-1])
+    y = array_ops.concat([y0, yk], axis=-1)
+    return y
+
+  def _inverse(self, y):
+    x0 = y[..., 0, array_ops.newaxis]
+    xk = math_ops.exp(y[..., 1:])
+    x = array_ops.concat([x0, xk], axis=-1)
+    return math_ops.cumsum(x, axis=-1)
+
+  def _inverse_log_det_jacobian(self, y):
+    # The Jacobian of the inverse mapping is lower
+    # triangular, with the diagonal elements being:
+    # J[i,i] = 1 if i=1, and
+    #          exp(y_i) if 1<i<=K
+    # which gives the absolute Jacobian determinant:
+    # |det(Jac)| = prod_{i=1}^{K} exp(y[i]).
+    # (1) - Stan Modeling Language User's Guide and Reference Manual
+    #       Version 2.17.0 session 35.2
+    return math_ops.reduce_sum(y[..., 1:], axis=-1)
+
+  def _forward_log_det_jacobian(self, x):
+    x = self._maybe_assert_valid_x(x)
+    return -math_ops.reduce_sum(
+        math_ops.log(x[..., 1:] - x[..., :-1]),
+        axis=-1)
+
+  def _maybe_assert_valid_x(self, x):
+    if not self.validate_args:
+      return x
+    is_valid = check_ops.assert_positive(
+        x[..., 1:] - x[..., :-1],
+        message="Forward transformation input must be strictly increasing.")
+    return control_flow_ops.with_dependencies([is_valid], x)
diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/permute.py b/tensorflow/contrib/distributions/python/ops/bijectors/permute.py
index 4978167803..12a16a3f2b 100644
--- a/tensorflow/contrib/distributions/python/ops/bijectors/permute.py
+++ b/tensorflow/contrib/distributions/python/ops/bijectors/permute.py
@@ -28,7 +28,7 @@ from tensorflow.python.ops import check_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import nn_ops
-from tensorflow.python.ops.distributions import bijector as bijector_lib
+from tensorflow.python.ops.distributions import bijector
 
 
 __all__ = [
@@ -36,7 +36,7 @@ __all__ = [
 ]
 
 
-class Permute(bijector_lib.Bijector):
+class Permute(bijector.Bijector):
   """Permutes the rightmost dimension of a `Tensor`.
 
   ```python
diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py b/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py
index f09ab21bce..66e8a5b9b3 100644
--- a/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py
+++ b/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py
@@ -25,7 +25,7 @@ from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import nn_ops
 from tensorflow.python.ops import template as template_ops
-from tensorflow.python.ops.distributions import bijector as bijector_lib
+from tensorflow.python.ops.distributions import bijector
 
 
 __all__ = [
@@ -34,7 +34,7 @@ __all__ = [
 ]
 
 
-class RealNVP(bijector_lib.Bijector):
+class RealNVP(bijector.Bijector):
   """RealNVP "affine coupling layer" for vector-valued events.
 
   Real NVP models a normalizing flow on a `D`-dimensional distribution via a
diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/reshape.py b/tensorflow/contrib/distributions/python/ops/bijectors/reshape.py
index f21b982ba6..5497c422e4 100644
--- a/tensorflow/contrib/distributions/python/ops/bijectors/reshape.py
+++ b/tensorflow/contrib/distributions/python/ops/bijectors/reshape.py
@@ -28,7 +28,7 @@ from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import check_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import math_ops
-from tensorflow.python.ops.distributions import bijector as bijector_lib
+from tensorflow.python.ops.distributions import bijector
 
 
 __all__ = [
@@ -44,7 +44,7 @@ def _ndims_from_shape(shape):
   return array_ops.shape(shape)[0]
 
 
-class Reshape(bijector_lib.Bijector):
+class Reshape(bijector.Bijector):
   """Reshapes the `event_shape` of a `Tensor`.
 
   The semantics generally follow that of `tf.reshape()`, with
diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/weibull.py b/tensorflow/contrib/distributions/python/ops/bijectors/weibull.py
index 39129cd22c..a22560fe80 100644
--- a/tensorflow/contrib/distributions/python/ops/bijectors/weibull.py
+++ b/tensorflow/contrib/distributions/python/ops/bijectors/weibull.py
@@ -128,7 +128,7 @@ class Weibull(bijector.Bijector):
       return x
     is_valid = check_ops.assert_non_negative(
         x,
-        message="Forward transformation input must be at least {}.".format(0))
+        message="Forward transformation input must be at least 0.")
     return control_flow_ops.with_dependencies([is_valid], x)
 
   def _maybe_assert_valid_y(self, y):
diff --git a/tensorflow/contrib/distributions/python/ops/shape.py b/tensorflow/contrib/distributions/python/ops/shape.py
index bac0b79d59..6a7f28713a 100644
--- a/tensorflow/contrib/distributions/python/ops/shape.py
+++ b/tensorflow/contrib/distributions/python/ops/shape.py
@@ -439,7 +439,7 @@ class _DistributionShape(object):
           if self._batch_ndims_is_0 and expand_batch_dim:
             squeeze_dims += [1]
           if squeeze_dims:
-            x = array_ops.squeeze(x, squeeze_dims=squeeze_dims)
+            x = array_ops.squeeze(x, axis=squeeze_dims)
             # x.shape: [prod(S)]+B+E
         _, batch_shape, event_shape = self.get_shape(x)
       else:
diff --git a/tensorflow/contrib/factorization/python/ops/gmm_ops.py b/tensorflow/contrib/factorization/python/ops/gmm_ops.py
index ccdd679d6a..e076631bc1 100644
--- a/tensorflow/contrib/factorization/python/ops/gmm_ops.py
+++ b/tensorflow/contrib/factorization/python/ops/gmm_ops.py
@@ -397,7 +397,7 @@ class GmmAlgorithm(object):
     # Compute the effective number of data points assigned to component k.
     with ops.control_dependencies(self._w):
       points_in_k = array_ops.squeeze(
-          math_ops.add_n(self._points_in_k), squeeze_dims=[0])
+          math_ops.add_n(self._points_in_k), axis=[0])
       # Update alpha.
       if 'w' in self._params:
         final_points_in_k = points_in_k / num_batches
diff --git a/tensorflow/contrib/layers/python/layers/layers.py b/tensorflow/contrib/layers/python/layers/layers.py
index 25c3b1e7ea..2f3e57653c 100644
--- a/tensorflow/contrib/layers/python/layers/layers.py
+++ b/tensorflow/contrib/layers/python/layers/layers.py
@@ -932,7 +932,8 @@ def convolution(inputs,
                 variables_collections=None,
                 outputs_collections=None,
                 trainable=True,
-                scope=None):
+                scope=None,
+                conv_dims=None):
   """Adds an N-D convolution followed by an optional batch_norm layer.
 
   It is required that 1 <= N <= 3.
@@ -993,6 +994,10 @@ def convolution(inputs,
     trainable: If `True` also add variables to the graph collection
       `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
     scope: Optional scope for `variable_scope`.
+    conv_dims: Optional convolution dimensionality, when set it would use the
+      corresponding convolution (e.g. 2 for Conv 2D, 3 for Conv 3D, ..). When
+      leaved to None it would select the convolution dimensionality based on
+      the input rank (i.e. Conv ND, with N = input_rank - 2).
 
   Returns:
     A tensor representing the output of the operation.
@@ -1015,6 +1020,9 @@ def convolution(inputs,
     inputs = ops.convert_to_tensor(inputs)
     input_rank = inputs.get_shape().ndims
 
+    if conv_dims is not None and conv_dims + 2 != input_rank:
+      raise ValueError('Convolution expects input with rank %d, got %d' %
+                       (conv_dims + 2, input_rank))
     if input_rank == 3:
       layer_class = convolutional_layers.Convolution1D
     elif input_rank == 4:
@@ -1061,10 +1069,134 @@ def convolution(inputs,
       outputs = activation_fn(outputs)
     return utils.collect_named_outputs(outputs_collections, sc.name, outputs)
 
+@add_arg_scope
+def convolution1d(inputs,
+                  num_outputs,
+                  kernel_size,
+                  stride=1,
+                  padding='SAME',
+                  data_format=None,
+                  rate=1,
+                  activation_fn=nn.relu,
+                  normalizer_fn=None,
+                  normalizer_params=None,
+                  weights_initializer=initializers.xavier_initializer(),
+                  weights_regularizer=None,
+                  biases_initializer=init_ops.zeros_initializer(),
+                  biases_regularizer=None,
+                  reuse=None,
+                  variables_collections=None,
+                  outputs_collections=None,
+                  trainable=True,
+                  scope=None):
+  return convolution(inputs,
+                     num_outputs,
+                     kernel_size,
+                     stride,
+                     padding,
+                     data_format,
+                     rate,
+                     activation_fn,
+                     normalizer_fn,
+                     normalizer_params,
+                     weights_initializer,
+                     weights_regularizer,
+                     biases_initializer,
+                     biases_regularizer,
+                     reuse,
+                     variables_collections,
+                     outputs_collections,
+                     trainable,
+                     scope,
+                     conv_dims=1)
+
+convolution1d.__doc__ = convolution.__doc__
 
-convolution2d = convolution
-convolution3d = convolution
+@add_arg_scope
+def convolution2d(inputs,
+                  num_outputs,
+                  kernel_size,
+                  stride=1,
+                  padding='SAME',
+                  data_format=None,
+                  rate=1,
+                  activation_fn=nn.relu,
+                  normalizer_fn=None,
+                  normalizer_params=None,
+                  weights_initializer=initializers.xavier_initializer(),
+                  weights_regularizer=None,
+                  biases_initializer=init_ops.zeros_initializer(),
+                  biases_regularizer=None,
+                  reuse=None,
+                  variables_collections=None,
+                  outputs_collections=None,
+                  trainable=True,
+                  scope=None):
+  return convolution(inputs,
+                     num_outputs,
+                     kernel_size,
+                     stride,
+                     padding,
+                     data_format,
+                     rate,
+                     activation_fn,
+                     normalizer_fn,
+                     normalizer_params,
+                     weights_initializer,
+                     weights_regularizer,
+                     biases_initializer,
+                     biases_regularizer,
+                     reuse,
+                     variables_collections,
+                     outputs_collections,
+                     trainable,
+                     scope,
+                     conv_dims=2)
+
+convolution2d.__doc__ = convolution.__doc__
 
+@add_arg_scope
+def convolution3d(inputs,
+                  num_outputs,
+                  kernel_size,
+                  stride=1,
+                  padding='SAME',
+                  data_format=None,
+                  rate=1,
+                  activation_fn=nn.relu,
+                  normalizer_fn=None,
+                  normalizer_params=None,
+                  weights_initializer=initializers.xavier_initializer(),
+                  weights_regularizer=None,
+                  biases_initializer=init_ops.zeros_initializer(),
+                  biases_regularizer=None,
+                  reuse=None,
+                  variables_collections=None,
+                  outputs_collections=None,
+                  trainable=True,
+                  scope=None):
+  return convolution(inputs,
+                     num_outputs,
+                     kernel_size,
+                     stride,
+                     padding,
+                     data_format,
+                     rate,
+                     activation_fn,
+                     normalizer_fn,
+                     normalizer_params,
+                     weights_initializer,
+                     weights_regularizer,
+                     biases_initializer,
+                     biases_regularizer,
+                     reuse,
+                     variables_collections,
+                     outputs_collections,
+                     trainable,
+                     scope,
+                     conv_dims=3)
+
+convolution3d.__doc__ = convolution.__doc__
 
 @add_arg_scope
 def convolution2d_in_plane(
@@ -1411,7 +1543,7 @@ def dense_to_sparse(tensor, eos_token=0, outputs_collections=None, scope=None):
   Args:
      tensor: An `int` `Tensor` to be converted to a `Sparse`.
      eos_token: An integer.
-       It is part of the target label that signfies the end of a sentence.
+       It is part of the target label that signifies the end of a sentence.
      outputs_collections: Collection to add the outputs.
      scope: Optional scope for name_scope.
   """
@@ -1555,7 +1687,7 @@ def _inner_flatten(inputs, new_rank, output_collections=None, scope=None):
     output_collections: Collection to which the outputs will be added.
     scope: Optional scope for `name_scope`.
   Returns:
-    A `Tensor` or `SparseTensor` conataining the same values as `inputs`, but
+    A `Tensor` or `SparseTensor` containing the same values as `inputs`, but
     with innermost dimensions flattened to obtain rank `new_rank`.
 
   Raises:
diff --git a/tensorflow/contrib/layers/python/layers/layers_test.py b/tensorflow/contrib/layers/python/layers/layers_test.py
index 997f910a2a..b01fd5d5c9 100644
--- a/tensorflow/contrib/layers/python/layers/layers_test.py
+++ b/tensorflow/contrib/layers/python/layers/layers_test.py
@@ -310,6 +310,17 @@ class BiasAddTest(test.TestCase):
 
 class ConvolutionTest(test.TestCase):
 
+  def testInvalidShape(self):
+    with self.test_session():
+      images_2d = random_ops.random_uniform((5, 7, 9, 3), seed=1)
+      with self.assertRaisesRegexp(
+          ValueError, 'Convolution expects input with rank 5, got 4'):
+        layers_lib.convolution3d(images_2d, 32, 3)
+      images_3d = random_ops.random_uniform((5, 6, 7, 9, 3), seed=1)
+      with self.assertRaisesRegexp(
+          ValueError, 'Convolution expects input with rank 4, got 5'):
+        layers_lib.convolution2d(images_3d, 32, 3)
+
   def testInvalidDataFormat(self):
     height, width = 7, 9
     with self.test_session():
@@ -3155,7 +3166,7 @@ class RepeatTests(test.TestCase):
     with self.test_session():
       images = np.random.uniform(size=(5, height, width, 3)).astype(np.float32)
       output = _layers.repeat(images, 3, layers_lib.conv2d, 32, [3, 3])
-      self.assertEqual(output.op.name, 'Repeat/convolution_3/Relu')
+      self.assertEqual(output.op.name, 'Repeat/convolution2d_3/Relu')
       self.assertListEqual(output.get_shape().as_list(), [5, 3, 3, 32])
 
   def testRepeatWithScope(self):
@@ -3749,7 +3760,7 @@ class StackTests(test.TestCase):
           layers_lib.convolution2d, [10, 20, 30],
           kernel_size=[3, 3],
           padding='SAME')
-      self.assertEqual(output.op.name, 'Stack/convolution_3/Relu')
+      self.assertEqual(output.op.name, 'Stack/convolution2d_3/Relu')
       self.assertListEqual(output.get_shape().as_list(), [5, 3, 3, 30])
 
   def testStackWithScope(self):
diff --git a/tensorflow/contrib/layers/python/layers/target_column.py b/tensorflow/contrib/layers/python/layers/target_column.py
index 3e639a180e..69bb6be814 100644
--- a/tensorflow/contrib/layers/python/layers/target_column.py
+++ b/tensorflow/contrib/layers/python/layers/target_column.py
@@ -270,7 +270,7 @@ class _RegressionTargetColumn(_TargetColumn):
 
   def logits_to_predictions(self, logits, proba=False):
     if self.num_label_columns == 1:
-      return array_ops.squeeze(logits, squeeze_dims=[1])
+      return array_ops.squeeze(logits, axis=[1])
     return logits
 
   def get_eval_ops(self, features, logits, labels, metrics=None):
@@ -418,7 +418,7 @@ def _softmax_cross_entropy_loss(logits, target):
                      "Instead got %s." % target.dtype)
   # sparse_softmax_cross_entropy_with_logits requires [batch_size] target.
   if len(target.get_shape()) == 2:
-    target = array_ops.squeeze(target, squeeze_dims=[1])
+    target = array_ops.squeeze(target, axis=[1])
   loss_vec = nn.sparse_softmax_cross_entropy_with_logits(
       labels=target, logits=logits)
   return loss_vec
diff --git a/tensorflow/contrib/learn/python/learn/estimators/head.py b/tensorflow/contrib/learn/python/learn/estimators/head.py
index 2b4b6eff39..e28e6854a5 100644
--- a/tensorflow/contrib/learn/python/learn/estimators/head.py
+++ b/tensorflow/contrib/learn/python/learn/estimators/head.py
@@ -777,7 +777,7 @@ class _RegressionHead(_SingleHead):
     key = prediction_key.PredictionKey.SCORES
     with ops.name_scope(None, "predictions", (logits,)):
       if self.logits_dimension == 1:
-        logits = array_ops.squeeze(logits, squeeze_dims=(1,), name=key)
+        logits = array_ops.squeeze(logits, axis=(1,), name=key)
       return {key: self._link_fn(logits)}
 
   def _metrics(self, eval_loss, predictions, labels, weights):
@@ -974,7 +974,7 @@ def _softmax_cross_entropy_loss(labels, logits, weights=None):
     is_squeezed_labels = False
     # TODO(ptucker): This will break for dynamic shapes.
     if len(labels.get_shape()) == 2:
-      labels = array_ops.squeeze(labels, squeeze_dims=(1,))
+      labels = array_ops.squeeze(labels, axis=(1,))
       is_squeezed_labels = True
 
     loss = nn.sparse_softmax_cross_entropy_with_logits(
diff --git a/tensorflow/contrib/learn/python/learn/ops/losses_ops.py b/tensorflow/contrib/learn/python/learn/ops/losses_ops.py
index 92976d1539..9f2cadb017 100644
--- a/tensorflow/contrib/learn/python/learn/ops/losses_ops.py
+++ b/tensorflow/contrib/learn/python/learn/ops/losses_ops.py
@@ -40,7 +40,7 @@ def mean_squared_error_regressor(tensor_in, labels, weights, biases, name=None):
                       [tensor_in, labels]):
     predictions = nn.xw_plus_b(tensor_in, weights, biases)
     if len(labels.get_shape()) == 1 and len(predictions.get_shape()) == 2:
-      predictions = array_ops_.squeeze(predictions, squeeze_dims=[1])
+      predictions = array_ops_.squeeze(predictions, axis=[1])
     return predictions, losses.mean_squared_error(labels, predictions)
 
 
diff --git a/tensorflow/contrib/lite/examples/label_image/label_image.cc b/tensorflow/contrib/lite/examples/label_image/label_image.cc
index a91467d345..456c5c6dc7 100644
--- a/tensorflow/contrib/lite/examples/label_image/label_image.cc
+++ b/tensorflow/contrib/lite/examples/label_image/label_image.cc
@@ -17,6 +17,7 @@ limitations under the License.
 #include <cstdio>
 #include <cstdlib>
 #include <fstream>
+#include <iomanip>
 #include <iostream>
 #include <memory>
 #include <sstream>
@@ -70,6 +71,23 @@ TfLiteStatus ReadLabelsFile(const string& file_name,
   return kTfLiteOk;
 }
 
+void PrintProfilingInfo(const profiling::ProfileEvent* e, uint32_t op_index,
+                        TfLiteRegistration registration) {
+  // output something like
+  // time (ms) , Node xxx, OpCode xxx, symblic name
+  //      5.352, Node   5, OpCode   4, DEPTHWISE_CONV_2D
+
+
+  LOG(INFO) << std::fixed << std::setw(10) << std::setprecision(3)
+            << (e->end_timestamp_us - e->begin_timestamp_us) / 1000.0
+            << ", Node " << std::setw(3) << std::setprecision(3) << op_index
+            << ", OpCode " << std::setw(3) << std::setprecision(3)
+            << registration.builtin_code << ", "
+            << EnumNameBuiltinOperator(
+                   (BuiltinOperator)registration.builtin_code)
+            << "\n";
+}
+
 void RunInference(Settings* s) {
   if (!s->model_name.c_str()) {
     LOG(ERROR) << "no model file name\n";
@@ -166,6 +184,11 @@ void RunInference(Settings* s) {
       exit(-1);
   }
 
+  profiling::Profiler* profiler = new profiling::Profiler();
+  interpreter->SetProfiler(profiler);
+
+  if (s->profiling) profiler->StartProfiling();
+
   struct timeval start_time, stop_time;
   gettimeofday(&start_time, NULL);
   for (int i = 0; i < s->loop_count; i++) {
@@ -179,6 +202,18 @@ void RunInference(Settings* s) {
             << (get_us(stop_time) - get_us(start_time)) / (s->loop_count * 1000)
             << " ms \n";
 
+  if (s->profiling) {
+    profiler->StopProfiling();
+    auto profile_events = profiler->GetProfileEvents();
+    for (int i = 0; i < profile_events.size(); i++) {
+      auto op_index = profile_events[i]->event_metadata;
+      const auto node_and_registration =
+          interpreter->node_and_registration(op_index);
+      const TfLiteRegistration registration = node_and_registration->second;
+      PrintProfilingInfo(profile_events[i], op_index, registration);
+    }
+  }
+
   const int output_size = 1000;
   const size_t num_results = 5;
   const float threshold = 0.001f;
@@ -217,13 +252,14 @@ void RunInference(Settings* s) {
 
 void display_usage() {
   LOG(INFO) << "label_image\n"
-            << "--accelerated, -a: [0|1], use Android NNAPI or note\n"
+            << "--accelerated, -a: [0|1], use Android NNAPI or not\n"
             << "--count, -c: loop interpreter->Invoke() for certain times\n"
             << "--input_mean, -b: input mean\n"
             << "--input_std, -s: input standard deviation\n"
             << "--image, -i: image_name.bmp\n"
             << "--labels, -l: labels for the model\n"
             << "--tflite_model, -m: model_name.tflite\n"
+            << "--profiling, -p: [0|1], profiling or not\n"
             << "--threads, -t: number of threads\n"
             << "--verbose, -v: [0|1] print more information\n"
             << "\n";
@@ -241,6 +277,7 @@ int Main(int argc, char** argv) {
         {"image", required_argument, 0, 'i'},
         {"labels", required_argument, 0, 'l'},
         {"tflite_model", required_argument, 0, 'm'},
+        {"profiling", required_argument, 0, 'p'},
         {"threads", required_argument, 0, 't'},
         {"input_mean", required_argument, 0, 'b'},
         {"input_std", required_argument, 0, 's'},
@@ -249,7 +286,7 @@ int Main(int argc, char** argv) {
     /* getopt_long stores the option index here. */
     int option_index = 0;
 
-    c = getopt_long(argc, argv, "a:b:c:f:i:l:m:s:t:v:", long_options,
+    c = getopt_long(argc, argv, "a:b:c:f:i:l:m:p:s:t:v:", long_options,
                     &option_index);
 
     /* Detect the end of the options. */
@@ -276,6 +313,10 @@ int Main(int argc, char** argv) {
       case 'm':
         s.model_name = optarg;
         break;
+      case 'p':
+        s.profiling = strtol(  // NOLINT(runtime/deprecated_fn)
+            optarg, (char**)NULL, 10);
+        break;
       case 's':
         s.input_std = strtod(optarg, NULL);
         break;
diff --git a/tensorflow/contrib/lite/examples/label_image/label_image.h b/tensorflow/contrib/lite/examples/label_image/label_image.h
index 4de32e33fb..4b48014e1c 100644
--- a/tensorflow/contrib/lite/examples/label_image/label_image.h
+++ b/tensorflow/contrib/lite/examples/label_image/label_image.h
@@ -25,6 +25,7 @@ struct Settings {
   bool verbose = false;
   bool accel = false;
   bool input_floating = false;
+  bool profiling = false;
   int loop_count = 1;
   float input_mean = 127.5f;
   float input_std = 127.5f;
diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml b/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml
index 2c4ce84473..d12435d5ab 100644
--- a/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml
+++ b/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml
@@ -84,4 +84,32 @@
             android:visibility="visible" />
     </RelativeLayout>
 
+    <RelativeLayout
+        android:id="@+id/control2"
+        android:layout_width="match_parent"
+        android:layout_height="135dp"
+        android:layout_alignParentLeft="true"
+        android:layout_alignParentStart="true"
+        android:layout_alignTop="@+id/control"
+        android:layout_marginLeft="300dp"
+        android:layout_marginStart="300dp"
+        android:background="@color/control_background">
+
+        <ToggleButton
+            android:id="@+id/button"
+            android:textOff="@string/tflite"
+            android:textOn="@string/nnapi"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:layout_alignParentLeft="true"
+            android:layout_alignParentStart="true" />
+
+        <NumberPicker
+            android:id="@+id/np"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:layout_below="@+id/button"
+            android:visibility="visible" />
+    </RelativeLayout>
+
 </RelativeLayout>
diff --git a/tensorflow/contrib/lite/kernels/topk_v2.cc b/tensorflow/contrib/lite/kernels/topk_v2.cc
index 807e84609f..ad9b744f1a 100644
--- a/tensorflow/contrib/lite/kernels/topk_v2.cc
+++ b/tensorflow/contrib/lite/kernels/topk_v2.cc
@@ -25,8 +25,8 @@ namespace builtin {
 namespace topk_v2 {
 constexpr int kInputTensor = 0;
 constexpr int kInputTopK = 1;
-constexpr int kOutputIndexes = 0;
-constexpr int kOutputValues = 1;
+constexpr int kOutputValues = 0;
+constexpr int kOutputIndexes = 1;
 
 namespace {
 TfLiteStatus ResizeOutput(TfLiteContext* context, TfLiteNode* node) {
diff --git a/tensorflow/contrib/lite/kernels/topk_v2_test.cc b/tensorflow/contrib/lite/kernels/topk_v2_test.cc
index 29f2a057cd..212f8acc76 100644
--- a/tensorflow/contrib/lite/kernels/topk_v2_test.cc
+++ b/tensorflow/contrib/lite/kernels/topk_v2_test.cc
@@ -31,8 +31,8 @@ class TopKV2OpModel : public SingleOpModel {
                 int top_k) {
     input_ = AddInput(input_type);
     top_k_ = AddInput(TensorType_INT32);
-    output_indexes_ = AddOutput(TensorType_INT32);
     output_values_ = AddOutput(input_type);
+    output_indexes_ = AddOutput(TensorType_INT32);
     SetBuiltinOp(BuiltinOperator_TOPK_V2, BuiltinOptions_TopKV2Options, 0);
     BuildInterpreter({input_shape, {1}});
     PopulateTensor<int32_t>(top_k_, {top_k});
diff --git a/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h b/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h
index ace4827d8c..4a648e4283 100644
--- a/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h
+++ b/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h
@@ -609,7 +609,7 @@ enum {
    * Long short-term memory unit (LSTM) recurrent network layer.
    *
    * The default non-peephole implementation is based on:
-   * http://deeplearning.cs.cmu.edu/pdfs/Hochreiter97_lstm.pdf
+   * http://www.bioinf.jku.at/publications/older/2604.pdf
    * S. Hochreiter and J. Schmidhuber. "Long Short-Term Memory". Neural
    * Computation, 9(8):1735-1780, 1997.
    *
diff --git a/tensorflow/contrib/lite/profiling/profile_buffer.h b/tensorflow/contrib/lite/profiling/profile_buffer.h
index b2f565376c..299b2a9cad 100644
--- a/tensorflow/contrib/lite/profiling/profile_buffer.h
+++ b/tensorflow/contrib/lite/profiling/profile_buffer.h
@@ -37,9 +37,9 @@ struct ProfileEvent {
   // Label of the event. This usually describes the event.
   const char* tag;
   // Timestamp in microseconds when the event began.
-  int64_t begin_timestamp_us;
+  uint64_t begin_timestamp_us;
   // Timestamp in microseconds when the event ended.
-  int64_t end_timestamp_us;
+  uint64_t end_timestamp_us;
   // The field containing the type of event. This must be one of the event types
   // in EventType.
   EventType event_type;
@@ -74,7 +74,7 @@ class ProfileBuffer {
     if (!enabled_) {
       return kInvalidEventHandle;
     }
-    int64_t timestamp = NowMicros();
+    uint64_t timestamp = NowMicros();
     int index = current_index_ % event_buffer_.size();
     event_buffer_[index].tag = tag;
     event_buffer_[index].event_type = event_type;
@@ -134,7 +134,7 @@ class ProfileBuffer {
   }
 
  private:
-  static int64_t NowMicros() {
+  static uint64_t NowMicros() {
     // TODO(shashishekhar): Refactor this to a separate file.
     struct timeval tv;
     gettimeofday(&tv, nullptr);
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc b/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc
index 89ad58f887..c1cf79f626 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc
@@ -124,6 +124,15 @@ bool PropagateArrayDataTypes::Run(Model* model, std::size_t op_index) {
       SetDataTypeForAllOutputs(model, op, rand_op->dtype);
       break;
     }
+    case OperatorType::kTopK_V2: {
+      // topk(values: T, k: int32) -> values: T, indices: int32
+      CHECK_EQ(op->inputs.size(), 2);
+      CHECK_EQ(op->outputs.size(), 2);
+      CHECK(model->GetArray(op->inputs[1]).data_type == ArrayDataType::kInt32);
+      model->GetArray(op->outputs[0]).data_type = model->GetArray(op->inputs[0]).data_type;
+      model->GetArray(op->outputs[1]).data_type = ArrayDataType ::kInt32;
+      break;
+    }
     case OperatorType::kTensorFlowUnsupported: {
       auto* unsupported_op = static_cast<TensorFlowUnsupportedOperator*>(op);
       // Some output tensors from the op could be eliminated by optimization.
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
index 19037bc503..4923f83d91 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
@@ -1087,8 +1087,8 @@ void ProcessGatherOperator(Model* model, GatherOperator* op) {
 void ProcessTopkV2Operator(Model* model, TopKV2Operator* op) {
   const auto& input_values = model->GetArray(op->inputs[0]);
   const auto& input_k = model->GetArray(op->inputs[1]);
-  auto& output_indexes = model->GetArray(op->outputs[0]);
-  auto& output_values = model->GetArray(op->outputs[1]);
+  auto& output_values = model->GetArray(op->outputs[0]);
+  auto& output_indexes = model->GetArray(op->outputs[1]);
 
   // Bail if we already know the output shape.
   if (output_indexes.has_shape()) {
diff --git a/tensorflow/contrib/lite/toco/import_tensorflow.cc b/tensorflow/contrib/lite/toco/import_tensorflow.cc
index 61e4c9d542..fa8b26bce0 100644
--- a/tensorflow/contrib/lite/toco/import_tensorflow.cc
+++ b/tensorflow/contrib/lite/toco/import_tensorflow.cc
@@ -1991,7 +1991,7 @@ void ConvertTopKV2Operator(const NodeDef& node,
     op->inputs.push_back(node.input(1));
   }
   // The op has two outputs.
-  op->outputs.push_back(node.name() + ":0");
+  op->outputs.push_back(node.name());
   op->outputs.push_back(node.name() + ":1");
   model->operators.emplace_back(op.release());
 }
diff --git a/tensorflow/contrib/lite/toco/tooling_util.cc b/tensorflow/contrib/lite/toco/tooling_util.cc
index 5a341294db..f334c51bbb 100644
--- a/tensorflow/contrib/lite/toco/tooling_util.cc
+++ b/tensorflow/contrib/lite/toco/tooling_util.cc
@@ -825,11 +825,6 @@ void FixNoOrphanedArray(Model* model) {
 void CheckEachArray(const Model& model) {
   for (const auto& array_entry : model.GetArrayMap()) {
     const auto& array = array_entry.second;
-    if (array->has_shape()) {
-      for (int d : array->shape().dims()) {
-        CHECK_GE(d, 1);
-      }
-    }
     // It's OK to have a buffer or an alloc, but not both.
     // (Since allocs are for transient arrays without a buffer).
     CHECK(!array->buffer || !array->alloc);
@@ -839,6 +834,10 @@ void CheckEachArray(const Model& model) {
       // The presence of a fixed buffer should imply the presence of a fixed
       // shape.
       CHECK(array->has_shape());
+      // Constant buffer should has a valid shape.
+      for (int d : array->shape().dims()) {
+        CHECK_GE(d, 1);
+      }
       // The shape flat-size should agree with the buffer length.
       CHECK_EQ(array->buffer->Length(),
                RequiredBufferSizeForShape(array->shape()));
diff --git a/tensorflow/contrib/mpi/mpi_utils.h b/tensorflow/contrib/mpi/mpi_utils.h
index df055ff567..4091925fc0 100644
--- a/tensorflow/contrib/mpi/mpi_utils.h
+++ b/tensorflow/contrib/mpi/mpi_utils.h
@@ -22,6 +22,7 @@ limitations under the License.
 #include <string>
 #include <vector>
 
+#include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 
 // Skip MPI C++ bindings support, this matches the usage in other places
diff --git a/tensorflow/contrib/opt/python/training/lazy_adam_optimizer.py b/tensorflow/contrib/opt/python/training/lazy_adam_optimizer.py
index aeca900bc8..72117c1e81 100644
--- a/tensorflow/contrib/opt/python/training/lazy_adam_optimizer.py
+++ b/tensorflow/contrib/opt/python/training/lazy_adam_optimizer.py
@@ -56,21 +56,21 @@ class LazyAdamOptimizer(adam.AdamOptimizer):
     epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype)
     lr = (lr_t * math_ops.sqrt(1 - beta2_power) / (1 - beta1_power))
 
-    # m := beta1 * m + (1 - beta1) * g_t
+    # \\(m := beta1 * m + (1 - beta1) * g_t\\)
     m = self.get_slot(var, "m")
     m_t = state_ops.scatter_update(m, grad.indices,
                                    beta1_t * array_ops.gather(m, grad.indices) +
                                    (1 - beta1_t) * grad.values,
                                    use_locking=self._use_locking)
 
-    # v := beta2 * v + (1 - beta2) * (g_t * g_t)
+    # \\(v := beta2 * v + (1 - beta2) * (g_t * g_t)\\)
     v = self.get_slot(var, "v")
     v_t = state_ops.scatter_update(v, grad.indices,
                                    beta2_t * array_ops.gather(v, grad.indices) +
                                    (1 - beta2_t) * math_ops.square(grad.values),
                                    use_locking=self._use_locking)
 
-    # variable -= learning_rate * m_t / (epsilon_t + sqrt(v_t))
+    # \\(variable -= learning_rate * m_t / (epsilon_t + sqrt(v_t))\\)
     m_t_slice = array_ops.gather(m_t, grad.indices)
     v_t_slice = array_ops.gather(v_t, grad.indices)
     denominator_slice = math_ops.sqrt(v_t_slice) + epsilon_t
diff --git a/tensorflow/contrib/optimizer_v2/adam.py b/tensorflow/contrib/optimizer_v2/adam.py
index 42b7f92a76..d538ad0fb0 100644
--- a/tensorflow/contrib/optimizer_v2/adam.py
+++ b/tensorflow/contrib/optimizer_v2/adam.py
@@ -40,23 +40,19 @@ class AdamOptimizer(optimizer_v2.OptimizerV2):
 
     Initialization:
 
-    ```
-    m_0 <- 0 (Initialize initial 1st moment vector)
-    v_0 <- 0 (Initialize initial 2nd moment vector)
-    t <- 0 (Initialize timestep)
-    ```
+    $$m_0 := 0 (Initialize initial 1st moment vector)$$
+    $$v_0 := 0 (Initialize initial 2nd moment vector)$$
+    $$t := 0 (Initialize timestep)$$
 
     The update rule for `variable` with gradient `g` uses an optimization
     described at the end of section2 of the paper:
 
-    ```
-    t <- t + 1
-    lr_t <- learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
+    $$t := t + 1$$
+    $$lr_t := \text{learning_rate} * \sqrt{(1 - beta_2^t) / (1 - beta_1^t)}$$
 
-    m_t <- beta1 * m_{t-1} + (1 - beta1) * g
-    v_t <- beta2 * v_{t-1} + (1 - beta2) * g * g
-    variable <- variable - lr_t * m_t / (sqrt(v_t) + epsilon)
-    ```
+    $$m_t := beta_1 * m_{t-1} + (1 - beta_1) * g$$
+    $$v_t := beta_2 * v_{t-1} + (1 - beta_2) * g * g$$
+    $$variable := variable - lr_t * m_t / (\sqrt{v_t} + \epsilon)$$
 
     The default value of 1e-8 for epsilon might not be a good default in
     general. For example, when training an Inception network on ImageNet a
diff --git a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py
index de5df91292..ba4933ddf7 100644
--- a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py
+++ b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py
@@ -307,6 +307,21 @@ class LSTMTest(test.TestCase):
     self._seed = 23489
     np.random.seed(self._seed)
 
+  def testDType(self):
+    # Test case for GitHub issue 16228
+    # Not passing dtype in constructor results in default float32
+    lstm = rnn_cell.LSTMCell(10)
+    input_tensor = array_ops.ones([10, 50])
+    lstm.build(input_tensor.get_shape())
+    self.assertEqual(lstm._bias.dtype, dtypes.float32_ref)
+
+    # Explicitly pass dtype in constructor
+    for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
+      lstm = rnn_cell.LSTMCell(10, dtype=dtype)
+      input_tensor = array_ops.ones([10, 50])
+      lstm.build(input_tensor.get_shape())
+      self.assertEqual(lstm._bias.dtype, dtype._as_ref)
+
   def testNoProjNoSharding(self):
     num_units = 3
     input_size = 5
diff --git a/tensorflow/contrib/tensor_forest/client/eval_metrics.py b/tensorflow/contrib/tensor_forest/client/eval_metrics.py
index 90033015eb..e893e1d1c8 100644
--- a/tensorflow/contrib/tensor_forest/client/eval_metrics.py
+++ b/tensorflow/contrib/tensor_forest/client/eval_metrics.py
@@ -37,7 +37,7 @@ def _top_k_generator(k):
   def _top_k(probabilities, targets):
     targets = math_ops.to_int32(targets)
     if targets.get_shape().ndims > 1:
-      targets = array_ops.squeeze(targets, squeeze_dims=[1])
+      targets = array_ops.squeeze(targets, axis=[1])
     return metric_ops.streaming_mean(nn.in_top_k(probabilities, targets, k))
   return _top_k
 
@@ -57,7 +57,7 @@ def _r2(probabilities, targets, weights=None):
 
 
 def _squeeze_and_onehot(targets, depth):
-  targets = array_ops.squeeze(targets, squeeze_dims=[1])
+  targets = array_ops.squeeze(targets, axis=[1])
   return array_ops.one_hot(math_ops.to_int32(targets), depth)
 
 
diff --git a/tensorflow/contrib/tensor_forest/hybrid/python/layers/fully_connected.py b/tensorflow/contrib/tensor_forest/hybrid/python/layers/fully_connected.py
index ff3ab21eaa..745a5b1caf 100644
--- a/tensorflow/contrib/tensor_forest/hybrid/python/layers/fully_connected.py
+++ b/tensorflow/contrib/tensor_forest/hybrid/python/layers/fully_connected.py
@@ -55,7 +55,7 @@ class ManyToOneLayer(hybrid_layer.HybridLayer):
 
       # There is always one activation per instance by definition, so squeeze
       # away the extra dimension.
-      return array_ops.squeeze(nn_activations, squeeze_dims=[1])
+      return array_ops.squeeze(nn_activations, axis=[1])
 
 
 class FlattenedFullyConnectedLayer(hybrid_layer.HybridLayer):
diff --git a/tensorflow/contrib/tensor_forest/python/tensor_forest.py b/tensorflow/contrib/tensor_forest/python/tensor_forest.py
index b9bcbb170b..7a35a70bbe 100644
--- a/tensorflow/contrib/tensor_forest/python/tensor_forest.py
+++ b/tensorflow/contrib/tensor_forest/python/tensor_forest.py
@@ -445,7 +445,7 @@ class RandomForestGraphs(object):
           mask = math_ops.less(
               r, array_ops.ones_like(r) * self.params.bagging_fraction)
           gather_indices = array_ops.squeeze(
-              array_ops.where(mask), squeeze_dims=[1])
+              array_ops.where(mask), axis=[1])
           # TODO(thomaswc): Calculate out-of-bag data and labels, and store
           # them for use in calculating statistics later.
           tree_data = array_ops.gather(processed_dense_features, gather_indices)
diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.cc b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
index b412b296e0..0774027711 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_graph.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
@@ -111,20 +111,22 @@ void GetSubGraphOutgoingEdges(const tensorflow::Graph& graph,
   }
 }
 
-std::pair<string, int> ParseTensorName(string name, int default_idx = 0) {
+std::pair<string, int> ParseTensorName(const string& name,
+                                       int default_idx = 0) {
+  string name_no_idx = name;
   int idx = default_idx;
-  size_t sep = name.find_last_of(':');
+  const size_t sep = name_no_idx.find_last_of(':');
   if (sep != string::npos) {
-    name = name.substr(0, sep);
+    name_no_idx = name_no_idx.substr(0, sep);
     idx = std::stoi(name.substr(sep + 1));
   }
-  return std::make_pair(name, idx);
+  return std::make_pair(name_no_idx, idx);
 }
 
 std::unordered_map<string, std::vector<int>> BuildTensorNameMap(
     const std::vector<string>& tensor_names) {
   std::unordered_map<string, std::vector<int>> result;
-  for (string const& tensor_name : tensor_names) {
+  for (const string& tensor_name : tensor_names) {
     string node_name;
     int index;
     std::tie(node_name, index) = ParseTensorName(tensor_name);
@@ -132,6 +134,7 @@ std::unordered_map<string, std::vector<int>> BuildTensorNameMap(
   }
   return result;
 }
+
 // TODO(sami): convert references to pointers
 struct ConvertGraphParams {
   ConvertGraphParams(
diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_management_test.py b/tensorflow/contrib/timeseries/python/timeseries/state_management_test.py
index d5dce30fda..5f7e3da2db 100644
--- a/tensorflow/contrib/timeseries/python/timeseries/state_management_test.py
+++ b/tensorflow/contrib/timeseries/python/timeseries/state_management_test.py
@@ -78,7 +78,7 @@ class StubTimeSeriesModel(model.TimeSeriesModel):
     batch_end_values = array_ops.squeeze(
         array_ops.slice(values, [0, array_ops.shape(times)[1] - 1, 0],
                         [-1, 1, -1]),
-        squeeze_dims=[1, 2])
+        axis=[1, 2])
     # A pretty odd but easy to think about loss: L1 loss on the batch end
     # values.
     loss = math_ops.reduce_sum(
diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/kalman_filter.py b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/kalman_filter.py
index 1fcd3e391b..a614386121 100644
--- a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/kalman_filter.py
+++ b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/kalman_filter.py
@@ -170,7 +170,7 @@ class KalmanFilter(object):
         math_ops.matmul(
             transition_matrices,
             prior_state[..., None]),
-        squeeze_dims=[-1])
+        axis=[-1])
     return advanced_state
 
   def predict_state_var(
@@ -254,7 +254,7 @@ class KalmanFilter(object):
             kalman_gain_transposed,
             array_ops.expand_dims(residual, -1),
             adjoint_a=True),
-        squeeze_dims=[-1])
+        axis=[-1])
     gain_obs = math_ops.matmul(
         kalman_gain_transposed, observation_model, adjoint_a=True)
     identity_extradim = linalg_ops.eye(
@@ -332,7 +332,7 @@ class KalmanFilter(object):
             array_ops.expand_dims(state_mean, 1),
             observation_model,
             adjoint_b=True),
-        squeeze_dims=[1])
+        axis=[1])
     observed_var = math_ops.matmul(
         math_ops.matmul(observation_model, state_var),
         observation_model,
diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index 2a849a3019..76ff372cd0 100644
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -2292,7 +2292,9 @@ tf_cuda_library(
 
 CORE_CPU_BASE_HDRS = GRAPH_HDRS + [
     "common_runtime/device.h",
+    "common_runtime/device_factory.h",
     "common_runtime/device_mgr.h",
+    "common_runtime/device_set.h",
     "common_runtime/eval_const_tensor.h",
     "common_runtime/graph_runner.h",
     "common_runtime/shape_refiner.h",
@@ -2350,9 +2352,7 @@ CORE_CPU_LIB_HEADERS = CORE_CPU_BASE_HDRS + [
     "common_runtime/copy_tensor.h",
     "common_runtime/costmodel_manager.h",
     "common_runtime/debugger_state_interface.h",
-    "common_runtime/device_factory.h",
     "common_runtime/device_resolver_local.h",
-    "common_runtime/device_set.h",
     "common_runtime/dma_helper.h",
     "common_runtime/eigen_thread_pool.h",
     "common_runtime/executor.h",
diff --git a/tensorflow/core/api_def/base_api/api_def_ApplyAdam.pbtxt b/tensorflow/core/api_def/base_api/api_def_ApplyAdam.pbtxt
index c2858a1bfb..b90f5473c8 100644
--- a/tensorflow/core/api_def/base_api/api_def_ApplyAdam.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_ApplyAdam.pbtxt
@@ -82,9 +82,9 @@ END
   }
   summary: "Update \'*var\' according to the Adam algorithm."
   description: <<END
-lr_t <- learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
-m_t <- beta1 * m_{t-1} + (1 - beta1) * g_t
-v_t <- beta2 * v_{t-1} + (1 - beta2) * g_t * g_t
-variable <- variable - lr_t * m_t / (sqrt(v_t) + epsilon)
+$$lr_t := \text{learning_rate} * \sqrt{(1 - beta_2^t) / (1 - beta_1^t)}$$
+$$m_t := beta_1 * m_{t-1} + (1 - beta_1) * g$$
+$$v_t := beta_2 * v_{t-1} + (1 - beta_2) * g * g$$
+$$variable := variable - lr_t * m_t / (\sqrt{v_t} + \epsilon)$$
 END
 }
diff --git a/tensorflow/core/api_def/base_api/api_def_Pad.pbtxt b/tensorflow/core/api_def/base_api/api_def_Pad.pbtxt
index e45e2375eb..ee4aad7899 100644
--- a/tensorflow/core/api_def/base_api/api_def_Pad.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_Pad.pbtxt
@@ -24,5 +24,6 @@ pad(t, paddings) ==> [[0, 0, 0, 0, 0, 0]
                       [0, 0, 2, 2, 0, 0]
                       [0, 0, 0, 0, 0, 0]]
 ```
+
 END
 }
diff --git a/tensorflow/core/api_def/base_api/api_def_QuantizeV2.pbtxt b/tensorflow/core/api_def/base_api/api_def_QuantizeV2.pbtxt
index b9e75caf02..37ac10dddb 100644
--- a/tensorflow/core/api_def/base_api/api_def_QuantizeV2.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_QuantizeV2.pbtxt
@@ -44,6 +44,7 @@ In 'MIN_COMBINED' mode, each value of the tensor will undergo the following:
 out[i] = (in[i] - min_range) * range(T) / (max_range - min_range)
 if T == qint8, out[i] -= (range(T) + 1) / 2.0
 ```
+
 here `range(T) = numeric_limits<T>::max() - numeric_limits<T>::min()`
 
 *MIN_COMBINED Mode Example*
@@ -87,6 +88,7 @@ choosing to elide the lowest possible value for symmetry (e.g., output range is
 
 We first find the range of values in our tensor. The
 range we use is always centered on 0, so we find m such that
+
 ```c++
   m = max(abs(input_min), abs(input_max))
 ```
@@ -95,6 +97,7 @@ Our input tensor range is then `[-m, m]`.
 
 Next, we choose our fixed-point quantization buckets, `[min_fixed, max_fixed]`.
 If T is signed, this is
+
 ```
   num_bits = sizeof(T) * 8
   [min_fixed, max_fixed] =
@@ -102,16 +105,19 @@ If T is signed, this is
 ```
 
 Otherwise, if T is unsigned, the fixed-point range is
+
 ```
   [min_fixed, max_fixed] = [0, (1 << num_bits) - 1]
 ```
 
 From this we compute our scaling factor, s:
+
 ```c++
   s = (max_fixed - min_fixed) / (2 * m)
 ```
 
 Now we can quantize the elements of our tensor:
+
 ```c++
 result = round(input * s)
 ```
diff --git a/tensorflow/core/api_def/base_api/api_def_ResourceApplyAdam.pbtxt b/tensorflow/core/api_def/base_api/api_def_ResourceApplyAdam.pbtxt
index bea1fd6762..ad0aeac004 100644
--- a/tensorflow/core/api_def/base_api/api_def_ResourceApplyAdam.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_ResourceApplyAdam.pbtxt
@@ -76,9 +76,9 @@ END
   }
   summary: "Update \'*var\' according to the Adam algorithm."
   description: <<END
-lr_t <- learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
-m_t <- beta1 * m_{t-1} + (1 - beta1) * g_t
-v_t <- beta2 * v_{t-1} + (1 - beta2) * g_t * g_t
-variable <- variable - lr_t * m_t / (sqrt(v_t) + epsilon)
+$$lr_t := \text{learning_rate} * \sqrt{(1 - beta_2^t) / (1 - beta_1^t)}$$
+$$m_t := beta_1 * m_{t-1} + (1 - beta_1) * g$$
+$$v_t := beta_2 * v_{t-1} + (1 - beta_2) * g * g$$
+$$variable := variable - lr_t * m_t / (\sqrt{v_t} + \epsilon)$$
 END
 }
diff --git a/tensorflow/core/api_def/base_api/api_def_ScatterNd.pbtxt b/tensorflow/core/api_def/base_api/api_def_ScatterNd.pbtxt
index 4cb8c064fc..58753a651a 100644
--- a/tensorflow/core/api_def/base_api/api_def_ScatterNd.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_ScatterNd.pbtxt
@@ -25,12 +25,12 @@ A new tensor with the given shape and updates applied according
 to the indices.
 END
   }
-  summary: "Scatter `updates` into a new (initially zero) tensor according to `indices`."
+  summary: "Scatter `updates` into a new tensor according to `indices`."
   description: <<END
-Creates a new tensor by applying sparse `updates` to individual
-values or slices within a zero tensor of the given `shape` according to
-indices.  This operator is the inverse of the @{tf.gather_nd} operator which
-extracts values or slices from a given tensor.
+Creates a new tensor by applying sparse `updates` to individual values or
+slices within a tensor (initially zero for numeric, empty for string) of
+the given `shape` according to indices.  This operator is the inverse of the
+@{tf.gather_nd} operator which extracts values or slices from a given tensor.
 
 **WARNING**: The order in which updates are applied is nondeterministic, so the
 output will be nondeterministic if `indices` contains duplicates.
diff --git a/tensorflow/core/common_runtime/graph_execution_state.cc b/tensorflow/core/common_runtime/graph_execution_state.cc
index 49b1df38dc..eb710bdbc5 100644
--- a/tensorflow/core/common_runtime/graph_execution_state.cc
+++ b/tensorflow/core/common_runtime/graph_execution_state.cc
@@ -490,7 +490,7 @@ Status GraphExecutionState::OptimizeGraph(
         cpu_device = device;
       }
     }
-    grappler::VirtualCluster cluster(device_map);
+    grappler::VirtualCluster cluster(device_map, device_set_);
     GraphDef new_graph;
     TF_RETURN_IF_ERROR(grappler::RunMetaOptimizer(
         item, rewrite_options, cpu_device, &cluster, &new_graph));
diff --git a/tensorflow/core/graph/mkl_layout_pass.cc b/tensorflow/core/graph/mkl_layout_pass.cc
index 5368774f2d..72a13d4da7 100644
--- a/tensorflow/core/graph/mkl_layout_pass.cc
+++ b/tensorflow/core/graph/mkl_layout_pass.cc
@@ -547,14 +547,14 @@ class MklLayoutRewritePass : public GraphOptimizationPass {
 
     // If Op has been specifically assigned to a non-CPU device, then No.
     if (!n->assigned_device_name().empty() &&
-        !StringPiece(n->assigned_device_name()).contains(kCPUDeviceSubStr)) {
+        !str_util::StrContains(n->assigned_device_name(),kCPUDeviceSubStr)) {
       result = false;
       reason = "Op has been assigned a runtime device that is not CPU.";
     }
 
     // If user has specifically assigned this op to a non-CPU device, then No.
     if (!n->def().device().empty() &&
-        !StringPiece(n->def().device()).contains(kCPUDeviceSubStr)) {
+        !str_util::StrContains(n->def().device(),kCPUDeviceSubStr)) {
       result = false;
       reason = "User has assigned a device that is not CPU.";
     }
@@ -2691,14 +2691,14 @@ class MklLayoutRewritePass : public GraphOptimizationPass {
 
     // If Op has been specifically assigned to a non-CPU device, then No.
     if (!n->assigned_device_name().empty() &&
-        !StringPiece(n->assigned_device_name()).contains(kCPUDeviceSubStr)) {
+        !str_util::StrContains(n->assigned_device_name(),kCPUDeviceSubStr)) {
       result = false;
       reason = "Op has been assigned a runtime device that is not CPU.";
     }
 
     // If user has specifically assigned this op to a non-CPU device, then No.
     if (!n->def().device().empty() &&
-        !StringPiece(n->def().device()).contains(kCPUDeviceSubStr)) {
+        !str_util::StrContains(n->def().device(),kCPUDeviceSubStr)) {
       result = false;
       reason = "User has assigned a device that is not CPU.";
     }
diff --git a/tensorflow/core/grappler/clusters/BUILD b/tensorflow/core/grappler/clusters/BUILD
index 9ecf5a6cf7..30c6126fbb 100644
--- a/tensorflow/core/grappler/clusters/BUILD
+++ b/tensorflow/core/grappler/clusters/BUILD
@@ -56,6 +56,7 @@ cc_library(
     ],
     visibility = ["//visibility:public"],
     deps = [
+        "//tensorflow/core:core_cpu_base",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
@@ -73,6 +74,7 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         ":cluster",
+        "//tensorflow/core:core_cpu_base",
         "//tensorflow/core:framework",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/grappler/costs:op_level_cost_estimator",
diff --git a/tensorflow/core/grappler/clusters/cluster.h b/tensorflow/core/grappler/clusters/cluster.h
index 0796ba65ec..d33aaa7e4c 100644
--- a/tensorflow/core/grappler/clusters/cluster.h
+++ b/tensorflow/core/grappler/clusters/cluster.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "tensorflow/core/common_runtime/device_set.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/grappler/grappler_item.h"
 #include "tensorflow/core/lib/core/status.h"
@@ -92,6 +93,10 @@ class Cluster {
   // sorted alphabetically.
   const std::vector<string> GetDeviceNames() const;
 
+  // The DeviceSet is not always available, but when it is it contains a
+  // superset of the devices listed in GetDevices/GetDeviceNames().
+  const DeviceSet* GetDeviceSet() const { return device_set_; }
+
   // Enables collecting the allocator stats. Call with enable=true must be made
   // before Provision().
   virtual Status EnablePeakMemoryStats(bool enable) {
@@ -119,6 +124,7 @@ class Cluster {
 
  protected:
   std::unordered_map<string, DeviceProperties> devices_;
+  const DeviceSet* device_set_ = nullptr;  // Not owned
   const int timeout_s_;
   SessionOptions options_;
   RunOptions run_options_;
diff --git a/tensorflow/core/grappler/clusters/virtual_cluster.cc b/tensorflow/core/grappler/clusters/virtual_cluster.cc
index abfa7bc48e..5c9b2320b5 100644
--- a/tensorflow/core/grappler/clusters/virtual_cluster.cc
+++ b/tensorflow/core/grappler/clusters/virtual_cluster.cc
@@ -37,6 +37,14 @@ VirtualCluster::VirtualCluster(
     : Cluster(0), node_estimator_(node_estimator), node_manager_(node_manager) {
   devices_ = devices;
 }
+
+VirtualCluster::VirtualCluster(
+    const std::unordered_map<string, DeviceProperties>& devices,
+    const DeviceSet* device_set)
+    : VirtualCluster(devices) {
+  device_set_ = device_set;
+}
+
 VirtualCluster::~VirtualCluster() {}
 
 Status VirtualCluster::Provision() { return Status::OK(); }
diff --git a/tensorflow/core/grappler/clusters/virtual_cluster.h b/tensorflow/core/grappler/clusters/virtual_cluster.h
index e5967bac3d..eebac68e1b 100644
--- a/tensorflow/core/grappler/clusters/virtual_cluster.h
+++ b/tensorflow/core/grappler/clusters/virtual_cluster.h
@@ -17,6 +17,8 @@ limitations under the License.
 #define TENSORFLOW_CORE_GRAPPLER_CLUSTERS_VIRTUAL_CLUSTER_H_
 
 #include <unordered_map>
+
+#include "tensorflow/core/common_runtime/device_set.h"
 #include "tensorflow/core/grappler/clusters/cluster.h"
 #include "tensorflow/core/grappler/costs/op_level_cost_estimator.h"
 #include "tensorflow/core/grappler/costs/virtual_scheduler.h"
@@ -34,6 +36,8 @@ class VirtualCluster : public Cluster {
   VirtualCluster(const std::unordered_map<string, DeviceProperties>& devices,
                  OpLevelCostEstimator* node_estimator,
                  ReadyNodeManager* node_manager);
+  VirtualCluster(const std::unordered_map<string, DeviceProperties>& devices,
+                 const DeviceSet* device_set);
 
   ~VirtualCluster() override;
 
diff --git a/tensorflow/core/grappler/costs/virtual_scheduler.h b/tensorflow/core/grappler/costs/virtual_scheduler.h
index 5116c8183c..67bf1e6980 100644
--- a/tensorflow/core/grappler/costs/virtual_scheduler.h
+++ b/tensorflow/core/grappler/costs/virtual_scheduler.h
@@ -199,7 +199,7 @@ class FirstReadyManager : public ReadyNodeManager {
   // current node.
   std::vector<const NodeDef*> nodes_;
   // Newly added nodes are added to waiting_queue_. That way, GetCurrNode(),
-  // wihch returns the front of the nodes_, always returns the same node,
+  // which returns the front of the nodes_, always returns the same node,
   // even if any of new nodes has time_ready smaller than the current node's.
   std::vector<const NodeDef*> waiting_queue_;
   // Comparator functor for heap; stl heap is max heap, so we use "greater than"
@@ -212,7 +212,7 @@ class FirstReadyManager : public ReadyNodeManager {
 };
 
 // CompositeNodeManager has a few other NodeManagers: per-device LIFO for normal
-// ops (neither _Send nor _Recv) and FirstyReadyManagers for _Send ops and _Recv
+// ops (neither _Send nor _Recv) and FirstReadyManagers for _Send ops and _Recv
 // ops, and then it chooses FirstReady among the ops chosen from each
 // internal NodeManagers. The objective is to maximize producer-consumer
 // locality within device, while processing nodes across devices, including
diff --git a/tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h b/tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h
index 796da91373..3148a5f809 100644
--- a/tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h
+++ b/tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h
@@ -33,7 +33,7 @@ class CustomGraphOptimizerRegistry {
   static std::vector<string> GetRegisteredOptimizers();
 
   typedef std::function<CustomGraphOptimizer*()> Creator;
-  // Regsiter graph optimizer which can be called during program initialization.
+  // Register graph optimizer which can be called during program initialization.
   // This class is not thread-safe.
   static void RegisterOptimizerOrDie(const Creator& optimizer_creator,
                                      const string& name);
diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer.cc b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
index 2edc4da9dc..5230177dca 100644
--- a/tensorflow/core/grappler/optimizers/meta_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
@@ -160,13 +160,26 @@ Status MetaOptimizer::InitializeOptimizersByName(
       VLOG(2) << "Can't register an optimizer by name: " << optimizer_name;
     }
   }
+  for (const auto& optimizer_config : cfg_.custom_optimizers()) {
+    auto custom_optimizer = CustomGraphOptimizerRegistry::CreateByNameOrNull(
+        optimizer_config.name());
+    if (custom_optimizer) {
+      VLOG(2) << "Registered custom configurable graph optimizer: "
+              << optimizer_config.name();
+      TF_RETURN_IF_ERROR(custom_optimizer->Init(&optimizer_config));
+      optimizers->push_back(std::move(custom_optimizer));
+    } else {
+      VLOG(2) << "Can't register an optimizer by name: "
+              << optimizer_config.name();
+    }
+  }
   return Status::OK();
 }
 
 Status MetaOptimizer::OptimizeGraph(Cluster* cluster, const GrapplerItem& item,
                                     GraphDef* optimized_graph) {
   std::vector<std::unique_ptr<GraphOptimizer>> optimizers;
-  if (cfg_.optimizers().empty()) {
+  if (cfg_.optimizers().empty() && cfg_.custom_optimizers().empty()) {
     TF_RETURN_IF_ERROR(InitializeOptimizers(&optimizers));
   } else {
     TF_RETURN_IF_ERROR(InitializeOptimizersByName(&optimizers));
@@ -337,7 +350,7 @@ bool MetaOptimizerEnabled(const RewriterConfig& cfg) {
          cfg.auto_parallel().enable() ||
          cfg.memory_optimization() != RewriterConfig::NO_MEM_OPT ||
          cfg.debug_stripper() == RewriterConfig::ON ||
-         !cfg.optimizers().empty();
+         !cfg.optimizers().empty() || !cfg.custom_optimizers().empty();
 }
 
 Status RunMetaOptimizer(const GrapplerItem& item, const RewriterConfig& cfg,
diff --git a/tensorflow/core/kernels/batch_util.cc b/tensorflow/core/kernels/batch_util.cc
index 52be1ab8d0..1182ed42e7 100644
--- a/tensorflow/core/kernels/batch_util.cc
+++ b/tensorflow/core/kernels/batch_util.cc
@@ -134,6 +134,8 @@ Status CopyElementToSlice(Tensor element, Tensor* parent, int64 index) {
   switch (element.dtype()) {
     TF_CALL_ALL_TYPES(HANDLE_TYPE);
     TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
+    TF_CALL_uint32(HANDLE_TYPE);
+    TF_CALL_uint64(HANDLE_TYPE);
 #undef HANDLE_TYPE
     default:
       return errors::Unimplemented("CopyElementToSlice Unhandled data type: ",
diff --git a/tensorflow/core/kernels/cwise_op_floor_div.cc b/tensorflow/core/kernels/cwise_op_floor_div.cc
index fecbf85989..24da61fdf6 100644
--- a/tensorflow/core/kernels/cwise_op_floor_div.cc
+++ b/tensorflow/core/kernels/cwise_op_floor_div.cc
@@ -16,8 +16,8 @@ limitations under the License.
 #include "tensorflow/core/kernels/cwise_ops_common.h"
 
 namespace tensorflow {
-REGISTER5(BinaryOp, CPU, "FloorDiv", functor::safe_floor_div, uint8, uint16,
-          int16, int32, int64);
+REGISTER6(BinaryOp, CPU, "FloorDiv", functor::safe_floor_div, uint8, uint16,
+          int8, int16, int32, int64);
 REGISTER3(BinaryOp, CPU, "FloorDiv", functor::floor_div_real, float,
           Eigen::half, double);
 
diff --git a/tensorflow/core/kernels/mkl_conv_ops.cc b/tensorflow/core/kernels/mkl_conv_ops.cc
index f0818eb96d..f2b14f1278 100644
--- a/tensorflow/core/kernels/mkl_conv_ops.cc
+++ b/tensorflow/core/kernels/mkl_conv_ops.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <map>
 #include <string>
 #include <vector>
+#include <memory>
 
 #include "tensorflow/core/framework/numeric_op.h"
 #include "tensorflow/core/framework/op_kernel.h"
@@ -42,14 +43,13 @@ limitations under the License.
 #include "tensorflow/core/util/mkl_util.h"
 
 #ifndef INTEL_MKL_ML
-
 #include "mkldnn.hpp"
 
 using mkldnn::prop_kind;
 using mkldnn::stream;
-
-using mkldnn::convolution_direct;
 using mkldnn::convolution_forward;
+using mkldnn::convolution_direct;
+
 #else
 #include "mkl_dnn.h"
 #include "mkl_dnn_types.h"
@@ -57,11 +57,232 @@ using mkldnn::convolution_forward;
 
 namespace tensorflow {
 
+#ifndef INTEL_MKL_ML
+
+struct ConvFwdDimensions {
+  memory::dims src_dims;
+  memory::dims filter_dims;
+  memory::dims bias_dims;
+  memory::dims dst_dims;
+  memory::dims strides;
+  memory::dims dilations;
+  memory::dims padding_left;
+  memory::dims padding_right;
+
+  ConvFwdDimensions(memory::dims src_dims,
+    memory::dims filter_dims, memory::dims bias_dims,
+    memory::dims dst_dims, memory::dims strides,
+    memory::dims dilations, memory::dims padding_left,
+    memory::dims padding_right) :
+      src_dims(src_dims), filter_dims(filter_dims),
+      bias_dims(bias_dims), dst_dims(dst_dims),
+      strides(strides), dilations(dilations),
+      padding_left(padding_left), padding_right(padding_right) {
+  }
+};
+
+template <typename T>
+class Conv2DFwd : public DnnOp {
+ public:
+  explicit Conv2DFwd(const ConvFwdDimensions& convFwdDims) {
+    fwd_stream_.reset(new stream(stream::kind::eager));
+    // create conv primitive
+    if (conv_fwd_ == nullptr) {
+      Setup(convFwdDims);
+    }
+  }
+
+  ~Conv2DFwd() {}
+
+  // Convolution forward execute with bias
+  //   src_data:    input data buffer of src
+  //   filter_data: input data buffer of filter (weights)
+  //   bias_data:   input data buffer of bias
+  //   dst_data:    output data buffer of dst
+  void Execute(T* src_data, T* filter_data, T* bias_data, T* dst_data) {
+    src_mem_->set_data_handle(static_cast<void*>(src_data));
+    filter_mem_->set_data_handle(static_cast<void*>(filter_data));
+    bias_mem_->set_data_handle(static_cast<void*>(bias_data));
+    dst_mem_->set_data_handle(static_cast<void*>(dst_data));
+    fwd_stream_->submit(fwd_primitives_);
+
+    // after exec, set data handle back
+    src_mem_->set_data_handle(DummyData);
+    filter_mem_->set_data_handle(DummyData);
+    bias_mem_->set_data_handle(DummyData);
+    dst_mem_->set_data_handle(DummyData);
+
+    return;
+  }
+
+  // Convolution forward execute without bias
+  //   src_data:    input data buffer of src
+  //   filter_data: input data buffer of filter (weights)
+  //   dst_data:    output data buffer of dst
+  void Execute(T* src_data, T* filter_data, T* dst_data) {
+    src_mem_->set_data_handle(static_cast<void*>(src_data));
+    filter_mem_->set_data_handle(static_cast<void*>(filter_data));
+    dst_mem_->set_data_handle(static_cast<void*>(dst_data));
+    fwd_stream_->submit(fwd_primitives_);
+
+    // after exec, set data handle back
+    src_mem_->set_data_handle(DummyData);
+    filter_mem_->set_data_handle(DummyData);
+    dst_mem_->set_data_handle(DummyData);
+
+    return;
+  }
+
+  // expected memory format for this primitive instance
+  memory::format src_fmt_;
+  memory::format filter_fmt_;
+
+  // convolution primitive
+  std::shared_ptr<mkldnn::convolution_forward::primitive_desc> fwd_pd_;
+  std::shared_ptr<mkldnn::primitive> conv_fwd_;
+
+ private:
+  void Setup(const ConvFwdDimensions& convFwdDims) {
+    // create memory descriptors for convolution data w/ no specified format
+    src_md_.reset(new memory::desc({convFwdDims.src_dims},
+        MklDnnType<T>(), memory::format::any));
+
+    filter_md_.reset(new memory::desc({convFwdDims.filter_dims},
+        MklDnnType<T>(), memory::format::any));
+
+    dst_md_.reset(new memory::desc({convFwdDims.dst_dims},
+        MklDnnType<T>(), memory::format::any));
+
+    if (!convFwdDims.bias_dims.empty())
+        bias_md_.reset(new memory::desc({convFwdDims.bias_dims},
+            MklDnnType<T>(), memory::format::any));
+
+    // create a convolution
+    if (!convFwdDims.bias_dims.empty()) {
+      fwd_desc_.reset(new convolution_forward::desc(prop_kind::forward,
+          convolution_direct, *src_md_, *filter_md_, *bias_md_, *dst_md_,
+          convFwdDims.strides, convFwdDims.dilations, convFwdDims.padding_left,
+          convFwdDims.padding_right, padding_kind::zero));
+    } else {
+      fwd_desc_.reset(new convolution_forward::desc(prop_kind::forward,
+          convolution_direct, *src_md_, *filter_md_, *dst_md_,
+          convFwdDims.strides, convFwdDims.dilations, convFwdDims.padding_left,
+          convFwdDims.padding_right, padding_kind::zero));
+    }
+
+    fwd_pd_.reset(new convolution_forward::primitive_desc(
+        *fwd_desc_, cpu_engine_));
+
+    // store the expected memory format
+    src_fmt_ = static_cast<mkldnn::memory::format>(
+        fwd_pd_.get()->src_primitive_desc().desc().data.format);
+
+    filter_fmt_ = static_cast<mkldnn::memory::format>(
+        fwd_pd_.get()->weights_primitive_desc().desc().data.format);
+
+    // create memory primitive based on dummy data
+    src_mem_.reset(new memory(fwd_pd_.get()->src_primitive_desc(), DummyData));
+    filter_mem_.reset(new memory(fwd_pd_.get()->weights_primitive_desc(),
+                      DummyData));
+    dst_mem_.reset(new memory(fwd_pd_.get()->dst_primitive_desc(), DummyData));
+
+    // create convolution primitive and add it to net
+    if (!convFwdDims.bias_dims.empty()) {
+        bias_mem_.reset(new memory({{{convFwdDims.bias_dims}, MklDnnType<T>(),
+                        memory::format::x}, cpu_engine_}, DummyData));
+        conv_fwd_.reset(new convolution_forward(*fwd_pd_, *src_mem_,
+                        *filter_mem_, *bias_mem_, *dst_mem_));
+    } else {
+        conv_fwd_.reset(new convolution_forward(*fwd_pd_, *src_mem_,
+                        *filter_mem_, *dst_mem_));
+    }
+
+    fwd_primitives_.push_back(*conv_fwd_);
+    return;
+  }
+
+  // MKLDNN memory
+  std::shared_ptr<mkldnn::memory> src_mem_;
+  std::shared_ptr<mkldnn::memory> filter_mem_;
+  std::shared_ptr<mkldnn::memory> bias_mem_;
+  std::shared_ptr<mkldnn::memory> dst_mem_;
+
+  std::shared_ptr<mkldnn::stream> fwd_stream_;
+  std::vector<mkldnn::primitive> fwd_primitives_;
+
+  // desc & prmitive desc
+  std::shared_ptr<mkldnn::convolution_forward::desc> fwd_desc_;
+
+  // memory desc
+  std::shared_ptr<mkldnn::memory::desc> src_md_;
+  std::shared_ptr<mkldnn::memory::desc> filter_md_;
+  std::shared_ptr<mkldnn::memory::desc> bias_md_;
+  std::shared_ptr<mkldnn::memory::desc> dst_md_;
+
+  engine cpu_engine_ = engine(engine::cpu, 0);
+};
+
+template <typename T>
+class Conv2DFwdFactory : public DnnOpFactory<T> {
+ public:
+  static Conv2DFwd<T>* Get(const ConvFwdDimensions& convFwdDims) {
+     Conv2DFwd<T>* conv2d_fwd = nullptr;
+
+     // try to find a suitable one in pool
+     conv2d_fwd = dynamic_cast<Conv2DFwd<T>*> (
+       Conv2DFwdFactory<T>::GetInstance().GetConv2DFwd(convFwdDims));
+
+     if (conv2d_fwd == nullptr) {
+       conv2d_fwd = new Conv2DFwd<T>(convFwdDims);
+       Conv2DFwdFactory<T>::GetInstance().SetConv2DFwd(
+           convFwdDims, conv2d_fwd);
+     }
+     return conv2d_fwd;
+  }
+
+ private:
+  Conv2DFwdFactory() {}
+  ~Conv2DFwdFactory() {}
+
+  static const int kDilationH = 0, kDilationW = 1;
+
+  static Conv2DFwdFactory& GetInstance() {
+    static Conv2DFwdFactory instance_;
+    return instance_;
+  }
+
+  static std::string CreateKey(const ConvFwdDimensions& convFwdDims) {
+    std::string prefix = "conv2d_fwd_";
+    FactoryKeyCreator key_creator;
+    key_creator.AddAsKey(prefix);
+    key_creator.AddAsKey(convFwdDims.src_dims);
+    key_creator.AddAsKey(convFwdDims.filter_dims);
+    key_creator.AddAsKey(convFwdDims.bias_dims);
+    key_creator.AddAsKey(convFwdDims.dst_dims);
+    key_creator.AddAsKey(convFwdDims.strides);
+    key_creator.AddAsKey(convFwdDims.dilations);
+    key_creator.AddAsKey(convFwdDims.padding_left);
+    key_creator.AddAsKey(convFwdDims.padding_right);
+    return key_creator.GetKey();
+  }
+
+  DnnOp* GetConv2DFwd(const ConvFwdDimensions& convFwdDims) {
+    std::string key = CreateKey(convFwdDims);
+    return this->GetOp(key);
+  }
+
+  void SetConv2DFwd(const ConvFwdDimensions& convFwdDims, DnnOp *op) {
+    std::string key = CreateKey(convFwdDims);
+    this->SetOp(key, op);
+  }
+};
+
+#endif
+
 typedef Eigen::ThreadPoolDevice CPUDevice;
 
-// MKL-DNN is now default. MKL-ML must be specified explicitly.
+// For now, MKL-ML is default. So making MKL-DNN not a default choice.
 #ifdef INTEL_MKL_ML
-
 template <typename Device, typename T, bool biasEnabled>
 class MklConv2DOp : public OpKernel {
  public:
@@ -528,8 +749,6 @@ class MklConv2DOp : public OpKernel {
 
   void Compute(OpKernelContext* context) override {
     try {
-      auto cpu_engine = engine(engine::cpu, 0);
-
       // Input tensors
       const Tensor& src_tensor = MklGetInput(context, kInputIndex_Src);
       const Tensor& filter_tensor = MklGetInput(context, kInputIndex_Filter);
@@ -538,16 +757,16 @@ class MklConv2DOp : public OpKernel {
       GetMklShape(context, kInputIndex_Src, &src_mkl_shape);
       GetMklShape(context, kInputIndex_Filter, &filter_mkl_shape);
       OP_REQUIRES(context, filter_mkl_shape.IsMklTensor() == false,
-                  errors::InvalidArgument("Filter should not be in "
-                                          "Mkl Layout"));
+            errors::InvalidArgument("Filter should not be in "
+            "Mkl Layout"));
 
       MklDnnData<T> src(&cpu_engine);
       MklDnnData<T> filter(&cpu_engine);
-      MklDnnData<T> output(&cpu_engine);
+      MklDnnData<T> dst(&cpu_engine);  // output
 
-      memory::dims src_dims, filter_dims, padding_l, padding_r,
+      memory::dims src_dims, filter_dims, padding_left, padding_right,
                    dilations, strides;
-      memory::dims output_dims_tf_order, output_dims_mkl_order;
+      memory::dims dst_dims_tf_order, dst_dims_mkl_order;
 
       // Get shapes of input tensors in MKL-DNN order
       MklDnnConvUtil conv_utl(context, strides_, padding_, data_format_,
@@ -555,31 +774,29 @@ class MklConv2DOp : public OpKernel {
       auto src_tf_shape = GetTfShape(context, kInputIndex_Src);
       auto filter_tf_shape = GetTfShape(context, kInputIndex_Filter);
       conv_utl.GetConvFwdSizesInMklOrder(
-          src_tf_shape, filter_tf_shape, &src_dims, &filter_dims, &strides,
-          &dilations, &output_dims_tf_order, &output_dims_mkl_order,
-          &padding_l, &padding_r);
+          src_tf_shape, filter_tf_shape, &src_dims, &filter_dims,
+          &strides, &dilations, &dst_dims_tf_order, &dst_dims_mkl_order,
+          &padding_left, &padding_right);
       if (!context->status().ok()) return;
 
       // Check for corner case - if there is nothing to compute, return.
-      TensorShape output_tf_shape = MklDnnDimsToTFShape(output_dims_tf_order);
+      TensorShape dst_tf_shape = MklDnnDimsToTFShape(dst_dims_tf_order);
 
       // Corner cases: output with 0 elements and 0 batch size.
-      Tensor* output_tensor = nullptr;
-      if (output_tf_shape.num_elements() == 0 || output_dims_tf_order[0] == 0) {
-        // TODO(jbobba): Verify correctness here
-        //               Need semantics for Null MKL tensor
-        MklDnnShape output_mkl_shape;
-        output_mkl_shape.SetMklTensor(false);
-
-        AllocateOutputSetMklShape(context, kOutputIndex_Dst, &output_tensor,
-                                  src_tf_shape, output_mkl_shape);
+      Tensor* dst_tensor = nullptr;
+      if (dst_tf_shape.num_elements() == 0 ||
+          dst_dims_tf_order[0] == 0) {
+        MklDnnShape dst_mkl_shape;
+        dst_mkl_shape.SetMklTensor(false);
+        AllocateOutputSetMklShape(context, kOutputIndex_Dst,
+                    &dst_tensor, src_tf_shape, dst_mkl_shape);
 
         // MklConv2D also outputs converted filter as 2nd output of Conv2D.
         filter_mkl_shape.SetMklTensor(false);
         Tensor* output_filter_tensor = nullptr;
         AllocateOutputSetMklShape(context, kOutputIndex_Filter,
-                                  &output_filter_tensor, filter_tf_shape,
-                                  filter_mkl_shape);
+                                  &output_filter_tensor,
+                                  filter_tf_shape, filter_mkl_shape);
         return;
       }
 
@@ -587,6 +804,7 @@ class MklConv2DOp : public OpKernel {
       // Describe how the inputs and outputs of Convolution look like. Also
       // specify buffers containing actual input and output data.
       auto tf_fmt = TFDataFormatToMklDnnDataFormat(data_format_);
+
       // If input is in MKL layout, then simply grab input layout; otherwise,
       // construct input Tf layout. For TF layout, although input shape
       // (src_dims) required is in MKL-DNN order, the layout is Tensorflow's
@@ -595,6 +813,7 @@ class MklConv2DOp : public OpKernel {
                         ? src_mkl_shape.GetMklLayout()
                         : memory::desc(src_dims, MklDnnType<T>(), tf_fmt);
       src.SetUsrMem(src_md, &src_tensor);
+
       // Although filter shape (filter_dims) required is in MKL-DNN order,
       // the layout is Tensorflow's layout (HWIO).
       auto filter_md = filter_mkl_shape.IsMklTensor()  // Should NEVER be true
@@ -603,98 +822,70 @@ class MklConv2DOp : public OpKernel {
                                           memory::format::hwio);
       filter.SetUsrMem(filter_md, &filter_tensor);
 
-      // Set output shape (output_dims) required in MKL-DNN order.
-      // Currently, we set output layout as Tensorflow's layout (NHWC or NCHW
-      // depending on data format). But later we propagate Mkl layout of the
-      // output to the next op directly.
-      output.SetUsrMem(output_dims_mkl_order, tf_fmt);
-
-      // Create memory descriptors for convolution data w/ no specified format.
-      src.SetOpMemDesc(src_dims, memory::format::any);
-      filter.SetOpMemDesc(filter_dims, memory::format::any);
-      output.SetOpMemDesc(output_dims_mkl_order, memory::format::any);
-
       // MKLDNN dilation starts from 0.
       dilations[kDilationH] -= 1;
       dilations[kDilationW] -= 1;
 
+      // get a conv2d fwd from primitive pool
+      Conv2DFwd<T> *conv2d_fwd = nullptr;
+      if (biasEnabled) {
+        memory::dims bias_dims = {};
+        conv_utl.GetBiasSizeInMklOrder(kInputIndex_Bias, &bias_dims);
+        ConvFwdDimensions convFwdDims(src_dims, filter_dims, bias_dims,
+          dst_dims_mkl_order, strides, dilations, padding_left, padding_right);
+        conv2d_fwd = Conv2DFwdFactory<T>::Get(convFwdDims);
+      } else {
+        ConvFwdDimensions convFwdDims(src_dims, filter_dims, NONE_DIMS,
+          dst_dims_mkl_order, strides, dilations, padding_left, padding_right);
+        conv2d_fwd = Conv2DFwdFactory<T>::Get(convFwdDims);
+      }
+
+      // allocate output tensors output_tensor and filter_out_tensor
+      std::shared_ptr<mkldnn::convolution_forward::primitive_desc>
+      conv_fwd_pd = conv2d_fwd->fwd_pd_;
+      AllocateOutputTensor(context, *conv_fwd_pd,
+                       dst_dims_mkl_order, tf_fmt, &dst_tensor);
+      Tensor* filter_out_tensor = nullptr;
+      AllocateFilterOutputTensor(context, *conv_fwd_pd,
+                                 TFShapeToMklDnnDims(filter_tf_shape),
+                                 &filter_out_tensor);
+
+      T* dst_data = static_cast<T*>(dst_tensor->flat<T>().data());
+
+      // check whether src/filter need reorder
+      std::vector<primitive> net;
+      if (src_md.data.format != conv2d_fwd->src_fmt_)
+          src.CheckReorderToOpMem(
+              conv_fwd_pd.get()->src_primitive_desc(), &net);
+
+      if (filter_md.data.format != conv2d_fwd->filter_fmt_)
+          filter.CheckReorderToOpMem(
+              conv_fwd_pd.get()->weights_primitive_desc(),
+              filter.GetTensorBuffer(filter_out_tensor), &net);
+      stream(stream::kind::eager).submit(net).wait();
+
+      T* src_data = static_cast<T*>(
+                src.GetOpMem().get_data_handle());
+      T* filter_data = static_cast<T*>(
+                filter.GetOpMem().get_data_handle());
+
+      // execute convolution
       if (biasEnabled) {
-          // Create convolution primitive with Bias.
-          MklDnnData<T> bias(&cpu_engine);
-          memory::dims bias_size;
-          conv_utl.GetBiasSizeInMklOrder(kInputIndex_Bias, &bias_size);
-          const Tensor& bias_tensor = MklGetInput(context, kInputIndex_Bias);
-          bias.SetUsrMem(bias_size, memory::format::x, &bias_tensor);
-          bias.SetOpMemDesc(bias_size, memory::format::any);
-
-          // Create convolution primitive with Bias.
-          // Use MKLDNN dilated convolution in case of dilated rate (>0).
-          auto conv_desc = (dilations[kDilationH] > 0 ||
-              dilations[kDilationW] > 0) ?
-              convolution_forward::desc(prop_kind::forward,
-                      convolution_direct, src.GetOpMemDesc(),
-                      filter.GetOpMemDesc(), bias.GetOpMemDesc(),
-                      output.GetOpMemDesc(), strides, dilations,
-                      padding_l, padding_r,
-                      TFPaddingToMklDnnPadding(padding_)):
-              convolution_forward::desc(prop_kind::forward,
-                      convolution_direct, src.GetOpMemDesc(),
-                      filter.GetOpMemDesc(), bias.GetOpMemDesc(),
-                      output.GetOpMemDesc(), strides,
-                      padding_l, padding_r,
-                      TFPaddingToMklDnnPadding(padding_));
-
-          auto conv_prim_desc = convolution_forward::primitive_desc(conv_desc,
-                                                                  cpu_engine);
-          AllocateOutputTensor(context, conv_prim_desc,
-                               output_dims_mkl_order, tf_fmt, &output_tensor);
-          // Set data handle for output.
-          output.SetUsrMemDataHandle(output_tensor);
-
-          Tensor* filter_out_tensor = nullptr;
-          AllocateFilterOutputTensor(context, conv_prim_desc,
-                TFShapeToMklDnnDims(filter_tf_shape),
-                &filter_out_tensor);
-
-          PrepareAndExecuteNet(conv_prim_desc, &src, &filter, &bias, &output,
-                               filter_out_tensor);
+        const Tensor& bias_tensor = MklGetInput(context, kInputIndex_Bias);
+        T* bias_data = static_cast<T*>(const_cast<T*>(
+            bias_tensor.flat<T>().data()));
+
+        conv2d_fwd->Execute(src_data, filter_data, bias_data, dst_data);
       } else {
-          // Create convolution primitive without Bias.
-          // Use MKLDNN dilated convolution in case of dilated rate (>0).
-          auto conv_desc = (dilations[kDilationH] > 0 ||
-            dilations[kDilationW] > 0) ?
-            convolution_forward::desc(prop_kind::forward,
-              convolution_direct, src.GetOpMemDesc(),
-              filter.GetOpMemDesc(), output.GetOpMemDesc(),
-              strides, dilations, padding_l, padding_r,
-              TFPaddingToMklDnnPadding(padding_)):
-          convolution_forward::desc(prop_kind::forward,
-              convolution_direct, src.GetOpMemDesc(),
-              filter.GetOpMemDesc(), output.GetOpMemDesc(),
-              strides, padding_l, padding_r,
-              TFPaddingToMklDnnPadding(padding_));
-
-          auto conv_prim_desc = convolution_forward::primitive_desc(conv_desc,
-                                                                  cpu_engine);
-          AllocateOutputTensor(context, conv_prim_desc, output_dims_mkl_order,
-                               tf_fmt, &output_tensor);
-          // Set data handle for output.
-          output.SetUsrMemDataHandle(output_tensor);
-
-          Tensor* filter_out_tensor = nullptr;
-          AllocateFilterOutputTensor(context, conv_prim_desc,
-                TFShapeToMklDnnDims(filter_tf_shape),
-                &filter_out_tensor);
-          PrepareAndExecuteNet(conv_prim_desc, &src, &filter,
-                              nullptr, &output, filter_out_tensor);
+        conv2d_fwd->Execute(src_data, filter_data, dst_data);
       }
-    } catch (mkldnn::error& e) {
+    } catch (mkldnn::error &e) {
       string error_msg = "Status: " + std::to_string(e.status) +
-                         ", message: " + std::string(e.message) + ", in file " +
-                         std::string(__FILE__) + ":" + std::to_string(__LINE__);
-      OP_REQUIRES_OK(
-          context,
-          errors::Aborted("Operation received an exception:", error_msg));
+                       ", message: " + std::string(e.message) +
+                       ", in file " + std::string(__FILE__) + ":" +
+                       std::to_string(__LINE__);
+      OP_REQUIRES_OK(context,
+        errors::Aborted("Operation received an exception:", error_msg));
     }
   }
 
@@ -706,6 +897,7 @@ class MklConv2DOp : public OpKernel {
   const int kInputIndex_Src = 0, kInputIndex_Filter = 1, kInputIndex_Bias = 2;
   const int kOutputIndex_Dst = 0, kOutputIndex_Filter = 1;
   const int kDilationH = 0, kDilationW = 1;
+  engine cpu_engine = engine(engine::cpu, 0);
 
   // Allocate output tensor.
   void AllocateOutputTensor(
diff --git a/tensorflow/core/kernels/scatter_nd_op.cc b/tensorflow/core/kernels/scatter_nd_op.cc
index 3a95dd1773..0caa7bd317 100644
--- a/tensorflow/core/kernels/scatter_nd_op.cc
+++ b/tensorflow/core/kernels/scatter_nd_op.cc
@@ -241,6 +241,7 @@ class ScatterNdUpdateOp : public OpKernel {
 TF_CALL_NUMBER_TYPES(REGISTER_SCATTER_ND_ADD_SUB_CPU);
 TF_CALL_NUMBER_TYPES(REGISTER_SCATTER_ND_UPDATE_CPU);
 TF_CALL_NUMBER_TYPES(REGISTER_SCATTER_ND_CPU);
+TF_CALL_string(REGISTER_SCATTER_ND_CPU);
 
 // Registers GPU kernels.
 #if GOOGLE_CUDA
diff --git a/tensorflow/core/kernels/scatter_nd_op_cpu_impl.h b/tensorflow/core/kernels/scatter_nd_op_cpu_impl.h
index e82660dcc1..7cfffa20c5 100644
--- a/tensorflow/core/kernels/scatter_nd_op_cpu_impl.h
+++ b/tensorflow/core/kernels/scatter_nd_op_cpu_impl.h
@@ -160,6 +160,7 @@ struct ScatterNdFunctor<CPUDevice, T, Index, OP, IXDIM> {
   REGISTER_SCATTER_ND_INDEX(type, scatter_nd_op::UpdateOp::SUB);
 
 TF_CALL_ALL_TYPES(REGISTER_SCATTER_ND_UPDATE);
+REGISTER_SCATTER_ND_INDEX(string, scatter_nd_op::UpdateOp::ADD);
 TF_CALL_NUMBER_TYPES(REGISTER_SCATTER_ND_MATH)
 
 #undef REGISTER_SCATTER_ND_MATH
diff --git a/tensorflow/core/kernels/segment_reduction_ops.h b/tensorflow/core/kernels/segment_reduction_ops.h
index bedd965966..4abfbfb1a6 100644
--- a/tensorflow/core/kernels/segment_reduction_ops.h
+++ b/tensorflow/core/kernels/segment_reduction_ops.h
@@ -16,35 +16,6 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_KERNELS_SEGMENT_REDUCTION_OPS_H_
 #define TENSORFLOW_CORE_KERNELS_SEGMENT_REDUCTION_OPS_H_
 
-
-// This file requires the following include because it uses CudaAtomicMax:
-// #include "tensorflow/core/util/cuda_kernel_helper.h"
-
-// Unfortunately we can't add the #include, since it breaks compilation for
-// non-GPU targets. This only breaks in clang, because it's more strict for
-// template code and CudaAtomicMax is used in template context.
-
-// This file requires the following include because it uses CudaAtomicMax:
-// #include "tensorflow/core/util/cuda_kernel_helper.h"
-
-// Unfortunately we can't add the #include, since it breaks compilation for
-// non-GPU targets. This only breaks in clang, because it's more strict for
-// template code and CudaAtomicMax is used in template context.
-
-// This file requires the following include because it uses CudaAtomicMax:
-// #include "tensorflow/core/util/cuda_kernel_helper.h"
-
-// Unfortunately we can't add the #include, since it breaks compilation for
-// non-GPU targets. This only breaks in clang, because it's more strict for
-// template code and CudaAtomicMax is used in template context.
-
-// This file requires the following include because it uses CudaAtomicMax:
-// #include "tensorflow/core/util/cuda_kernel_helper.h"
-
-// Unfortunately we can't add the #include, since it breaks compilation for
-// non-GPU targets. This only breaks in clang, because it's more strict for
-// template code and CudaAtomicMax is used in template context.
-
 // This file requires the following include because it uses CudaAtomicMax:
 // #include "tensorflow/core/util/cuda_kernel_helper.h"
 
diff --git a/tensorflow/core/platform/default/gpu/cupti_wrapper.h b/tensorflow/core/platform/default/gpu/cupti_wrapper.h
index acd889e474..e3ebe6ca1d 100644
--- a/tensorflow/core/platform/default/gpu/cupti_wrapper.h
+++ b/tensorflow/core/platform/default/gpu/cupti_wrapper.h
@@ -23,7 +23,7 @@ limitations under the License.
 #if defined(WIN32)
 #include "extras/CUPTI/include/cupti.h"
 #else
-#include "cuda/extras/CUPTI/include/cupti.h"
+#include "cupti.h"
 #endif
 namespace perftools {
 namespace gputools {
diff --git a/tensorflow/core/public/version.h b/tensorflow/core/public/version.h
index 0ca7d8475f..ba69efb289 100644
--- a/tensorflow/core/public/version.h
+++ b/tensorflow/core/public/version.h
@@ -24,7 +24,7 @@ limitations under the License.
 
 // TF_VERSION_SUFFIX is non-empty for pre-releases (e.g. "-alpha", "-alpha.1",
 // "-beta", "-rc", "-rc.1")
-#define TF_VERSION_SUFFIX "-rc0"
+#define TF_VERSION_SUFFIX "-rc1"
 
 #define TF_STR_HELPER(x) #x
 #define TF_STR(x) TF_STR_HELPER(x)
diff --git a/tensorflow/core/util/mkl_util.h b/tensorflow/core/util/mkl_util.h
index bc6d2d77a4..50a8e30574 100644
--- a/tensorflow/core/util/mkl_util.h
+++ b/tensorflow/core/util/mkl_util.h
@@ -19,6 +19,8 @@ limitations under the License.
 
 #include <string>
 #include <vector>
+#include <unordered_map>
+#include <utility>
 
 #include "mkl_dnn.h"
 #include "mkl_dnn_types.h"
@@ -1759,7 +1761,90 @@ class MklDnnData {
   }
 };
 
-#endif  // INTEL_MKL_ML
+/// Base class for operations with reuse of DNN primitives
+///
+class DnnOp {
+ public:
+  virtual ~DnnOp() {}
+
+  // Dummy data. Its size, hard-coded as 256 here, does
+  // not matter since MKL should never operate on this buffer.
+  unsigned char DummyData[256];
+};
+
+const mkldnn::memory::dims NONE_DIMS = {};
+// This constant is used to declare dummy buffer (size), for MKL primitives
+template <typename T>
+class DnnOpFactory {
+ public:
+  DnnOpFactory() {}
+  ~DnnOpFactory() {}
+
+  DnnOp* GetOp(const std::string& key) {
+    auto stream_iter = DnnOpFactory<T>::GetHashMap().find(key);
+    if (stream_iter == DnnOpFactory<T>::GetHashMap().end()) {
+      return nullptr;
+    } else {
+      return stream_iter->second;
+    }
+  }
+
+  void SetOp(const std::string& key, DnnOp* op) {
+    auto stream_iter = DnnOpFactory<T>::GetHashMap().find(key);
+
+    CHECK(stream_iter == DnnOpFactory<T>::GetHashMap().end());
+
+    DnnOpFactory<T>::GetHashMap()[key] = op;
+  }
+
+ private:
+  static inline std::unordered_map<std::string, DnnOp*> &GetHashMap() {
+    static thread_local std::unordered_map<std::string, DnnOp*> map_;
+    return map_;
+  }
+};
+
+// utility class for creating keys of MKL primitive pool.
+class FactoryKeyCreator {
+ public:
+  FactoryKeyCreator() {
+    key_.reserve(kMaxKeyLength);
+  }
+
+  ~FactoryKeyCreator() {}
+
+  void AddAsKey(const string &str) {
+    auto buffer = reinterpret_cast<const char *>(str.c_str());
+    Append(buffer, str.length());
+  }
+
+  void AddAsKey(const mkldnn::memory::dims &dims) {
+    for (unsigned int i = 0; i < dims.size(); i++) {
+      AddAsKey<int>(dims[i]);
+    }
+  }
+
+  template <typename T>
+  void AddAsKey(const T data) {
+    auto buffer = reinterpret_cast<const char *>(&data);
+    Append(buffer, sizeof(T));
+  }
+
+  std::string GetKey() {
+    return key_;
+  }
+
+ private:
+  string key_;
+  const char delimiter = 'x';
+  const int kMaxKeyLength = 256;
+  void Append(const char* data, int len) {
+    key_.append(data, len);
+    key_.append(1, delimiter);
+  }
+};
+
+#endif  // INTEL_MKL_DNN
 
 }  // namespace tensorflow
 #endif  // INTEL_MKL
diff --git a/tensorflow/docs_src/community/roadmap.md b/tensorflow/docs_src/community/roadmap.md
index a3170a10f2..0463ca05fe 100644
--- a/tensorflow/docs_src/community/roadmap.md
+++ b/tensorflow/docs_src/community/roadmap.md
@@ -1,5 +1,5 @@
 # Roadmap
-**Last updated: Feb 15, 2018**
+**Last updated: Apr 27, 2018**
 
 TensorFlow is a rapidly moving, community supported project. This document is intended 
 to provide guidance about priorities and focus areas of the core set of TensorFlow 
@@ -14,12 +14,12 @@ expected in the next one to two releases.
 
 ### APIs
 #### High Level APIs:
-* Easy multi-GPU utilization with Estimators
+* Easy multi-GPU and TPU utilization with Estimators
 * Easy-to-use high-level pre-made estimators for Gradient Boosted Trees, Time Series, and other models
 
 #### Eager Execution:
 * Efficient utilization of multiple GPUs
-* Distributed training (multi-machine)
+* Distributed training support (multi-machine)
 * Performance improvements
 * Simpler export to a GraphDef/SavedModel 
 
@@ -31,14 +31,14 @@ to create Keras models Eager- style via Model subclassing)
 
 #### Official Models:
 * A set of 
-[reference models](https://github.com/tensorflow/models/tree/master/official) 
+[models](https://github.com/tensorflow/models/tree/master/official) 
 across image recognition, speech, object detection, and 
   translation that demonstrate best practices and serve as a starting point for 
   high-performance model development.
 
 #### Contrib:
-* Deprecation notices added to parts of tf.contrib where preferred implementations exist outside of tf.contrib.
-* As much as possible, large projects inside tf.contrib moved to separate repositories.
+* Deprecate parts of tf.contrib where preferred implementations exist outside of tf.contrib.
+* As much as possible, move large projects inside tf.contrib to separate repositories.
 * The tf.contrib module will eventually be discontinued in its current form, experimental development will in future happen in other repositories.
 
 
@@ -50,36 +50,72 @@ across image recognition, speech, object detection, and
 
 ### Platforms
 #### TensorFlow Lite:
-* Increased coverage of supported ops in TensorFlow Lite
+* Increase coverage of supported ops in TensorFlow Lite
 * Easier conversion of a trained TensorFlow graph for use on TensorFlow Lite
 * Support for GPU acceleration in TensorFlow Lite (iOS and Android)
 * Support for hardware accelerators via Android NeuralNets API 
-* Improved CPU performance by quantization and other network optimizations (eg. pruning, distillation)
-* Increased support for devices beyond Android and iOS (eg. RPi, Cortex-M)
+* Improve CPU performance by quantization and other network optimizations (eg. pruning, distillation)
+* Increase support for devices beyond Android and iOS (eg. RPi, Cortex-M)
+
+#### TensorFlow.js:
+* Release package for Node.js bindings to the TensorFlow C API through the TensorFlow.js backend interface
+* Expand support for importing TensorFlow SavedModels and Keras models into browser with unified APIs supporting retraining in browser
+* Improve Layers API and allow model exporting/saving
+* Release tfjs-data API for efficient data input pipelines
+
+#### TensorFlow with Swift:
+* Establish open source project including documentation, open design, and code availability.
+* Continue implementing and refining implementation and design through 2018.
+* Aim for implementation to be solid enough for general use later in 2018.
 
 ### Performance
 #### Distributed TensorFlow:
-* Multi-GPU support optimized for a variety of GPU topologies
-* Improved mechanisms for distributing computations on several machines
+* Optimize Multi-GPU support for a variety of GPU topologies
+* Improve mechanisms for distributing computations on several machines
+
+#### GPU Optimizations:
+* Simplify mixed precision API with initial example model and guide.
+* Finalize TensorRT API and move to core.
+* CUDA 9.2 and NCCL 2.x default in TensorFlow builds.
+* Optimizations for DGX-2.
+* Remove support for CUDA less than 8.x and cuDNN less than 6.x.
 
-#### Optimizations:
-* Mixed precision training support with initial example model and guide
-* Native TensorRT support
+
+#### CPU Optimizations
 * Int8 support for SkyLake via MKL
 * Dynamic loading of SIMD-optimized kernels
+* MKL for Linux and Windows
+
+### End-to-end ML systems:
+#### TensorFlow Hub:
+* Expand support for module-types in TF Hub with TF Eager integration, Keras layers integration, and TensorFlow.js integration
+* Accept variable-sized image input
+* Improve multi-GPU estimator support
+* Document and improve TPU integration
+
+#### TensorFlow Extended:
+* Open source more of the TensorFlow Extended platform to facilitate adoption of TensorFlow in production settings.
+* Release TFX libraries for Data Validation
+
+### Documentation and Resources:
+* Update documentation, tutorials and Getting Started guides on all features and APIs
+* Update [Youtube Tensorflow channel](https://youtube.com/tensorflow) weekly with new content:
+Coding TensorFlow - where we teach folks coding with tensorflow
+TensorFlow Meets - where we highlight community contributions
+Ask TensorFlow - where we answer community questions
+Guest and Showcase videos
+* Update [Official TensorFlow blog](https://blog.tensorflow.org) with regular articles from Google team and the Community
 
-### Documentation and Usability:
-* Updated documentation, tutorials and Getting Started guides
-* Process to enable external contributions to tutorials, documentation, and blogs showcasing best practice use-cases of TensorFlow and high-impact applications
 
 ### Community and Partner Engagement
 #### Special Interest Groups: 
-* Mobilizing the community to work together in focused domains
+* Mobilize the community to work together in focused domains
 * [tf-distribute](https://groups.google.com/a/tensorflow.org/forum/#!forum/tf-distribute): build and packaging of TensorFlow
-* More to be identified and launched
+* SIG TensorBoard, SIG Rust, and more to be identified and launched
 
 #### Community:
 * Incorporate public feedback on significant design decisions via a Request-for-Comment (RFC) process
 * Formalize process for external contributions to land in TensorFlow and associated projects 
 * Grow global TensorFlow communities and user groups
 * Collaborate with partners to co-develop and publish research papers
+* Process to enable external contributions to tutorials, documentation, and blogs showcasing best practice use-cases of TensorFlow and high-impact applications
diff --git a/tensorflow/docs_src/get_started/checkpoints.md b/tensorflow/docs_src/get_started/checkpoints.md
index 4aa07c7f2a..8dfd91e3c8 100644
--- a/tensorflow/docs_src/get_started/checkpoints.md
+++ b/tensorflow/docs_src/get_started/checkpoints.md
@@ -38,8 +38,10 @@ Estimators automatically write the following to disk:
     uses to create visualizations.
 
 To specify the top-level directory in which the Estimator stores its
-information, assign a value to the optional `model_dir` argument of any
-Estimator's constructor.  For example, the following code sets the `model_dir`
+information, assign a value to the optional `model_dir` argument of *any*
+`Estimator`'s constructor.
+Taking `DNNClassifier` as an example,
+the following code sets the `model_dir`
 argument to the `models/iris` directory:
 
 ```python
diff --git a/tensorflow/docs_src/get_started/feature_columns.md b/tensorflow/docs_src/get_started/feature_columns.md
index 9c777a0077..79c2667979 100644
--- a/tensorflow/docs_src/get_started/feature_columns.md
+++ b/tensorflow/docs_src/get_started/feature_columns.md
@@ -138,7 +138,7 @@ The model will represent the buckets as follows:
 |< 1960               | [1, 0, 0, 0] |
 |>= 1960 but < 1980   | [0, 1, 0, 0] |
 |>= 1980 but < 2000   | [0, 0, 1, 0] |
-|> 2000               | [0, 0, 0, 1] |
+|>= 2000              | [0, 0, 0, 1] |
 
 Why would you want to split a number—a perfectly valid input to your
 model—into a categorical value? Well, notice that the categorization splits a
diff --git a/tensorflow/docs_src/get_started/index.md b/tensorflow/docs_src/get_started/index.md
index b28cb9df75..746126c720 100644
--- a/tensorflow/docs_src/get_started/index.md
+++ b/tensorflow/docs_src/get_started/index.md
@@ -10,7 +10,7 @@ course prior to diving into TensorFlow documentation:
 TensorFlow is a tool for machine learning. While it contains a wide range of
 functionality, TensorFlow is mainly designed for deep neural network models.
 
-The easiest way to get started with tensorflow is using Eager Execution.
+The easiest way to get started with TensorFlow is using Eager Execution.
 
   * @{$get_started/eager}, is for anyone new to  machine learning or TensorFlow.
 
diff --git a/tensorflow/docs_src/install/install_c.md b/tensorflow/docs_src/install/install_c.md
index 995b8ae666..8c165aad52 100644
--- a/tensorflow/docs_src/install/install_c.md
+++ b/tensorflow/docs_src/install/install_c.md
@@ -38,7 +38,7 @@ enable TensorFlow for C:
          OS="linux" # Change to "darwin" for macOS
          TARGET_DIRECTORY="/usr/local"
          curl -L \
-           "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-${OS}-x86_64-1.8.0-rc0.tar.gz" |
+           "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-${OS}-x86_64-1.8.0-rc1.tar.gz" |
            sudo tar -C $TARGET_DIRECTORY -xz
 
      The `tar` command extracts the TensorFlow C library into the `lib`
diff --git a/tensorflow/docs_src/install/install_go.md b/tensorflow/docs_src/install/install_go.md
index 2938a8f7ee..26cbcc9a9b 100644
--- a/tensorflow/docs_src/install/install_go.md
+++ b/tensorflow/docs_src/install/install_go.md
@@ -38,7 +38,7 @@ steps to install this library and enable TensorFlow for Go:
          TF_TYPE="cpu" # Change to "gpu" for GPU support
          TARGET_DIRECTORY='/usr/local'
          curl -L \
-           "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-$(go env GOOS)-x86_64-1.8.0-rc0.tar.gz" |
+           "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-${TF_TYPE}-$(go env GOOS)-x86_64-1.8.0-rc1.tar.gz" |
          sudo tar -C $TARGET_DIRECTORY -xz
 
      The `tar` command extracts the TensorFlow C library into the `lib`
diff --git a/tensorflow/docs_src/install/install_java.md b/tensorflow/docs_src/install/install_java.md
index 05604d95c5..05b2878701 100644
--- a/tensorflow/docs_src/install/install_java.md
+++ b/tensorflow/docs_src/install/install_java.md
@@ -36,7 +36,7 @@ following to the project's `pom.xml` to use the TensorFlow Java APIs:
 <dependency>
   <groupId>org.tensorflow</groupId>
   <artifactId>tensorflow</artifactId>
-  <version>1.8.0-rc0</version>
+  <version>1.8.0-rc1</version>
 </dependency>
 ```
 
@@ -65,7 +65,7 @@ As an example, these steps will create a Maven project that uses TensorFlow:
                <dependency>
                  <groupId>org.tensorflow</groupId>
                  <artifactId>tensorflow</artifactId>
-                 <version>1.8.0-rc0</version>
+                 <version>1.8.0-rc1</version>
                </dependency>
              </dependencies>
          </project>
@@ -124,12 +124,12 @@ instead:
 <dependency>
   <groupId>org.tensorflow</groupId>
   <artifactId>libtensorflow</artifactId>
-  <version>1.8.0-rc0</version>
+  <version>1.8.0-rc1</version>
 </dependency>
 <dependency>
   <groupId>org.tensorflow</groupId>
   <artifactId>libtensorflow_jni_gpu</artifactId>
-  <version>1.8.0-rc0</version>
+  <version>1.8.0-rc1</version>
 </dependency>
 ```
 
@@ -148,7 +148,7 @@ refer to the simpler instructions above instead.
 Take the following steps to install TensorFlow for Java on Linux or macOS:
 
   1. Download
-     [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.8.0-rc0.jar),
+     [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.8.0-rc1.jar),
      which is the TensorFlow Java Archive (JAR).
 
   2. Decide whether you will run TensorFlow for Java on CPU(s) only or with
@@ -167,7 +167,7 @@ Take the following steps to install TensorFlow for Java on Linux or macOS:
          OS=$(uname -s | tr '[:upper:]' '[:lower:]')
          mkdir -p ./jni
          curl -L \
-           "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-${TF_TYPE}-${OS}-x86_64-1.8.0-rc0.tar.gz" |
+           "https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-${TF_TYPE}-${OS}-x86_64-1.8.0-rc1.tar.gz" |
            tar -xz -C ./jni
 
 ### Install on Windows
@@ -175,10 +175,10 @@ Take the following steps to install TensorFlow for Java on Linux or macOS:
 Take the following steps to install TensorFlow for Java on Windows:
 
   1. Download
-     [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.8.0-rc0.jar),
+     [libtensorflow.jar](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-1.8.0-rc1.jar),
      which is the TensorFlow Java Archive (JAR).
   2. Download the following Java Native Interface (JNI) file appropriate for
-     [TensorFlow for Java on Windows](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-cpu-windows-x86_64-1.8.0-rc0.zip).
+     [TensorFlow for Java on Windows](https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow_jni-cpu-windows-x86_64-1.8.0-rc1.zip).
   3. Extract this .zip file.
 
 
@@ -227,7 +227,7 @@ must be part of your `classpath`. For example, you can include the
 downloaded `.jar` in your `classpath` by using the `-cp` compilation flag
 as follows:
 
-<pre><b>javac -cp libtensorflow-1.8.0-rc0.jar HelloTF.java</b></pre>
+<pre><b>javac -cp libtensorflow-1.8.0-rc1.jar HelloTF.java</b></pre>
 
 
 ### Running
@@ -241,11 +241,11 @@ two files are available to the JVM:
 For example, the following command line executes the `HelloTF` program on Linux
 and macOS X:
 
-<pre><b>java -cp libtensorflow-1.8.0-rc0.jar:. -Djava.library.path=./jni HelloTF</b></pre>
+<pre><b>java -cp libtensorflow-1.8.0-rc1.jar:. -Djava.library.path=./jni HelloTF</b></pre>
 
 And the following command line executes the `HelloTF` program on Windows:
 
-<pre><b>java -cp libtensorflow-1.8.0-rc0.jar;. -Djava.library.path=jni HelloTF</b></pre>
+<pre><b>java -cp libtensorflow-1.8.0-rc1.jar;. -Djava.library.path=jni HelloTF</b></pre>
 
 If the program prints <tt>Hello from <i>version</i></tt>, you've successfully
 installed TensorFlow for Java and are ready to use the API.  If the program
diff --git a/tensorflow/docs_src/install/install_linux.md b/tensorflow/docs_src/install/install_linux.md
index e087b0c221..761555ca9a 100644
--- a/tensorflow/docs_src/install/install_linux.md
+++ b/tensorflow/docs_src/install/install_linux.md
@@ -1,350 +1,273 @@
 # Installing TensorFlow on Ubuntu
 
-This guide explains how to install TensorFlow on Ubuntu. Although these
-instructions might also work on other Linux variants, we have only
-tested (and we only support) these instructions on machines meeting the
-following requirements:
+This guide explains how to install TensorFlow on Ubuntu Linux. While these
+instructions may work on other Linux variants, they are tested and supported with
+the following system requirements:
 
-  * 64-bit desktops or laptops
-  * Ubuntu 16.04 or higher
+* 64-bit desktops or laptops
+* Ubuntu 16.04 or higher
 
 
-## Determine which TensorFlow to install
+## Choose which TensorFlow to install
 
-You must choose one of the following types of TensorFlow to install:
+The following TensorFlow variants are available for installation:
 
-  * **TensorFlow with CPU support only**. If your system does not have a
-    NVIDIA® GPU, you must install this version. Note that this version of
-    TensorFlow is typically much easier to install (typically,
-    in 5 or 10 minutes), so even if you have an NVIDIA GPU, we recommend
-    installing this version first.
-  * **TensorFlow with GPU support**. TensorFlow programs typically run
-    significantly faster on a GPU than on a CPU. Therefore, if your
-    system has a NVIDIA® GPU meeting the prerequisites shown below and you
-    need to run performance-critical applications, you should ultimately
-    install this version.
+* __TensorFlow with CPU support only__. If your system does not have a
+  NVIDIA®&nbsp;GPU, you must install this version. This version of TensorFlow is
+  usually easier to install, so even if you have an NVIDIA GPU, we recommend
+  installing this version first.
+* __TensorFlow with GPU support__. TensorFlow programs usually run much faster on
+  a GPU instead of a CPU. If you run performance-critical applications and your
+  system has an NVIDIA®&nbsp;GPU that meets the prerequisites, you should install
+  this version. See [TensorFlow GPU support](#NVIDIARequirements) for details.
 
-<a name="NVIDIARequirements"></a>
-### NVIDIA requirements to run TensorFlow with GPU support
-
-If you are installing TensorFlow with GPU support using one of the
-mechanisms described in this guide, then the following NVIDIA software
-must be installed on your system:
-
-  * [CUDA Toolkit 9.0](http://nvidia.com/cuda). For details, see
-    [NVIDIA's documentation](http://docs.nvidia.com/cuda/cuda-installation-guide-linux/).
-    Ensure that you append the relevant CUDA pathnames to the
-    `LD_LIBRARY_PATH` environment variable as described in the
-    NVIDIA documentation.
-  * [cuDNN SDK v7](http://developer.nvidia.com/cudnn). For details, see
-    [NVIDIA's documentation](http://docs.nvidia.com/deeplearning/sdk/cudnn-install/).
-    Ensure that you create the `CUDA_HOME` environment variable as
-    described in the NVIDIA documentation.
-  * GPU card with CUDA Compute Capability 3.0 or higher for building
-    from source and 3.5 or higher for our binaries. See
-    [NVIDIA documentation](https://developer.nvidia.com/cuda-gpus) for
-    a list of supported GPU cards.
-  * [GPU drivers](http://nvidia.com/driver) supporting your version of the CUDA
-    Toolkit.
-  * The libcupti-dev library, which is the NVIDIA CUDA Profile Tools Interface.
-    This library provides advanced profiling support. To install this library,
-    issue the following command for CUDA Toolkit >= 8.0:
-
-    <pre>
-    $ <b>sudo apt-get install cuda-command-line-tools</b>
-    </pre>
-
-    and add its path to your `LD_LIBRARY_PATH` environment variable:
-
-    <pre>
-    $ <b>export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:+${LD_LIBRARY_PATH}:}/usr/local/cuda/extras/CUPTI/lib64</b>
-    </pre>
-
-    For CUDA Toolkit <= 7.5 do:
-
-    <pre>
-    $ <b>sudo apt-get install libcupti-dev</b>
-    </pre>
-
-  * **[OPTIONAL]**  For optimized inferencing performance, you can also install
-    **NVIDIA TensorRT 3.0**. The minimal set of TensorRT runtime components needed
-    for use with the pre-built `tensorflow-gpu` package can be installed as follows:
-
-    <pre>
-    $ <b>wget https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1404/x86_64/nvinfer-runtime-trt-repo-ubuntu1404-3.0.4-ga-cuda9.0_1.0-1_amd64.deb</b>
-    $ <b>sudo dpkg -i nvinfer-runtime-trt-repo-ubuntu1404-3.0.4-ga-cuda9.0_1.0-1_amd64.deb</b>
-    $ <b>sudo apt-get update</b>
-    $ <b>sudo apt-get install -y --allow-downgrades libnvinfer-dev libcudnn7-dev=7.0.5.15-1+cuda9.0 libcudnn7=7.0.5.15-1+cuda9.0</b>
-    </pre>
-
-    **IMPORTANT:** For compatibility with the pre-built `tensorflow-gpu`
-    package, please use the Ubuntu **14.04** package of TensorRT as shown above,
-    even when installing onto an Ubuntu 16.04 system.<br/>
-    <br/>
-    To build the TensorFlow-TensorRT integration module from source rather than
-    using pre-built binaries, see the [module documentation](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/tensorrt#using-tensorrt-in-tensorflow).
-    For detailed TensorRT installation instructions, see [NVIDIA's TensorRT documentation](http://docs.nvidia.com/deeplearning/sdk/tensorrt-install-guide/index.html).<br/>
-    <br/>
-    To avoid cuDNN version conflicts during later system upgrades, you can hold
-    the cuDNN version at 7.0.5:
-
-    <pre>
-    $ <b> sudo apt-mark hold libcudnn7 libcudnn7-dev</b>
-    </pre>
-
-    To later allow upgrades, you can remove the hold:
-
-    <pre>
-    $ <b> sudo apt-mark unhold libcudnn7 libcudnn7-dev</b>
-    </pre>
-
-If you have an earlier version of the preceding packages, please upgrade to
-the specified versions. If upgrading is not possible, then you may still run
-TensorFlow with GPU support, if you @{$install_sources$install TensorFlow from Sources}.
-
-
-## Determine how to install TensorFlow
-
-You must pick the mechanism by which you install TensorFlow. The
-supported choices are as follows:
-
-  * [Virtualenv](#InstallingVirtualenv)
-  * ["native" pip](#InstallingNativePip)
-  * [Docker](#InstallingDocker)
-  * [Anaconda](#InstallingAnaconda)
-  * installing from sources, which is documented in
-    [a separate guide](https://www.tensorflow.org/install/install_sources).
-
-**We recommend the Virtualenv installation.**
-[Virtualenv](https://virtualenv.pypa.io/en/stable/)
-is a virtual Python environment isolated from other Python development,
-incapable of interfering with or being affected by other Python programs
-on the same machine.  During the Virtualenv installation process,
-you will install not only TensorFlow but also all the packages that
-TensorFlow requires.  (This is actually pretty easy.)
-To start working with TensorFlow, you simply need to "activate" the
-virtual environment.  All in all, Virtualenv provides a safe and
-reliable mechanism for installing and running TensorFlow.
-
-Native pip installs TensorFlow directly on your system without going
-through any container system. **We recommend the native pip install for
-system administrators aiming to make TensorFlow available to everyone on a
-multi-user system.** Since a native pip installation is not walled-off in
-a separate container, the pip installation might interfere with other
-Python-based installations on your system. However, if you understand pip
-and your Python environment, a native pip installation often entails only
-a single command.
-
-Docker completely isolates the TensorFlow installation
-from pre-existing packages on your machine. The Docker container contains
-TensorFlow and all its dependencies. Note that the Docker image can be quite
-large (hundreds of MBs). You might choose the Docker installation if you are
-incorporating TensorFlow into a larger application architecture that already
-uses Docker.
 
-In Anaconda, you may use conda to create a virtual environment.
-However, within Anaconda, we recommend installing TensorFlow with the
-`pip install` command, not with the `conda install` command.
+## How to install TensorFlow
 
-**NOTE:** The conda package is community supported, not officially supported.
-That is, the TensorFlow team neither tests nor maintains the conda package.
-Use that package at your own risk.
+There are a few options to install TensorFlow on your machine:
 
+* [Use pip in a virtual environment](#InstallingVirtualenv) *(recommended)*
+* [Use pip in your system environment](#InstallingNativePip)
+* [Configure a Docker container](#InstallingDocker)
+* [Use pip in Anaconda](#InstallingAnaconda)
+* [Install TensorFlow from source](/install/install_sources)
 
 <a name="InstallingVirtualenv"></a>
-## Installing with Virtualenv
+### Use `pip` in a virtual environment
 
-Take the following steps to install TensorFlow with Virtualenv:
+Key Point: Using a virtual environment is the recommended install method.
 
-  1. Install pip and Virtualenv by issuing one of the following commands:
+The [Virtualenv](https://virtualenv.pypa.io/en/stable/) tool creates virtual
+Python environments that are isolated from other Python development on the same
+machine. In this scenario, you install TensorFlow and its dependencies within a
+virtual environment that is available when *activated*. Virtualenv provides a
+reliable way to install and run TensorFlow while avoiding conflicts with the rest
+of the system.
 
-     <pre>$ <b>sudo apt-get install python-pip python-dev python-virtualenv</b> # for Python 2.7
-    $ <b>sudo apt-get install python3-pip python3-dev python-virtualenv</b> # for Python 3.n</pre>
+##### 1. Install Python, `pip`, and `virtualenv`.
 
-  2. Create a Virtualenv environment by issuing one of the following commands:
+On Ubuntu, Python is automatically installed and `pip` is *usually* installed.
+Confirm the `python` and `pip` versions:
 
-     <pre>$ <b>virtualenv --system-site-packages</b> <i>targetDirectory</i> # for Python 2.7
-    $ <b>virtualenv --system-site-packages -p python3</b> <i>targetDirectory</i> # for Python 3.n</pre>
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">python -V  # or: python3 -V</code>
+  <code class="devsite-terminal">pip -V     # or: pip3 -V</code>
+</pre>
+
+To install these packages on Ubuntu:
+
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">sudo apt-get install python-pip python-dev python-virtualenv   # for Python 2.7</code>
+  <code class="devsite-terminal">sudo apt-get install python3-pip python3-dev python-virtualenv # for Python 3.n</code>
+</pre>
 
-     where <code><em>targetDirectory</em></code> specifies the top of the
-     Virtualenv tree.  Our instructions assume that
-     <code><em>targetDirectory</em></code> is `~/tensorflow`, but you may
-     choose any directory.
+We *recommend* using `pip` version 8.1 or higher. If using a release before
+version 8.1,  upgrade `pip`:
 
-  3. Activate the Virtualenv environment by issuing one of the following
-     commands:
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">sudo pip install -U pip</code>
+</pre>
+
+If not using Ubuntu and [setuptools](https://pypi.org/project/setuptools/) is
+installed, use `easy_install` to install `pip`:
+
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">easy_install -U pip</code>
+</pre>
 
-     <pre>$ <b>source ~/tensorflow/bin/activate</b> # bash, sh, ksh, or zsh
-    $ <b>source ~/tensorflow/bin/activate.csh</b>  # csh or tcsh
-    $ <b>. ~/tensorflow/bin/activate.fish</b>  # fish</pre>
+##### 2. Create a directory for the virtual environment and choose a Python interpreter.
 
-     The preceding <tt>source</tt> command should change your prompt
-     to the following:
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">mkdir ~/tensorflow  # somewhere to work out of</code>
+  <code class="devsite-terminal">cd ~/tensorflow</code>
+  <code># Choose one of the following Python environments for the ./venv directory:</code>
+  <code class="devsite-terminal">virtualenv --system-site-packages <var>venv</var>            # Use python default (Python 2.7)</code>
+  <code class="devsite-terminal">virtualenv --system-site-packages -p python3 <var>venv</var> # Use Python 3.n</code>
+</pre>
 
-     <pre>(tensorflow)$ </pre>
+##### 3. Activate the Virtualenv environment.
 
-  4. Ensure pip ≥8.1 is installed:
+Use one of these shell-specific commands to activate the virtual environment:
 
-     <pre>(tensorflow)$ <b>easy_install -U pip</b></pre>
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">source ~/tensorflow/<var>venv</var>/bin/activate      # bash, sh, ksh, or zsh</code>
+  <code class="devsite-terminal">source ~/tensorflow/<var>venv</var>/bin/activate.csh  # csh or tcsh</code>
+  <code class="devsite-terminal">. ~/tensorflow/<var>venv</var>/bin/activate.fish      # fish</code>
+</pre>
 
-  5. Issue one of the following commands to install TensorFlow in the active
-     Virtualenv environment:
+When the Virtualenv is activated, the shell prompt displays as `(venv) $`.
 
-     <pre>(tensorflow)$ <b>pip install --upgrade tensorflow</b>      # for Python 2.7
-    (tensorflow)$ <b>pip3 install --upgrade tensorflow</b>     # for Python 3.n
-    (tensorflow)$ <b>pip install --upgrade tensorflow-gpu</b>  # for Python 2.7 and GPU
-    (tensorflow)$ <b>pip3 install --upgrade tensorflow-gpu</b> # for Python 3.n and GPU</pre>
+##### 4. Upgrade `pip` in the virtual environment.
 
-     If the above command succeeds, skip Step 6. If the preceding
-     command fails, perform Step 6.
+Within the active virtual environment, upgrade `pip`:
 
-  6. (Optional) If Step 5 failed (typically because you invoked a pip version
-     lower than 8.1), install TensorFlow in the active Virtualenv environment
-     by issuing a command of the following format:
+<pre class="prettyprint lang-bsh">
+(venv)$ pip install -U pip
+</pre>
 
-     <pre>(tensorflow)$ <b>pip install --upgrade</b> <i>tfBinaryURL</i>   # Python 2.7
-    (tensorflow)$ <b>pip3 install --upgrade</b> <i>tfBinaryURL</i>  # Python 3.n </pre>
+You can install other Python packages within the virtual environment without
+affecting packages outside the `virtualenv`.
 
-     where <code><em>tfBinaryURL</em></code> identifies the URL of the
-     TensorFlow Python package. The appropriate value of
-     <code><em>tfBinaryURL</em></code>depends on the operating system,
-     Python version, and GPU support. Find the appropriate value for
-     <code><em>tfBinaryURL</em></code> for your system
-     [here](#the_url_of_the_tensorflow_python_package).  For example, if you
-     are installing TensorFlow for Linux, Python 3.4, and CPU-only support,
-     issue the following command to install TensorFlow in the active
-     Virtualenv environment:
+##### 5. Install TensorFlow in the virtual environment.
 
-     <pre>(tensorflow)$ <b>pip3 install --upgrade \
-     https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc0-cp34-cp34m-linux_x86_64.whl</b></pre>
+Choose one of the available TensorFlow packages for installation:
 
-If you encounter installation problems, see
-[Common Installation Problems](#common_installation_problems).
+* `tensorflow` —Current release for CPU
+* `tensorflow-gpu` —Current release with GPU support
+* `tf-nightly` —Nightly build for CPU
+* `tf-nightly-gpu` —Nightly build with GPU support
 
+Within an active Virtualenv environment, use `pip` to install the package:
 
-### Next Steps
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">pip install -U tensorflow</code>
+</pre>
 
-After installing TensorFlow,
-[validate the installation](#ValidateYourInstallation).
+Use `pip list` to show the packages installed in the virtual environment.
+[Validate the install](#ValidateYourInstallation) and test the version:
 
-Note that you must activate the Virtualenv environment each time you
-use TensorFlow. If the Virtualenv environment is not currently active,
-invoke one of the following commands:
+<pre class="prettyprint lang-bsh">
+(venv)$ python -c "import tensorflow as tf; print(tf.__version__)"
+</pre>
 
-<pre> $ <b>source ~/tensorflow/bin/activate</b>      # bash, sh, ksh, or zsh
-$ <b>source ~/tensorflow/bin/activate.csh</b>  # csh or tcsh</pre>
+Success: TensorFlow is now installed.
 
-When the Virtualenv environment is active, you may run
-TensorFlow programs from this shell.  Your prompt will become
-the following to indicate that your tensorflow environment is active:
+Use the `deactivate` command to stop the Python virtual environment.
 
-<pre>(tensorflow)$ </pre>
+#### Problems
 
-When you are done using TensorFlow, you may deactivate the
-environment by invoking the `deactivate` function as follows:
+If the above steps failed, try installing the TensorFlow binary using the remote
+URL of the `pip` package:
 
-<pre>(tensorflow)$ <b>deactivate</b> </pre>
+<pre class="prettyprint lang-bsh">
+(venv)$ pip install --upgrade <var>remote-pkg-URL</var>   # Python 2.7
+(venv)$ pip3 install --upgrade <var>remote-pkg-URL</var>  # Python 3.n
+</pre>
 
-The prompt will revert back to your default prompt (as defined by the
-`PS1` environment variable).
+The <var>remote-pkg-URL</var> depends on the operating system, Python version,
+and GPU support. See [here](#the_url_of_the_tensorflow_python_package) for the
+URL naming scheme and location.
 
+See [Common Installation Problems](#common_installation_problems) if you
+encounter problems.
 
-### Uninstalling TensorFlow
+#### Uninstall TensorFlow
 
-To uninstall TensorFlow, simply remove the tree you created.
-For example:
+To uninstall TensorFlow, remove the Virtualenv directory you created in step 2:
 
-<pre>$ <b>rm -r</b> <i>targetDirectory</i> </pre>
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">deactivate  # stop the virtualenv</code>
+  <code class="devsite-terminal">rm -r ~/tensorflow/<var>venv</var></code>
+</pre>
 
 
 <a name="InstallingNativePip"></a>
-## Installing with native pip
+### Use `pip` in your system environment
+
+Use `pip` to install the TensorFlow package directly on your system without
+using a container or virtual environment for isolation. This method is
+recommended for system administrators that want a TensorFlow installation that is
+available to everyone on a multi-user system.
 
-You may install TensorFlow through pip, choosing between a simple
-installation procedure or a more complex one.
+Since a system install is not isolated, it could interfere with other
+Python-based installations. But if you understand `pip` and your Python
+environment, a system `pip` install is straightforward.
 
-**Note:** The
+See the
 [REQUIRED_PACKAGES section of setup.py](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/pip_package/setup.py)
-lists the TensorFlow packages that pip will install or upgrade.
+for a list of packages that TensorFlow installs.
 
+##### 1. Install Python, `pip`, and `virtualenv`.
 
-### Prerequisite: Python and Pip
+On Ubuntu, Python is automatically installed and `pip` is *usually* installed.
+Confirm the `python` and `pip` versions:
 
-Python is automatically installed on Ubuntu.  Take a moment to confirm
-(by issuing a `python -V` command) that one of the following Python
-versions is already installed on your system:
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">python -V  # or: python3 -V</code>
+  <code class="devsite-terminal">pip -V     # or: pip3 -V</code>
+</pre>
+
+To install these packages on Ubuntu:
 
-  * Python 2.7
-  * Python 3.4+
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">sudo apt-get install python-pip python-dev   # for Python 2.7</code>
+  <code class="devsite-terminal">sudo apt-get install python3-pip python3-dev # for Python 3.n</code>
+</pre>
 
-The pip or pip3 package manager is *usually* installed on Ubuntu.  Take a
-moment to confirm (by issuing a `pip -V` or `pip3 -V` command)
-that pip or pip3 is installed.  We strongly recommend version 8.1 or higher
-of pip or pip3.  If Version 8.1 or later is not installed, issue the
-following command, which will either install or upgrade to the latest
-pip version:
+We *recommend* using `pip` version 8.1 or higher. If using a release before
+version 8.1,  upgrade `pip`:
 
-<pre>$ <b>sudo apt-get install python-pip python-dev</b>   # for Python 2.7
-$ <b>sudo apt-get install python3-pip python3-dev</b> # for Python 3.n
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">sudo pip install -U pip</code>
 </pre>
 
+If not using Ubuntu and [setuptools](https://pypi.org/project/setuptools/) is
+installed, use `easy_install` to install `pip`:
 
-### Install TensorFlow
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">easy_install -U pip</code>
+</pre>
 
-Assuming the prerequisite software is installed on your Linux host,
-take the following steps:
+##### 2. Install TensorFlow on system.
 
-  1. Install TensorFlow by invoking **one** of the following commands:
+Choose one of the available TensorFlow packages for installation:
 
-     <pre>$ <b>pip install tensorflow</b>      # Python 2.7; CPU support (no GPU support)
-    $ <b>pip3 install tensorflow</b>     # Python 3.n; CPU support (no GPU support)
-    $ <b>pip install tensorflow-gpu</b>  # Python 2.7;  GPU support
-    $ <b>pip3 install tensorflow-gpu</b> # Python 3.n; GPU support </pre>
+* `tensorflow` —Current release for CPU
+* `tensorflow-gpu` —Current release with GPU support
+* `tf-nightly` —Nightly build for CPU
+* `tf-nightly-gpu` —Nightly build with GPU support
 
-     If the preceding command runs to completion, you should now
-     [validate your installation](#ValidateYourInstallation).
+And use `pip` to install the package for Python 2 or 3:
 
-  2. (Optional.) If Step 1 failed, install the latest version of TensorFlow
-     by issuing a command of the following format:
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">sudo pip install -U tensorflow   # Python 2.7</code>
+  <code class="devsite-terminal">sudo pip3 install -U tensorflow  # Python 3.n</code>
+</pre>
 
-     <pre>$ <b>sudo pip  install --upgrade</b> <i>tfBinaryURL</i>   # Python 2.7
-    $ <b>sudo pip3 install --upgrade</b> <i>tfBinaryURL</i>   # Python 3.n </pre>
+Use `pip list` to show the packages installed on the system.
+[Validate the install](#ValidateYourInstallation) and test the version:
 
-     where <code><em>tfBinaryURL</em></code> identifies the URL of the
-     TensorFlow Python package. The appropriate value of
-     <code><em>tfBinaryURL</em></code> depends on the operating system,
-     Python version, and GPU support. Find the appropriate value for
-     <code><em>tfBinaryURL</em></code>
-     [here](#the_url_of_the_tensorflow_python_package).  For example, to
-     install TensorFlow for Linux, Python 3.4, and CPU-only support, issue
-     the following command:
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">python -c "import tensorflow as tf; print(tf.__version__)"</code>
+</pre>
 
-     <pre>
-     $ <b>sudo pip3 install --upgrade \
-     https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc0-cp34-cp34m-linux_x86_64.whl</b>
-     </pre>
+Success: TensorFlow is now installed.
 
-     If this step fails, see
-     [Common Installation Problems](#common_installation_problems).
+#### Problems
 
+If the above steps failed, try installing the TensorFlow binary using the remote
+URL of the `pip` package:
 
-### Next Steps
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">sudo pip install --upgrade <var>remote-pkg-URL</var>   # Python 2.7</code>
+  <code class="devsite-terminal">sudo pip3 install --upgrade <var>remote-pkg-URL</var>  # Python 3.n</code>
+</pre>
 
-After installing TensorFlow, [validate your installation](#ValidateYourInstallation).
+The <var>remote-pkg-URL</var> depends on the operating system, Python version,
+and GPU support. See [here](#the_url_of_the_tensorflow_python_package) for the
+URL naming scheme and location.
 
+See [Common Installation Problems](#common_installation_problems) if you
+encounter problems.
 
-### Uninstalling TensorFlow
+#### Uninstall TensorFlow
 
-To uninstall TensorFlow, issue one of following commands:
+To uninstall TensorFlow on your system, use one of following commands:
 
-<pre>
-$ <b>sudo pip uninstall tensorflow</b>  # for Python 2.7
-$ <b>sudo pip3 uninstall tensorflow</b> # for Python 3.n
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">sudo pip uninstall tensorflow   # for Python 2.7</code>
+  <code class="devsite-terminal">sudo pip3 uninstall tensorflow  # for Python 3.n</code>
 </pre>
 
-
 <a name="InstallingDocker"></a>
-## Installing with Docker
+### Configure a Docker container
+
+Docker completely isolates the TensorFlow installation
+from pre-existing packages on your machine. The Docker container contains
+TensorFlow and all its dependencies. Note that the Docker image can be quite
+large (hundreds of MBs). You might choose the Docker installation if you are
+incorporating TensorFlow into a larger application architecture that already
+uses Docker.
 
 Take the following steps to install TensorFlow through Docker:
 
@@ -364,7 +287,7 @@ Take the following steps to install TensorFlow through Docker:
 The remainder of this section explains how to launch a Docker container.
 
 
-### CPU-only
+#### CPU-only
 
 To launch a Docker container with CPU-only support (that is, without
 GPU support), enter a command of the following format:
@@ -414,7 +337,7 @@ $ <b>docker run -it -p 8888:8888 tensorflow/tensorflow</b>
 Docker will download the TensorFlow binary image the first time you launch it.
 
 
-### GPU support
+#### GPU support
 
 Prior to installing TensorFlow with GPU support, ensure that your system meets all
 [NVIDIA software requirements](#NVIDIARequirements).  To launch a Docker container
@@ -470,14 +393,22 @@ For more details see the
 [TensorFlow docker readme](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/tools/docker).
 
 
-### Next Steps
+#### Next Steps
 
 You should now
 [validate your installation](#ValidateYourInstallation).
 
 
 <a name="InstallingAnaconda"></a>
-## Installing with Anaconda
+### Use `pip` in Anaconda
+
+Anaconda provides the `conda` utility to create a virtual environment. However,
+within Anaconda, we recommend installing TensorFlow using the `pip install`
+command and *not* with the `conda install` command.
+
+Caution: `conda` is a community supported package this is not officially
+maintained by the TensorFlow team. Use this package at your own risk since it is
+not tested on new TensorFlow releases.
 
 Take the following steps to install TensorFlow in an Anaconda environment:
 
@@ -507,7 +438,7 @@ Take the following steps to install TensorFlow in an Anaconda environment:
 
      <pre>
      (tensorflow)$ <b>pip install --ignore-installed --upgrade \
-     https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc0-cp34-cp34m-linux_x86_64.whl</b></pre>
+     https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc1-cp34-cp34m-linux_x86_64.whl</b></pre>
 
 <a name="ValidateYourInstallation"></a>
 ## Validate your installation
@@ -563,11 +494,89 @@ installation problems](#common_installation_problems).
 If you are new to machine learning, we recommend the following:
 
 *  [Machine Learning Crash Course](https://developers.google.com/machine-learning/crash-course)
-*  @{$get_started/get_started_for_beginners$Getting Started for ML Beginners}
+*  @{$get_started/eager}
 
 If you are experienced with machine learning but new to TensorFlow, see
 @{$get_started/eager}.
 
+<a name="NVIDIARequirements"></a>
+## TensorFlow GPU support
+
+To install TensorFlow with GPU support, configure the following NVIDIA® software
+on your system:
+
+* [CUDA Toolkit 9.0](http://nvidia.com/cuda). For details, see
+  [NVIDIA's documentation](http://docs.nvidia.com/cuda/cuda-installation-guide-linux/).
+  Append the relevant CUDA pathnames to the `LD_LIBRARY_PATH` environmental
+  variable as described in the NVIDIA documentation.
+* [cuDNN SDK v7](http://developer.nvidia.com/cudnn). For details, see
+  [NVIDIA's documentation](http://docs.nvidia.com/deeplearning/sdk/cudnn-install/).
+  Create the `CUDA_HOME` environment variable as described in the NVIDIA
+  documentation.
+* A GPU card with CUDA Compute Capability 3.0 or higher for building TensorFlow
+  from source. To use the TensorFlow binaries, version 3.5 or higher is required.
+  See the [NVIDIA documentation](https://developer.nvidia.com/cuda-gpus) for a
+  list of supported GPU cards.
+* [GPU drivers](http://nvidia.com/driver) that support your version of the CUDA
+  Toolkit.
+* The `libcupti-dev` library is the NVIDIA CUDA Profile Tools Interface. This
+  library provides advanced profiling support. To install this library,
+  use the following command for CUDA Toolkit >= 8.0:
+
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">sudo apt-get install cuda-command-line-tools</code>
+</pre>
+
+Add this path to the `LD_LIBRARY_PATH` environmental variable:
+
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:+${LD_LIBRARY_PATH}:}/usr/local/cuda/extras/CUPTI/lib64</code>
+</pre>
+
+For CUDA Toolkit <= 7.5 use:
+
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">sudo apt-get install libcupti-dev</code>
+</pre>
+
+* *OPTIONAL*:  For optimized performance during inference, install
+  *NVIDIA&nbsp;TensorRT&nbsp;3.0*. To install the minimal amount of TensorRT
+  runtime components required to use with the pre-built `tensorflow-gpu` package:
+
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">wget https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1404/x86_64/nvinfer-runtime-trt-repo-ubuntu1404-3.0.4-ga-cuda9.0_1.0-1_amd64.deb</code>
+  <code class="devsite-terminal">sudo dpkg -i nvinfer-runtime-trt-repo-ubuntu1404-3.0.4-ga-cuda9.0_1.0-1_amd64.deb</code>
+  <code class="devsite-terminal">sudo apt-get update</code>
+  <code class="devsite-terminal">sudo apt-get install -y --allow-downgrades libnvinfer-dev libcudnn7-dev=7.0.5.15-1+cuda9.0 libcudnn7=7.0.5.15-1+cuda9.0</code>
+</pre>
+
+Note: For compatibility with the pre-built `tensorflow-gpu` package, use the
+Ubuntu *14.04* package of TensorRT (shown above). Use this even when installing
+on an Ubuntu 16.04 system.
+
+To build the TensorFlow-TensorRT integration module from source instead of using
+the pre-built binaries, see the
+[module documentation](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/tensorrt#using-tensorrt-in-tensorflow).
+For detailed TensorRT installation instructions, see
+[NVIDIA's TensorRT documentation](http://docs.nvidia.com/deeplearning/sdk/tensorrt-install-guide/index.html).
+
+To avoid cuDNN version conflicts during later system upgrades, hold the cuDNN
+version at 7.0.5:
+
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">sudo apt-mark hold libcudnn7 libcudnn7-dev</code>
+</pre>
+
+To allow upgrades, remove the this hold:
+
+<pre class="prettyprint lang-bsh">
+  <code class="devsite-terminal">sudo apt-mark unhold libcudnn7 libcudnn7-dev</code>
+</pre>
+
+If you have an earlier version of the preceding packages, upgrade to the
+specified versions. If upgrading is not possible, you can still run TensorFlow
+with GPU support by @{$install_sources}.
+
 
 ## Common installation problems
 
@@ -581,7 +590,7 @@ ask a new question about it on Stack Overflow and specify
 the `tensorflow` tag.
 
 <table>
-<tr> <th>Stack Overflow Link</th> <th>Error Message</th> </tr>
+<tr> <th>Link to GitHub or Stack&nbsp;Overflow</th> <th>Error Message</th> </tr>
 
 <tr>
   <td><a href="https://stackoverflow.com/q/36159194">36159194</a></td>
@@ -681,14 +690,14 @@ This section documents the relevant values for Linux installations.
 CPU only:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc0-cp27-none-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc1-cp27-none-linux_x86_64.whl
 </pre>
 
 
 GPU support:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0rc0-cp27-none-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0rc1-cp27-none-linux_x86_64.whl
 </pre>
 
 Note that GPU support requires the NVIDIA hardware and software described in
@@ -700,14 +709,14 @@ Note that GPU support requires the NVIDIA hardware and software described in
 CPU only:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc0-cp34-cp34m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc1-cp34-cp34m-linux_x86_64.whl
 </pre>
 
 
 GPU support:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0rc0-cp34-cp34m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0rc1-cp34-cp34m-linux_x86_64.whl
 </pre>
 
 Note that GPU support requires the NVIDIA hardware and software described in
@@ -719,14 +728,14 @@ Note that GPU support requires the NVIDIA hardware and software described in
 CPU only:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc0-cp35-cp35m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc1-cp35-cp35m-linux_x86_64.whl
 </pre>
 
 
 GPU support:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0rc0-cp35-cp35m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0rc1-cp35-cp35m-linux_x86_64.whl
 </pre>
 
 
@@ -738,14 +747,14 @@ Note that GPU support requires the NVIDIA hardware and software described in
 CPU only:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc0-cp36-cp36m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.8.0rc1-cp36-cp36m-linux_x86_64.whl
 </pre>
 
 
 GPU support:
 
 <pre>
-https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0rc0-cp36-cp36m-linux_x86_64.whl
+https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow_gpu-1.8.0rc1-cp36-cp36m-linux_x86_64.whl
 </pre>
 
 
diff --git a/tensorflow/docs_src/install/install_mac.md b/tensorflow/docs_src/install/install_mac.md
index af24aaaca8..90d9ea0288 100644
--- a/tensorflow/docs_src/install/install_mac.md
+++ b/tensorflow/docs_src/install/install_mac.md
@@ -119,7 +119,7 @@ Take the following steps to install TensorFlow with Virtualenv:
      TensorFlow in the active Virtualenv is as follows:
 
      <pre> $ <b>pip3 install --upgrade \
-     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc0-py3-none-any.whl</b></pre>
+     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc1-py3-none-any.whl</b></pre>
 
 If you encounter installation problems, see
 [Common Installation Problems](#common-installation-problems).
@@ -242,7 +242,7 @@ take the following steps:
      issue the following command:
 
      <pre> $ <b>sudo pip3 install --upgrade \
-     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc0-py3-none-any.whl</b> </pre>
+     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc1-py3-none-any.whl</b> </pre>
 
      If the preceding command fails, see
      [installation problems](#common-installation-problems).
@@ -350,7 +350,7 @@ Take the following steps to install TensorFlow in an Anaconda environment:
      TensorFlow for Python 2.7:
 
      <pre> (<i>targetDirectory</i>)$ <b>pip install --ignore-installed --upgrade \
-     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc0-py2-none-any.whl</b></pre>
+     https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc1-py2-none-any.whl</b></pre>
 
 
 <a name="ValidateYourInstallation"></a>
@@ -524,7 +524,7 @@ The value you specify depends on your Python version.
 
 
 <pre>
-https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc0-py2-none-any.whl
+https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc1-py2-none-any.whl
 </pre>
 
 
@@ -532,5 +532,5 @@ https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc0-py2-none-a
 
 
 <pre>
-https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc0-py3-none-any.whl
+https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.8.0rc1-py3-none-any.whl
 </pre>
diff --git a/tensorflow/docs_src/install/install_sources.md b/tensorflow/docs_src/install/install_sources.md
index 649c5b4751..a4fec382f4 100644
--- a/tensorflow/docs_src/install/install_sources.md
+++ b/tensorflow/docs_src/install/install_sources.md
@@ -354,10 +354,10 @@ Invoke `pip install` to install that pip package.
 The filename of the `.whl` file depends on your platform.
 For example, the following command will install the pip package
 
-for TensorFlow 1.8.0rc0 on Linux:
+for TensorFlow 1.8.0rc1 on Linux:
 
 <pre>
-$ <b>sudo pip install /tmp/tensorflow_pkg/tensorflow-1.8.0rc0-py2-none-any.whl</b>
+$ <b>sudo pip install /tmp/tensorflow_pkg/tensorflow-1.8.0rc1-py2-none-any.whl</b>
 </pre>
 
 ## Validate your installation
diff --git a/tensorflow/docs_src/performance/xla/tfcompile.md b/tensorflow/docs_src/performance/xla/tfcompile.md
index f57ca3948d..8521d7eacb 100644
--- a/tensorflow/docs_src/performance/xla/tfcompile.md
+++ b/tensorflow/docs_src/performance/xla/tfcompile.md
@@ -86,7 +86,7 @@ code. `tf_library` utilizes `tfcompile` to compile the TensorFlow graph into
 executable code.
 
 ```build
-load("//third_party/tensorflow/compiler/aot:tfcompile.bzl", "tf_library")
+load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")
 
 # Use the tf_library macro to compile your graph into executable code.
 tf_library(
@@ -258,8 +258,8 @@ file.
 
 ```build
 # Example of linking your binary
-# Also see //third_party/tensorflow/compiler/aot/tests/BUILD
-load("//third_party/tensorflow/compiler/aot:tfcompile.bzl", "tf_library")
+# Also see //tensorflow/compiler/aot/tests/BUILD
+load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")
 
 # The same tf_library call from step 2 above.
 tf_library(
diff --git a/tensorflow/examples/tutorials/estimators/__init__.py b/tensorflow/examples/tutorials/estimators/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tensorflow/examples/tutorials/input_fn/__init__.py b/tensorflow/examples/tutorials/input_fn/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tensorflow/examples/tutorials/layers/__init__.py b/tensorflow/examples/tutorials/layers/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tensorflow/examples/tutorials/monitors/__init__.py b/tensorflow/examples/tutorials/monitors/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tensorflow/examples/tutorials/monitors/iris_monitors.py b/tensorflow/examples/tutorials/monitors/iris_monitors.py
index 850d105f7b..a2b7fe6023 100644
--- a/tensorflow/examples/tutorials/monitors/iris_monitors.py
+++ b/tensorflow/examples/tutorials/monitors/iris_monitors.py
@@ -32,9 +32,9 @@ IRIS_TEST = os.path.join(os.path.dirname(__file__), "iris_test.csv")
 def main(unused_argv):
   # Load datasets.
   training_set = tf.contrib.learn.datasets.base.load_csv_with_header(
-      filename=IRIS_TRAINING, target_dtype=np.int, features_dtype=np.float)
+      filename=IRIS_TRAINING, target_dtype=np.int, features_dtype=np.float32)
   test_set = tf.contrib.learn.datasets.base.load_csv_with_header(
-      filename=IRIS_TEST, target_dtype=np.int, features_dtype=np.float)
+      filename=IRIS_TEST, target_dtype=np.int, features_dtype=np.float32)
 
   validation_metrics = {
       "accuracy":
@@ -83,7 +83,7 @@ def main(unused_argv):
 
   # Classify two new flower samples.
   new_samples = np.array(
-      [[6.4, 3.2, 4.5, 1.5], [5.8, 3.1, 5.0, 1.7]], dtype=float)
+      [[6.4, 3.2, 4.5, 1.5], [5.8, 3.1, 5.0, 1.7]], dtype=np.float32)
   y = list(classifier.predict(new_samples))
   print("Predictions: {}".format(str(y)))
 
diff --git a/tensorflow/go/README.md b/tensorflow/go/README.md
index b1bd87eb0c..e251356ec8 100644
--- a/tensorflow/go/README.md
+++ b/tensorflow/go/README.md
@@ -5,7 +5,7 @@ Construct and execute TensorFlow graphs in Go.
 [![GoDoc](https://godoc.org/github.com/tensorflow/tensorflow/tensorflow/go?status.svg)](https://godoc.org/github.com/tensorflow/tensorflow/tensorflow/go)
 
 > *WARNING*: The API defined in this package is not stable and can change
-> without notice. The same goes for the awkward package path
+> without notice. The same goes for the package path:
 > (`github.com/tensorflow/tensorflow/tensorflow/go`).
 
 ## Quickstart
diff --git a/tensorflow/go/op/wrappers.go b/tensorflow/go/op/wrappers.go
index c12ea51563..2f1be51ada 100644
--- a/tensorflow/go/op/wrappers.go
+++ b/tensorflow/go/op/wrappers.go
@@ -21386,7 +21386,7 @@ func ImageSummaryBadColor(value tf.Tensor) ImageSummaryAttr {
 //    generated sequentially as '*tag*/image/0', '*tag*/image/1', etc.
 //
 // The `bad_color` argument is the color to use in the generated images for
-// non-finite input values.  It is a `unit8` 1-D tensor of length `channels`.
+// non-finite input values.  It is a `uint8` 1-D tensor of length `channels`.
 // Each element must be in the range `[0, 255]` (It represents the value of a
 // pixel in the output image).  Non-finite values in the input tensor are
 // replaced by this tensor in the output image.  The default value is the color
diff --git a/tensorflow/python/estimator/estimator.py b/tensorflow/python/estimator/estimator.py
index 63099b44bb..0970f00124 100644
--- a/tensorflow/python/estimator/estimator.py
+++ b/tensorflow/python/estimator/estimator.py
@@ -644,11 +644,9 @@ class Estimator(object):
               sharded=True)
           saver_for_restore.restore(session, checkpoint_path)
 
-          # pylint: disable=protected-access
           local_init_op = (
               estimator_spec.scaffold.local_init_op or
-              monitored_session.Scaffold._default_local_init_op())
-          # pylint: enable=protected-access
+              monitored_session.Scaffold.default_local_init_op())
 
           # Perform the export
           builder = saved_model_builder.SavedModelBuilder(temp_export_dir)
diff --git a/tensorflow/python/keras/_impl/keras/estimator.py b/tensorflow/python/keras/_impl/keras/estimator.py
index b922a6c683..c3c3fceb45 100644
--- a/tensorflow/python/keras/_impl/keras/estimator.py
+++ b/tensorflow/python/keras/_impl/keras/estimator.py
@@ -29,12 +29,14 @@ from tensorflow.python.estimator import run_config as run_config_lib
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import random_seed
 from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib
+from tensorflow.python.framework import tensor_util
 from tensorflow.python.keras._impl.keras import backend as K
 from tensorflow.python.keras._impl.keras import models
 from tensorflow.python.keras._impl.keras import optimizers
 from tensorflow.python.keras._impl.keras.engine.base_layer import Layer
 from tensorflow.python.keras._impl.keras.engine.network import Network
 from tensorflow.python.keras._impl.keras.utils.generic_utils import CustomObjectScope
+from tensorflow.python.ops import check_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import metrics as metrics_module
 from tensorflow.python.ops import variables as variables_module
@@ -55,6 +57,17 @@ def _cast_tensor_to_floatx(x):
     return math_ops.cast(x, K.floatx())
 
 
+def _convert_tensor(x):
+  """Create or cast tensor if needed."""
+  if not tensor_util.is_tensor(x):
+    # x is a numpy array
+    x = sparse_tensor_lib.convert_to_tensor_or_sparse_tensor(x)
+  if check_ops.is_numeric_tensor(x):
+    # is_numeric_tensor returns False if provided with a numpy array
+    x = _cast_tensor_to_floatx(x)
+  return x
+
+
 def _any_variable_initalized():
   """Check if any variable has been initialized in the Keras model.
 
@@ -86,7 +99,7 @@ def _create_ordered_io(keras_model, estimator_io, is_input=True):
   if isinstance(estimator_io, (list, tuple)):
     # Case currently not supported by most built-in input_fn,
     # but it's good to have for sanity
-    return [_cast_tensor_to_floatx(x) for x in estimator_io]
+    return [_convert_tensor(x) for x in estimator_io]
   elif isinstance(estimator_io, dict):
     if is_input:
       if keras_model._is_graph_network:
@@ -108,12 +121,12 @@ def _create_ordered_io(keras_model, estimator_io, is_input=True):
             'It needs to match one '
             'of the following: %s' % ('input' if is_input else 'output', key,
                                       ', '.join(keras_io_names)))
-      tensors = [_cast_tensor_to_floatx(estimator_io[io_name])
+      tensors = [_convert_tensor(estimator_io[io_name])
                  for io_name in keras_io_names]
     return tensors
   else:
     # Plain array.
-    return _cast_tensor_to_floatx(estimator_io)
+    return _convert_tensor(estimator_io)
 
 
 def _in_place_subclassed_model_reset(model):
@@ -274,8 +287,7 @@ def _clone_and_build_model(mode,
                                         is_input=False)
   else:
     target_tensors = [
-        _cast_tensor_to_floatx(
-            sparse_tensor_lib.convert_to_tensor_or_sparse_tensor(labels))
+        _convert_tensor(labels)
     ]
 
   if keras_model._is_graph_network:
diff --git a/tensorflow/python/keras/_impl/keras/estimator_test.py b/tensorflow/python/keras/_impl/keras/estimator_test.py
index 653cdc01e2..80fa87d041 100644
--- a/tensorflow/python/keras/_impl/keras/estimator_test.py
+++ b/tensorflow/python/keras/_impl/keras/estimator_test.py
@@ -30,6 +30,7 @@ from tensorflow.python.estimator.inputs import numpy_io
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.keras._impl import keras
+from tensorflow.python.keras._impl.keras import backend as K
 from tensorflow.python.keras._impl.keras import testing_utils
 from tensorflow.python.keras._impl.keras.applications import mobilenet
 from tensorflow.python.keras._impl.keras.optimizers import SGD
@@ -142,16 +143,20 @@ def randomize_io_type(array, name):
 
 
 def multi_inputs_multi_outputs_model():
-  # test multi-input layer
   a = keras.layers.Input(shape=(16,), name='input_a')
   b = keras.layers.Input(shape=(16,), name='input_b')
+  m = keras.layers.Input(shape=(8,), dtype='bool', name='input_m')
   dense = keras.layers.Dense(8, name='dense_1')
+
   a_2 = dense(a)
+  # Apply a mask
+  s_2 = keras.layers.Lambda(lambda k:
+                            K.switch(k[0], k[1], K.zeros_like(k[1])))([m, a_2])
   b_2 = dense(b)
-  merged = keras.layers.concatenate([a_2, b_2], name='merge')
+  merged = keras.layers.concatenate([s_2, b_2], name='merge')
   c = keras.layers.Dense(3, activation='softmax', name='dense_2')(merged)
   d = keras.layers.Dense(2, activation='softmax', name='dense_3')(merged)
-  model = keras.models.Model(inputs=[a, b], outputs=[c, d])
+  model = keras.models.Model(inputs=[a, b, m], outputs=[c, d])
   model.compile(
       loss='categorical_crossentropy',
       optimizer='rmsprop',
@@ -352,18 +357,27 @@ class TestKerasEstimator(test_util.TensorFlowTestCase):
         test_samples=50,
         input_shape=(16,),
         num_classes=2)
+    np.random.seed(_RANDOM_SEED)
+    (input_m_train, _), (input_m_test, _) = testing_utils.get_test_data(
+        train_samples=_TRAIN_SIZE,
+        test_samples=50,
+        input_shape=(8,),
+        num_classes=2)
+
     c_train = keras.utils.to_categorical(c_train)
     c_test = keras.utils.to_categorical(c_test)
     d_train = keras.utils.to_categorical(d_train)
     d_test = keras.utils.to_categorical(d_test)
 
     def train_input_fn():
-      input_dict = {'input_a': a_train, 'input_b': b_train}
+      input_dict = {'input_a': a_train, 'input_b': b_train,
+                    'input_m': input_m_train > 0}
       output_dict = {'dense_2': c_train, 'dense_3': d_train}
       return input_dict, output_dict
 
     def eval_input_fn():
-      input_dict = {'input_a': a_test, 'input_b': b_test}
+      input_dict = {'input_a': a_test, 'input_b': b_test,
+                    'input_m': input_m_test > 0}
       output_dict = {'dense_2': c_test, 'dense_3': d_test}
       return input_dict, output_dict
 
diff --git a/tensorflow/python/kernel_tests/division_past_test.py b/tensorflow/python/kernel_tests/division_past_test.py
index 2ff2f89407..9ddd62e63c 100644
--- a/tensorflow/python/kernel_tests/division_past_test.py
+++ b/tensorflow/python/kernel_tests/division_past_test.py
@@ -35,8 +35,7 @@ class DivisionTestCase(test.TestCase):
     """Test all the different ways to divide."""
     values = [1, 2, 7, 11]
     functions = (lambda x: x), constant_op.constant
-    # TODO(irving): Test int8, int16 once we support casts for those.
-    dtypes = np.int32, np.int64, np.float32, np.float64
+    dtypes = np.int8, np.int16, np.int32, np.int64, np.float32, np.float64
 
     tensors = []
     checks = []
diff --git a/tensorflow/python/kernel_tests/reduce_join_op_test.py b/tensorflow/python/kernel_tests/reduce_join_op_test.py
index 7f3049b9f8..fb9e5cc2a3 100644
--- a/tensorflow/python/kernel_tests/reduce_join_op_test.py
+++ b/tensorflow/python/kernel_tests/reduce_join_op_test.py
@@ -160,7 +160,7 @@ class ReduceJoinTest(UnicodeTestCase):
             separator=separator)
       if not reduction_indices:
         truth = constant_op.constant(truth)
-      truth_squeezed = array_ops.squeeze(truth, squeeze_dims=reduction_indices)
+      truth_squeezed = array_ops.squeeze(truth, axis=reduction_indices)
       output_array = output.eval()
       output_keep_dims_array = output_keep_dims.eval()
       truth_array = truth.eval()
diff --git a/tensorflow/python/kernel_tests/reduction_ops_test.py b/tensorflow/python/kernel_tests/reduction_ops_test.py
index 589ea54973..ea78b58d88 100644
--- a/tensorflow/python/kernel_tests/reduction_ops_test.py
+++ b/tensorflow/python/kernel_tests/reduction_ops_test.py
@@ -889,9 +889,9 @@ class AnyReductionTest(test.TestCase):
 
 class CountNonzeroReductionTest(test.TestCase):
 
-  def _compare(self, x, reduction_axes, keepdims, use_gpu=False,
+  def _compare(self, x, reduction_axes, keepdims, use_gpu=False, zero=0,
                feed_dict=None):
-    np_ans = (x != 0).astype(np.int32)
+    np_ans = (x != zero).astype(np.int32)
     if reduction_axes is None:
       np_ans = np.sum(np_ans, keepdims=keepdims)
     else:
@@ -958,6 +958,37 @@ class CountNonzeroReductionTest(test.TestCase):
           y = math_ops.count_nonzero(x, [0])
           self.assertAllEqual(y.eval(), np.zeros(9938))
 
+  def testStringReduce(self):
+    # Test case for GitHub issue 18712
+    with self.test_session() as sess:
+      v = math_ops.count_nonzero(constant_op.constant(["test"]))
+      self.assertAllClose(sess.run(v), 1)
+
+  def testStringReduce1D(self):
+    # Create a 1D array of strings
+    x = np.asarray(["", "", "a", "", "", "b"])
+    self._compare(x, None, keepdims=False, zero=np.str(""))
+    self._compare(x, [], keepdims=False, zero=np.str(""))
+    self._compare(x, [0], keepdims=False, zero=np.str(""))
+    self._compare(x, None, keepdims=True, zero=np.str(""))
+    self._compare(x, [], keepdims=True, zero=np.str(""))
+    self._compare(x, [0], keepdims=True, zero=np.str(""))
+
+  def testStringReduce2D(self):
+    # Create a 2D array of strings
+    x = np.asarray([["", "", "a", "", "", "b"],
+                    ["", "c", "", "d", "", ""],
+                    ["e", "", "f", "", "", ""]])
+    self._compare(x, None, keepdims=False, zero=np.str(""))
+    self._compare(x, [], keepdims=False, zero=np.str(""))
+    self._compare(x, [0], keepdims=False, zero=np.str(""))
+    self._compare(x, [1], keepdims=False, zero=np.str(""))
+    self._compare(x, [0, 1], keepdims=False, zero=np.str(""))
+    self._compare(x, None, keepdims=True, zero=np.str(""))
+    self._compare(x, [], keepdims=True, zero=np.str(""))
+    self._compare(x, [0], keepdims=True, zero=np.str(""))
+    self._compare(x, [0, 1], keepdims=True, zero=np.str(""))
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/python/kernel_tests/scatter_nd_ops_test.py b/tensorflow/python/kernel_tests/scatter_nd_ops_test.py
index 9f57949515..b7477a768a 100644
--- a/tensorflow/python/kernel_tests/scatter_nd_ops_test.py
+++ b/tensorflow/python/kernel_tests/scatter_nd_ops_test.py
@@ -364,6 +364,42 @@ class ScatterNdTest(test.TestCase):
     del input_  # input_ is not used in scatter_nd
     return array_ops.scatter_nd(indices, updates, shape)
 
+  def testString(self):
+    indices = constant_op.constant([[4], [3], [1], [7]],
+                                   dtype=dtypes.int32)
+    updates = constant_op.constant(["four", "three", "one", "seven"],
+                                   dtype=dtypes.string)
+    expected = np.array([b"", b"one", b"", b"three", b"four",
+                         b"", b"", b"seven"])
+    scatter = self.scatter_nd(indices, updates, shape=(8,))
+    with self.test_session() as sess:
+      result = sess.run(scatter)
+      self.assertAllEqual(expected, result)
+
+    # Same indice is updated twice by same value.
+    indices = constant_op.constant([[4], [3], [3], [7]],
+                                   dtype=dtypes.int32)
+    updates = constant_op.constant(["a", "b", "b", "c"],
+                                   dtype=dtypes.string)
+    expected = np.array([b"", b"", b"", b"bb", b"a", b"", b"", b"c"])
+    scatter = self.scatter_nd(indices, updates, shape=(8,))
+    with self.test_session() as sess:
+      result = sess.run(scatter)
+      self.assertAllEqual(expected, result)
+
+    # Same indice is updated twice by different value.
+    indices = constant_op.constant([[4], [3], [3], [7]],
+                                   dtype=dtypes.int32)
+    updates = constant_op.constant(["a", "b", "c", "d"],
+                                   dtype=dtypes.string)
+    expected = [np.array([b"", b"", b"", b"bc", b"a", b"", b"", b"d"]),
+                np.array([b"", b"", b"", b"cb", b"a", b"", b"", b"d"])]
+    scatter = self.scatter_nd(indices, updates, shape=(8,))
+    with self.test_session() as sess:
+      result = sess.run(scatter)
+      self.assertTrue(np.array_equal(result, expected[0]) or
+                      np.array_equal(result, expected[1]))
+
   def testRank3ValidShape(self):
     indices = array_ops.zeros([2, 2, 2], dtypes.int32)
     updates = array_ops.zeros([2, 2, 2], dtypes.int32)
@@ -584,6 +620,10 @@ class ScatterNdNonAliasingAddTest(ScatterNdTest):
         shape, dtype=updates.dtype))
     return array_ops.scatter_nd_non_aliasing_add(input_, indices, updates)
 
+  def testString(self):
+    # Not supported yet.
+    pass
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/python/ops/array_grad.py b/tensorflow/python/ops/array_grad.py
index 57d2657838..3678bd4c1f 100644
--- a/tensorflow/python/ops/array_grad.py
+++ b/tensorflow/python/ops/array_grad.py
@@ -196,7 +196,7 @@ def _ConcatGradHelper(op, grad, start_value_index, end_value_index, dim_index):
             array_ops.where(
                 math_ops.logical_and(grad.indices >= start,
                                      grad.indices < end)),
-            squeeze_dims=[1])
+            axis=[1])
         new_indices = array_ops.gather(grad.indices, indices_to_select) - start
         new_values = array_ops.gather(grad.values, indices_to_select)
         out_grads.append(ops.IndexedSlices(new_values, new_indices, size))
diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py
index 3c2593066a..e235047aff 100644
--- a/tensorflow/python/ops/array_ops.py
+++ b/tensorflow/python/ops/array_ops.py
@@ -994,9 +994,7 @@ def unstack(value, num=None, axis=0, name="unstack"):
     `value[:, i, :, :]` and each tensor in `output` will have shape `(A, C, D)`.
   Etc.
 
-  This is the opposite of stack.  The numpy equivalent is
-
-      tf.unstack(x, n) = np.unstack(x)
+  This is the opposite of stack.
 
   Args:
     value: A rank `R > 0` `Tensor` to be unstacked.
@@ -1720,8 +1718,10 @@ def placeholder(dtype, shape=None, name=None):
     print(sess.run(y, feed_dict={x: rand_array}))  # Will succeed.
   ```
 
-  @compatibility{eager} Placeholders are not compatible with eager execution.
-
+  @compatibility(eager)
+  Placeholders are not compatible with eager execution.
+  @end_compatibility
+  
   Args:
     dtype: The type of elements in the tensor to be fed.
     shape: The shape of the tensor to be fed (optional). If the shape is not
diff --git a/tensorflow/python/ops/image_ops_impl.py b/tensorflow/python/ops/image_ops_impl.py
index 601010bce9..bd5b2ae83b 100644
--- a/tensorflow/python/ops/image_ops_impl.py
+++ b/tensorflow/python/ops/image_ops_impl.py
@@ -652,7 +652,7 @@ def pad_to_bounding_box(image, offset_height, offset_width, target_height,
     padded.set_shape(padded_shape)
 
     if not is_batch:
-      padded = array_ops.squeeze(padded, squeeze_dims=[0])
+      padded = array_ops.squeeze(padded, axis=[0])
 
     return padded
 
@@ -732,7 +732,7 @@ def crop_to_bounding_box(image, offset_height, offset_width, target_height,
     cropped.set_shape(cropped_shape)
 
     if not is_batch:
-      cropped = array_ops.squeeze(cropped, squeeze_dims=[0])
+      cropped = array_ops.squeeze(cropped, axis=[0])
 
     return cropped
 
@@ -849,7 +849,7 @@ def resize_image_with_crop_or_pad(image, target_height, target_width):
     resized = control_flow_ops.with_dependencies(assert_ops, resized)
 
     if not is_batch:
-      resized = array_ops.squeeze(resized, squeeze_dims=[0])
+      resized = array_ops.squeeze(resized, axis=[0])
 
     return resized
 
@@ -942,7 +942,7 @@ def resize_images(images,
            for x in [new_width_const, width, new_height_const, height]) and (
                width == new_width_const and height == new_height_const):
       if not is_batch:
-        images = array_ops.squeeze(images, squeeze_dims=[0])
+        images = array_ops.squeeze(images, axis=[0])
       return images
 
     if method == ResizeMethod.BILINEAR:
@@ -965,7 +965,7 @@ def resize_images(images,
     images.set_shape([None, new_height_const, new_width_const, None])
 
     if not is_batch:
-      images = array_ops.squeeze(images, squeeze_dims=[0])
+      images = array_ops.squeeze(images, axis=[0])
     return images
 
 
diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py
index 57660578aa..7ac3bd8091 100644
--- a/tensorflow/python/ops/math_ops.py
+++ b/tensorflow/python/ops/math_ops.py
@@ -1338,8 +1338,18 @@ def count_nonzero(input_tensor,
   tf.count_nonzero(x, [0, 1])  # 3
   ```
 
+  **NOTE** Strings are compared against zero-length empty string `""`. Any
+  string with a size greater than zero is already considered as nonzero.
+
+  For example:
+  ```python
+  x = tf.constant(["", "a", "  ", "b", ""])
+  tf.count_nonzero(x) # 3, with "a", "  ", and "b" as nonzero strings.
+  ```
+
   Args:
-    input_tensor: The tensor to reduce. Should be of numeric type, or `bool`.
+    input_tensor: The tensor to reduce. Should be of numeric type, `bool`,
+      or `string`.
     axis: The dimensions to reduce. If `None` (the default),
       reduces all dimensions. Must be in the range
       `[-rank(input_tensor), rank(input_tensor))`.
@@ -1359,7 +1369,8 @@ def count_nonzero(input_tensor,
 
   with ops.name_scope(name, "count_nonzero", [input_tensor]):
     input_tensor = ops.convert_to_tensor(input_tensor, name="input_tensor")
-    zero = input_tensor.dtype.as_numpy_dtype()
+    # A scalar of 'zero' is enough as `not_equal` will broadcast.
+    zero = array_ops.zeros([], dtype=input_tensor.dtype)
     return cast(
         reduce_sum(
             # int64 reduction happens on GPU
diff --git a/tensorflow/python/ops/nn_impl.py b/tensorflow/python/ops/nn_impl.py
index d0d5ed07ce..576627e78e 100644
--- a/tensorflow/python/ops/nn_impl.py
+++ b/tensorflow/python/ops/nn_impl.py
@@ -765,9 +765,9 @@ def weighted_moments(x, axes, frequency_weights, name=None, keep_dims=False):
     weighted_variance = math_ops.multiply(weighted_distsq, divisor)
 
     if not keep_dims:
-      weighted_mean = array_ops.squeeze(weighted_mean, squeeze_dims=axes)
+      weighted_mean = array_ops.squeeze(weighted_mean, axis=axes)
       weighted_variance = array_ops.squeeze(
-          weighted_variance, squeeze_dims=axes)
+          weighted_variance, axis=axes)
 
     if needs_cast:
       weighted_mean = math_ops.cast(weighted_mean, dtypes.float16)
diff --git a/tensorflow/python/ops/rnn_cell_impl.py b/tensorflow/python/ops/rnn_cell_impl.py
index 86dc053c0f..67f753485b 100644
--- a/tensorflow/python/ops/rnn_cell_impl.py
+++ b/tensorflow/python/ops/rnn_cell_impl.py
@@ -785,10 +785,14 @@ class LSTMCell(LayerRNNCell):
         shape=[input_depth + h_depth, 4 * self._num_units],
         initializer=self._initializer,
         partitioner=maybe_partitioner)
+    if self.dtype is None:
+      initializer = init_ops.zeros_initializer
+    else:
+      initializer = init_ops.zeros_initializer(dtype=self.dtype)
     self._bias = self.add_variable(
         _BIAS_VARIABLE_NAME,
         shape=[4 * self._num_units],
-        initializer=init_ops.zeros_initializer(dtype=self.dtype))
+        initializer=initializer)
     if self._use_peepholes:
       self._w_f_diag = self.add_variable("w_f_diag", shape=[self._num_units],
                                          initializer=self._initializer)
diff --git a/tensorflow/python/training/adam.py b/tensorflow/python/training/adam.py
index 006e360389..6fa3ff6658 100644
--- a/tensorflow/python/training/adam.py
+++ b/tensorflow/python/training/adam.py
@@ -43,23 +43,19 @@ class AdamOptimizer(optimizer.Optimizer):
 
     Initialization:
 
-    ```
-    m_0 <- 0 (Initialize initial 1st moment vector)
-    v_0 <- 0 (Initialize initial 2nd moment vector)
-    t <- 0 (Initialize timestep)
-    ```
+    $$m_0 := 0 (Initialize initial 1st moment vector)$$
+    $$v_0 := 0 (Initialize initial 2nd moment vector)$$
+    $$t := 0 (Initialize timestep)$$
 
     The update rule for `variable` with gradient `g` uses an optimization
     described at the end of section2 of the paper:
 
-    ```
-    t <- t + 1
-    lr_t <- learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
+    $$t := t + 1$$
+    $$lr_t := \text{learning_rate} * \sqrt{(1 - beta_2^t) / (1 - beta_1^t)}$$
 
-    m_t <- beta1 * m_{t-1} + (1 - beta1) * g
-    v_t <- beta2 * v_{t-1} + (1 - beta2) * g * g
-    variable <- variable - lr_t * m_t / (sqrt(v_t) + epsilon)
-    ```
+    $$m_t := beta_1 * m_{t-1} + (1 - beta_1) * g$$
+    $$v_t := beta_2 * v_{t-1} + (1 - beta_2) * g * g$$
+    $$variable := variable - lr_t * m_t / (\sqrt{v_t} + \epsilon)$$
 
     The default value of 1e-8 for epsilon might not be a good default in
     general. For example, when training an Inception network on ImageNet a
diff --git a/tensorflow/python/training/input_test.py b/tensorflow/python/training/input_test.py
index 3a25bfe343..1b1e89cb26 100644
--- a/tensorflow/python/training/input_test.py
+++ b/tensorflow/python/training/input_test.py
@@ -497,6 +497,28 @@ class BatchTest(test_lib.TestCase):
   def testOneThreadDict(self):
     self._testOneThreadHelper(use_dict=True)
 
+  def testUint32DataTypes(self):
+    values = constant_op.constant([0, 1, 2, 3, 4, 5], dtype=dtypes.uint32)
+    batched = inp.batch([values], batch_size=2)
+    with self.test_session() as sess:
+      coord = coordinator.Coordinator()
+      threads = queue_runner_impl.start_queue_runners(sess=sess, coord=coord)
+      sess.run(batched)
+      coord.request_stop()
+      for thread in threads:
+        thread.join()
+
+  def testUint64DataTypes(self):
+    values = constant_op.constant([0, 1, 2, 3, 4, 5], dtype=dtypes.uint64)
+    batched = inp.batch([values], batch_size=2)
+    with self.test_session() as sess:
+      coord = coordinator.Coordinator()
+      threads = queue_runner_impl.start_queue_runners(sess=sess, coord=coord)
+      sess.run(batched)
+      coord.request_stop()
+      for thread in threads:
+        thread.join()
+
   def testOneThreadDynamicPad(self):
     with self.test_session() as sess:
       batch_size = 10
diff --git a/tensorflow/python/training/monitored_session.py b/tensorflow/python/training/monitored_session.py
index 4ce6f6d002..f584a009d9 100644
--- a/tensorflow/python/training/monitored_session.py
+++ b/tensorflow/python/training/monitored_session.py
@@ -202,7 +202,7 @@ class Scaffold(object):
     if self._local_init_op is None:
       self._local_init_op = Scaffold.get_or_default(
           'local_init_op', ops.GraphKeys.LOCAL_INIT_OP,
-          Scaffold._default_local_init_op)
+          Scaffold.default_local_init_op)
     if self._summary_op is None:
       self._summary_op = Scaffold.get_or_default('summary_op',
                                                  ops.GraphKeys.SUMMARY_OP,
@@ -267,7 +267,17 @@ class Scaffold(object):
     return op
 
   @staticmethod
-  def _default_local_init_op():
+  def default_local_init_op():
+    """Returns an op that groups the default local init ops.
+
+    This op is used during session initialization when a Scaffold is
+    initialized without specifying the local_init_op arg. It includes
+    `tf.local_variables_initializer`, `tf.tables_initializer`, and also
+    initializes local session resources.
+
+    Returns:
+      The default Scaffold local init op.
+    """
     return control_flow_ops.group(
         variables.local_variables_initializer(),
         lookup_ops.tables_initializer(),
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-scaffold.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-scaffold.pbtxt
index 62b956c5ef..38cc98b48e 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-scaffold.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-scaffold.pbtxt
@@ -38,6 +38,10 @@ tf_class {
     name: "__init__"
     argspec: "args=[\'self\', \'init_op\', \'init_feed_dict\', \'init_fn\', \'ready_op\', \'ready_for_local_init_op\', \'local_init_op\', \'summary_op\', \'saver\', \'copy_from_scaffold\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\', \'None\'], "
   }
+  member_method {
+    name: "default_local_init_op"
+    argspec: "args=[], varargs=None, keywords=None, defaults=None"
+  }
   member_method {
     name: "finalize"
     argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
diff --git a/tensorflow/tools/ci_build/ci_parameterized_build.sh b/tensorflow/tools/ci_build/ci_parameterized_build.sh
index 9d23b508aa..797e0a6db5 100755
--- a/tensorflow/tools/ci_build/ci_parameterized_build.sh
+++ b/tensorflow/tools/ci_build/ci_parameterized_build.sh
@@ -237,7 +237,7 @@ function get_cuda_capability_version() {
 CTYPE=${TF_BUILD_CONTAINER_TYPE}
 
 # Determine if the machine is a Mac
-OPT_FLAG=""
+OPT_FLAG="--test_output=errors"
 if [[ "$(uname -s)" == "Darwin" ]]; then
   DO_DOCKER=0
 
diff --git a/tensorflow/tools/ci_build/install/install_python3.5_pip_packages.sh b/tensorflow/tools/ci_build/install/install_python3.5_pip_packages.sh
index aefc49f604..204a82f647 100755
--- a/tensorflow/tools/ci_build/install/install_python3.5_pip_packages.sh
+++ b/tensorflow/tools/ci_build/install/install_python3.5_pip_packages.sh
@@ -39,6 +39,9 @@ if [[ -z $pip35_version ]]; then
 fi
 
 set -e
+pip3.5 install --upgrade setuptools
+pip3.5 install --upgrade pip
+
 pip3.5 install --upgrade virtualenv
 
 # Install six.
diff --git a/tensorflow/tools/ci_build/install/install_python3.6_pip_packages.sh b/tensorflow/tools/ci_build/install/install_python3.6_pip_packages.sh
index bfaa044c82..275abeb669 100755
--- a/tensorflow/tools/ci_build/install/install_python3.6_pip_packages.sh
+++ b/tensorflow/tools/ci_build/install/install_python3.6_pip_packages.sh
@@ -49,9 +49,13 @@ cd Python-3.6.1
 make altinstall
 ln -s /usr/local/bin/pip3.6 /usr/local/bin/pip3
 
+pip3 install --upgrade setuptools
+pip3 install --upgrade pip
+
 pip3 install --upgrade virtualenv
 
 set -e
+
 # Install six.
 pip3 install --upgrade absl-py
 pip3 install --upgrade six==1.10.0
diff --git a/tensorflow/tools/ci_build/windows/cpu/bazel/run_cc_test_windows.sh b/tensorflow/tools/ci_build/windows/cpu/bazel/run_cc_test_windows.sh
index 748a961e44..dc9af221ec 100644
--- a/tensorflow/tools/ci_build/windows/cpu/bazel/run_cc_test_windows.sh
+++ b/tensorflow/tools/ci_build/windows/cpu/bazel/run_cc_test_windows.sh
@@ -44,7 +44,7 @@ source "tensorflow/tools/ci_build/windows/bazel/bazel_test_lib.sh" \
 
 run_configure_for_cpu_build
 
-# Compliling the following test is extremely slow with -c opt
+# Compiling the following test is extremely slow with -c opt
 slow_compiling_test="//tensorflow/core/kernels:eigen_backward_spatial_convolutions_test"
 
 # Find all the passing cc_tests on Windows and store them in a variable
diff --git a/tensorflow/tools/ci_build/windows/gpu/bazel/run_cc_test_windows.sh b/tensorflow/tools/ci_build/windows/gpu/bazel/run_cc_test_windows.sh
index f26f8727e5..f1114f4ffa 100644
--- a/tensorflow/tools/ci_build/windows/gpu/bazel/run_cc_test_windows.sh
+++ b/tensorflow/tools/ci_build/windows/gpu/bazel/run_cc_test_windows.sh
@@ -46,7 +46,7 @@ clean_output_base
 
 run_configure_for_gpu_build
 
-# Compliling the following test is extremely slow with -c opt
+# Compiling the following test is extremely slow with -c opt
 slow_compiling_test="//tensorflow/core/kernels:eigen_backward_spatial_convolutions_test"
 
 # Find all the passing cc_tests on Windows and store them in a variable
diff --git a/tensorflow/tools/docker/Dockerfile b/tensorflow/tools/docker/Dockerfile
index 78cb4d250e..a3ff8211e3 100644
--- a/tensorflow/tools/docker/Dockerfile
+++ b/tensorflow/tools/docker/Dockerfile
@@ -7,6 +7,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
         build-essential \
         curl \
         libfreetype6-dev \
+        libhdf5-serial-dev \
         libpng12-dev \
         libzmq3-dev \
         pkg-config \
diff --git a/tensorflow/tools/docker/Dockerfile.devel b/tensorflow/tools/docker/Dockerfile.devel
index 390d7442c3..b9996395d0 100644
--- a/tensorflow/tools/docker/Dockerfile.devel
+++ b/tensorflow/tools/docker/Dockerfile.devel
@@ -8,6 +8,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
         git \
         libcurl3-dev \
         libfreetype6-dev \
+        libhdf5-serial-dev \
         libpng12-dev \
         libzmq3-dev \
         pkg-config \
@@ -28,9 +29,12 @@ RUN curl -fSsL -O https://bootstrap.pypa.io/get-pip.py && \
     rm get-pip.py
 
 RUN pip --no-cache-dir install \
+        Pillow \
+        h5py \
         ipykernel \
         jupyter \
         matplotlib \
+        mock \
         numpy \
         scipy \
         sklearn \
diff --git a/tensorflow/tools/docker/Dockerfile.devel-gpu b/tensorflow/tools/docker/Dockerfile.devel-gpu
index 293028d229..7e5e6ef2d5 100644
--- a/tensorflow/tools/docker/Dockerfile.devel-gpu
+++ b/tensorflow/tools/docker/Dockerfile.devel-gpu
@@ -17,6 +17,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
         libcudnn7-dev=7.0.5.15-1+cuda9.0 \
         libcurl3-dev \
         libfreetype6-dev \
+        libhdf5-serial-dev \
         libpng12-dev \
         libzmq3-dev \
         pkg-config \
@@ -37,9 +38,12 @@ RUN curl -fSsL -O https://bootstrap.pypa.io/get-pip.py && \
     rm get-pip.py
 
 RUN pip --no-cache-dir install \
+        Pillow \
+        h5py \
         ipykernel \
         jupyter \
         matplotlib \
+        mock \
         numpy \
         scipy \
         sklearn \
diff --git a/tensorflow/tools/docker/Dockerfile.gpu b/tensorflow/tools/docker/Dockerfile.gpu
index 9e1708662e..bff4a20392 100644
--- a/tensorflow/tools/docker/Dockerfile.gpu
+++ b/tensorflow/tools/docker/Dockerfile.gpu
@@ -14,6 +14,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
         curl \
         libcudnn7=7.0.5.15-1+cuda9.0 \
         libfreetype6-dev \
+        libhdf5-serial-dev \
         libpng12-dev \
         libzmq3-dev \
         pkg-config \
diff --git a/tensorflow/tools/docker/README.md b/tensorflow/tools/docker/README.md
index f46c56e11a..525f2995ce 100644
--- a/tensorflow/tools/docker/README.md
+++ b/tensorflow/tools/docker/README.md
@@ -16,12 +16,12 @@ quick links here:
 
 We currently maintain two Docker container images:
 
-* `gcr.io/tensorflow/tensorflow` - TensorFlow with all dependencies - CPU only!
+* `tensorflow/tensorflow` - TensorFlow with all dependencies - CPU only!
 
-* `gcr.io/tensorflow/tensorflow:latest-gpu` - TensorFlow with all dependencies
+* `tensorflow/tensorflow:latest-gpu` - TensorFlow with all dependencies
   and support for NVidia CUDA
 
-Note: We also publish the same containers into
+Note: We store all our containers on 
 [Docker Hub](https://hub.docker.com/r/tensorflow/tensorflow/tags/).
 
 
@@ -29,12 +29,12 @@ Note: We also publish the same containers into
 
 Run non-GPU container using
 
-    $ docker run -it -p 8888:8888 gcr.io/tensorflow/tensorflow
+    $ docker run -it -p 8888:8888 tensorflow/tensorflow
 
 For GPU support install NVidia drivers (ideally latest) and
 [nvidia-docker](https://github.com/NVIDIA/nvidia-docker). Run using
 
-    $ nvidia-docker run -it -p 8888:8888 gcr.io/tensorflow/tensorflow:latest-gpu
+    $ nvidia-docker run -it -p 8888:8888 tensorflow/tensorflow:latest-gpu
 
 
 Note: If you would have a problem running nvidia-docker you may try the old method
@@ -44,7 +44,7 @@ it there and try using nvidia-docker as described above.
     $ # The old, not recommended way to run docker with gpu support:
     $ export CUDA_SO=$(\ls /usr/lib/x86_64-linux-gnu/libcuda.* | xargs -I{} echo '-v {}:{}')
     $ export DEVICES=$(\ls /dev/nvidia* | xargs -I{} echo '--device {}:{}')
-    $ docker run -it -p 8888:8888 $CUDA_SO $DEVICES gcr.io/tensorflow/tensorflow:latest-gpu
+    $ docker run -it -p 8888:8888 $CUDA_SO $DEVICES tensorflow/tensorflow:latest-gpu
 
 
 ## More containers
diff --git a/tensorflow/tools/graph_transforms/fold_old_batch_norms.cc b/tensorflow/tools/graph_transforms/fold_old_batch_norms.cc
index d86f65325b..f1d361e07d 100644
--- a/tensorflow/tools/graph_transforms/fold_old_batch_norms.cc
+++ b/tensorflow/tools/graph_transforms/fold_old_batch_norms.cc
@@ -159,6 +159,9 @@ Status FuseScaleOffsetToConvWeights(const std::vector<float>& scale_values,
   NodeDef bias_add_node;
   bias_add_node.set_op("BiasAdd");
   bias_add_node.set_name(conv_output_name);
+  if (!conv_node.attr().count("data_format")) {
+    CopyNodeAttr(conv_node, "data_format", "data_format", &bias_add_node);
+  }
   CopyNodeAttr(conv_node, "T", "T", &bias_add_node);
   AddNodeInput(conv_node.name(), &bias_add_node);
   AddNodeInput(bias_offset_node.name(), &bias_add_node);
diff --git a/tensorflow/tools/pip_package/setup.py b/tensorflow/tools/pip_package/setup.py
index f84a91d009..937d41c36c 100644
--- a/tensorflow/tools/pip_package/setup.py
+++ b/tensorflow/tools/pip_package/setup.py
@@ -31,7 +31,7 @@ from setuptools.dist import Distribution
 # This version string is semver compatible, but incompatible with pip.
 # For pip, we will remove all '-' characters from this string, and use the
 # result for pip.
-_VERSION = '1.8.0-rc0'
+_VERSION = '1.8.0-rc1'
 
 REQUIRED_PACKAGES = [
     'absl-py >= 0.1.6',
@@ -69,7 +69,7 @@ else:
 if 'tf_nightly' in project_name:
   for i, pkg in enumerate(REQUIRED_PACKAGES):
     if 'tensorboard' in pkg:
-      REQUIRED_PACKAGES[i] = 'tb-nightly >= 1.8.0a0, < 1.9.0a0'
+      REQUIRED_PACKAGES[i] = 'tb-nightly >= 1.9.0a0, < 1.10.0a0'
       break
 
 # weakref.finalize and enum were introduced in Python 3.4
diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 152da547c1..16da59c5cf 100644
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -50,31 +50,31 @@ def tf_workspace(path_prefix="", tf_repo_name=""):
   mkl_repository(
       name = "mkl_linux",
       urls = [
-          "https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.12/mklml_lnx_2018.0.1.20171227.tgz",
-          "https://github.com/intel/mkl-dnn/releases/download/v0.12/mklml_lnx_2018.0.1.20171227.tgz",
+          "https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.13/mklml_lnx_2018.0.2.20180127.tgz",
+          "https://github.com/intel/mkl-dnn/releases/download/v0.13/mklml_lnx_2018.0.2.20180127.tgz",
       ],
-      sha256 = "feacc3d82565c1231470359b42c696236fae873704e0b013436afba5fd4fd30f",
-      strip_prefix = "mklml_lnx_2018.0.1.20171227",
+      sha256 = "74844bd77294742bf2396ff040369d1aa4cdd9e826fcd38cf8398ae83564d146",
+      strip_prefix = "mklml_lnx_2018.0.2.20180127",
       build_file = clean_dep("//third_party/mkl:mkl.BUILD")
   )
   mkl_repository(
       name = "mkl_windows",
       urls = [
-          "https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.12/mklml_win_2018.0.1.20171227.zip",
-          "https://github.com/intel/mkl-dnn/releases/download/v0.12/mklml_win_2018.0.1.20171227.zip"
+          "https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.13/mklml_win_2018.0.2.20180127.zip",
+          "https://github.com/intel/mkl-dnn/releases/download/v0.13/mklml_win_2018.0.2.20180127.zip"
       ],
-      sha256 = "24bae8d7b22b431a654acadea43f2243c46ae6b1e5a73a4a936825f31d284ee4",
-      strip_prefix = "mklml_win_2018.0.1.20171227",
+      sha256 = "d8fbf0faa0684bffa3548005d05fe5cfe56ff9dbc0e15e7612d7ac01055a6ded",
+      strip_prefix = "mklml_win_2018.0.2.20180127",
       build_file = clean_dep("//third_party/mkl:mkl.BUILD")
   )
   mkl_repository(
       name = "mkl_darwin",
       urls = [
-          "https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.12/mklml_mac_2018.0.1.20171227.tgz",
-          "https://github.com/intel/mkl-dnn/releases/download/v0.12/mklml_mac_2018.0.1.20171227.tgz"
+          "https://mirror.bazel.build/github.com/intel/mkl-dnn/releases/download/v0.13/mklml_mac_2018.0.2.20180127.tgz",
+          "https://github.com/intel/mkl-dnn/releases/download/v0.13/mklml_mac_2018.0.2.20180127.tgz"
       ],
-      sha256 = "0e954ec6fd3dc5e37f64c4043f6b5613dd687558da3df1028b3b7c29ff5cf77f",
-      strip_prefix = "mklml_mac_2018.0.1.20171227",
+      sha256 = "aa740d71e14562bfea56e6829e6dc186e7487cbcf6748a88dec73826b7ec1943",
+      strip_prefix = "mklml_mac_2018.0.2.20180127",
       build_file = clean_dep("//third_party/mkl:mkl.BUILD")
   )
 
@@ -85,11 +85,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""):
   tf_http_archive(
       name = "mkl_dnn",
       urls = [
-          "https://mirror.bazel.build/github.com/intel/mkl-dnn/archive/v0.12.tar.gz",
-          "https://github.com/intel/mkl-dnn/archive/v0.12.tar.gz",
+          "https://mirror.bazel.build/github.com/intel/mkl-dnn/archive/v0.13.tar.gz",
+          "https://github.com/intel/mkl-dnn/archive/v0.13.tar.gz",
       ],
-      sha256 = "86fa2a8c12a56e3b725945acedeaa82492746be02545aba6d710f097e013e19e",
-      strip_prefix = "mkl-dnn-0.12",
+      sha256 = "d2cfd93a70cfe86ebe054477c530c9b5c1218b70f75856eb6d1956c68ee89e8f",
+      strip_prefix = "mkl-dnn-0.13",
       build_file = clean_dep("//third_party/mkl_dnn:mkldnn.BUILD"),
   )
 
diff --git a/third_party/eigen3/unsupported/Eigen/CXX11/src/NeuralNetworks/BackwardSpatialConvolutions.h b/third_party/eigen3/unsupported/Eigen/CXX11/src/NeuralNetworks/BackwardSpatialConvolutions.h
index 188dc75bf6..0f4ada246c 100644
--- a/third_party/eigen3/unsupported/Eigen/CXX11/src/NeuralNetworks/BackwardSpatialConvolutions.h
+++ b/third_party/eigen3/unsupported/Eigen/CXX11/src/NeuralNetworks/BackwardSpatialConvolutions.h
@@ -280,9 +280,9 @@ SpatialConvolutionBackwardKernel(const Input& input, const OutputBackward& outpu
     eigen_assert(input_dims[0] == pre_contract_dims[0]);
   }
 
-  // We will contract along dimensions (1, 2) in in and (1, 3) in out, if
+  // We will contract along dimensions (1, 2) in and (1, 3) in out, if
   // this is col-major.
-  // For row-major, it's dimensions (0, 1) in in and (0, 2) in out.
+  // For row-major, it's dimensions (0, 1) in and (0, 2) in out.
   array<IndexPair<TensorIndex>, 2> contract_dims;
   if (isColMajor) {
     // col-major: in.contract(output.patches)
-- 
GitLab


From fdcdf752dca18d479932119a2445e0129fcd54a9 Mon Sep 17 00:00:00 2001
From: Mark Heffernan <meheff@google.com>
Date: Tue, 1 May 2018 14:52:40 -0700
Subject: [PATCH 112/395] Fix bug in peak buffer accounting in buffer
 assignment. Buffer assignment keeps track of the set of logical buffers which
 are live at the point of peak memory usage for each allocation. Previously
 colocated buffers were not properly accounted for. This CL addresses this
 problem.

PiperOrigin-RevId: 195001567
---
 .../compiler/xla/service/buffer_assignment.cc | 196 ++++++++----------
 .../compiler/xla/service/buffer_assignment.h  |  23 +-
 .../xla/service/buffer_assignment_test.cc     |  77 ++++++-
 3 files changed, 169 insertions(+), 127 deletions(-)

diff --git a/tensorflow/compiler/xla/service/buffer_assignment.cc b/tensorflow/compiler/xla/service/buffer_assignment.cc
index dbe45e932c..94ccfedf62 100644
--- a/tensorflow/compiler/xla/service/buffer_assignment.cc
+++ b/tensorflow/compiler/xla/service/buffer_assignment.cc
@@ -292,112 +292,6 @@ BufferAllocationProto BufferAllocation::ToProto() const {
   return proto;
 }
 
-std::pair<int64, std::vector<const LogicalBuffer*>>
-BufferAllocation::ComputePeakMemoryLogicalBuffers() const {
-  if (HeapTraces().empty()) {
-    // Just return the largest LogicalBuffer in the allocation.
-    const LogicalBuffer* largest_buffer = nullptr;
-    int64 largest_size = 0;
-    for (const auto& pair : assigned_buffers()) {
-      const LogicalBuffer* buffer = pair.first;
-      int64 size = pair.second.size;
-      if (largest_buffer == nullptr) {
-        largest_buffer = buffer;
-        largest_size = size;
-        continue;
-      }
-      // Tie-break with LogicalBuffer::Id so the return value is stable relative
-      // to changing addresses.
-      if (size > largest_size ||
-          ((size == largest_size) && (largest_buffer->id() > buffer->id()))) {
-        largest_buffer = buffer;
-        largest_size = size;
-      }
-    }
-    CHECK(largest_buffer != nullptr)
-        << "No logical buffers in allocation: " << ToString();
-    return {largest_size, {largest_buffer}};
-  }
-
-  // Create a map from LogicalBuffer::Id to LogicalBuffer* for the logical
-  // buffers in this allocation.
-  tensorflow::gtl::FlatMap<LogicalBuffer::Id, const LogicalBuffer*>
-      id_to_buffer;
-  tensorflow::gtl::FlatMap<const LogicalBuffer*, int64> buffer_sizes;
-  for (const auto& pair : assigned_buffers()) {
-    const LogicalBuffer* buffer = pair.first;
-    const OffsetSize& offset_size = pair.second;
-    id_to_buffer[buffer->id()] = buffer;
-    buffer_sizes[buffer] = offset_size.size;
-  }
-
-  // Returns how much the given event increases the total size of live
-  // buffers. Can be negative.
-  auto memory_delta = [this, &id_to_buffer, &buffer_sizes](
-                          const HeapSimulatorTrace::Event& event) -> int64 {
-    const LogicalBuffer* buffer = id_to_buffer.at(event.buffer_id());
-    const int64 buffer_size = buffer_sizes.at(buffer);
-    if (event.kind() == HeapSimulatorTrace::Event::ALLOC) {
-      return buffer_size;
-    } else if (event.kind() == HeapSimulatorTrace::Event::SHARE_WITH) {
-      // Sharing a buffer does not change the live set size for the purposes of
-      // the heap simulator. Even though the shared-with buffer may be smaller,
-      // the entire allocation remains live.
-      return 0;
-    } else if (event.kind() == HeapSimulatorTrace::Event::FREE) {
-      return -1 * buffer_size;
-    }
-    LOG(FATAL) << "Unknown event kind: " << event.kind();
-  };
-
-  int64 total_max_live_size = 0;
-  std::vector<const LogicalBuffer*> live_buffers_vector;
-  for (const HeapSimulatorTrace& heap_trace : HeapTraces()) {
-    // First compute the size of the maximal live set.
-    int64 max_live_size = 0;
-    int64 live_size = 0;
-    for (const auto& event : heap_trace.events()) {
-      live_size += memory_delta(event);
-      if (max_live_size < live_size) {
-        max_live_size = live_size;
-      }
-    }
-
-    // Next gather the set of logical buffers live at the earliest point of
-    // maximal live set size.
-    tensorflow::gtl::FlatSet<const LogicalBuffer*> live_buffers;
-    live_size = 0;
-    for (const auto& event : heap_trace.events()) {
-      const LogicalBuffer* buffer = id_to_buffer.at(event.buffer_id());
-      if (event.kind() == HeapSimulatorTrace::Event::ALLOC) {
-        InsertOrDie(&live_buffers, buffer);
-      } else if (event.kind() == HeapSimulatorTrace::Event::SHARE_WITH) {
-        // Nothing to do.
-      } else if (event.kind() == HeapSimulatorTrace::Event::FREE) {
-        CHECK(ContainsKey(live_buffers, buffer));
-        live_buffers.erase(buffer);
-      }
-
-      live_size += memory_delta(event);
-      if (live_size == max_live_size) {
-        break;
-      }
-    }
-    CHECK_EQ(live_size, max_live_size);
-    total_max_live_size += max_live_size;
-
-    live_buffers_vector.insert(live_buffers_vector.end(), live_buffers.begin(),
-                               live_buffers.end());
-  }
-
-  // Stabily sort the live buffers.
-  std::sort(live_buffers_vector.begin(), live_buffers_vector.end(),
-            [](const LogicalBuffer* a, const LogicalBuffer* b) {
-              return a->id() < b->id();
-            });
-  return {total_max_live_size, live_buffers_vector};
-}
-
 string BufferAllocation::ToString() const {
   string output;
   Appendf(&output, "allocation %lld: %p, size %lld", index_, this, size());
@@ -610,6 +504,7 @@ BufferAllocation* BufferAssignment::NewAllocation(const LogicalBuffer& buffer,
   BufferAllocation* allocation =
       NewEmptyAllocation(size, is_thread_local, is_reusable, buffer.color());
   AddAssignment(allocation, buffer, /*offset=*/0, size);
+  allocation->peak_buffers_.push_back(&buffer);
   return allocation;
 }
 
@@ -680,6 +575,10 @@ void BufferAssignment::CombineTempAllocations() {
         CHECK_EQ(temp_allocation.HeapTraces().size(), 1);
         combined_allocation->AddHeapTrace(temp_allocation.HeapTraces().front());
       }
+      combined_allocation->peak_buffers_.insert(
+          combined_allocation->peak_buffers_.end(),
+          temp_allocation.peak_buffers_.begin(),
+          temp_allocation.peak_buffers_.end());
     }
     // Replace all existing temporary allocations with the new combined
     // allocations.
@@ -1228,6 +1127,89 @@ Status BufferAssigner::AssignBuffersWithSequentialOrdering(
   return Status::OK();
 }
 
+namespace {
+
+// Computes and returns the set of logical buffers live at the point of maximal
+// liveness in the given heap trace. LogicalBuffers are (stabily) sorted by id.
+std::vector<const LogicalBuffer*> ComputePeakMemoryLogicalBuffers(
+    const BufferAllocation& allocation, const HeapSimulatorTrace& heap_trace) {
+  // Create a map from LogicalBuffer::Id to LogicalBuffer* for the logical
+  // buffers in this allocation.
+  tensorflow::gtl::FlatMap<LogicalBuffer::Id, const LogicalBuffer*>
+      id_to_buffer;
+  tensorflow::gtl::FlatMap<const LogicalBuffer*, int64> buffer_sizes;
+  for (const auto& pair : allocation.assigned_buffers()) {
+    const LogicalBuffer* buffer = pair.first;
+    const BufferAllocation::OffsetSize& offset_size = pair.second;
+    id_to_buffer[buffer->id()] = buffer;
+    buffer_sizes[buffer] = offset_size.size;
+  }
+
+  // Returns how much the given event increases the total size of live
+  // buffers. Can be negative.
+  auto memory_delta = [&id_to_buffer, &buffer_sizes](
+                          const HeapSimulatorTrace::Event& event) -> int64 {
+    const LogicalBuffer* buffer = id_to_buffer.at(event.buffer_id());
+    const int64 buffer_size = buffer_sizes.at(buffer);
+    if (event.kind() == HeapSimulatorTrace::Event::ALLOC) {
+      return buffer_size;
+    } else if (event.kind() == HeapSimulatorTrace::Event::SHARE_WITH) {
+      // Sharing a buffer does not change the live set size for the purposes of
+      // the heap simulator. Even though the shared-with buffer may be smaller,
+      // the entire allocation remains live.
+      return 0;
+    } else if (event.kind() == HeapSimulatorTrace::Event::FREE) {
+      return -1 * buffer_size;
+    }
+    LOG(FATAL) << "Unknown event kind: " << event.kind();
+  };
+
+  // First compute the size of the maximal live set.
+  int64 max_live_size = 0;
+  int64 live_size = 0;
+  for (const auto& event : heap_trace.events()) {
+    live_size += memory_delta(event);
+    if (max_live_size < live_size) {
+      max_live_size = live_size;
+    }
+  }
+
+  // Next gather the set of logical buffers live at the earliest point of
+  // maximal live set size.
+  tensorflow::gtl::FlatSet<const LogicalBuffer*> live_buffers;
+  live_size = 0;
+  for (const auto& event : heap_trace.events()) {
+    const LogicalBuffer* buffer = id_to_buffer.at(event.buffer_id());
+    if (event.kind() == HeapSimulatorTrace::Event::ALLOC) {
+      InsertOrDie(&live_buffers, buffer);
+    } else if (event.kind() == HeapSimulatorTrace::Event::SHARE_WITH) {
+      // Nothing to do.
+    } else if (event.kind() == HeapSimulatorTrace::Event::FREE) {
+      CHECK(ContainsKey(live_buffers, buffer));
+      live_buffers.erase(buffer);
+    }
+
+    live_size += memory_delta(event);
+    if (live_size == max_live_size) {
+      break;
+    }
+  }
+  CHECK_EQ(live_size, max_live_size);
+
+  std::vector<const LogicalBuffer*> live_buffers_vector;
+  live_buffers_vector.insert(live_buffers_vector.end(), live_buffers.begin(),
+                             live_buffers.end());
+
+  // Stabily sort the live buffers.
+  std::sort(live_buffers_vector.begin(), live_buffers_vector.end(),
+            [](const LogicalBuffer* a, const LogicalBuffer* b) {
+              return a->id() < b->id();
+            });
+  return live_buffers_vector;
+}
+
+}  // namespace
+
 void BufferAssigner::AssignBuffersFromHeapSimulator(
     const HeapSimulator::Result& result, BufferAssignment* assignment,
     LogicalBuffer::Color color) {
@@ -1246,6 +1228,8 @@ void BufferAssigner::AssignBuffersFromHeapSimulator(
     const HeapSimulator::Chunk& chunk = buffer_chunk.second;
     assignment->AddAssignment(allocation, buffer, chunk.offset, chunk.size);
   }
+  allocation->peak_buffers_ =
+      ComputePeakMemoryLogicalBuffers(*allocation, result.debug_trace);
 
   VLOG(1) << "Ran heap simulation for allocation: " << allocation->ToString();
   allocation->AddHeapTrace(result.debug_trace);
diff --git a/tensorflow/compiler/xla/service/buffer_assignment.h b/tensorflow/compiler/xla/service/buffer_assignment.h
index 3086d0e2ca..15fd905e8d 100644
--- a/tensorflow/compiler/xla/service/buffer_assignment.h
+++ b/tensorflow/compiler/xla/service/buffer_assignment.h
@@ -206,17 +206,15 @@ class BufferAllocation {
     return heap_traces_;
   }
 
-  // Compute and return the LogicalBuffers which are live at the point of peak
-  // memory usage for the given allocation. The point of peak memory usage is
-  // the point at which the total size of all live logical buffers is
-  // maximal. If peak memory is reached at multiple points, the set of logical
-  // buffers live at the earliest maximal point is returned. The vector is
-  // stabily asserted by LogicalBuffer::Index.
-  //
-  // The return value is a pair of total size of the logical buffers at peak,
-  // and the buffers themselves.
-  std::pair<int64, std::vector<const LogicalBuffer*>>
-  ComputePeakMemoryLogicalBuffers() const;
+  // Returns the LogicalBuffers which are live at the point of peak memory usage
+  // for this allocation. The point of peak memory usage is the point at which
+  // the total size of all live logical buffers is maximal. If peak memory is
+  // reached at multiple points, the set of logical buffers live at the earliest
+  // maximal point is returned. The vector is stabily sorted by
+  // LogicalBuffer::Index.
+  const std::vector<const LogicalBuffer*>& PeakMemoryLogicalBuffers() const {
+    return peak_buffers_;
+  }
 
   // Get the number of bytes lost to fragmentation. This is equal to the
   // difference between the size of the allocation and the size of the maximal
@@ -291,6 +289,9 @@ class BufferAllocation {
 
   int64 fragmentation_bytes_ = 0;
   std::vector<HeapSimulatorTrace> heap_traces_;
+
+  // Set of buffers live at the point of peak memory usage for this allocation.
+  std::vector<const LogicalBuffer*> peak_buffers_;
 };
 
 // Add stream operators for nicer output of CHECK/RET_CHECK failures.
diff --git a/tensorflow/compiler/xla/service/buffer_assignment_test.cc b/tensorflow/compiler/xla/service/buffer_assignment_test.cc
index 3ec9795a65..40cf6483aa 100644
--- a/tensorflow/compiler/xla/service/buffer_assignment_test.cc
+++ b/tensorflow/compiler/xla/service/buffer_assignment_test.cc
@@ -1519,12 +1519,8 @@ TEST_F(BufferAssignmentTest, TrivialPeakBuffers) {
   // single logical buffer should be exactly the logical buffer in that
   // allocation.
   const BufferAllocation& mul_buffer = GetTopLevelAllocation(*buffers, mul);
-  int64 peak_size;
-  std::vector<const LogicalBuffer*> peak_buffers;
-
-  std::tie(peak_size, peak_buffers) =
-      mul_buffer.ComputePeakMemoryLogicalBuffers();
-  EXPECT_EQ(peak_size, ShapeUtil::ByteSizeOf(f32vec100_));
+  const std::vector<const LogicalBuffer*>& peak_buffers =
+      mul_buffer.PeakMemoryLogicalBuffers();
   ASSERT_EQ(peak_buffers.size(), 1);
   EXPECT_EQ(peak_buffers[0]->instruction(), mul);
 }
@@ -1555,6 +1551,7 @@ TEST_F(BufferAssignmentTest, PeakBuffers) {
       HloInstruction::CreateConcatenate(concat_shape, {rev, neg}, 0));
   // Make the root tiny so no interior nodes can share its buffer.
   auto root = builder.AddInstruction(HloInstruction::CreateSlice(
+
       ShapeUtil::MakeShape(F32, {1}), concat, {0}, {1}, {1}));
 
   auto module = CreateNewModule();
@@ -1569,12 +1566,10 @@ TEST_F(BufferAssignmentTest, PeakBuffers) {
   EXPECT_TRUE(buffer.IsPreallocatedTempBuffer());
   ASSERT_EQ(buffer.assigned_buffers().size(), 4);
 
-  int64 peak_size;
-  std::vector<const LogicalBuffer*> peak_buffers;
-  std::tie(peak_size, peak_buffers) = buffer.ComputePeakMemoryLogicalBuffers();
+  const std::vector<const LogicalBuffer*>& peak_buffers =
+      buffer.PeakMemoryLogicalBuffers();
 
   // The peak live set should be concat and its inputs.
-  EXPECT_EQ(peak_size, ShapeUtil::ByteSizeOf(ShapeUtil::MakeShape(F32, {400})));
   ASSERT_EQ(peak_buffers.size(), 3);
   std::vector<const HloInstruction*> peak_instructions;
   for (const LogicalBuffer* logical_buffer : peak_buffers) {
@@ -1583,6 +1578,68 @@ TEST_F(BufferAssignmentTest, PeakBuffers) {
   EXPECT_THAT(peak_instructions, UnorderedElementsAre(rev, neg, concat));
 }
 
+TEST_F(BufferAssignmentTest, PeakBuffersWhile) {
+  auto module = CreateNewModule();
+  const Shape shape = ShapeUtil::MakeShape(F32, {123, 123});
+  HloComputation* condition;
+  {
+    auto b = HloComputation::Builder(TestName() + ".cond");
+    b.AddInstruction(HloInstruction::CreateParameter(0, shape, "x"));
+    b.AddInstruction(
+        HloInstruction::CreateConstant(Literal::CreateR0<bool>(true)));
+    condition = module->AddEmbeddedComputation(b.Build());
+  }
+  HloComputation* body;
+  {
+    auto b = HloComputation::Builder(TestName() + ".body");
+    auto param =
+        b.AddInstruction(HloInstruction::CreateParameter(0, shape, "x"));
+    b.AddInstruction(
+        HloInstruction::CreateUnary(shape, HloOpcode::kNegate, param));
+    body = module->AddEmbeddedComputation(b.Build());
+  }
+  auto builder = HloComputation::Builder(TestName());
+  auto param =
+      builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0"));
+  auto copy = builder.AddInstruction(
+      HloInstruction::CreateUnary(shape, HloOpcode::kCopy, param));
+  auto while_op = builder.AddInstruction(
+      HloInstruction::CreateWhile(shape, condition, body, copy));
+  // This broadcast should get a temporary allocation which is merged with the
+  // allocation for the while. Peak buffers should include the while and the
+  // broadcast.
+  auto bcast = builder.AddInstruction(HloInstruction::CreateBroadcast(
+      ShapeUtil::MakeShape(F32, {123, 123, 123}), while_op, {0, 1}));
+  builder.AddInstruction(HloInstruction::CreateReverse(
+      ShapeUtil::MakeShape(F32, {123, 123, 123}), bcast, {0}));
+  module->AddEntryComputation(builder.Build());
+
+  auto buffers = RunBufferAssignment(module.get());
+  const BufferAllocation& buffer = GetTopLevelAllocation(*buffers, bcast);
+  const std::vector<const LogicalBuffer*>& peak_buffers =
+      buffer.PeakMemoryLogicalBuffers();
+  ASSERT_EQ(peak_buffers.size(), 2);
+
+  // The peak buffers should include the broadcast and one of the colocated
+  // buffers of the while (body param, condition param, body root, or the while
+  // itself).
+  const LogicalBuffer* bcast_buffer;
+  const LogicalBuffer* nonbcast_buffer;
+  if (peak_buffers[0]->instruction() == bcast) {
+    bcast_buffer = peak_buffers[0];
+    nonbcast_buffer = peak_buffers[1];
+  } else {
+    bcast_buffer = peak_buffers[1];
+    nonbcast_buffer = peak_buffers[0];
+  }
+  EXPECT_EQ(bcast_buffer->instruction(), bcast);
+  EXPECT_TRUE(
+      nonbcast_buffer->instruction() == while_op ||
+      nonbcast_buffer->instruction() == body->parameter_instruction(0) ||
+      nonbcast_buffer->instruction() == body->root_instruction() ||
+      nonbcast_buffer->instruction() == condition->parameter_instruction(0));
+}
+
 class WhileBufferAssignmentTest : public HloTestBase {
  protected:
   std::unique_ptr<HloComputation> BuildWhileConditionComputation(
-- 
GitLab


From 19ad98e8393547076706285b922bd801d763033f Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Wed, 2 May 2018 00:56:43 +0300
Subject: [PATCH 113/395] [tf.data] Fix debug output.

---
 tensorflow/contrib/data/python/ops/interleave_ops.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/contrib/data/python/ops/interleave_ops.py b/tensorflow/contrib/data/python/ops/interleave_ops.py
index 0852fc6be8..2a9c5b45f8 100644
--- a/tensorflow/contrib/data/python/ops/interleave_ops.py
+++ b/tensorflow/contrib/data/python/ops/interleave_ops.py
@@ -240,9 +240,9 @@ def sample_from_datasets(datasets, weights=None, seed=None):
   selector_input = dataset_ops.Dataset.zip(
       (logits_ds, random_ops.RandomDataset(seed).batch(2))).map(select_dataset)
 
-  logging.warn('selector_input.output_types: %s', selector_input.output_types)
-  logging.warn('selector_input.output_shapes: %s', selector_input.output_shapes)
+  logging.warn('selector_input.output_types: %s', str(selector_input.output_types))
+  logging.warn('selector_input.output_shapes: %s', str(selector_input.output_shapes))
   for i, dataset in enumerate(datasets):
-    logging.warn('dataset %i output_types: %s' % (i, dataset.output_types))
-    logging.warn('dataset %i output_shapes: %s' % (i, dataset.output_shapes))
+    logging.warn('dataset %i output_types: %s' % (i, str(dataset.output_types)))
+    logging.warn('dataset %i output_shapes: %s' % (i, str(dataset.output_shapes)))
   return DirectedInterleaveDataset(selector_input, datasets)
-- 
GitLab


From 415ea7360d3f57249fc18e068852a8b8ce6d7f77 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 15:00:20 -0700
Subject: [PATCH 114/395] Making ids unique in nn.embedding_lookup_sparse. This
 helps to reduce RPC calls for looking up the embeddings when there are
 repeated ids in the batch.

PiperOrigin-RevId: 195002785
---
 tensorflow/python/ops/embedding_ops.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensorflow/python/ops/embedding_ops.py b/tensorflow/python/ops/embedding_ops.py
index 6f2a34c731..bcc717b043 100644
--- a/tensorflow/python/ops/embedding_ops.py
+++ b/tensorflow/python/ops/embedding_ops.py
@@ -385,7 +385,7 @@ def embedding_lookup_sparse(params,
       ```
 
   Raises:
-    TypeError: If `sp_ids` is not a `SparseTensor`, or if `sp_weights` is 
+    TypeError: If `sp_ids` is not a `SparseTensor`, or if `sp_weights` is
       neither `None` nor `SparseTensor`.
     ValueError: If `combiner` is not one of {"mean", "sqrtn", "sum"}.
   """
@@ -421,10 +421,7 @@ def embedding_lookup_sparse(params,
       segment_ids = math_ops.cast(segment_ids, dtypes.int32)
 
     ids = sp_ids.values
-    if ignore_weights:
-      ids, idx = array_ops.unique(ids)
-    else:
-      idx = None
+    ids, idx = array_ops.unique(ids)
 
     embeddings = embedding_lookup(
         params, ids, partition_strategy=partition_strategy, max_norm=max_norm)
@@ -433,6 +430,8 @@ def embedding_lookup_sparse(params,
       if weights.dtype != embeddings.dtype:
         weights = math_ops.cast(weights, embeddings.dtype)
 
+      embeddings = array_ops.gather(embeddings, idx)
+
       # Reshape weights to allow broadcast
       ones = array_ops.fill(
           array_ops.expand_dims(array_ops.rank(embeddings) - 1, 0), 1)
-- 
GitLab


From 707b0c9cc4d5335d04fce4addb8ed2f158cbd1c0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 15:01:22 -0700
Subject: [PATCH 115/395] Minor JNI performance improvement.

PiperOrigin-RevId: 195002949
---
 tensorflow/contrib/lite/java/src/main/native/tensor_jni.cc | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/tensorflow/contrib/lite/java/src/main/native/tensor_jni.cc b/tensorflow/contrib/lite/java/src/main/native/tensor_jni.cc
index 17f4be09c6..005dca0253 100644
--- a/tensorflow/contrib/lite/java/src/main/native/tensor_jni.cc
+++ b/tensorflow/contrib/lite/java/src/main/native/tensor_jni.cc
@@ -238,10 +238,6 @@ Java_org_tensorflow_lite_Tensor_shape(JNIEnv* env, jclass clazz, jlong handle) {
   if (tensor == nullptr) return nullptr;
   int num_dims = tensor->dims->size;
   jintArray result = env->NewIntArray(num_dims);
-  jint* dims = env->GetIntArrayElements(result, nullptr);
-  for (int i = 0; i < num_dims; ++i) {
-    dims[i] = static_cast<jint>(tensor->dims->data[i]);
-  }
-  env->ReleaseIntArrayElements(result, dims, 0);
+  env->SetIntArrayRegion(result, 0, num_dims, tensor->dims->data);
   return result;
 }
-- 
GitLab


From 6f10fb5b583cb7b883a41a45a69b22fd84eeb10e Mon Sep 17 00:00:00 2001
From: Benoit Steiner <bsteiner@google.com>
Date: Tue, 1 May 2018 15:19:42 -0700
Subject: [PATCH 116/395] Fixed some outdated comments

PiperOrigin-RevId: 195006088
---
 .../core/grappler/costs/graph_properties.cc      | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/tensorflow/core/grappler/costs/graph_properties.cc b/tensorflow/core/grappler/costs/graph_properties.cc
index a12d9b932b..431efb08cb 100644
--- a/tensorflow/core/grappler/costs/graph_properties.cc
+++ b/tensorflow/core/grappler/costs/graph_properties.cc
@@ -382,7 +382,7 @@ class TopoQueue {
   std::set<const NodeDef*, CompareNodes> queue_;
 };
 
-// Merge and relax symbolic shapes.
+// Processes symbolic shapes.
 // Each symbolic shape or dimension is represented by a handle. Unlike the TF
 // shape refiner which creates new handles every time it processes an unknown
 // shape/dimension, the symbolic shape refiner assigns a specific handle to each
@@ -864,11 +864,8 @@ Status GraphProperties::RelaxEnqueueShapesAndMergeTypes(
   return Status::OK();
 }
 
-// If a Merge node has a NextIteration node as an input then that input will
-// try to forward an UnknownShape at graph construction time. However, the
-// Merge shape function will always propagate an UnknownShape if any of its
-// inputs are UnknownShapes. So we need to ignore the input from NextIteration
-// nodes to propagate any known shape from the Merge node.
+// Compute the output shape of the merge node as the union of the available
+// input shapes.
 Status GraphProperties::UpdateMergeNode(SymbolicShapeRefiner* shape_refiner,
                                         const NodeDef* node,
                                         bool* new_shapes) const {
@@ -914,8 +911,7 @@ Status GraphProperties::UpdateMergeNode(SymbolicShapeRefiner* shape_refiner,
   return Status::OK();
 }
 
-// Manually propagate the input shape for Enter nodes and update any Merge node
-// outputs.
+// Manually propagate the input shape for Enter nodes.
 Status GraphProperties::UpdateEnter(SymbolicShapeRefiner* shape_refiner,
                                     const NodeDef* node, bool* new_shapes) {
   auto enter_ctx = shape_refiner->GetContext(node);
@@ -955,6 +951,8 @@ Status GraphProperties::UpdateShapes(
     // Properly handle merge nodes.
     TF_RETURN_IF_ERROR(UpdateMergeNode(shape_refiner, n, new_shapes));
   } else if (IsEnqueue(*n)) {
+    // Make sure the shapes of enqueued tensors are propagated to the queue
+    // itself.
     TF_RETURN_IF_ERROR(
         UpdateEnqueue(n, resource_handles, shape_refiner, new_shapes));
   } else {
@@ -1209,7 +1207,6 @@ Status GraphProperties::InferStatically(bool assume_valid_feeds) {
 
     // Fill input properties.
     {
-      // CHECK_EQ(ctx->num_inputs(), node.num_inputs());
       auto& input_properties = input_properties_[node.name()];
 
       // Should always be empty, node names in graph are supposed to be unique.
@@ -1233,7 +1230,6 @@ Status GraphProperties::InferStatically(bool assume_valid_feeds) {
 
     // Fill output properties.
     {
-      // CHECK_EQ(ctx->num_outputs(), node->num_outputs());
       auto& output_properties = output_properties_[node.name()];
 
       // Should always be empty, node names in graph are supposed to be unique.
-- 
GitLab


From 33978e881c8b0aed71e26858641736313a486c12 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 15:47:26 -0700
Subject: [PATCH 117/395] Go: Update generated wrapper functions for TensorFlow
 ops. PiperOrigin-RevId: 195010310

---
 tensorflow/go/op/wrappers.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/go/op/wrappers.go b/tensorflow/go/op/wrappers.go
index 2f1be51ada..c12ea51563 100644
--- a/tensorflow/go/op/wrappers.go
+++ b/tensorflow/go/op/wrappers.go
@@ -21386,7 +21386,7 @@ func ImageSummaryBadColor(value tf.Tensor) ImageSummaryAttr {
 //    generated sequentially as '*tag*/image/0', '*tag*/image/1', etc.
 //
 // The `bad_color` argument is the color to use in the generated images for
-// non-finite input values.  It is a `uint8` 1-D tensor of length `channels`.
+// non-finite input values.  It is a `unit8` 1-D tensor of length `channels`.
 // Each element must be in the range `[0, 255]` (It represents the value of a
 // pixel in the output image).  Non-finite values in the input tensor are
 // replaced by this tensor in the output image.  The default value is the color
-- 
GitLab


From b25e6fe32cccd29ec4cb4014bbb45d62b75835b4 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 15:47:27 -0700
Subject: [PATCH 118/395] Implementation of the fully-connected TFLite Op using
 the symmetric quantization.

PiperOrigin-RevId: 195010312
---
 tensorflow/contrib/lite/kernels/BUILD         |   2 +
 .../contrib/lite/kernels/fully_connected.cc   | 117 ++++++++++++++-
 .../lite/kernels/fully_connected_test.cc      | 141 ++++++++++++++++--
 tensorflow/contrib/lite/kernels/test_util.h   |  17 +++
 4 files changed, 255 insertions(+), 22 deletions(-)

diff --git a/tensorflow/contrib/lite/kernels/BUILD b/tensorflow/contrib/lite/kernels/BUILD
index 689f9bfa71..57b3136cce 100644
--- a/tensorflow/contrib/lite/kernels/BUILD
+++ b/tensorflow/contrib/lite/kernels/BUILD
@@ -31,6 +31,7 @@ cc_library(
         "//tensorflow/contrib/lite:framework",
         "//tensorflow/contrib/lite:schema_fbs_version",
         "//tensorflow/contrib/lite:string_util",
+        "//tensorflow/contrib/lite/kernels/internal:tensor_utils",
         "//tensorflow/contrib/lite/testing:util",
         "//tensorflow/core:tflite_portable_logging",
         "@com_google_googletest//:gtest",
@@ -672,6 +673,7 @@ tf_cc_test(
         ":builtin_ops",
         "//tensorflow/contrib/lite:framework",
         "//tensorflow/contrib/lite/kernels:test_util",
+        "//tensorflow/contrib/lite/kernels/internal:tensor_utils",
         "@com_google_absl//absl/memory",
         "@com_google_googletest//:gtest",
     ],
diff --git a/tensorflow/contrib/lite/kernels/fully_connected.cc b/tensorflow/contrib/lite/kernels/fully_connected.cc
index 888e67966c..c5bf50da5f 100644
--- a/tensorflow/contrib/lite/kernels/fully_connected.cc
+++ b/tensorflow/contrib/lite/kernels/fully_connected.cc
@@ -55,19 +55,24 @@ struct OpData {
   // uint8_t these would be 0 and 255.
   int32_t output_activation_min;
   int32_t output_activation_max;
+  // The index of the temporary tensor where the quantized inputs are cached.
+  int input_quantized_index;
 };
 
 constexpr int kInputTensor = 0;
 constexpr int kWeightsTensor = 1;
 constexpr int kBiasTensor = 2;
 constexpr int kOutputTensor = 0;
+constexpr int kScratchBufferTensor = 1;
 
 void* Init(TfLiteContext* context, const char* buffer, size_t length) {
   // This is a builtin op, so we don't use the contents in 'buffer', if any.
   // Instead, we allocate a new object to carry information from Prepare() to
   // Eval().
   gemm_support::IncrementUsageCounter(context);
-  return new OpData;
+  auto* op_data = new OpData;
+  context->AddTensors(context, 1, &op_data->input_quantized_index);
+  return op_data;
 }
 
 void Free(TfLiteContext* context, void* buffer) {
@@ -121,6 +126,27 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                                   &data->output_activation_max);
   }
 
+  // If we have to perform on-the-fly quantization (with quantized weights and
+  // float inputs) first we need to quantize the inputs. Allocate a temporary
+  // buffer to store the intermediate quantized values.
+  if (input->type == kTfLiteFloat32 && filter->type == kTfLiteUInt8) {
+    TfLiteIntArrayFree(node->temporaries);
+    node->temporaries = TfLiteIntArrayCreate(1);
+    node->temporaries->data[0] = data->input_quantized_index;
+
+    TfLiteTensor* input_quantized =
+        &context->tensors[node->temporaries->data[0]];
+    input_quantized->type = kTfLiteUInt8;
+    input_quantized->allocation_type = kTfLiteArenaRw;
+
+    // TODO(raziel): add this logic to ResizeTensor.
+    if (!TfLiteIntArrayEqual(input_quantized->dims, input->dims)) {
+      TfLiteIntArray* input_quantized_size = TfLiteIntArrayCopy(input->dims);
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, input_quantized,
+                                                       input_quantized_size));
+    }
+  }
+
   // Resize output.
   TfLiteIntArray* output_size_array = TfLiteIntArrayCreate(2);
   output_size_array->data[0] = batch_size;
@@ -163,6 +189,74 @@ TfLiteStatus EvalPie(TfLiteContext* context, TfLiteNode* node,
   return kTfLiteOk;
 }
 
+TfLiteStatus EvalPieQuantized(TfLiteContext* context, TfLiteNode* node,
+                              TfLiteFullyConnectedParams* params, OpData* data,
+                              TfLiteTensor* input, TfLiteTensor* filter,
+                              TfLiteTensor* bias, TfLiteTensor* input_quantized,
+                              TfLiteTensor* output) {
+  // Check the types for this hybrid Op.
+  TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_EQ(context, filter->type, kTfLiteUInt8);
+  TF_LITE_ENSURE_EQ(context, bias->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32);
+
+  int total_input_size = 1;
+  for (int i = 0; i < input->dims->size; i++) {
+    total_input_size *= input->dims->data[i];
+  }
+
+  const int input_size = filter->dims->data[1];
+  const int batch_size = total_input_size / filter->dims->data[1];
+  const int num_units = filter->dims->data[0];
+
+  // Output = bias if bias tensor exists.
+  if (bias) {
+    tensor_utils::VectorBatchVectorAssign(bias->data.f, num_units, batch_size,
+                                          output->data.f);
+  } else {
+    tensor_utils::ZeroVector(output->data.f, batch_size * num_units);
+  }
+
+  // TODO(mirkov): change std::minmax_element with a vectorized call.
+  auto minmax_element =
+      std::minmax_element(input->data.f, input->data.f + total_input_size);
+  // Save matrix multiplication computation for all zero input.
+  if (*minmax_element.first == 0.0 && *minmax_element.second == 0.0) {
+    tensor_utils::ApplyActivationToVector(output->data.f,
+                                          batch_size * num_units,
+                                          params->activation, output->data.f);
+    return kTfLiteOk;
+  }
+
+  // Quantize input from float to uint8 + quantization params (scaling factor).
+  float min, max;
+  float* scaling_factors = new float[batch_size];
+
+  // Quantize each batch independently.
+  for (int b = 0; b < batch_size; ++b) {
+    const int offset = b * input_size;
+    tensor_utils::SymmetricQuantizeFloats(
+        input->data.f + offset, input_size,
+        reinterpret_cast<int8_t*>(input_quantized->data.uint8) + offset, &min,
+        &max, &scaling_factors[b]);
+    // Incorporate scaling of the filter.
+    scaling_factors[b] *= filter->params.scale;
+  }
+
+  // Compute output += weight * quantized_input
+  tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+      reinterpret_cast<int8_t*>(filter->data.uint8), num_units, input_size,
+      reinterpret_cast<int8_t*>(input_quantized->data.uint8), scaling_factors,
+      batch_size, output->data.f, /*result_stride=*/1);
+
+  // Apply activation function to floats.
+  tensor_utils::ApplyActivationToVector(output->data.f, batch_size * num_units,
+                                        params->activation, output->data.f);
+  delete[] scaling_factors;
+
+  return kTfLiteOk;
+}
+
 #define TF_LITE_MACRO_DISPATCH(macro_name, params, target_namespace) \
   if (params->activation == kTfLiteActNone) {                        \
     macro_name(target_namespace, kNone);                             \
@@ -178,7 +272,8 @@ template <KernelType kernel_type>
 TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
                            TfLiteFullyConnectedParams* params, OpData* data,
                            TfLiteTensor* input, TfLiteTensor* filter,
-                           TfLiteTensor* bias, TfLiteTensor* output) {
+                           TfLiteTensor* bias, TfLiteTensor* input_quantized,
+                           TfLiteTensor* output) {
   gemmlowp::GemmContext* gemm_context = gemm_support::GetFromContext(context);
 
   int32_t input_offset = -input->params.zero_point;
@@ -195,9 +290,15 @@ TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
   if (kernel_type == kReference) {
     TF_LITE_FULLY_CONNECTED(reference_ops);
   } else if (kernel_type == kPie) {
-    // TODO(ahentz): we don't have a quantized version of the PIE kernels, so
-    // we just defer to the MINI ones.
-    TF_LITE_FULLY_CONNECTED(optimized_ops);
+    if (input->type == kTfLiteFloat32) {
+      // Pie currently only supports quantized models and float inputs/outputs.
+      return EvalPieQuantized(context, node, params, data, input, filter, bias,
+                              input_quantized, output);
+    } else {
+      // TODO(ahentz): we don't have a quantized version of the PIE kernels, so
+      // we just defer to the MINI ones.
+      TF_LITE_FULLY_CONNECTED(optimized_ops);
+    }
   } else {
     TF_LITE_FULLY_CONNECTED(optimized_ops);
   }
@@ -245,13 +346,15 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   TfLiteTensor* bias = GetOptionalInputTensor(context, node, kBiasTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  switch (input->type) {  // Already know in/out types are same.
+  TfLiteTensor* input_quantized = &context->tensors[node->temporaries->data[0]];
+
+  switch (filter->type) {  // Already know in/out types are same.
     case kTfLiteFloat32:
       return EvalFloat<kernel_type>(context, node, params, data, input, filter,
                                     bias, output);
     case kTfLiteUInt8:
       return EvalQuantized<kernel_type>(context, node, params, data, input,
-                                        filter, bias, output);
+                                        filter, bias, input_quantized, output);
     default:
       context->ReportError(context, "Type not currently supported.");
       return kTfLiteError;
diff --git a/tensorflow/contrib/lite/kernels/fully_connected_test.cc b/tensorflow/contrib/lite/kernels/fully_connected_test.cc
index 87413000a9..05dd028b48 100644
--- a/tensorflow/contrib/lite/kernels/fully_connected_test.cc
+++ b/tensorflow/contrib/lite/kernels/fully_connected_test.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <gtest/gtest.h>
 #include "absl/memory/memory.h"
 #include "tensorflow/contrib/lite/interpreter.h"
+#include "tensorflow/contrib/lite/kernels/internal/tensor_utils.h"
 #include "tensorflow/contrib/lite/kernels/register.h"
 #include "tensorflow/contrib/lite/kernels/test_util.h"
 #include "tensorflow/contrib/lite/model.h"
@@ -224,6 +225,60 @@ class QuantizedFullyConnectedOpModel : public BaseFullyConnectedOpModel {
   }
 };
 
+// In the hybrid model the weights are quantized (to uint8). But the bias,
+// input (and output) are expected to be in float precision.
+class HybridFullyConnectedOpModel : public SingleOpModel {
+ public:
+  HybridFullyConnectedOpModel(int units, int batches, const TensorData& input,
+                              const TensorData& weights,
+                              const TensorData& output = {TensorType_FLOAT32})
+      : batches_(batches), units_(units) {
+    int total_input_size = 1;
+    for (int i = 0; i < input.shape.size(); ++i) {
+      total_input_size *= input.shape[i];
+    }
+    input_size_ = total_input_size / batches_;
+
+    input_ = AddInput(input);
+    weights_ = AddInput(weights);
+
+    TensorData bias{TensorType_FLOAT32, {units_}};
+    bias_ = AddInput(bias);
+
+    output_ = AddOutput(output);
+
+    SetBuiltinOp(
+        BuiltinOperator_FULLY_CONNECTED, BuiltinOptions_FullyConnectedOptions,
+        CreateFullyConnectedOptions(builder_, ActivationFunctionType_RELU)
+            .Union());
+    resolver_ = absl::make_unique<SingleOpResolver>(
+        BuiltinOperator_FULLY_CONNECTED,
+        ops::builtin::Register_FULLY_CONNECTED_PIE());
+    BuildInterpreter({GetShape(input_), GetShape(weights_), GetShape(bias_)});
+  }
+  void SetBias(std::initializer_list<float> f) { PopulateTensor(bias_, f); }
+  void SetWeights(std::initializer_list<float> data) {
+    SymmetricQuantizeAndPopulate(weights_, data);
+  }
+
+  void SetInput(std::initializer_list<float> f) { PopulateTensor(input_, f); }
+  std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
+
+  int input_size() { return input_size_; }
+  int num_units() { return units_; }
+  int num_batches() { return batches_; }
+
+ protected:
+  int input_;
+  int weights_;
+  int bias_;
+  int output_;
+
+  int batches_;
+  int units_;
+  int input_size_;
+};
+
 const auto kKernelMap = new std::map<string, TfLiteRegistration*>({
     {"Reference", ops::builtin::Register_FULLY_CONNECTED_REF()},
     {"NeonOptimized", ops::builtin::Register_FULLY_CONNECTED_NEON_OPT()},
@@ -231,18 +286,43 @@ const auto kKernelMap = new std::map<string, TfLiteRegistration*>({
     {"Pie", ops::builtin::Register_FULLY_CONNECTED_PIE()},
 });
 
-class FullyConnectedOpTest : public SingleOpTest {
+class FloatFullyConnectedOpTest : public SingleOpTest {
  protected:
   const std::map<string, TfLiteRegistration*>& GetKernelMap() override {
     return *kKernelMap;
   }
 };
 
+const auto kKernelMapNoPie = new std::map<string, TfLiteRegistration*>({
+    {"Reference", ops::builtin::Register_FULLY_CONNECTED_REF()},
+    {"NeonOptimized", ops::builtin::Register_FULLY_CONNECTED_NEON_OPT()},
+    {"GenericOptimized", ops::builtin::Register_FULLY_CONNECTED_GENERIC_OPT()},
+});
+
+class QuantizedFullyConnectedOpTest : public SingleOpTest {
+ protected:
+  const std::map<string, TfLiteRegistration*>& GetKernelMap() override {
+    return *kKernelMapNoPie;
+  }
+};
+
+const auto kKernelMapPie = new std::map<string, TfLiteRegistration*>({
+    {"Pie", ops::builtin::Register_FULLY_CONNECTED_PIE()},
+});
+
+// Hybrid mode is used by the Pie quantized kernel.
+class HybridFullyConnectedOpTest : public SingleOpTest {
+ protected:
+  const std::map<string, TfLiteRegistration*>& GetKernelMap() override {
+    return *kKernelMapPie;
+  }
+};
+
 // TODO(ahentz): add more small tests like this one, focused on making sure the
 // calculations are correct.
-TEST_P(FullyConnectedOpTest, SimpleTest) {
-  FloatFullyConnectedOpModel m(GetRegistration(), 3, 2,
-                               {TensorType_FLOAT32, {2, 10}});
+TEST_P(FloatFullyConnectedOpTest, SimpleTest) {
+  FloatFullyConnectedOpModel m(GetRegistration(), /*units=*/3, /*batches=*/2,
+                               /*input=*/{TensorType_FLOAT32, {2, 10}});
   m.SetWeights({
       1, 2, 3, 4, 5, 6, 7, 8, 9, 10,  // u = 0
       1, 2, 3, 4, 5, 6, 7, 8, 9, 10,  // u = 1
@@ -260,9 +340,9 @@ TEST_P(FullyConnectedOpTest, SimpleTest) {
   EXPECT_THAT(m.GetOutput(), ElementsAre(24, 25, 26, 58, 59, 60));
 }
 
-TEST_P(FullyConnectedOpTest, SimpleTestQuantized) {
+TEST_P(QuantizedFullyConnectedOpTest, SimpleTestQuantized) {
   QuantizedFullyConnectedOpModel m(
-      GetRegistration(), 3, 2,
+      GetRegistration(), /*units=*/3, /*batches*/ 2,
       /*input=*/{TensorType_UINT8, {2, 10}, -63.5, 64},
       /*output=*/{TensorType_UINT8, {}, -127, 128});
 
@@ -288,13 +368,40 @@ TEST_P(FullyConnectedOpTest, SimpleTestQuantized) {
   EXPECT_THAT(m.GetOutput(), ElementsAre(151, 152, 153, 185, 186, 187));
 }
 
-TEST(FullyConnectedOpTest, SimpleTest4DInput) {
+TEST(HybridFullyConnectedOpTest, SimpleTestQuantized) {
+  HybridFullyConnectedOpModel m(
+      /*units=*/3, /*batches=*/2,
+      /*input=*/{TensorType_FLOAT32, {2, 10}},
+      /*weights=*/{TensorType_UINT8, {3, 10}, -63.5, 64});  // PIE
+
+  m.SetWeights({
+      1, 2, 3, 4, 5, 6, 7, 8, 9, 10,  // u = 0
+      1, 2, 3, 4, 5, 6, 7, 8, 9, 10,  // u = 1
+      1, 2, 3, 4, 5, 6, 7, 8, 9, 10,  // u = 1
+  });
+  m.SetBias({1, 2, 3});
+
+  m.SetInput({
+      1, 2, 3, 4, 5, 6, 7, 8,  -9, -10,  // b = 0
+      1, 2, 3, 4, 5, 6, 7, -8, 9,  -10,  // b = 1
+  });
+
+  m.Invoke();
+
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear(
+                                 {
+                                     24, 25, 26,  //
+                                     58, 59, 60,  //
+                                 },
+                                 /*max_abs_error=*/1.3f)));
+}
+
+TEST(FloatFullyConnectedOpTest, SimpleTest4DInput) {
   // Note that it is not required that the first dimension be the number of
   // batches. All we care is that the input can be evenly distributed in
   // batches. In this case, we need the input to have multiples of '2'.
   FloatFullyConnectedOpModel m(ops::builtin::Register_FULLY_CONNECTED_PIE(),
-                               /*units=*/3,
-                               /*batches=*/2,
+                               /*units=*/3, /*batches=*/2,
                                /*input=*/{TensorType_FLOAT32, {4, 1, 5, 1}});
   m.SetWeights({
       1, 2, 3, 4, 5, 6, 7, 8, 9, 10,  // u = 0
@@ -316,9 +423,9 @@ TEST(FullyConnectedOpTest, SimpleTest4DInput) {
                              }));
 }
 
-TEST_P(FullyConnectedOpTest, SimpleTest4dInputQuantized) {
+TEST_P(QuantizedFullyConnectedOpTest, SimpleTest4dInputQuantized) {
   QuantizedFullyConnectedOpModel m(
-      GetRegistration(), 3, 2,
+      GetRegistration(), /*units=*/3, /*batches=*/2,
       /*input=*/{TensorType_UINT8, {4, 1, 5, 1}, -63.5, 64},
       /*output=*/{TensorType_UINT8, {}, -127, 128});
 
@@ -345,14 +452,18 @@ TEST_P(FullyConnectedOpTest, SimpleTest4dInputQuantized) {
 }
 
 INSTANTIATE_TEST_CASE_P(
-    FullyConnectedOpTest, FullyConnectedOpTest,
+    FloatFullyConnectedOpTest, FloatFullyConnectedOpTest,
     ::testing::ValuesIn(SingleOpTest::GetKernelTags(*kKernelMap)));
 
+INSTANTIATE_TEST_CASE_P(
+    QuantizedFullyConnectedOpTest, QuantizedFullyConnectedOpTest,
+    ::testing::ValuesIn(SingleOpTest::GetKernelTags(*kKernelMapNoPie)));
+
 // TODO(ahentz): Reconsider this test. Having arbitrary weights makes it hard
 // to debug errors and doesn't necessarily test all the important details.
-TEST_P(FullyConnectedOpTest, BlackBoxTest) {
-  FloatFullyConnectedOpModel m(GetRegistration(), 16, 2,
-                               {TensorType_FLOAT32, {2, 8}});
+TEST_P(FloatFullyConnectedOpTest, BlackBoxTest) {
+  FloatFullyConnectedOpModel m(GetRegistration(), /*units=*/16, /*batches=*/2,
+                               /*input=*/{TensorType_FLOAT32, {2, 8}});
   m.SetWeights(
       {0.091327,  0.103366,  -0.316505, -0.083120, 0.149366,  -0.196636,
        -0.123672, 0.062800,  0.063031,  0.191670,  -0.062001, -0.061504,
diff --git a/tensorflow/contrib/lite/kernels/test_util.h b/tensorflow/contrib/lite/kernels/test_util.h
index a9064d54e7..6fb6fe27eb 100644
--- a/tensorflow/contrib/lite/kernels/test_util.h
+++ b/tensorflow/contrib/lite/kernels/test_util.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include <gtest/gtest.h>
 
 #include "tensorflow/contrib/lite/interpreter.h"
+#include "tensorflow/contrib/lite/kernels/internal/tensor_utils.h"
 #include "tensorflow/contrib/lite/kernels/register.h"
 #include "tensorflow/contrib/lite/model.h"
 #include "tensorflow/contrib/lite/string_util.h"
@@ -133,6 +134,22 @@ class SingleOpModel {
     PopulateTensor(index, 0, q.data(), q.data() + q.size());
   }
 
+  void SymmetricQuantizeAndPopulate(int index,
+                                    std::initializer_list<float> data) {
+    TfLiteTensor* t = interpreter_->tensor(index);
+    std::vector<float> values(data);
+    const int length = values.size();
+    std::vector<int8_t> q(length);
+    float min, max, scaling_factor;
+    tensor_utils::SymmetricQuantizeFloats(values.data(), length, q.data(), &min,
+                                          &max, &scaling_factor);
+    // Update quantization params.
+    t->params.scale = scaling_factor;
+    t->params.zero_point = 0;
+    PopulateTensor(index, /*offset=*/0, reinterpret_cast<uint8_t*>(q.data()),
+                   reinterpret_cast<uint8_t*>(q.data() + q.size()));
+  }
+
   const std::vector<int>& GetShape(int id) { return tensor_data_.at(id).shape; }
 
   float GetScale(int id) { return tensor_data_.at(id).scale; }
-- 
GitLab


From b2aebe0721f630a2cbc4769d1d5b9eb5b1691824 Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Wed, 2 May 2018 02:01:14 +0300
Subject: [PATCH 119/395] [tf.data] Try fixing the Windows build by adding the
 directed interleave kernel to this cmake file.

---
 tensorflow/contrib/cmake/tf_core_kernels.cmake | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/contrib/cmake/tf_core_kernels.cmake b/tensorflow/contrib/cmake/tf_core_kernels.cmake
index f38c9e0513..1505d3e208 100644
--- a/tensorflow/contrib/cmake/tf_core_kernels.cmake
+++ b/tensorflow/contrib/cmake/tf_core_kernels.cmake
@@ -68,6 +68,7 @@ if(tensorflow_BUILD_CONTRIB_KERNELS)
       "${tensorflow_source_dir}/tensorflow/contrib/coder/kernels/range_coder_ops.cc"
       "${tensorflow_source_dir}/tensorflow/contrib/coder/kernels/range_coder_ops_util.cc"
       "${tensorflow_source_dir}/tensorflow/contrib/coder/ops/coder_ops.cc"
+      "${tensorflow_source_dir}/tensorflow/contrib/data/kernels/directed_interleave_dataset_op.cc"
       "${tensorflow_source_dir}/tensorflow/contrib/data/kernels/ignore_errors_dataset_op.cc"
       "${tensorflow_source_dir}/tensorflow/contrib/data/kernels/prefetching_kernels.cc"
       "${tensorflow_source_dir}/tensorflow/contrib/data/kernels/threadpool_dataset_op.cc"
-- 
GitLab


From 29cd3f96322f3d5326a2dbe6a9c502919159c9fc Mon Sep 17 00:00:00 2001
From: joel-shor <joelshor@google.com>
Date: Wed, 2 May 2018 02:14:30 +0300
Subject: [PATCH 120/395] [tf.data] Remove debug code.

---
 tensorflow/contrib/data/python/ops/BUILD             | 1 -
 tensorflow/contrib/data/python/ops/interleave_ops.py | 6 ------
 2 files changed, 7 deletions(-)

diff --git a/tensorflow/contrib/data/python/ops/BUILD b/tensorflow/contrib/data/python/ops/BUILD
index 9959ccc005..7a3e42cc72 100644
--- a/tensorflow/contrib/data/python/ops/BUILD
+++ b/tensorflow/contrib/data/python/ops/BUILD
@@ -184,7 +184,6 @@ py_library(
         "//tensorflow/python/data/ops:readers",
         "//tensorflow/python/data/util:nest",
         "//tensorflow/python/data/util:sparse",
-        "//tensorflow/python:platform",
     ],
 )
 
diff --git a/tensorflow/contrib/data/python/ops/interleave_ops.py b/tensorflow/contrib/data/python/ops/interleave_ops.py
index 2a9c5b45f8..812a50ecbf 100644
--- a/tensorflow/contrib/data/python/ops/interleave_ops.py
+++ b/tensorflow/contrib/data/python/ops/interleave_ops.py
@@ -30,7 +30,6 @@ from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.util import deprecation
-from tensorflow.python.platform import tf_logging as logging
 
 
 def parallel_interleave(map_func,
@@ -240,9 +239,4 @@ def sample_from_datasets(datasets, weights=None, seed=None):
   selector_input = dataset_ops.Dataset.zip(
       (logits_ds, random_ops.RandomDataset(seed).batch(2))).map(select_dataset)
 
-  logging.warn('selector_input.output_types: %s', str(selector_input.output_types))
-  logging.warn('selector_input.output_shapes: %s', str(selector_input.output_shapes))
-  for i, dataset in enumerate(datasets):
-    logging.warn('dataset %i output_types: %s' % (i, str(dataset.output_types)))
-    logging.warn('dataset %i output_shapes: %s' % (i, str(dataset.output_shapes)))
   return DirectedInterleaveDataset(selector_input, datasets)
-- 
GitLab


From 210abebd3febdd2c44ab5021bcebf8f1f5d451c4 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 16:17:21 -0700
Subject: [PATCH 121/395] [TF:XLA] Separate on-host and on-device shape and
 layout in HloModule.

Previously, only one layout was stored with an HLO module. This CL allows
HLO passes to modify the on-device layouts without affecting the on-host
layout (provided by the client)

PiperOrigin-RevId: 195014875
---
 .../compiler/xla/client/local_client.cc       | 36 ++++++++++---
 tensorflow/compiler/xla/client/local_client.h |  9 ----
 .../compiler/xla/service/cpu/cpu_compiler.cc  |  2 +-
 .../xla/service/cpu/cpu_executable.cc         |  5 +-
 .../xla/service/cpu/cpu_layout_assignment.h   |  3 +-
 .../service/cpu/cpu_layout_assignment_test.cc |  4 +-
 tensorflow/compiler/xla/service/executable.h  |  4 +-
 .../compiler/xla/service/gpu/gpu_compiler.cc  |  2 +-
 .../xla/service/gpu/gpu_layout_assignment.h   |  3 +-
 .../service/gpu/gpu_layout_assignment_test.cc |  8 +--
 tensorflow/compiler/xla/service/hlo_module.cc | 17 +++---
 tensorflow/compiler/xla/service/hlo_module.h  | 16 ++++--
 .../compiler/xla/service/hlo_module_config.cc | 17 ++++--
 .../compiler/xla/service/hlo_module_config.h  | 43 ++++++++++-----
 .../xla/service/interpreter/compiler.cc       |  2 +-
 .../compiler/xla/service/layout_assignment.cc | 15 +++---
 .../compiler/xla/service/layout_assignment.h  |  4 +-
 .../xla/service/layout_assignment_test.cc     |  8 +--
 tensorflow/compiler/xla/service/service.cc    | 52 +++++++++++++++----
 tensorflow/compiler/xla/service/service.h     |  3 ++
 tensorflow/compiler/xla/tests/BUILD           |  1 +
 tensorflow/compiler/xla/tests/hlo_test_base.h | 20 +++++--
 .../compiler/xla/tools/parser/hlo_parser.cc   | 10 +++-
 .../xla/tools/parser/hlo_parser_test.cc       |  2 +-
 24 files changed, 195 insertions(+), 91 deletions(-)

diff --git a/tensorflow/compiler/xla/client/local_client.cc b/tensorflow/compiler/xla/client/local_client.cc
index 1c12705903..1acc6f8686 100644
--- a/tensorflow/compiler/xla/client/local_client.cc
+++ b/tensorflow/compiler/xla/client/local_client.cc
@@ -51,27 +51,49 @@ LocalExecutable::LocalExecutable(std::unique_ptr<Executable> executable,
 tensorflow::Status LocalExecutable::ValidateExecutionOptions(
     const tensorflow::gtl::ArraySlice<const ShapedBuffer*> arguments,
     const ExecutableRunOptions& run_options, const Backend& backend) {
-  const ComputationLayout& computation_layout =
-      executable_->module_config().entry_computation_layout();
+  const ComputationLayout& host_computation_layout =
+      executable_->module_config().host_entry_computation_layout();
+  const ComputationLayout& device_computation_layout =
+      executable_->module_config().device_entry_computation_layout();
 
   // Check argument number, shapes, and layouts.
-  if (arguments.size() != computation_layout.parameter_count()) {
+  if (arguments.size() != host_computation_layout.parameter_count()) {
     return InvalidArgument(
         "invalid number of arguments for computation: expected %d, got %zu",
-        computation_layout.parameter_count(), arguments.size());
+        host_computation_layout.parameter_count(), arguments.size());
+  }
+  if (arguments.size() != device_computation_layout.parameter_count()) {
+    return InvalidArgument(
+        "invalid number of arguments for computation: expected %d, got %zu",
+        device_computation_layout.parameter_count(), arguments.size());
   }
   for (int i = 0; i < arguments.size(); ++i) {
-    if (!computation_layout.parameter_layout(i).MatchesLayoutInShape(
+    if (!host_computation_layout.parameter_layout(i).MatchesLayoutInShape(
             arguments[i]->on_host_shape())) {
       return InvalidParameterArgument(
           executable_.get(), i,
-          "Argument does not match shape or layout of computation parameter "
+          "Argument does not match host shape or layout of computation "
+          "parameter "
           "%d: want %s, got %s",
           i,
-          ShapeUtil::HumanString(computation_layout.parameter_layout(i).shape())
+          ShapeUtil::HumanString(
+              host_computation_layout.parameter_layout(i).shape())
               .c_str(),
           ShapeUtil::HumanString(arguments[i]->on_host_shape()).c_str());
     }
+    if (!device_computation_layout.parameter_layout(i).MatchesLayoutInShape(
+            arguments[i]->on_device_shape())) {
+      return InvalidParameterArgument(
+          executable_.get(), i,
+          "Argument does not match device shape or layout of computation "
+          "parameter "
+          "%d: want %s, got %s",
+          i,
+          ShapeUtil::HumanString(
+              device_computation_layout.parameter_layout(i).shape())
+              .c_str(),
+          ShapeUtil::HumanString(arguments[i]->on_device_shape()).c_str());
+    }
   }
 
   if (run_options.stream() != nullptr) {
diff --git a/tensorflow/compiler/xla/client/local_client.h b/tensorflow/compiler/xla/client/local_client.h
index 4ce7059f7e..d8fd7a5623 100644
--- a/tensorflow/compiler/xla/client/local_client.h
+++ b/tensorflow/compiler/xla/client/local_client.h
@@ -43,15 +43,6 @@ class LocalExecutable {
       const tensorflow::gtl::ArraySlice<const ShapedBuffer*> arguments,
       ExecutableRunOptions run_options);
 
-  // Return the layout (contained in a shape) of the result produced by the
-  // computation.
-  const Shape& result_layout() const {
-    return executable_->module_config()
-        .entry_computation_layout()
-        .result_layout()
-        .shape();
-  }
-
   // Return the options used to build the executable.
   const ExecutableBuildOptions& build_options() const { return build_options_; }
 
diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
index ec2bb6c762..d8ba289f29 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
+++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
@@ -294,7 +294,7 @@ Status CpuCompiler::RunHloPasses(HloModule* module, bool is_aot_compile) {
       ReducePrecisionInsertion::PassTiming::AFTER_FUSION);
 
   pipeline.AddPass<CpuLayoutAssignment>(
-      module->mutable_entry_computation_layout());
+      module->device_entry_computation_layout());
   // The LayoutAssignment pass may leave behind kCopy instructions which are
   // duplicate or NOPs, so remove them with algebraic simplification and CSE.
   pipeline.AddPass<HloPassFix<AlgebraicSimplifier>>(
diff --git a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc
index aabf4d5161..32613b8690 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc
+++ b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc
@@ -249,8 +249,9 @@ StatusOr<ScopedShapedBuffer> CpuExecutable::CreateResultShapedBuffer(
     std::vector<bool>* buffers_in_result) {
   se::Stream* stream = run_options->stream();
   ScopedShapedBuffer result_buffer(
-      /*on_host_shape=*/result_shape(), /*on_device_shape=*/result_shape(),
-      run_options->allocator(), stream->parent()->device_ordinal());
+      /*on_host_shape=*/host_result_shape(),
+      /*on_device_shape=*/host_result_shape(), run_options->allocator(),
+      stream->parent()->device_ordinal());
 
   // Copy DeviceMemoryBase values which contain the array(s) of the result into
   // the respective location in ShapedBuffer which is returned to the caller.
diff --git a/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.h b/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.h
index c8edbb9e15..09adb5cb02 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.h
+++ b/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.h
@@ -27,7 +27,8 @@ namespace cpu {
 // layout constraints for operands and results of library calls.
 class CpuLayoutAssignment : public LayoutAssignment {
  public:
-  explicit CpuLayoutAssignment(ComputationLayout* entry_computation_layout)
+  explicit CpuLayoutAssignment(
+      const ComputationLayout& entry_computation_layout)
       : LayoutAssignment(entry_computation_layout) {}
   ~CpuLayoutAssignment() override {}
 
diff --git a/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment_test.cc b/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment_test.cc
index 6ba030fff3..ba4c5a23d3 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment_test.cc
+++ b/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment_test.cc
@@ -49,7 +49,7 @@ class CpuLayoutAssignmentTest : public HloTestBase {
  protected:
   void AssignLayouts(HloModule* module,
                      ComputationLayout* entry_computation_layout) {
-    cpu::CpuLayoutAssignment layout_assignment(entry_computation_layout);
+    cpu::CpuLayoutAssignment layout_assignment(*entry_computation_layout);
     EXPECT_IS_OK(layout_assignment.Run(module).status());
   }
 };
@@ -311,7 +311,7 @@ static StatusOr<DotOutputFusionLayoutAssignmentResult> RunDotOutputFusion(
   result.addend_fusion_param = fusion_instruction->operand(
       fused_add->operand(1 - dot_operand_idx_in_add)->parameter_number());
 
-  cpu::CpuLayoutAssignment layout_assignment(&computation_layout);
+  cpu::CpuLayoutAssignment layout_assignment(computation_layout);
   TF_ASSIGN_OR_RETURN(result.layout_assignment_changed_something,
                       layout_assignment.Run(module));
 
diff --git a/tensorflow/compiler/xla/service/executable.h b/tensorflow/compiler/xla/service/executable.h
index 99762f4586..4f0466c544 100644
--- a/tensorflow/compiler/xla/service/executable.h
+++ b/tensorflow/compiler/xla/service/executable.h
@@ -140,8 +140,8 @@ class Executable {
 
   // The shape (including layout) that results from this execution. This is the
   // shape of the DeviceMemoryBase result value in ExecuteOnStream above.
-  const Shape& result_shape() const {
-    return hlo_module_->config().entry_computation_layout().result_shape();
+  const Shape& host_result_shape() const {
+    return hlo_module_->config().host_entry_computation_layout().result_shape();
   }
 
   // TODO(b/74197823): Delete the session module dumping helpers.
diff --git a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
index 796c3070f2..4fdc4c8961 100644
--- a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
+++ b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
@@ -248,7 +248,7 @@ tensorflow::Status OptimizeHloModule(HloModule* hlo_module,
   {
     HloPassPipeline pipeline("layout_assignment");
     pipeline.AddPass<GpuLayoutAssignment>(
-        hlo_module->mutable_entry_computation_layout());
+        hlo_module->device_entry_computation_layout());
 
     // The LayoutAssignment pass may leave behind kCopy instructions which are
     // duplicate or NOPs, so remove them with algebraic simplification and CSE.
diff --git a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.h b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.h
index 86a3a7111f..51aae79c3d 100644
--- a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.h
+++ b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.h
@@ -27,7 +27,8 @@ namespace gpu {
 // layout constraints for operands and results of library calls.
 class GpuLayoutAssignment : public LayoutAssignment {
  public:
-  explicit GpuLayoutAssignment(ComputationLayout* entry_computation_layout)
+  explicit GpuLayoutAssignment(
+      const ComputationLayout& entry_computation_layout)
       : LayoutAssignment(entry_computation_layout) {}
   ~GpuLayoutAssignment() override {}
 
diff --git a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment_test.cc b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment_test.cc
index 4c45d2e94a..7c80195594 100644
--- a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment_test.cc
+++ b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment_test.cc
@@ -69,7 +69,7 @@ TEST_F(LayoutAssignmentTest, Elementwise) {
         *computation_layout.mutable_result_layout() =
             ShapeLayout(result_shape_with_layout);
 
-        GpuLayoutAssignment layout_assignment(&computation_layout);
+        GpuLayoutAssignment layout_assignment(computation_layout);
         EXPECT_TRUE(layout_assignment.Run(module.get()).ValueOrDie());
 
         for (const HloInstruction* operand : add->operands()) {
@@ -156,7 +156,7 @@ TEST_F(LayoutAssignmentTest, BatchNormInference) {
         *computation_layout.mutable_result_layout() = ShapeLayout(result_shape);
       }
 
-      GpuLayoutAssignment layout_assignment(&computation_layout);
+      GpuLayoutAssignment layout_assignment(computation_layout);
       EXPECT_TRUE(layout_assignment.Run(module.get()).ValueOrDie());
 
       // The first operand to batchnorm should have the same layout as the
@@ -225,7 +225,7 @@ TEST_F(LayoutAssignmentTest, BatchNormTraining) {
                 {result_shape, offset_scale_shape, offset_scale_shape}));
       }
 
-      GpuLayoutAssignment layout_assignment(&computation_layout);
+      GpuLayoutAssignment layout_assignment(computation_layout);
       EXPECT_TRUE(layout_assignment.Run(module.get()).ValueOrDie());
 
       // The first operand to batchnorm should have the same layout as the
@@ -305,7 +305,7 @@ TEST_F(LayoutAssignmentTest, BatchNormGrad) {
                   {result_shape, scale_shape, scale_shape}));
         }
 
-        GpuLayoutAssignment layout_assignment(&computation_layout);
+        GpuLayoutAssignment layout_assignment(computation_layout);
         EXPECT_TRUE(layout_assignment.Run(module.get()).ValueOrDie());
 
         // The first and fourth operands to the batchnorm call should have the
diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc
index d4bad16f79..987c4b2719 100644
--- a/tensorflow/compiler/xla/service/hlo_module.cc
+++ b/tensorflow/compiler/xla/service/hlo_module.cc
@@ -55,7 +55,7 @@ HloComputation* HloModule::AddComputationInternal(
 
     // If the module configuration has no entry layout computation set, create a
     // default one based on the program shape.
-    if (!config_.has_entry_computation_layout()) {
+    if (!config_.has_host_entry_computation_layout()) {
       config_.SetDefaultComputationLayout(
           entry_computation_->ComputeProgramShape());
     }
@@ -229,11 +229,14 @@ StatusOr<std::unique_ptr<HloModule>> HloModule::CreateFromProto(
   TF_RET_CHECK(proto.has_program_shape())
       << "No program shape found in the proto";
   const auto& expected_program_shape = proto.program_shape();
-  TF_RET_CHECK(expected_program_shape.parameters_size() ==
-               module_config.entry_computation_layout().parameter_count());
+  TF_RET_CHECK(
+      expected_program_shape.parameters_size() ==
+      module_config.device_entry_computation_layout().parameter_count());
   for (int i = 0; i < expected_program_shape.parameters_size(); ++i) {
     const Shape& parameter_shape =
-        module_config.entry_computation_layout().parameter_layout(i).shape();
+        module_config.device_entry_computation_layout()
+            .parameter_layout(i)
+            .shape();
     TF_RET_CHECK(ShapeUtil::Compatible(expected_program_shape.parameters(i),
                                        parameter_shape))
         << "HloModuleConfig has different shape for parameter " << i
@@ -243,7 +246,7 @@ StatusOr<std::unique_ptr<HloModule>> HloModule::CreateFromProto(
         << ", actual: " << ShapeUtil::HumanStringWithLayout(parameter_shape);
   }
   const Shape& result_shape =
-      module_config.entry_computation_layout().result_layout().shape();
+      module_config.device_entry_computation_layout().result_layout().shape();
   TF_RET_CHECK(
       ShapeUtil::Compatible(expected_program_shape.result(), result_shape))
       << "HloModuleConfig has different result shape than the HLO module. "
@@ -303,7 +306,7 @@ StatusOr<HloModuleConfig> HloModule::CreateModuleConfigFromProto(
   // The module config is constructed with default layouts regardless of what is
   // passed in via the ProgramShape. Set the layouts to the appropriate values.
   ComputationLayout* entry_layout =
-      module_config.mutable_entry_computation_layout();
+      module_config.mutable_host_entry_computation_layout();
   for (int64 i = 0; i < entry_layout->parameter_count(); ++i) {
     TF_RETURN_IF_ERROR(
         entry_layout->mutable_parameter_layout(i)->CopyLayoutFromShape(
@@ -311,6 +314,8 @@ StatusOr<HloModuleConfig> HloModule::CreateModuleConfigFromProto(
   }
   TF_RETURN_IF_ERROR(entry_layout->mutable_result_layout()->CopyLayoutFromShape(
       program_shape.result()));
+  *module_config.mutable_device_entry_computation_layout() =
+      module_config.host_entry_computation_layout();
 
   return module_config;
 }
diff --git a/tensorflow/compiler/xla/service/hlo_module.h b/tensorflow/compiler/xla/service/hlo_module.h
index aa843ead51..82d790ec3b 100644
--- a/tensorflow/compiler/xla/service/hlo_module.h
+++ b/tensorflow/compiler/xla/service/hlo_module.h
@@ -98,12 +98,20 @@ class HloModule {
     return entry_computation_;
   }
 
-  ComputationLayout* mutable_entry_computation_layout() {
-    return config_.mutable_entry_computation_layout();
+  ComputationLayout* mutable_host_entry_computation_layout() {
+    return config_.mutable_host_entry_computation_layout();
   }
 
-  const ComputationLayout& entry_computation_layout() const {
-    return config_.entry_computation_layout();
+  const ComputationLayout& host_entry_computation_layout() const {
+    return config_.host_entry_computation_layout();
+  }
+
+  ComputationLayout* mutable_device_entry_computation_layout() {
+    return config_.mutable_device_entry_computation_layout();
+  }
+
+  const ComputationLayout& device_entry_computation_layout() const {
+    return config_.device_entry_computation_layout();
   }
 
   const VersionedComputationHandle& entry_computation_handle() const {
diff --git a/tensorflow/compiler/xla/service/hlo_module_config.cc b/tensorflow/compiler/xla/service/hlo_module_config.cc
index 4205b0402c..dae5578a31 100644
--- a/tensorflow/compiler/xla/service/hlo_module_config.cc
+++ b/tensorflow/compiler/xla/service/hlo_module_config.cc
@@ -31,11 +31,13 @@ using tensorflow::strings::StrAppend;
 HloModuleConfig::HloModuleConfig() {}
 
 HloModuleConfig::HloModuleConfig(const ProgramShape& program_shape)
-    : entry_computation_layout_(program_shape) {}
+    : host_entry_computation_layout_(program_shape),
+      device_entry_computation_layout_(program_shape) {}
 
 void HloModuleConfig::SetDefaultComputationLayout(
     const ProgramShape& program_shape) {
-  entry_computation_layout_ = ComputationLayout(program_shape);
+  host_entry_computation_layout_ = ComputationLayout(program_shape);
+  device_entry_computation_layout_ = ComputationLayout(program_shape);
 }
 
 string HloModuleConfig::compilation_cache_key() const {
@@ -44,11 +46,18 @@ string HloModuleConfig::compilation_cache_key() const {
   StrAppend(&key, "::(");
   std::vector<string> params;
   for (const ShapeLayout& param_layout :
-       entry_computation_layout_->parameter_layouts()) {
+       host_entry_computation_layout_->parameter_layouts()) {
     params.push_back(param_layout.shape().DebugString());
   }
   StrAppend(&key, tensorflow::str_util::Join(params, ", "), ") => ",
-            entry_computation_layout_->result_shape().SerializeAsString());
+            host_entry_computation_layout_->result_shape().SerializeAsString());
+  for (const ShapeLayout& param_layout :
+       device_entry_computation_layout_->parameter_layouts()) {
+    params.push_back(param_layout.shape().DebugString());
+  }
+  StrAppend(
+      &key, tensorflow::str_util::Join(params, ", "), ") => ",
+      device_entry_computation_layout_->result_shape().SerializeAsString());
   if (seed() != 0) {
     // TODO(b/32083678): force recompilation to reset global state.
     static std::atomic<int> counter{0};
diff --git a/tensorflow/compiler/xla/service/hlo_module_config.h b/tensorflow/compiler/xla/service/hlo_module_config.h
index 586a03d412..cdb0b29a23 100644
--- a/tensorflow/compiler/xla/service/hlo_module_config.h
+++ b/tensorflow/compiler/xla/service/hlo_module_config.h
@@ -41,26 +41,44 @@ class HloModuleConfig {
   explicit HloModuleConfig(const ProgramShape& program_shape);
 
   // Checks if this config has an entry computation layout already.
-  bool has_entry_computation_layout() const {
-    return entry_computation_layout_.has_value();
+  bool has_host_entry_computation_layout() const {
+    return host_entry_computation_layout_.has_value();
+  }
+
+  bool has_device_entry_computation_layout() const {
+    return device_entry_computation_layout_.has_value();
   }
 
   // Sets the entry computation layout for this config. If the entry computation
   // layout already exists, it is silently replaced.
   void SetDefaultComputationLayout(const ProgramShape& program_shape);
 
-  // Returns a constant reference to the layout of the entry computation.
+  // Returns a constant reference to the on-host layout of the entry
+  // computation. Assumes the layout was set.
+  const ComputationLayout& host_entry_computation_layout() const {
+    CHECK(host_entry_computation_layout_.has_value());
+    return *host_entry_computation_layout_;
+  }
+
+  // Returns a mutable pointer to the layout of the on-host entry computation.
   // Assumes the layout was set.
-  const ComputationLayout& entry_computation_layout() const {
-    CHECK(entry_computation_layout_.has_value());
-    return *entry_computation_layout_;
+  ComputationLayout* mutable_host_entry_computation_layout() {
+    CHECK(host_entry_computation_layout_.has_value());
+    return &(*host_entry_computation_layout_);
   }
 
-  // Returns a mutable pointer to the layout of the entry computation. Assumes
-  // the layout was set.
-  ComputationLayout* mutable_entry_computation_layout() {
-    CHECK(entry_computation_layout_.has_value());
-    return &(*entry_computation_layout_);
+  // Returns a constant reference to the on-device layout of the entry
+  // computation. Assumes the layout was set.
+  const ComputationLayout& device_entry_computation_layout() const {
+    CHECK(device_entry_computation_layout_.has_value());
+    return *device_entry_computation_layout_;
+  }
+
+  // Returns a mutable pointer to the layout of the on-device entry computation.
+  // Assumes the layout was set.
+  ComputationLayout* mutable_device_entry_computation_layout() {
+    CHECK(device_entry_computation_layout_.has_value());
+    return &(*device_entry_computation_layout_);
   }
 
   // Returns whether to enable HLO-level profiling.
@@ -109,7 +127,8 @@ class HloModuleConfig {
  private:
   // If you add new members, be sure to update compilation_cache_key.
 
-  tensorflow::gtl::optional<ComputationLayout> entry_computation_layout_;
+  tensorflow::gtl::optional<ComputationLayout> host_entry_computation_layout_;
+  tensorflow::gtl::optional<ComputationLayout> device_entry_computation_layout_;
 
   // Whether this is a 'host module'.
   bool is_host_module_ = false;
diff --git a/tensorflow/compiler/xla/service/interpreter/compiler.cc b/tensorflow/compiler/xla/service/interpreter/compiler.cc
index 76b3ecad26..eecbbcb93d 100644
--- a/tensorflow/compiler/xla/service/interpreter/compiler.cc
+++ b/tensorflow/compiler/xla/service/interpreter/compiler.cc
@@ -45,7 +45,7 @@ Status InterpreterCompiler::RunHloOptimization(HloModule* hlo_module) {
   HloPassPipeline pipeline("Interpreter");
 
   pipeline.AddPass<LayoutAssignment>(
-      hlo_module->mutable_entry_computation_layout());
+      hlo_module->device_entry_computation_layout());
 
   return pipeline.Run(hlo_module).status();
 }
diff --git a/tensorflow/compiler/xla/service/layout_assignment.cc b/tensorflow/compiler/xla/service/layout_assignment.cc
index 2494569db5..cfa7ba5e81 100644
--- a/tensorflow/compiler/xla/service/layout_assignment.cc
+++ b/tensorflow/compiler/xla/service/layout_assignment.cc
@@ -909,22 +909,19 @@ Status LayoutAssignment::CheckLayouts(HloModule* module) {
 }
 
 LayoutAssignment::LayoutAssignment(
-    ComputationLayout* entry_computation_layout,
+    const ComputationLayout& entry_computation_layout,
     ChannelLayoutConstraints* channel_constraints)
     : entry_computation_layout_(entry_computation_layout),
       channel_layout_constraints_(channel_constraints) {
   VLOG(1) << "entry computation layout given to layout assignment: "
-          << entry_computation_layout_->ToString();
+          << entry_computation_layout_.ToString();
   // Layouts of all parameter instructions must be set.
   for (const ShapeLayout& parameter_layout :
-       entry_computation_layout_->parameter_layouts()) {
+       entry_computation_layout_.parameter_layouts()) {
     CHECK(parameter_layout.LayoutIsSet());
   }
-  // If the result layout is not set, then choose the default.
-  // TODO(b/29118294): Choose a better layout in this case.
-  if (!entry_computation_layout_->result_layout().LayoutIsSet()) {
-    entry_computation_layout_->mutable_result_layout()->SetToDefaultLayout();
-  }
+  // TODO(b/29118294): Choose a better layout if the result layout is not set.
+  CHECK(entry_computation_layout_.result_layout().LayoutIsSet());
 }
 
 std::unique_ptr<Layout> LayoutAssignment::ChooseOperandLayoutFromOutputLayout(
@@ -1597,7 +1594,7 @@ StatusOr<bool> LayoutAssignment::Run(HloModule* module) {
     }
     if (computation == module->entry_computation()) {
       TF_RETURN_IF_ERROR(RunOnComputation(
-          *entry_computation_layout_, *points_to_analysis,
+          entry_computation_layout_, *points_to_analysis,
           module->entry_computation(), channel_layout_constraints_));
     } else {
       ComputationLayout computation_layout(computation->ComputeProgramShape());
diff --git a/tensorflow/compiler/xla/service/layout_assignment.h b/tensorflow/compiler/xla/service/layout_assignment.h
index ae4986d6ad..c83ae0388b 100644
--- a/tensorflow/compiler/xla/service/layout_assignment.h
+++ b/tensorflow/compiler/xla/service/layout_assignment.h
@@ -288,7 +288,7 @@ class LayoutAssignment : public HloPassInterface {
   // If channel_constraints is nullptr, no kSend or kRecvs must be contained
   // within any module passed to `Run`.
   explicit LayoutAssignment(
-      ComputationLayout* entry_computation_layout,
+      const ComputationLayout& entry_computation_layout,
       ChannelLayoutConstraints* channel_constraints = nullptr);
   ~LayoutAssignment() override {}
   tensorflow::StringPiece name() const override { return "layout-assignment"; }
@@ -402,7 +402,7 @@ class LayoutAssignment : public HloPassInterface {
   // necessary conditions.
   Status CheckLayouts(HloModule* module);
 
-  ComputationLayout* entry_computation_layout_;
+  const ComputationLayout& entry_computation_layout_;
 
  protected:
   // Sets up the copy instruction according to the characteristic (sharding,
diff --git a/tensorflow/compiler/xla/service/layout_assignment_test.cc b/tensorflow/compiler/xla/service/layout_assignment_test.cc
index 4b1c9bad41..7e1bb11eaa 100644
--- a/tensorflow/compiler/xla/service/layout_assignment_test.cc
+++ b/tensorflow/compiler/xla/service/layout_assignment_test.cc
@@ -53,7 +53,7 @@ class LayoutAssignmentTest : public HloTestBase {
  protected:
   void AssignLayouts(HloModule* module,
                      ComputationLayout* entry_computation_layout) {
-    LayoutAssignment layout_assignment(entry_computation_layout);
+    LayoutAssignment layout_assignment(*entry_computation_layout);
     EXPECT_IS_OK(layout_assignment.Run(module).status());
   }
 };
@@ -285,7 +285,7 @@ TEST_F(LayoutAssignmentTest, ConflictingLayoutTuple) {
   TF_CHECK_OK(computation_layout.mutable_result_layout()->CopyLayoutFromShape(
       result_shape));
 
-  LayoutAssignment layout_assignment(&computation_layout);
+  LayoutAssignment layout_assignment(computation_layout);
   AssignLayouts(module.get(), &computation_layout);
 
   // Layout assignment should have deep copied the result of the computation to
@@ -488,7 +488,7 @@ class OperandsMustBeTheSameLayoutAssignment : public LayoutAssignment {
  public:
   explicit OperandsMustBeTheSameLayoutAssignment(
       ComputationLayout* entry_computation_layout)
-      : LayoutAssignment(entry_computation_layout) {}
+      : LayoutAssignment(*entry_computation_layout) {}
 
  protected:
   Status PropagateBufferConstraint(
@@ -808,7 +808,7 @@ TEST_F(LayoutAssignmentTest, InternalErrorOnBitcast) {
 
   ComputationLayout computation_layout(
       module->entry_computation()->ComputeProgramShape());
-  LayoutAssignment layout_assignment(&computation_layout);
+  LayoutAssignment layout_assignment(computation_layout);
   Status error_status = layout_assignment.Run(module.get()).status();
   EXPECT_FALSE(error_status.ok());
   EXPECT_THAT(
diff --git a/tensorflow/compiler/xla/service/service.cc b/tensorflow/compiler/xla/service/service.cc
index 175ee96bbc..6ce03ab39d 100644
--- a/tensorflow/compiler/xla/service/service.cc
+++ b/tensorflow/compiler/xla/service/service.cc
@@ -296,8 +296,10 @@ StatusOr<std::unique_ptr<HloModuleConfig>> Service::CreateModuleConfig(
     const ExecutionOptions* execution_options,
     const UserComputation* user_computation) {
   auto config = MakeUnique<HloModuleConfig>(program_shape);
-  auto* computation_layout = config->mutable_entry_computation_layout();
-
+  ComputationLayout* host_computation_layout =
+      config->mutable_host_entry_computation_layout();
+  ComputationLayout* device_computation_layout =
+      config->mutable_device_entry_computation_layout();
   if (program_shape.parameters_size() != argument_shapes.size()) {
     return InvalidArgument("computation takes %d parameters, but %zu given",
                            program_shape.parameters_size(),
@@ -322,9 +324,10 @@ StatusOr<std::unique_ptr<HloModuleConfig>> Service::CreateModuleConfig(
           i, ShapeUtil::HumanString(program_shape.parameters(i)).c_str(),
           ShapeUtil::HumanString(*argument_shapes[i]).c_str());
     }
-    TF_RETURN_IF_ERROR(
-        computation_layout->mutable_parameter_layout(i)->CopyLayoutFromShape(
-            *argument_shapes[i]));
+    TF_RETURN_IF_ERROR(host_computation_layout->mutable_parameter_layout(i)
+                           ->CopyLayoutFromShape(*argument_shapes[i]));
+    TF_RETURN_IF_ERROR(device_computation_layout->mutable_parameter_layout(i)
+                           ->CopyLayoutFromShape(*argument_shapes[i]));
   }
   if (execution_options != nullptr &&
       execution_options->has_shape_with_output_layout()) {
@@ -333,10 +336,17 @@ StatusOr<std::unique_ptr<HloModuleConfig>> Service::CreateModuleConfig(
     TF_RETURN_IF_ERROR(ValidateResultShapeWithLayout(shape_with_output_layout,
                                                      program_shape.result()));
     TF_RETURN_IF_ERROR(
-        computation_layout->mutable_result_layout()->CopyLayoutFromShape(
+        host_computation_layout->mutable_result_layout()->CopyLayoutFromShape(
+            shape_with_output_layout));
+    TF_RETURN_IF_ERROR(
+        device_computation_layout->mutable_result_layout()->CopyLayoutFromShape(
             shape_with_output_layout));
   } else {
-    computation_layout->mutable_result_layout()->Clear();
+    // If the result layout is not set, then choose the default.
+    // TODO(b/29118294): Allow the compiler to choose a better layout in this
+    // case.
+    host_computation_layout->mutable_result_layout()->SetToDefaultLayout();
+    device_computation_layout->mutable_result_layout()->SetToDefaultLayout();
   }
 
   config->set_replica_count(options_.number_of_replicas());
@@ -488,6 +498,22 @@ StatusOr<std::vector<std::unique_ptr<Executable>>> Service::BuildExecutables(
   return std::move(executables);
 }
 
+Status Service::ValidateEntryComputationLayout(HloModule* module) {
+  const ComputationLayout& on_device =
+      module->device_entry_computation_layout();
+  for (int64 i = 0; i < on_device.parameter_count(); ++i) {
+    TF_RET_CHECK(ShapeUtil::Equal(
+        on_device.parameter_shape(i),
+        execute_backend_->transfer_manager()->HostShapeToDeviceShape(
+            module->host_entry_computation_layout().parameter_shape(i))));
+  }
+  TF_RET_CHECK(ShapeUtil::Equal(
+      module->device_entry_computation_layout().result_shape(),
+      execute_backend_->transfer_manager()->HostShapeToDeviceShape(
+          module->host_entry_computation_layout().result_shape())));
+  return tensorflow::Status::OK();
+}
+
 StatusOr<std::unique_ptr<Executable>> Service::BuildExecutable(
     const VersionedComputationHandle& versioned_handle,
     std::unique_ptr<HloModuleConfig> module_config, Backend* backend,
@@ -526,6 +552,8 @@ StatusOr<std::unique_ptr<Executable>> Service::BuildExecutable(
   TF_ASSIGN_OR_RETURN(
       module, backend->compiler()->RunHloPasses(std::move(module), executor,
                                                 device_allocator));
+  // Check that on-host and on-device shapes are consistent.
+  TF_RETURN_IF_ERROR(ValidateEntryComputationLayout(module.get()));
 
   TF_ASSIGN_OR_RETURN(std::unique_ptr<Executable> executable,
                       backend->compiler()->RunBackend(
@@ -889,7 +917,7 @@ tensorflow::Status Service::ExecuteParallel(const ExecuteParallelRequest* arg,
         CreateModuleConfig(*program_shape, replicated_arguments.front(),
                            request.execution_options(), user_computation));
     VLOG(3) << "ExecuteParallel created HloModuleConfig computation layout: "
-            << module_config->entry_computation_layout().ToString();
+            << module_config->host_entry_computation_layout().ToString();
 
     // Adds to the vectors to build and execute the computations after the loop.
     all_arguments.push_back(replicated_arguments);
@@ -992,7 +1020,7 @@ tensorflow::Status Service::ExecuteGraphParallel(
                            /*user_computation=*/nullptr));
     VLOG(3)
         << "ExecuteGraphParallel created HloModuleConfig computation layout: "
-        << module_config->entry_computation_layout().ToString();
+        << module_config->host_entry_computation_layout().ToString();
 
     // Adds to the vectors to build and execute the computations after the loop.
     all_arguments.push_back(replicated_arguments);
@@ -1142,7 +1170,7 @@ tensorflow::Status Service::Execute(const ExecuteRequest* arg,
                          arg->execution_options(), user_computation));
 
   VLOG(3) << "Execute created HloModuleConfig computation layout: "
-          << module_config->entry_computation_layout().ToString();
+          << module_config->host_entry_computation_layout().ToString();
 
   TF_ASSIGN_OR_RETURN(
       std::shared_ptr<Executable> executable,
@@ -1212,6 +1240,8 @@ StatusOr<std::unique_ptr<Executable>> Service::BuildExecutable(
   TF_ASSIGN_OR_RETURN(
       module, backend->compiler()->RunHloPasses(std::move(module), executor,
                                                 device_allocator));
+  // Check that on-host and on-device shapes are consistent.
+  TF_RETURN_IF_ERROR(ValidateEntryComputationLayout(module.get()));
 
   TF_ASSIGN_OR_RETURN(std::unique_ptr<Executable> executable,
                       backend->compiler()->RunBackend(
@@ -1313,7 +1343,7 @@ tensorflow::Status Service::ExecuteAsync(const ExecuteAsyncRequest* arg,
                          arg->execution_options(), user_computation));
 
   VLOG(3) << "ExecuteAsync created HloModuleConfig computation layout: "
-          << module_config->entry_computation_layout().ToString();
+          << module_config->host_entry_computation_layout().ToString();
 
   ExecutionProfile profile;
 
diff --git a/tensorflow/compiler/xla/service/service.h b/tensorflow/compiler/xla/service/service.h
index 476bd0597d..f84fe407e0 100644
--- a/tensorflow/compiler/xla/service/service.h
+++ b/tensorflow/compiler/xla/service/service.h
@@ -295,6 +295,9 @@ class Service : public ServiceInterface {
       const ExecutionOptions& execution_options,
       tensorflow::gtl::ArraySlice<const GlobalDataHandle*> arguments);
 
+  // Assert that host- and device-shapes are in a consistent state.
+  Status ValidateEntryComputationLayout(HloModule* module);
+
  protected:
   friend class LocalExecutable;
 
diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD
index 840292010d..54cf0543b8 100644
--- a/tensorflow/compiler/xla/tests/BUILD
+++ b/tensorflow/compiler/xla/tests/BUILD
@@ -632,6 +632,7 @@ xla_test(
         "//tensorflow/compiler/xla/client/xla_client:xla_builder",
         "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/compiler/xla/tests:client_library_test_base",
+        "//tensorflow/compiler/xla/tests:hlo_test_base",
         "//tensorflow/compiler/xla/tests:literal_test_util",
         "//tensorflow/compiler/xla/tests:xla_internal_test_main",
         "//tensorflow/core:lib",
diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.h b/tensorflow/compiler/xla/tests/hlo_test_base.h
index 6491208895..9539ae0680 100644
--- a/tensorflow/compiler/xla/tests/hlo_test_base.h
+++ b/tensorflow/compiler/xla/tests/hlo_test_base.h
@@ -177,9 +177,13 @@ class HloTestBase : public ::testing::Test {
   // 'layout'.
   void ForceParameterLayout(HloModule* module, int64 param_no,
                             const Layout& layout) {
-    ASSERT_LT(param_no,
-              module->mutable_entry_computation_layout()->parameter_count());
-    module->mutable_entry_computation_layout()
+    ASSERT_LT(
+        param_no,
+        module->mutable_host_entry_computation_layout()->parameter_count());
+    module->mutable_host_entry_computation_layout()
+        ->mutable_parameter_layout(param_no)
+        ->ResetLayout(layout);
+    module->mutable_device_entry_computation_layout()
         ->mutable_parameter_layout(param_no)
         ->ResetLayout(layout);
   }
@@ -187,7 +191,10 @@ class HloTestBase : public ::testing::Test {
   // Convenience method to force the layout of the computation result in a
   // module. The result layout of 'module' is set to 'layout'.
   void ForceResultLayout(HloModule* module, const Layout& layout) {
-    module->mutable_entry_computation_layout()
+    module->mutable_host_entry_computation_layout()
+        ->mutable_result_layout()
+        ->ResetLayout(layout);
+    module->mutable_device_entry_computation_layout()
         ->mutable_result_layout()
         ->ResetLayout(layout);
   }
@@ -195,7 +202,10 @@ class HloTestBase : public ::testing::Test {
   // Convenience method to clear the layout of the computation result in
   // 'module'.
   void ForceClearResultLayout(HloModule* module) {
-    module->mutable_entry_computation_layout()
+    module->mutable_host_entry_computation_layout()
+        ->mutable_result_layout()
+        ->Clear();
+    module->mutable_device_entry_computation_layout()
         ->mutable_result_layout()
         ->Clear();
   }
diff --git a/tensorflow/compiler/xla/tools/parser/hlo_parser.cc b/tensorflow/compiler/xla/tools/parser/hlo_parser.cc
index fdbfc0210e..1bb31ddb7b 100644
--- a/tensorflow/compiler/xla/tools/parser/hlo_parser.cc
+++ b/tensorflow/compiler/xla/tools/parser/hlo_parser.cc
@@ -303,12 +303,18 @@ bool HloParser::ParseComputations() {
     // set the layouts to what the hlo text says.
     for (int p = 0; p < computation->num_parameters(); p++) {
       const Shape& param_shape = computation->parameter_instruction(p)->shape();
-      TF_CHECK_OK(module_->mutable_entry_computation_layout()
+      TF_CHECK_OK(module_->mutable_host_entry_computation_layout()
+                      ->mutable_parameter_layout(p)
+                      ->CopyLayoutFromShape(param_shape));
+      TF_CHECK_OK(module_->mutable_device_entry_computation_layout()
                       ->mutable_parameter_layout(p)
                       ->CopyLayoutFromShape(param_shape));
     }
     const Shape& result_shape = computation->root_instruction()->shape();
-    TF_CHECK_OK(module_->mutable_entry_computation_layout()
+    TF_CHECK_OK(module_->mutable_host_entry_computation_layout()
+                    ->mutable_result_layout()
+                    ->CopyLayoutFromShape(result_shape));
+    TF_CHECK_OK(module_->mutable_device_entry_computation_layout()
                     ->mutable_result_layout()
                     ->CopyLayoutFromShape(result_shape));
   }
diff --git a/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc b/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc
index adc8b1d620..4e085bc89c 100644
--- a/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc
+++ b/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc
@@ -1239,7 +1239,7 @@ ENTRY %Reduce (input: f32[8,16,256]) -> f32[8,16] {
 
   auto module = Parse(original);
   TF_ASSERT_OK(module.status());
-  auto program_layout = module.ValueOrDie()->entry_computation_layout();
+  auto program_layout = module.ValueOrDie()->host_entry_computation_layout();
   ASSERT_EQ(program_layout.parameter_count(), 1);
   auto param_layout = program_layout.parameter_layout(0).layout();
   auto result_layout = program_layout.result_layout().layout();
-- 
GitLab


From fb8f040f2a927c6df149238da7c4278cf781d081 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 16:20:47 -0700
Subject: [PATCH 122/395] Allow `warm_start_from` argument to be a SavedModel
 path.

PiperOrigin-RevId: 195015356
---
 tensorflow/python/estimator/estimator.py      | 28 ++++++++++-----
 tensorflow/python/estimator/estimator_test.py | 35 +++++++++++++++++++
 2 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/tensorflow/python/estimator/estimator.py b/tensorflow/python/estimator/estimator.py
index 0970f00124..3691c99dda 100644
--- a/tensorflow/python/estimator/estimator.py
+++ b/tensorflow/python/estimator/estimator.py
@@ -155,12 +155,12 @@ class Estimator(object):
       config: Configuration object.
       params: `dict` of hyper parameters that will be passed into `model_fn`.
               Keys are names of parameters, values are basic python types.
-      warm_start_from: Optional string filepath to a checkpoint to warm-start
-                       from, or a `tf.estimator.WarmStartSettings` object to
-                       fully configure warm-starting.  If the string filepath is
-                       provided instead of a `WarmStartSettings`, then all
-                       variables are warm-started, and it is assumed that
-                       vocabularies and Tensor names are unchanged.
+      warm_start_from: Optional string filepath to a checkpoint or SavedModel to
+                       warm-start from, or a `tf.estimator.WarmStartSettings`
+                       object to fully configure warm-starting.  If the string
+                       filepath is provided instead of a `WarmStartSettings`,
+                       then all variables are warm-started, and it is assumed
+                       that vocabularies and Tensor names are unchanged.
 
     Raises:
       ValueError: parameters of `model_fn` don't match `params`.
@@ -1502,7 +1502,7 @@ def _get_default_warm_start_settings(warm_start_from):
 
   Args:
     warm_start_from: Either a string representing the filepath of a checkpoint
-      to initialize from, or an instance of WarmStartSettings.
+      or SavedModel to initialize from, or an instance of WarmStartSettings.
 
   Returns:
     Either None or an instance of WarmStartSettings.
@@ -1513,9 +1513,19 @@ def _get_default_warm_start_settings(warm_start_from):
   """
   if warm_start_from is None:
     return None
-  if isinstance(warm_start_from, six.string_types):
+  if isinstance(warm_start_from, (six.string_types, six.binary_type)):
+    # Infer that this is a SavedModel if export_path +
+    # 'variables/variables.index' exists, and if so, construct the
+    # WarmStartSettings pointing to export_path + 'variables/variables'.
+    if gfile.Exists(os.path.join(compat.as_bytes(warm_start_from),
+                                 compat.as_bytes('variables/variables.index'))):
+      logging.info('Warm-starting from a SavedModel')
+      return WarmStartSettings(ckpt_to_initialize_from=os.path.join(
+          compat.as_bytes(warm_start_from),
+          compat.as_bytes('variables/variables')))
     return WarmStartSettings(ckpt_to_initialize_from=warm_start_from)
   elif isinstance(warm_start_from, WarmStartSettings):
     return warm_start_from
   else:
-    raise ValueError('warm_start_from must be a string or a WarmStartSettings')
+    raise ValueError('warm_start_from must be a string or a WarmStartSettings, '
+                     'instead got {}'.format(type(warm_start_from)))
diff --git a/tensorflow/python/estimator/estimator_test.py b/tensorflow/python/estimator/estimator_test.py
index 74114fab3b..4d958f8b43 100644
--- a/tensorflow/python/estimator/estimator_test.py
+++ b/tensorflow/python/estimator/estimator_test.py
@@ -658,6 +658,41 @@ class EstimatorTrainTest(test.TestCase):
         5, estimator._load_global_step_from_checkpoint_dir(
             warm_started_est.model_dir))
 
+  def test_warm_starts_from_savedmodel(self):
+    def _make_model_fn(x):
+      def _variable_creating_and_export_model_fn(features, labels, mode):
+        _, _ = features, labels
+        variable_scope.get_variable('x', initializer=x)
+        global_step = training.get_global_step()
+        return model_fn_lib.EstimatorSpec(
+            mode,
+            predictions={'y': constant_op.constant(1.0)},
+            loss=constant_op.constant(1.),
+            train_op=state_ops.assign_add(global_step, 1),
+            export_outputs={'test': export_output.ClassificationOutput(
+                constant_op.constant([4.2]), constant_op.constant(['label']))})
+      return _variable_creating_and_export_model_fn
+
+    est = estimator.Estimator(model_fn=_make_model_fn(42.))
+    est.train(dummy_input_fn, steps=10)
+    feature_spec = {'x': parsing_ops.VarLenFeature(dtype=dtypes.int64),
+                    'y': parsing_ops.VarLenFeature(dtype=dtypes.int64)}
+    serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
+        feature_spec)
+    tmpdir = tempfile.mkdtemp()
+    export_dir_base = os.path.join(
+        compat.as_bytes(tmpdir), compat.as_bytes('export'))
+    export_dir = est.export_savedmodel(
+        export_dir_base, serving_input_receiver_fn)
+
+    warm_started_est = estimator.Estimator(
+        model_fn=_make_model_fn(36.),
+        warm_start_from=export_dir)
+    warm_started_est.train(dummy_input_fn, steps=5)
+    # warm_start is called after the model_fn, so x should have the value
+    # from the SavedModel.
+    self.assertEqual(42., warm_started_est.get_variable_value('x'))
+
   def test_max_step(self):
     est = estimator.Estimator(model_fn=model_fn_global_step_incrementer)
     est.train(dummy_input_fn, max_steps=5)
-- 
GitLab


From f5dbc1e16622f433f41f195bb33f56d674a004ce Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 16:33:03 -0700
Subject: [PATCH 123/395] Check for overflow in shape calculation.

PiperOrigin-RevId: 195017114
---
 tensorflow/contrib/lite/toco/BUILD            |  12 +
 .../contrib/lite/toco/import_tensorflow.cc    | 505 ++++++++++--------
 .../lite/toco/import_tensorflow_test.cc       | 160 ++++++
 tensorflow/contrib/lite/toco/toco_port.h      |   5 +
 tensorflow/contrib/lite/toco/tooling_util.h   |  29 +
 .../contrib/lite/toco/tooling_util_test.cc    |  81 +++
 6 files changed, 562 insertions(+), 230 deletions(-)
 create mode 100644 tensorflow/contrib/lite/toco/import_tensorflow_test.cc

diff --git a/tensorflow/contrib/lite/toco/BUILD b/tensorflow/contrib/lite/toco/BUILD
index f92e546ab8..f16225fd66 100644
--- a/tensorflow/contrib/lite/toco/BUILD
+++ b/tensorflow/contrib/lite/toco/BUILD
@@ -364,6 +364,18 @@ cc_library(
     }),
 )
 
+tf_cc_test(
+    name = "import_tensorflow_test",
+    srcs = ["import_tensorflow_test.cc"],
+    deps = [
+        ":toco_tooling",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:graph",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "tooling_util",
     srcs = [
diff --git a/tensorflow/contrib/lite/toco/import_tensorflow.cc b/tensorflow/contrib/lite/toco/import_tensorflow.cc
index fa8b26bce0..453ff29b0d 100644
--- a/tensorflow/contrib/lite/toco/import_tensorflow.cc
+++ b/tensorflow/contrib/lite/toco/import_tensorflow.cc
@@ -62,6 +62,9 @@ using tensorflow::TensorProto;
 using tensorflow::TensorShapeProto;
 
 namespace toco {
+
+using port::Status;
+
 namespace {
 bool HasAttr(const NodeDef& node, const string& attr_name) {
   return node.attr().count(attr_name) > 0;
@@ -113,7 +116,7 @@ const TensorShapeProto& GetShapeAttr(const NodeDef& node,
 }
 
 const TensorProto& GetTensorAttr(const NodeDef& node, const string& attr_name) {
-  CHECK(HasAttr(node, attr_name));
+  CHECK(HasAttr(node, attr_name)) << "No attr named '" << attr_name << "'";
   const auto& attr = node.attr().at(attr_name);
   CHECK_EQ(attr.value_case(), AttrValue::kTensor);
   return attr.tensor();
@@ -145,9 +148,9 @@ ArrayDataType ConvertDataType(tensorflow::DataType dtype) {
   return ArrayDataType::kNone;
 }
 
-void ImportShape(const TFLITE_PROTO_NS::RepeatedPtrField<
-                     tensorflow::TensorShapeProto_Dim>& input_dims,
-                 Shape* shape) {
+Status ImportShape(const TFLITE_PROTO_NS::RepeatedPtrField<
+                       tensorflow::TensorShapeProto_Dim>& input_dims,
+                   int* input_flat_size, Shape* shape) {
   std::vector<int> input_dims_only_sizes;
   for (auto& d : input_dims) {
     if (d.size() == 0) {
@@ -155,23 +158,33 @@ void ImportShape(const TFLITE_PROTO_NS::RepeatedPtrField<
       // them of flat size 0 even though they have other nonzero dims.
       // This breaks our invariant, that array dims can't be 0.
       // For now, tweaking this to record a 0-D shape instead.
-      input_dims_only_sizes.clear();
-      break;
+      shape->mutable_dims()->clear();
+      if (input_flat_size != nullptr) *input_flat_size = 0;
+      return Status::OK();
+    }
+    // TensorFlow's shapes use int64s, while TOCO uses ints.
+    if (d.size() > std::numeric_limits<int>::max()) {
+      return Status(false, "Shape element overflows");
     }
+
     input_dims_only_sizes.push_back(d.size());
   }
   *shape->mutable_dims() = input_dims_only_sizes;
+
+  if (input_flat_size == nullptr) return Status::OK();
+
+  return NumElements(input_dims_only_sizes, input_flat_size);
 }
 
-void ImportFloatArray(const TensorProto& input_tensor, Array* output_array) {
+Status ImportFloatArray(const TensorProto& input_tensor, Array* output_array) {
   CHECK_EQ(input_tensor.dtype(), DT_FLOAT);
   const auto& input_shape = input_tensor.tensor_shape();
   CHECK_LE(input_shape.dim_size(), 4);
-  ImportShape(input_shape.dim(), output_array->mutable_shape());
-  int input_flat_size = 1;
-  for (int k = 0; k < input_shape.dim_size(); k++) {
-    input_flat_size *= input_shape.dim(k).size();
-  }
+  int input_flat_size;
+  auto status = ImportShape(input_shape.dim(), &input_flat_size,
+                            output_array->mutable_shape());
+  if (!status.ok()) return status;
+
   auto& output_float_data =
       output_array->GetMutableBuffer<ArrayDataType::kFloat>().data;
   output_float_data.resize(RequiredBufferSizeForShape(output_array->shape()),
@@ -189,20 +202,22 @@ void ImportFloatArray(const TensorProto& input_tensor, Array* output_array) {
     toco::port::CopyToBuffer(input_tensor.tensor_content(),
                              reinterpret_cast<char*>(output_float_data.data()));
   } else {
-    LOG(FATAL) << "Neither input_content nor float_val have the right "
-                  "dimensions for this float tensor.";
+    return Status(false,
+                  "Neither input_content nor float_val have the right "
+                  "dimensions for this float tensor");
   }
+  return Status::OK();
 }
 
-void ImportQuint8Array(const TensorProto& input_tensor, Array* output_array) {
+Status ImportQuint8Array(const TensorProto& input_tensor, Array* output_array) {
   CHECK_EQ(input_tensor.dtype(), DT_QUINT8);
   const auto& input_shape = input_tensor.tensor_shape();
   CHECK_LE(input_shape.dim_size(), 4);
-  ImportShape(input_shape.dim(), output_array->mutable_shape());
-  int input_flat_size = 1;
-  for (int k = 0; k < input_shape.dim_size(); k++) {
-    input_flat_size *= input_shape.dim(k).size();
-  }
+  int input_flat_size;
+  auto status = ImportShape(input_shape.dim(), &input_flat_size,
+                            output_array->mutable_shape());
+  if (!status.ok()) return status;
+
   auto& output_int_data =
       output_array->GetMutableBuffer<ArrayDataType::kUint8>().data;
   output_int_data.resize(RequiredBufferSizeForShape(output_array->shape()), 0);
@@ -215,20 +230,22 @@ void ImportQuint8Array(const TensorProto& input_tensor, Array* output_array) {
     toco::port::CopyToBuffer(input_tensor.tensor_content(),
                              reinterpret_cast<char*>(output_int_data.data()));
   } else {
-    LOG(FATAL) << "Neither input_content nor int_val have the right "
-                  "dimensions for this uint8 tensor.";
+    return Status(false,
+                  "Neither input_content nor int_val have the right dimensions "
+                  "for this uint8 tensor");
   }
+  return Status::OK();
 }
 
-void ImportInt32Array(const TensorProto& input_tensor, Array* output_array) {
+Status ImportInt32Array(const TensorProto& input_tensor, Array* output_array) {
   CHECK_EQ(input_tensor.dtype(), DT_INT32);
   const auto& input_shape = input_tensor.tensor_shape();
   CHECK_LE(input_shape.dim_size(), 4);
-  ImportShape(input_shape.dim(), output_array->mutable_shape());
-  int input_flat_size = 1;
-  for (int k = 0; k < input_shape.dim_size(); k++) {
-    input_flat_size *= input_shape.dim(k).size();
-  }
+  int input_flat_size;
+  auto status = ImportShape(input_shape.dim(), &input_flat_size,
+                            output_array->mutable_shape());
+  if (!status.ok()) return status;
+
   auto& output_int_data =
       output_array->GetMutableBuffer<ArrayDataType::kInt32>().data;
   output_int_data.resize(RequiredBufferSizeForShape(output_array->shape()), 0);
@@ -241,20 +258,22 @@ void ImportInt32Array(const TensorProto& input_tensor, Array* output_array) {
     toco::port::CopyToBuffer(input_tensor.tensor_content(),
                              reinterpret_cast<char*>(output_int_data.data()));
   } else {
-    LOG(FATAL) << "Neither input_content nor int_val have the right "
-                  "dimensions for this int32 tensor.";
+    return Status(false,
+                  "Neither input_content nor int_val have the right dimensions "
+                  "for this int32 tensor");
   }
+  return Status::OK();
 }
 
-void ImportInt64Array(const TensorProto& input_tensor, Array* output_array) {
+Status ImportInt64Array(const TensorProto& input_tensor, Array* output_array) {
   CHECK_EQ(input_tensor.dtype(), DT_INT64);
   const auto& input_shape = input_tensor.tensor_shape();
   CHECK_LE(input_shape.dim_size(), 4);
-  ImportShape(input_shape.dim(), output_array->mutable_shape());
-  int input_flat_size = 1;
-  for (int k = 0; k < input_shape.dim_size(); k++) {
-    input_flat_size *= input_shape.dim(k).size();
-  }
+  int input_flat_size;
+  auto status = ImportShape(input_shape.dim(), &input_flat_size,
+                            output_array->mutable_shape());
+  if (!status.ok()) return status;
+
   auto& output_int_data =
       output_array->GetMutableBuffer<ArrayDataType::kInt64>().data;
   output_int_data.resize(RequiredBufferSizeForShape(output_array->shape()), 0);
@@ -267,20 +286,22 @@ void ImportInt64Array(const TensorProto& input_tensor, Array* output_array) {
     toco::port::CopyToBuffer(input_tensor.tensor_content(),
                              reinterpret_cast<char*>(output_int_data.data()));
   } else {
-    LOG(FATAL) << "Neither input_content nor int64_val have the right "
-                  "dimensions for this int64 tensor.";
+    return Status(false,
+                  "Neither input_content nor int64_val have the right "
+                  "dimensions for this int64 tensor");
   }
+  return Status::OK();
 }
 
-void ImportBoolArray(const TensorProto& input_tensor, Array* output_array) {
+Status ImportBoolArray(const TensorProto& input_tensor, Array* output_array) {
   CHECK_EQ(input_tensor.dtype(), DT_BOOL);
   const auto& input_shape = input_tensor.tensor_shape();
   CHECK_LE(input_shape.dim_size(), 4);
-  ImportShape(input_shape.dim(), output_array->mutable_shape());
-  int input_flat_size = 1;
-  for (int k = 0; k < input_shape.dim_size(); k++) {
-    input_flat_size *= input_shape.dim(k).size();
-  }
+  int input_flat_size;
+  auto status = ImportShape(input_shape.dim(), &input_flat_size,
+                            output_array->mutable_shape());
+  if (!status.ok()) return status;
+
   auto& output_bool_data =
       output_array->GetMutableBuffer<ArrayDataType::kBool>().data;
   output_bool_data.resize(RequiredBufferSizeForShape(output_array->shape()),
@@ -300,20 +321,25 @@ void ImportBoolArray(const TensorProto& input_tensor, Array* output_array) {
     // assuming that 'false' is implied.
     // So far only encountered that in an array with 1 entry, let's
     // require that until we encounter a graph where that's not the case.
-    CHECK_EQ(output_bool_data.size(), 1);
+    if (output_bool_data.size() != 1) {
+      return Status(false,
+                    "Neither input_content nor bool_val have the right "
+                    "dimensions for this bool tensor");
+    }
     output_bool_data[0] = false;
   }
+  return Status::OK();
 }
 
-void ImportStringArray(const TensorProto& input_tensor, Array* output_array) {
+Status ImportStringArray(const TensorProto& input_tensor, Array* output_array) {
   CHECK_EQ(input_tensor.dtype(), DT_STRING);
   const auto& input_shape = input_tensor.tensor_shape();
   CHECK_LE(input_shape.dim_size(), 4);
-  ImportShape(input_shape.dim(), output_array->mutable_shape());
-  int input_flat_size = 1;
-  for (int k = 0; k < input_shape.dim_size(); k++) {
-    input_flat_size *= input_shape.dim(k).size();
-  }
+  int input_flat_size;
+  auto status = ImportShape(input_shape.dim(), &input_flat_size,
+                            output_array->mutable_shape());
+  if (!status.ok()) return status;
+
   auto& output_string_data =
       output_array->GetMutableBuffer<ArrayDataType::kString>().data;
   output_string_data.resize(RequiredBufferSizeForShape(output_array->shape()));
@@ -324,6 +350,7 @@ void ImportStringArray(const TensorProto& input_tensor, Array* output_array) {
   for (int i = 0; i < input_flat_size; ++i) {
     output_string_data[i] = input_tensor.string_val(i);
   }
+  return Status::OK();
 }
 
 // Count the number of inputs of a given node. If
@@ -363,38 +390,40 @@ string CreateConstArray(Model* model, string const& name,
   return array_name;
 }
 
-void ConvertConstOperator(const NodeDef& node,
-                          const TensorFlowImportFlags& tf_import_flags,
-                          Model* model) {
+Status ConvertConstOperator(const NodeDef& node,
+                            const TensorFlowImportFlags& tf_import_flags,
+                            Model* model) {
   CHECK_EQ(node.op(), "Const");
   const auto& tensor = GetTensorAttr(node, "value");
   const auto dtype = GetDataTypeAttr(node, "dtype");
 
+  Status status = Status::OK();
+
   auto& array = model->GetOrCreateArray(node.name());
   switch (dtype) {
     case DT_FLOAT:
       array.data_type = ArrayDataType::kFloat;
-      ImportFloatArray(tensor, &array);
+      status = ImportFloatArray(tensor, &array);
       break;
     case DT_INT32:
       array.data_type = ArrayDataType::kInt32;
-      ImportInt32Array(tensor, &array);
+      status = ImportInt32Array(tensor, &array);
       break;
     case DT_QUINT8:
       array.data_type = ArrayDataType::kUint8;
-      ImportQuint8Array(tensor, &array);
+      status = ImportQuint8Array(tensor, &array);
       break;
     case DT_INT64:
       array.data_type = ArrayDataType::kInt64;
-      ImportInt64Array(tensor, &array);
+      status = ImportInt64Array(tensor, &array);
       break;
     case DT_STRING:
       array.data_type = ArrayDataType::kString;
-      ImportStringArray(tensor, &array);
+      status = ImportStringArray(tensor, &array);
       break;
     case DT_BOOL:
       array.data_type = ArrayDataType::kBool;
-      ImportBoolArray(tensor, &array);
+      status = ImportBoolArray(tensor, &array);
       break;
     default:
       array.data_type = ArrayDataType::kNone;
@@ -404,6 +433,10 @@ void ConvertConstOperator(const NodeDef& node,
       array.GetMutableBuffer<ArrayDataType::kNone>();
       break;
   }
+  if (!status.ok()) {
+    status.AppendMessage(" (while processing node '" + node.name() + "')");
+  }
+  return status;
 }
 
 void ConvertConvOperator(const NodeDef& node,
@@ -2033,6 +2066,186 @@ void ConvertDynamicStitchOperator(const NodeDef& node,
 
 }  // namespace
 
+namespace internal {
+Status ImportTensorFlowNode(const tensorflow::NodeDef& node,
+                            const TensorFlowImportFlags& tf_import_flags,
+                            Model* model) {
+  // TODO(ahentz): Historically these functions all CHECK-fail on error. We've
+  // been slowly converting them to return Status.
+  if (node.op() == "Const") {
+    return ConvertConstOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Conv2D") {
+    ConvertConvOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Conv2DBackpropInput") {
+    ConvertTransposeConvOperator(node, tf_import_flags, model);
+  } else if (node.op() == "DepthwiseConv2dNative") {
+    ConvertDepthwiseConvOperator(node, tf_import_flags, model);
+  } else if (node.op() == "DepthToSpace") {
+    ConvertDepthToSpaceOperator(node, tf_import_flags, model);
+  } else if (node.op() == "SpaceToDepth") {
+    ConvertSpaceToDepthOperator(node, tf_import_flags, model);
+  } else if (node.op() == "BiasAdd") {
+    ConvertBiasAddOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Relu") {
+    ConvertReluOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Relu6") {
+    ConvertRelu6Operator(node, tf_import_flags, model);
+  } else if (node.op() == "Sigmoid") {
+    ConvertLogisticOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Tanh") {
+    ConvertTanhOperator(node, tf_import_flags, model);
+  } else if (node.op() == "MaxPool") {
+    ConvertMaxPoolOperator(node, tf_import_flags, model);
+  } else if (node.op() == "AvgPool") {
+    ConvertAvgPoolOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Reshape") {
+    ConvertReshapeOperator(node, tf_import_flags, model);
+  } else if (node.op() == "BatchMatMul") {
+    ConvertBatchMatMulOperator(node, tf_import_flags, model);
+  } else if (node.op() == "MatMul") {
+    ConvertMatMulOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Div" || node.op() == "RealDiv") {
+    ConvertDivOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Identity" || node.op() == "CheckNumerics" ||
+             node.op() == "StopGradient") {
+    ConvertIdentityOperator(node, tf_import_flags, model);
+  } else if (node.op() == "FakeQuantWithMinMaxVars") {
+    ConvertFakeQuantWithMinMaxVars(node, tf_import_flags, model);
+  } else if (node.op() == "FakeQuantWithMinMaxArgs") {
+    ConvertFakeQuantWithMinMaxArgs(node, tf_import_flags, model);
+  } else if (node.op() == "Neg") {
+    ConvertNegOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Rsqrt") {
+    ConvertRsqrtOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Squeeze") {
+    ConvertSqueezeOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Sqrt") {
+    ConvertSqrtOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Square") {
+    ConvertSquareOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Add") {
+    ConvertAddOperator(node, tf_import_flags, model);
+  } else if (node.op() == "AddN") {
+    ConvertAddNOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Mul") {
+    ConvertMulOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Sub") {
+    ConvertSubOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Sum") {
+    ConvertSumOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Tile") {
+    ConvertTileOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Concat" || node.op() == "ConcatV2") {
+    ConvertConcatOperator(node, tf_import_flags, model);
+  } else if (node.op() == "LRN") {
+    ConvertLRNOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Softmax") {
+    ConvertSoftmaxOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Log") {
+    ConvertLogOperator(node, tf_import_flags, model);
+  } else if (node.op() == "LogSoftmax") {
+    ConvertLogSoftmaxOperator(node, tf_import_flags, model);
+  } else if (node.op() == "All") {
+    ConvertAllOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Assert") {
+    ConvertAssertOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Less") {
+    ConvertLessOperator(node, tf_import_flags, model);
+  } else if (node.op() == "LessEqual") {
+    ConvertLessEqualOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Greater") {
+    ConvertGreaterOperator(node, tf_import_flags, model);
+  } else if (node.op() == "GreaterEqual") {
+    ConvertGreaterEqualOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Max") {
+    ConvertMaxOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Min") {
+    ConvertMinOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Maximum") {
+    ConvertMaximumOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Minimum") {
+    ConvertMinimumOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Merge") {
+    ConvertMergeOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Pad") {
+    ConvertPadOperator(node, tf_import_flags, model);
+  } else if (node.op() == "StridedSlice") {
+    ConvertStridedSliceOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Shape") {
+    ConvertShapeOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Slice") {
+    ConvertSliceOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Split") {
+    ConvertSplitOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Switch") {
+    ConvertSwitchOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Placeholder") {
+    ConvertPlaceholderOperator(node, tf_import_flags, model);
+  } else if (node.op() == "PlaceholderWithDefault") {
+    ConvertIdentityOperator(node, tf_import_flags, model);
+  } else if (node.op() == "LegacyFedInput") {
+    ConvertPlaceholderOperator(node, tf_import_flags, model);
+  } else if (node.op() == "NoOp") {
+    ConvertNoOpOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Cast") {
+    ConvertCastOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Floor") {
+    ConvertFloorOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Gather" || node.op() == "GatherV2") {
+    ConvertGatherOperator(node, tf_import_flags, model);
+  } else if (node.op() == "ResizeBilinear") {
+    ConvertResizeBilinearOperator(node, tf_import_flags, model);
+  } else if (node.op() == "BatchNormWithGlobalNormalization") {
+    ConvertBatchNormWithGlobalNormalizationOperator(node, tf_import_flags,
+                                                    model);
+  } else if (node.op() == "FusedBatchNorm") {
+    ConvertFusedBatchNormOperator(node, tf_import_flags, model);
+  } else if (node.op() == "SpaceToBatchND") {
+    ConvertSpaceToBatchNDOperator(node, tf_import_flags, model);
+  } else if (node.op() == "BatchToSpaceND") {
+    ConvertBatchToSpaceNDOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Mean") {
+    ConvertMeanOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Svdf") {
+    ConvertSvdfOperator(node, tf_import_flags, model);
+  } else if (node.op() == "NextIteration") {
+    ConvertOperatorSpecialCasedAsRNNBackEdge(node, tf_import_flags, model);
+  } else if (node.op() == "ExpandDims") {
+    ConvertExpandDimsOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Fill") {
+    ConvertFillOperator(node, tf_import_flags, model);
+  } else if (node.op() == "FloorDiv") {
+    ConvertFloorDivOperator(node, tf_import_flags, model);
+  } else if (node.op() == "FloorMod") {
+    ConvertFloorModOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Range") {
+    ConvertRangeOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Rank") {
+    ConvertRankOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Stack" || node.op() == "Pack") {
+    ConvertStackOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Transpose") {
+    ConvertTransposeOperator(node, tf_import_flags, model);
+  } else if (node.op() == "ArgMax") {
+    ConvertArgMaxOperator(node, tf_import_flags, model);
+  } else if (node.op() == "Exp") {
+    ConvertExpOperator(node, tf_import_flags, model);
+  } else if (node.op() == "TopK" || node.op() == "TopKV2") {
+    ConvertTopKV2Operator(node, tf_import_flags, model);
+  } else if (node.op() == "DynamicPartition") {
+    ConvertDynamicPartitionOperator(node, tf_import_flags, model);
+  } else if (node.op() == "DynamicStitch" ||
+             node.op() == "ParallelDynamicStitch") {
+    ConvertDynamicStitchOperator(node, tf_import_flags, model);
+  } else if (node.op() == "RandomUniform") {
+    ConvertRandomUniform(node, tf_import_flags, model);
+  } else {
+    ConvertUnsupportedOperator(node, tf_import_flags, model);
+  }
+  return Status::OK();
+}
+}  // namespace internal
+
 std::unique_ptr<Model> ImportTensorFlowGraphDef(
     const ModelFlags& model_flags, const TensorFlowImportFlags& tf_import_flags,
     const GraphDef& tf_graph) {
@@ -2058,176 +2271,8 @@ std::unique_ptr<Model> ImportTensorFlowGraphDef(
 
   for (auto node : inlined_graph.node()) {
     StripZeroOutputIndexFromInputs(&node);
-    if (node.op() == "Const") {
-      ConvertConstOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Conv2D") {
-      ConvertConvOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Conv2DBackpropInput") {
-      ConvertTransposeConvOperator(node, tf_import_flags, model);
-    } else if (node.op() == "DepthwiseConv2dNative") {
-      ConvertDepthwiseConvOperator(node, tf_import_flags, model);
-    } else if (node.op() == "DepthToSpace") {
-      ConvertDepthToSpaceOperator(node, tf_import_flags, model);
-    } else if (node.op() == "SpaceToDepth") {
-      ConvertSpaceToDepthOperator(node, tf_import_flags, model);
-    } else if (node.op() == "BiasAdd") {
-      ConvertBiasAddOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Relu") {
-      ConvertReluOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Relu6") {
-      ConvertRelu6Operator(node, tf_import_flags, model);
-    } else if (node.op() == "Sigmoid") {
-      ConvertLogisticOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Tanh") {
-      ConvertTanhOperator(node, tf_import_flags, model);
-    } else if (node.op() == "MaxPool") {
-      ConvertMaxPoolOperator(node, tf_import_flags, model);
-    } else if (node.op() == "AvgPool") {
-      ConvertAvgPoolOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Reshape") {
-      ConvertReshapeOperator(node, tf_import_flags, model);
-    } else if (node.op() == "BatchMatMul") {
-      ConvertBatchMatMulOperator(node, tf_import_flags, model);
-    } else if (node.op() == "MatMul") {
-      ConvertMatMulOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Div" || node.op() == "RealDiv") {
-      ConvertDivOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Identity" || node.op() == "CheckNumerics" ||
-               node.op() == "StopGradient") {
-      ConvertIdentityOperator(node, tf_import_flags, model);
-    } else if (node.op() == "FakeQuantWithMinMaxVars") {
-      ConvertFakeQuantWithMinMaxVars(node, tf_import_flags, model);
-    } else if (node.op() == "FakeQuantWithMinMaxArgs") {
-      ConvertFakeQuantWithMinMaxArgs(node, tf_import_flags, model);
-    } else if (node.op() == "Neg") {
-      ConvertNegOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Rsqrt") {
-      ConvertRsqrtOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Squeeze") {
-      ConvertSqueezeOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Sqrt") {
-      ConvertSqrtOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Square") {
-      ConvertSquareOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Add") {
-      ConvertAddOperator(node, tf_import_flags, model);
-    } else if (node.op() == "AddN") {
-      ConvertAddNOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Mul") {
-      ConvertMulOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Sub") {
-      ConvertSubOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Sum") {
-      ConvertSumOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Tile") {
-      ConvertTileOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Concat" || node.op() == "ConcatV2") {
-      ConvertConcatOperator(node, tf_import_flags, model);
-    } else if (node.op() == "LRN") {
-      ConvertLRNOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Softmax") {
-      ConvertSoftmaxOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Log") {
-      ConvertLogOperator(node, tf_import_flags, model);
-    } else if (node.op() == "LogSoftmax") {
-      ConvertLogSoftmaxOperator(node, tf_import_flags, model);
-    } else if (node.op() == "All") {
-      ConvertAllOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Assert") {
-      ConvertAssertOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Less") {
-      ConvertLessOperator(node, tf_import_flags, model);
-    } else if (node.op() == "LessEqual") {
-      ConvertLessEqualOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Greater") {
-      ConvertGreaterOperator(node, tf_import_flags, model);
-    } else if (node.op() == "GreaterEqual") {
-      ConvertGreaterEqualOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Max") {
-      ConvertMaxOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Min") {
-      ConvertMinOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Maximum") {
-      ConvertMaximumOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Minimum") {
-      ConvertMinimumOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Merge") {
-      ConvertMergeOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Pad") {
-      ConvertPadOperator(node, tf_import_flags, model);
-    } else if (node.op() == "StridedSlice") {
-      ConvertStridedSliceOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Shape") {
-      ConvertShapeOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Slice") {
-      ConvertSliceOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Split") {
-      ConvertSplitOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Switch") {
-      ConvertSwitchOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Placeholder") {
-      ConvertPlaceholderOperator(node, tf_import_flags, model);
-    } else if (node.op() == "PlaceholderWithDefault") {
-      ConvertIdentityOperator(node, tf_import_flags, model);
-    } else if (node.op() == "LegacyFedInput") {
-      ConvertPlaceholderOperator(node, tf_import_flags, model);
-    } else if (node.op() == "NoOp") {
-      ConvertNoOpOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Cast") {
-      ConvertCastOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Floor") {
-      ConvertFloorOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Gather" || node.op() == "GatherV2") {
-      ConvertGatherOperator(node, tf_import_flags, model);
-    } else if (node.op() == "ResizeBilinear") {
-      ConvertResizeBilinearOperator(node, tf_import_flags, model);
-    } else if (node.op() == "BatchNormWithGlobalNormalization") {
-      ConvertBatchNormWithGlobalNormalizationOperator(node, tf_import_flags,
-                                                      model);
-    } else if (node.op() == "FusedBatchNorm") {
-      ConvertFusedBatchNormOperator(node, tf_import_flags, model);
-    } else if (node.op() == "SpaceToBatchND") {
-      ConvertSpaceToBatchNDOperator(node, tf_import_flags, model);
-    } else if (node.op() == "BatchToSpaceND") {
-      ConvertBatchToSpaceNDOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Mean") {
-      ConvertMeanOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Svdf") {
-      ConvertSvdfOperator(node, tf_import_flags, model);
-    } else if (node.op() == "NextIteration") {
-      ConvertOperatorSpecialCasedAsRNNBackEdge(node, tf_import_flags, model);
-    } else if (node.op() == "ExpandDims") {
-      ConvertExpandDimsOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Fill") {
-      ConvertFillOperator(node, tf_import_flags, model);
-    } else if (node.op() == "FloorDiv") {
-      ConvertFloorDivOperator(node, tf_import_flags, model);
-    } else if (node.op() == "FloorMod") {
-      ConvertFloorModOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Range") {
-      ConvertRangeOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Rank") {
-      ConvertRankOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Stack" || node.op() == "Pack") {
-      ConvertStackOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Transpose") {
-      ConvertTransposeOperator(node, tf_import_flags, model);
-    } else if (node.op() == "ArgMax") {
-      ConvertArgMaxOperator(node, tf_import_flags, model);
-    } else if (node.op() == "Exp") {
-      ConvertExpOperator(node, tf_import_flags, model);
-    } else if (node.op() == "TopK" || node.op() == "TopKV2") {
-      ConvertTopKV2Operator(node, tf_import_flags, model);
-    } else if (node.op() == "DynamicPartition") {
-      ConvertDynamicPartitionOperator(node, tf_import_flags, model);
-    } else if (node.op() == "DynamicStitch" ||
-               node.op() == "ParallelDynamicStitch") {
-      ConvertDynamicStitchOperator(node, tf_import_flags, model);
-    } else if (node.op() == "RandomUniform") {
-      ConvertRandomUniform(node, tf_import_flags, model);
-    } else {
-      ConvertUnsupportedOperator(node, tf_import_flags, model);
-    }
+    auto status = internal::ImportTensorFlowNode(node, tf_import_flags, model);
+    CHECK(status.ok()) << status.error_message();
   }
 
   ResolveModelFlags(model_flags, model);
diff --git a/tensorflow/contrib/lite/toco/import_tensorflow_test.cc b/tensorflow/contrib/lite/toco/import_tensorflow_test.cc
new file mode 100644
index 0000000000..5dc78f73ad
--- /dev/null
+++ b/tensorflow/contrib/lite/toco/import_tensorflow_test.cc
@@ -0,0 +1,160 @@
+/* Copyright 2018 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.
+==============================================================================*/
+#include "tensorflow/contrib/lite/toco/import_tensorflow.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/attr_value_util.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+
+namespace toco {
+
+using port::Status;
+using tensorflow::AttrValue;
+using tensorflow::DT_BOOL;
+using tensorflow::DT_FLOAT;
+using tensorflow::DT_INT32;
+using tensorflow::DT_INT64;
+using tensorflow::DT_QUINT8;
+using tensorflow::DT_STRING;
+using tensorflow::NodeDef;
+
+namespace internal {
+Status ImportTensorFlowNode(const NodeDef&, const TensorFlowImportFlags&,
+                            Model*);
+}  // namespace internal
+
+namespace {
+
+class ShapeImportTest : public ::testing::TestWithParam<tensorflow::DataType> {
+ protected:
+  ShapeImportTest() {}
+
+  void BuildConstNode(std::initializer_list<int64_t> shape,
+                      tensorflow::DataType dtype, int64_t num_elements,
+                      NodeDef* node) {
+    node->set_op("Const");
+    node->set_name("Node1");
+
+    // An attribute describing the type of this const node.
+    AttrValue dtype_attr;
+    SetAttrValue(dtype, &dtype_attr);
+    (*node->mutable_attr())["dtype"] = dtype_attr;
+
+    // An attribute describing the content of this const node.
+    tensorflow::TensorProto t;
+    t.set_dtype(dtype);
+    auto* s = t.mutable_tensor_shape();
+    for (auto d : shape) {
+      s->add_dim()->set_size(d);
+    }
+
+    // TODO(ahentz): also need to test via tensor_content()
+    switch (dtype) {
+      case DT_FLOAT:
+        for (int64_t i = 0; i < num_elements; ++i) {
+          t.add_float_val(i / 10000.0);
+        }
+        break;
+      case DT_INT32:
+        for (int64_t i = 0; i < num_elements; ++i) {
+          t.add_int_val(i % std::numeric_limits<int>::max());
+        }
+        break;
+      case DT_QUINT8:
+        for (int64_t i = 0; i < num_elements; ++i) {
+          t.add_int_val(i % std::numeric_limits<uint8_t>::max());
+        }
+        break;
+      case DT_INT64:
+        for (int64_t i = 0; i < num_elements; ++i) {
+          t.add_int64_val(i);
+        }
+        break;
+      case DT_STRING:
+        break;
+      case DT_BOOL:
+        for (int64_t i = 0; i < num_elements; ++i) {
+          t.add_bool_val(i % 2);
+        }
+        break;
+      default:
+        break;
+    }
+
+    AttrValue value_attr;
+    SetAttrValue(t, &value_attr);
+    (*node->mutable_attr())["value"] = value_attr;
+  }
+
+  Status ImportNode(const NodeDef& node) {
+    Model model;
+    return internal::ImportTensorFlowNode(node, TensorFlowImportFlags(),
+                                          &model);
+  }
+};
+
+std::vector<tensorflow::DataType> TestTypes() {
+  return {DT_FLOAT, DT_INT32, DT_INT64, DT_BOOL, DT_QUINT8};
+}
+
+TEST_P(ShapeImportTest, ShapeElementIsNegative) {
+  NodeDef node;
+  BuildConstNode({1, -2, 10}, GetParam(), 0, &node);
+  auto status = ImportNode(node);
+  EXPECT_EQ(status.error_message(),
+            "Tensor shape should not include negative values (while processing "
+            "node 'Node1')");
+}
+INSTANTIATE_TEST_CASE_P(ShapeElementIsNegative, ShapeImportTest,
+                        ::testing::ValuesIn(TestTypes()));
+
+TEST_P(ShapeImportTest, ShapeElementTooLarge) {
+  NodeDef node;
+  BuildConstNode({3000000000}, GetParam(), 0, &node);
+  auto status = ImportNode(node);
+  EXPECT_EQ(status.error_message(),
+            "Shape element overflows (while processing node 'Node1')");
+}
+INSTANTIATE_TEST_CASE_P(ShapeElementTooLarge, ShapeImportTest,
+                        ::testing::ValuesIn(TestTypes()));
+
+TEST_P(ShapeImportTest, ShapeTooLarge) {
+  NodeDef node;
+  BuildConstNode({1000000, 2000000, 2000000, 2000000}, GetParam(), 0, &node);
+  auto status = ImportNode(node);
+  EXPECT_EQ(status.error_message(),
+            "Tensor shape is too large (while processing node 'Node1')");
+}
+INSTANTIATE_TEST_CASE_P(ShapeTooLarge, ShapeImportTest,
+                        ::testing::ValuesIn(TestTypes()));
+
+TEST_P(ShapeImportTest, ValidShapeButZeroElements) {
+  NodeDef node;
+  BuildConstNode({1, 2, 2, 2}, GetParam(), 0, &node);
+  auto status = ImportNode(node);
+  EXPECT_THAT(status.error_message(),
+              ::testing::MatchesRegex(
+                  "Neither input_content nor .*_val have the right dimensions "
+                  "for this .* tensor .while processing node 'Node1'."));
+}
+INSTANTIATE_TEST_CASE_P(ValidShapeButZeroElements, ShapeImportTest,
+                        ::testing::ValuesIn(TestTypes()));
+
+}  // namespace
+}  // namespace toco
diff --git a/tensorflow/contrib/lite/toco/toco_port.h b/tensorflow/contrib/lite/toco/toco_port.h
index 2d5c231bef..906792ef56 100644
--- a/tensorflow/contrib/lite/toco/toco_port.h
+++ b/tensorflow/contrib/lite/toco/toco_port.h
@@ -38,10 +38,15 @@ namespace port {
 
 class Status {
  public:
+  static Status OK() { return Status(true, ""); }
+
+  // Create a failed status with no message.
   Status() {}
 
   Status(bool ok, const string& message) : ok_(ok), message_(message) {}
 
+  void AppendMessage(const string& message) { message_ += message; }
+
   bool ok() const { return ok_; }
 
   const string error_message() const { return message_; }
diff --git a/tensorflow/contrib/lite/toco/tooling_util.h b/tensorflow/contrib/lite/toco/tooling_util.h
index 5cc15fa57b..f5b596df0f 100644
--- a/tensorflow/contrib/lite/toco/tooling_util.h
+++ b/tensorflow/contrib/lite/toco/tooling_util.h
@@ -294,6 +294,35 @@ void FinishBuildingRNNStates(Model* model);
 
 void UseArraysExtraInfo(Model* model, bool quantize_output);
 
+// Calculates the number of elements in tensor given a shape. Shape elements
+// are assumed to be of type T, while the result total is of type U. If U
+// doesn't have enough range to represent the sum of elements, an error is
+// returned.
+template <typename T, typename U>
+port::Status NumElements(const std::vector<T>& shape, U* num_elements) {
+  static_assert(
+      std::numeric_limits<T>::max() <= std::numeric_limits<uint64_t>::max(),
+      "vector type exceed capabilities of NumElements");
+
+  *num_elements = 1;
+  for (const T& dim : shape) {
+    if (dim < 0) {
+      // TensorFlow's shapes sometimes include -1 to represent an "unknown"
+      // size but TOCO isn't able to create arrays of unknown sizes and will
+      // crash in RequiredBufferSizeForShape().
+      return port::Status(false,
+                          "Tensor shape should not include negative values");
+    }
+    if (static_cast<uint64_t>(dim) >
+        std::numeric_limits<U>::max() / *num_elements) {
+      *num_elements = 0;
+      return port::Status(false, "Tensor shape is too large");
+    }
+    *num_elements *= dim;
+  }
+  return port::Status::OK();
+}
+
 }  // namespace toco
 
 #endif  // TENSORFLOW_CONTRIB_LITE_TOCO_TOOLING_UTIL_H_
diff --git a/tensorflow/contrib/lite/toco/tooling_util_test.cc b/tensorflow/contrib/lite/toco/tooling_util_test.cc
index 22955ce956..87fd30db2c 100644
--- a/tensorflow/contrib/lite/toco/tooling_util_test.cc
+++ b/tensorflow/contrib/lite/toco/tooling_util_test.cc
@@ -93,4 +93,85 @@ TEST_P(ShapeTest, Agrees) {
 INSTANTIATE_TEST_CASE_P(AgreeBroadcast, ShapeTest,
                         ::testing::ValuesIn(CreateShapePairs()));
 
+static const char kNegativeValuesMessage[] =
+    "Tensor shape should not include negative values";
+static const char kLargeTensorMessage[] = "Tensor shape is too large";
+
+TEST(NumElementsTest, Int) {
+  int count;
+  port::Status status = port::Status::OK();
+
+  status = NumElements(std::vector<int>{1024, 1024, 2047}, &count);
+  EXPECT_TRUE(status.ok());
+  EXPECT_EQ(count, 2146435072);
+
+  status = NumElements(std::vector<int>{1, 2, -3}, &count);
+  EXPECT_EQ(status.error_message(), kNegativeValuesMessage);
+
+  status = NumElements(std::vector<int>{1024, 1024, 2048}, &count);
+  EXPECT_EQ(status.error_message(), kLargeTensorMessage);
+}
+
+TEST(NumElementsTest, Int32) {
+  int32_t count;
+  port::Status status = port::Status::OK();
+
+  status = NumElements(std::vector<int32_t>{1024, 1024, 2047}, &count);
+  EXPECT_TRUE(status.ok());
+  EXPECT_EQ(count, 2146435072);
+
+  status = NumElements(std::vector<int32_t>{1, 2, -3}, &count);
+  EXPECT_EQ(status.error_message(), kNegativeValuesMessage);
+
+  status = NumElements(std::vector<int32_t>{1024, 1024, 2048}, &count);
+  EXPECT_EQ(status.error_message(), kLargeTensorMessage);
+}
+
+TEST(NumElementsTest, Int64) {
+  int64_t count;
+  port::Status status = port::Status::OK();
+
+  status = NumElements(std::vector<int64_t>{16777216, 16777216, 32767}, &count);
+  EXPECT_TRUE(status.ok());
+  EXPECT_EQ(count, 9223090561878065152LL);
+
+  status = NumElements(std::vector<int64_t>{1, 2, -3}, &count);
+  EXPECT_EQ(status.error_message(), kNegativeValuesMessage);
+
+  status = NumElements(std::vector<int64_t>{16777216, 16777216, 32768}, &count);
+  EXPECT_EQ(status.error_message(), kLargeTensorMessage);
+}
+
+TEST(NumElementsTest, UnsignedInt32) {
+  uint32_t count;
+  port::Status status = port::Status::OK();
+
+  status = NumElements(std::vector<uint32_t>{1024, 2048, 2047}, &count);
+  EXPECT_TRUE(status.ok());
+  EXPECT_EQ(count, 4292870144);
+
+  status = NumElements(std::vector<int>{1, 2, -3}, &count);
+  EXPECT_EQ(status.error_message(), kNegativeValuesMessage);
+
+  status = NumElements(std::vector<uint32_t>{1024, 2048, 2048}, &count);
+  EXPECT_EQ(status.error_message(), kLargeTensorMessage);
+}
+
+TEST(NumElementsTest, UnsignedInt64) {
+  uint64_t count;
+  port::Status status = port::Status::OK();
+
+  status =
+      NumElements(std::vector<uint64_t>{16777216, 16777216, 65535}, &count);
+  EXPECT_TRUE(status.ok());
+  EXPECT_EQ(count, 18446462598732840960ULL);
+
+  status = NumElements(std::vector<int>{1, 2, -3}, &count);
+  EXPECT_EQ(status.error_message(), kNegativeValuesMessage);
+
+  status =
+      NumElements(std::vector<uint64_t>{16777216, 16777216, 65536}, &count);
+  EXPECT_EQ(status.error_message(), kLargeTensorMessage);
+}
+
 }  // namespace toco
-- 
GitLab


From 1aa7aaa731ad8b64345fbfec2f53b49a77a9b94d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 16:38:19 -0700
Subject: [PATCH 124/395] Adds logistic_regression_head.

PiperOrigin-RevId: 195017830
---
 tensorflow/contrib/estimator/__init__.py      |   1 +
 .../estimator/python/estimator/head.py        |  69 +++++++++-
 .../estimator/python/estimator/head_test.py   | 119 ++++++++++++++++++
 3 files changed, 187 insertions(+), 2 deletions(-)

diff --git a/tensorflow/contrib/estimator/__init__.py b/tensorflow/contrib/estimator/__init__.py
index be20d1b777..f66d844660 100644
--- a/tensorflow/contrib/estimator/__init__.py
+++ b/tensorflow/contrib/estimator/__init__.py
@@ -38,6 +38,7 @@ _allowed_symbols = [
     'binary_classification_head',
     'clip_gradients_by_norm',
     'forward_features',
+    'logistic_regression_head',
     'multi_class_head',
     'multi_head',
     'multi_label_head',
diff --git a/tensorflow/contrib/estimator/python/estimator/head.py b/tensorflow/contrib/estimator/python/estimator/head.py
index 3dcf0374c8..2a6d17e81b 100644
--- a/tensorflow/contrib/estimator/python/estimator/head.py
+++ b/tensorflow/contrib/estimator/python/estimator/head.py
@@ -205,8 +205,9 @@ def regression_head(weight_column=None,
   shape `[D0, D1, ... DN, label_dimension]`.
 
   Also supports custom `inverse_link_fn`, also known as 'mean function'.
-  `inverse_link_fn` takes `logits` as argument and returns predicted values.
-  This function is the inverse of the link function defined in
+  `inverse_link_fn` is only used in `PREDICT` mode. It takes `logits` as
+  argument and returns predicted values. This function is the inverse of the
+  link function defined in
   https://en.wikipedia.org/wiki/Generalized_linear_model#Link_function
   Namely, for poisson regression, set `inverse_link_fn=tf.exp`.
 
@@ -305,6 +306,70 @@ def poisson_regression_head(
       name=name)
 
 
+def logistic_regression_head(
+    weight_column=None,
+    loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE,
+    name=None):
+  """Creates a `_Head` for logistic regression.
+
+  Uses `sigmoid_cross_entropy_with_logits` loss, which is the same as
+  `binary_classification_head`. The differences compared to
+  `binary_classification_head` are:
+
+  * Does not support `label_vocabulary`. Instead, labels must be float in the
+    range [0, 1].
+  * Does not calculate some metrics that do not make sense, such as AUC.
+  * In `PREDICT` mode, only returns logits and predictions
+    (`=tf.sigmoid(logits)`), whereas `binary_classification_head` also returns
+    probabilities, classes, and class_ids.
+  * Export output defaults to `RegressionOutput`, whereas
+    `binary_classification_head` defaults to `PredictOutput`.
+
+  The head expects `logits` with shape `[D0, D1, ... DN, 1]`.
+  In many applications, the shape is `[batch_size, 1]`.
+
+  The `labels` shape must match `logits`, namely
+  `[D0, D1, ... DN]` or `[D0, D1, ... DN, 1]`.
+
+  If `weight_column` is specified, weights must be of shape
+  `[D0, D1, ... DN]` or `[D0, D1, ... DN, 1]`.
+
+  This is implemented as a generalized linear model, see
+  https://en.wikipedia.org/wiki/Generalized_linear_model.
+
+  Args:
+    weight_column: A string or a `_NumericColumn` created by
+      `tf.feature_column.numeric_column` defining feature column representing
+      weights. It is used to down weight or boost examples during training. It
+      will be multiplied by the loss of the example.
+    loss_reduction: One of `tf.losses.Reduction` except `NONE`. Describes how to
+      reduce training loss over batch and label dimension. Defaults to
+      `SUM_OVER_BATCH_SIZE`, namely weighted sum of losses divided by
+      `batch size * label_dimension`. See `tf.losses.Reduction`.
+    name: name of the head. If provided, summary and metrics keys will be
+      suffixed by `"/" + name`. Also used as `name_scope` when creating ops.
+
+  Returns:
+    An instance of `_Head` for logistic regression.
+
+  Raises:
+    ValueError: If `loss_reduction` is invalid.
+  """
+  def _logistic_loss(labels, logits):
+    labels = head_lib._assert_range(  # pylint:disable=protected-access
+        labels, n_classes=2, message='Labels must be in range [0, 1]')
+    return nn.sigmoid_cross_entropy_with_logits(
+        labels=labels, logits=logits)
+  # TODO(roumposg): Rename to _regression_head, since it supports loss_fn arg.
+  return head_lib._regression_head_with_mean_squared_error_loss(  # pylint:disable=protected-access
+      weight_column=weight_column,
+      label_dimension=1,
+      loss_reduction=loss_reduction,
+      loss_fn=_logistic_loss,
+      inverse_link_fn=math_ops.sigmoid,
+      name=name)
+
+
 def multi_label_head(n_classes,
                      weight_column=None,
                      thresholds=None,
diff --git a/tensorflow/contrib/estimator/python/estimator/head_test.py b/tensorflow/contrib/estimator/python/estimator/head_test.py
index 98962ca427..19b86df556 100644
--- a/tensorflow/contrib/estimator/python/estimator/head_test.py
+++ b/tensorflow/contrib/estimator/python/estimator/head_test.py
@@ -1211,5 +1211,124 @@ class PoissonRegressionHead(test.TestCase):
       self.assertAllClose(logits, spec.predictions[keys.LOGITS].eval())
 
 
+class LogisticRegressionHead(test.TestCase):
+
+  def setUp(self):
+    ops.reset_default_graph()
+
+  def test_train(self):
+    head = head_lib.logistic_regression_head()
+
+    # Create estimator spec.
+    logits = np.array([[0], [-1], [1]], dtype=np.float32)
+    labels = np.array([[.4], [.6], [.8]], dtype=np.float32)
+    # Following the documentation in
+    # tf.nn.sigmoid_cross_entropy_with_logits:
+    # With x = logits, z = labels.
+    # loss  = max(x, 0) - x * z + log(1 + exp(-abs(x)))
+    # loss = [0 - 0 * 0.4 + ln(1 + exp(-0)),
+    #         0 + 1 * 0.6 + ln(1 + exp(-1)),
+    #         1 - 1 * 0.8 + ln(1 + exp(-1))]
+    #      = [0.6931, 0.9133, 0.5133]
+    # training_loss = (0.6931 + 0.9133 + 0.5133) / 3
+    expected_loss = 0.7066
+    atol = 0.001
+    expected_train_result = b'my_train_op'
+    def _train_op_fn(loss):
+      with ops.control_dependencies((check_ops.assert_near(
+          math_ops.to_float(expected_loss), math_ops.to_float(loss),
+          atol=atol, name='assert_loss'),)):
+        return constant_op.constant(expected_train_result)
+
+    spec = head.create_estimator_spec(
+        features={'x': np.array(((42.,),), dtype=np.int32)},
+        mode=model_fn.ModeKeys.TRAIN,
+        logits=logits,
+        labels=labels,
+        train_op_fn=_train_op_fn)
+
+    with self.test_session() as sess:
+      _initialize_variables(self, spec.scaffold)
+      loss, train_result = sess.run([spec.loss, spec.train_op])
+      self.assertAlmostEqual(expected_loss, loss, delta=atol)
+      self.assertEqual(expected_train_result, train_result)
+
+  def test_train_labels_too_large(self):
+    head = head_lib.logistic_regression_head()
+
+    # Create estimator spec.
+    logits = np.array([[0], [-1], [1]], dtype=np.float32)
+    labels = np.array([[.4], [1.2], [.8]], dtype=np.float32)
+    expected_train_result = b'my_train_op'
+    def _train_op_fn(loss):
+      del loss
+      return constant_op.constant(expected_train_result)
+
+    spec = head.create_estimator_spec(
+        features={'x': np.array(((42.,),), dtype=np.int32)},
+        mode=model_fn.ModeKeys.TRAIN,
+        logits=logits,
+        labels=labels,
+        train_op_fn=_train_op_fn)
+
+    with self.test_session() as sess:
+      _initialize_variables(self, spec.scaffold)
+      with self.assertRaisesRegexp(
+          errors.InvalidArgumentError,
+          r'\[Labels must be in range \[0, 1\]\] .* \[\[0.4\]\[1.2\]\[0.8\]\]'):
+        _ = sess.run(spec.loss)
+
+  def test_train_labels_negative(self):
+    head = head_lib.logistic_regression_head()
+
+    # Create estimator spec.
+    logits = np.array([[0], [-1], [1]], dtype=np.float32)
+    labels = np.array([[.4], [-0.2], [.8]], dtype=np.float32)
+    expected_train_result = b'my_train_op'
+    def _train_op_fn(loss):
+      del loss
+      return constant_op.constant(expected_train_result)
+
+    spec = head.create_estimator_spec(
+        features={'x': np.array(((42.,),), dtype=np.int32)},
+        mode=model_fn.ModeKeys.TRAIN,
+        logits=logits,
+        labels=labels,
+        train_op_fn=_train_op_fn)
+
+    with self.test_session() as sess:
+      _initialize_variables(self, spec.scaffold)
+      with self.assertRaisesRegexp(
+          errors.InvalidArgumentError,
+          r'\[Labels must be in range \[0, 1\]\] .* \[\[0.4\]\[-0.2\]\[0.8\]\]'
+      ):
+        _ = sess.run(spec.loss)
+
+  def test_predict(self):
+    head = head_lib.logistic_regression_head()
+
+    # Create estimator spec.
+    logits = np.array([[0], [-1], [1]], dtype=np.float32)
+    expected_predictions = 1. / (1. + np.exp(-logits))
+    spec = head.create_estimator_spec(
+        features={'x': np.array(((42.,),), dtype=np.int32)},
+        mode=model_fn.ModeKeys.PREDICT,
+        logits=logits)
+
+    # Assert spec contains expected tensors.
+    keys = prediction_keys.PredictionKeys
+    self.assertItemsEqual(
+        (keys.PREDICTIONS, keys.LOGITS), spec.predictions.keys())
+    self.assertEqual(dtypes.float32, spec.predictions[keys.PREDICTIONS].dtype)
+    self.assertEqual(dtypes.float32, spec.predictions[keys.LOGITS].dtype)
+
+    # Assert predictions.
+    with self.test_session():
+      _initialize_variables(self, spec.scaffold)
+      self.assertAllClose(
+          expected_predictions, spec.predictions[keys.PREDICTIONS].eval())
+      self.assertAllClose(logits, spec.predictions[keys.LOGITS].eval())
+
+
 if __name__ == '__main__':
   test.main()
-- 
GitLab


From 2b547749bccd73bb95d05b6e73f26ea8ae9f3be6 Mon Sep 17 00:00:00 2001
From: Benoit Steiner <bsteiner@google.com>
Date: Tue, 1 May 2018 16:38:23 -0700
Subject: [PATCH 125/395] Avoid making a copy of the graph needlessly

PiperOrigin-RevId: 195017837
---
 tensorflow/core/grappler/costs/graph_properties.cc | 2 +-
 tensorflow/core/grappler/costs/graph_properties.h  | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/grappler/costs/graph_properties.cc b/tensorflow/core/grappler/costs/graph_properties.cc
index 431efb08cb..2c7b57971a 100644
--- a/tensorflow/core/grappler/costs/graph_properties.cc
+++ b/tensorflow/core/grappler/costs/graph_properties.cc
@@ -1074,7 +1074,7 @@ Status GraphProperties::InferStatically(bool assume_valid_feeds) {
     }
   }
 
-  GraphView graph_view(&item_.graph);
+  GraphView graph_view(const_cast<GraphDef*>(&item_.graph));
 
   // List the resources and the nodes using them. Also collect the Merge nodes,
   // fed nodes, and primary inputs.
diff --git a/tensorflow/core/grappler/costs/graph_properties.h b/tensorflow/core/grappler/costs/graph_properties.h
index ecc10fddb8..8703613a12 100644
--- a/tensorflow/core/grappler/costs/graph_properties.h
+++ b/tensorflow/core/grappler/costs/graph_properties.h
@@ -38,6 +38,7 @@ class TopoQueue;
 // and data type properties.
 class GraphProperties {
  public:
+  // The item must outlive the properties
   explicit GraphProperties(const GrapplerItem& item) : item_(item) {}
 
   // Infer the shapes through abstract interpretation. Feed information can be
@@ -112,7 +113,7 @@ class GraphProperties {
       int num_loops) const;
 
   // Data members
-  GrapplerItem item_;
+  const GrapplerItem& item_;
   std::map<string, std::vector<OpInfo::TensorProperties>> input_properties_;
   std::map<string, std::vector<OpInfo::TensorProperties>> output_properties_;
   const std::vector<OpInfo::TensorProperties> missing_properties_;
-- 
GitLab


From 833803d76093bfcd738136694e2d78db8856b5ae Mon Sep 17 00:00:00 2001
From: Rachel Lim <rachelim@google.com>
Date: Tue, 1 May 2018 16:52:24 -0700
Subject: [PATCH 126/395] Fix wrongly ordered lines

PiperOrigin-RevId: 195019769
---
 tensorflow/core/ops/dataset_ops.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/ops/dataset_ops.cc b/tensorflow/core/ops/dataset_ops.cc
index 5f10ad24b6..73174c184c 100644
--- a/tensorflow/core/ops/dataset_ops.cc
+++ b/tensorflow/core/ops/dataset_ops.cc
@@ -478,11 +478,11 @@ REGISTER_OP("TextLineDataset")
       shape_inference::ShapeHandle unused;
       // `filenames` must be a scalar or a vector.
       TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 1, &unused));
-      return shape_inference::ScalarShape(c);
       // `compression_type` could only be a scalar.
       TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
       // `buffer_size` could only be a scalar.
       TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
+      return shape_inference::ScalarShape(c);
     });
 
 REGISTER_OP("SqlDataset")
-- 
GitLab


From 448b2ca80f91782f50c50f72bf7feafedcd744b6 Mon Sep 17 00:00:00 2001
From: Allen Lavoie <allenl@google.com>
Date: Tue, 1 May 2018 16:53:51 -0700
Subject: [PATCH 127/395] Sharding for
 tensorflow/contrib/timeseries/python/timeseries/state_space_models:structural_ensemble_test

PiperOrigin-RevId: 195019968
---
 .../timeseries/python/timeseries/state_space_models/BUILD      | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD
index 5d33e23a42..3c07a74ed8 100644
--- a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD
+++ b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD
@@ -176,8 +176,9 @@ py_library(
 
 py_test(
     name = "structural_ensemble_test",
-    timeout = "long",  # Moderate but for asan/tsan timeouts
+    timeout = "long",  # Moderate but for asan/tsan/msan timeouts
     srcs = ["structural_ensemble_test.py"],
+    shard_count = 4,
     srcs_version = "PY2AND3",
     deps = [
         ":state_space_model",
-- 
GitLab


From f453b62176cf57659ca0485e3b37b8f08c05b966 Mon Sep 17 00:00:00 2001
From: Mustafa Ispir <ispir@google.com>
Date: Tue, 1 May 2018 17:21:24 -0700
Subject: [PATCH 128/395] test fix

PiperOrigin-RevId: 195023740
---
 .../contrib/layers/python/layers/embedding_ops_test.py      | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensorflow/contrib/layers/python/layers/embedding_ops_test.py b/tensorflow/contrib/layers/python/layers/embedding_ops_test.py
index bf25144982..dd2395f8c9 100644
--- a/tensorflow/contrib/layers/python/layers/embedding_ops_test.py
+++ b/tensorflow/contrib/layers/python/layers/embedding_ops_test.py
@@ -31,6 +31,7 @@ from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors_impl
 from tensorflow.python.framework import random_seed
 from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib
+from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import gradient_checker
 from tensorflow.python.ops import init_ops
@@ -691,11 +692,12 @@ class EmbeddingLookupSparseWithDistributedAggregationTest(test.TestCase):
       index += num_val
     return grouped_vals
 
+  @test_util.enable_c_shapes
   def testEmbeddingLookupSparse(self):
     vocab_size = 13
     batch_size = 10
     param_shape = [2, 5]
-    expected_lookup_result_shape = [None] + param_shape
+    expected_lookup_result_shape = param_shape
 
     sp_ids, sp_weights, ids, weights, vals_per_batch_entry = (
         self._RandomIdsAndWeights(batch_size, vocab_size))
@@ -719,7 +721,7 @@ class EmbeddingLookupSparseWithDistributedAggregationTest(test.TestCase):
                 None if ignore_weights else sp_weights,
                 combiner=combiner)
 
-        self.assertEqual(embedding_sum.get_shape().as_list(),
+        self.assertEqual(embedding_sum.get_shape().as_list()[1:],
                          expected_lookup_result_shape)
 
         tf_embedding_sum = embedding_sum.eval(feed_dict=feed_dict)
-- 
GitLab


From 62356ad4fde3cab5eb3b565b802badb02b4ab835 Mon Sep 17 00:00:00 2001
From: Patrick Nguyen <drpng@google.com>
Date: Tue, 1 May 2018 17:48:36 -0700
Subject: [PATCH 129/395] Re-apply CL 194140820, which reverts #18251
 (convolution change).

PiperOrigin-RevId: 195027049
---
 .../contrib/layers/python/layers/layers.py    | 142 +-----------------
 .../layers/python/layers/layers_test.py       |  15 +-
 2 files changed, 7 insertions(+), 150 deletions(-)

diff --git a/tensorflow/contrib/layers/python/layers/layers.py b/tensorflow/contrib/layers/python/layers/layers.py
index 2f3e57653c..25c3b1e7ea 100644
--- a/tensorflow/contrib/layers/python/layers/layers.py
+++ b/tensorflow/contrib/layers/python/layers/layers.py
@@ -932,8 +932,7 @@ def convolution(inputs,
                 variables_collections=None,
                 outputs_collections=None,
                 trainable=True,
-                scope=None,
-                conv_dims=None):
+                scope=None):
   """Adds an N-D convolution followed by an optional batch_norm layer.
 
   It is required that 1 <= N <= 3.
@@ -994,10 +993,6 @@ def convolution(inputs,
     trainable: If `True` also add variables to the graph collection
       `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
     scope: Optional scope for `variable_scope`.
-    conv_dims: Optional convolution dimensionality, when set it would use the
-      corresponding convolution (e.g. 2 for Conv 2D, 3 for Conv 3D, ..). When
-      leaved to None it would select the convolution dimensionality based on
-      the input rank (i.e. Conv ND, with N = input_rank - 2).
 
   Returns:
     A tensor representing the output of the operation.
@@ -1020,9 +1015,6 @@ def convolution(inputs,
     inputs = ops.convert_to_tensor(inputs)
     input_rank = inputs.get_shape().ndims
 
-    if conv_dims is not None and conv_dims + 2 != input_rank:
-      raise ValueError('Convolution expects input with rank %d, got %d' %
-                       (conv_dims + 2, input_rank))
     if input_rank == 3:
       layer_class = convolutional_layers.Convolution1D
     elif input_rank == 4:
@@ -1069,134 +1061,10 @@ def convolution(inputs,
       outputs = activation_fn(outputs)
     return utils.collect_named_outputs(outputs_collections, sc.name, outputs)
 
-@add_arg_scope
-def convolution1d(inputs,
-                  num_outputs,
-                  kernel_size,
-                  stride=1,
-                  padding='SAME',
-                  data_format=None,
-                  rate=1,
-                  activation_fn=nn.relu,
-                  normalizer_fn=None,
-                  normalizer_params=None,
-                  weights_initializer=initializers.xavier_initializer(),
-                  weights_regularizer=None,
-                  biases_initializer=init_ops.zeros_initializer(),
-                  biases_regularizer=None,
-                  reuse=None,
-                  variables_collections=None,
-                  outputs_collections=None,
-                  trainable=True,
-                  scope=None):
-  return convolution(inputs,
-                     num_outputs,
-                     kernel_size,
-                     stride,
-                     padding,
-                     data_format,
-                     rate,
-                     activation_fn,
-                     normalizer_fn,
-                     normalizer_params,
-                     weights_initializer,
-                     weights_regularizer,
-                     biases_initializer,
-                     biases_regularizer,
-                     reuse,
-                     variables_collections,
-                     outputs_collections,
-                     trainable,
-                     scope,
-                     conv_dims=1)
-
-convolution1d.__doc__ = convolution.__doc__
 
-@add_arg_scope
-def convolution2d(inputs,
-                  num_outputs,
-                  kernel_size,
-                  stride=1,
-                  padding='SAME',
-                  data_format=None,
-                  rate=1,
-                  activation_fn=nn.relu,
-                  normalizer_fn=None,
-                  normalizer_params=None,
-                  weights_initializer=initializers.xavier_initializer(),
-                  weights_regularizer=None,
-                  biases_initializer=init_ops.zeros_initializer(),
-                  biases_regularizer=None,
-                  reuse=None,
-                  variables_collections=None,
-                  outputs_collections=None,
-                  trainable=True,
-                  scope=None):
-  return convolution(inputs,
-                     num_outputs,
-                     kernel_size,
-                     stride,
-                     padding,
-                     data_format,
-                     rate,
-                     activation_fn,
-                     normalizer_fn,
-                     normalizer_params,
-                     weights_initializer,
-                     weights_regularizer,
-                     biases_initializer,
-                     biases_regularizer,
-                     reuse,
-                     variables_collections,
-                     outputs_collections,
-                     trainable,
-                     scope,
-                     conv_dims=2)
-
-convolution2d.__doc__ = convolution.__doc__
+convolution2d = convolution
+convolution3d = convolution
 
-@add_arg_scope
-def convolution3d(inputs,
-                  num_outputs,
-                  kernel_size,
-                  stride=1,
-                  padding='SAME',
-                  data_format=None,
-                  rate=1,
-                  activation_fn=nn.relu,
-                  normalizer_fn=None,
-                  normalizer_params=None,
-                  weights_initializer=initializers.xavier_initializer(),
-                  weights_regularizer=None,
-                  biases_initializer=init_ops.zeros_initializer(),
-                  biases_regularizer=None,
-                  reuse=None,
-                  variables_collections=None,
-                  outputs_collections=None,
-                  trainable=True,
-                  scope=None):
-  return convolution(inputs,
-                     num_outputs,
-                     kernel_size,
-                     stride,
-                     padding,
-                     data_format,
-                     rate,
-                     activation_fn,
-                     normalizer_fn,
-                     normalizer_params,
-                     weights_initializer,
-                     weights_regularizer,
-                     biases_initializer,
-                     biases_regularizer,
-                     reuse,
-                     variables_collections,
-                     outputs_collections,
-                     trainable,
-                     scope,
-                     conv_dims=3)
-
-convolution3d.__doc__ = convolution.__doc__
 
 @add_arg_scope
 def convolution2d_in_plane(
@@ -1543,7 +1411,7 @@ def dense_to_sparse(tensor, eos_token=0, outputs_collections=None, scope=None):
   Args:
      tensor: An `int` `Tensor` to be converted to a `Sparse`.
      eos_token: An integer.
-       It is part of the target label that signifies the end of a sentence.
+       It is part of the target label that signfies the end of a sentence.
      outputs_collections: Collection to add the outputs.
      scope: Optional scope for name_scope.
   """
@@ -1687,7 +1555,7 @@ def _inner_flatten(inputs, new_rank, output_collections=None, scope=None):
     output_collections: Collection to which the outputs will be added.
     scope: Optional scope for `name_scope`.
   Returns:
-    A `Tensor` or `SparseTensor` containing the same values as `inputs`, but
+    A `Tensor` or `SparseTensor` conataining the same values as `inputs`, but
     with innermost dimensions flattened to obtain rank `new_rank`.
 
   Raises:
diff --git a/tensorflow/contrib/layers/python/layers/layers_test.py b/tensorflow/contrib/layers/python/layers/layers_test.py
index b01fd5d5c9..997f910a2a 100644
--- a/tensorflow/contrib/layers/python/layers/layers_test.py
+++ b/tensorflow/contrib/layers/python/layers/layers_test.py
@@ -310,17 +310,6 @@ class BiasAddTest(test.TestCase):
 
 class ConvolutionTest(test.TestCase):
 
-  def testInvalidShape(self):
-    with self.test_session():
-      images_2d = random_ops.random_uniform((5, 7, 9, 3), seed=1)
-      with self.assertRaisesRegexp(
-          ValueError, 'Convolution expects input with rank 5, got 4'):
-        layers_lib.convolution3d(images_2d, 32, 3)
-      images_3d = random_ops.random_uniform((5, 6, 7, 9, 3), seed=1)
-      with self.assertRaisesRegexp(
-          ValueError, 'Convolution expects input with rank 4, got 5'):
-        layers_lib.convolution2d(images_3d, 32, 3)
-
   def testInvalidDataFormat(self):
     height, width = 7, 9
     with self.test_session():
@@ -3166,7 +3155,7 @@ class RepeatTests(test.TestCase):
     with self.test_session():
       images = np.random.uniform(size=(5, height, width, 3)).astype(np.float32)
       output = _layers.repeat(images, 3, layers_lib.conv2d, 32, [3, 3])
-      self.assertEqual(output.op.name, 'Repeat/convolution2d_3/Relu')
+      self.assertEqual(output.op.name, 'Repeat/convolution_3/Relu')
       self.assertListEqual(output.get_shape().as_list(), [5, 3, 3, 32])
 
   def testRepeatWithScope(self):
@@ -3760,7 +3749,7 @@ class StackTests(test.TestCase):
           layers_lib.convolution2d, [10, 20, 30],
           kernel_size=[3, 3],
           padding='SAME')
-      self.assertEqual(output.op.name, 'Stack/convolution2d_3/Relu')
+      self.assertEqual(output.op.name, 'Stack/convolution_3/Relu')
       self.assertListEqual(output.get_shape().as_list(), [5, 3, 3, 30])
 
   def testStackWithScope(self):
-- 
GitLab


From 92939c55e47b10c6d1ccd82bb31d877efca12235 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 17:57:02 -0700
Subject: [PATCH 130/395] Internal change.

PiperOrigin-RevId: 195027918
---
 tensorflow/contrib/lite/kernels/fully_connected.cc | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensorflow/contrib/lite/kernels/fully_connected.cc b/tensorflow/contrib/lite/kernels/fully_connected.cc
index c5bf50da5f..470b52b7bc 100644
--- a/tensorflow/contrib/lite/kernels/fully_connected.cc
+++ b/tensorflow/contrib/lite/kernels/fully_connected.cc
@@ -272,8 +272,7 @@ template <KernelType kernel_type>
 TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
                            TfLiteFullyConnectedParams* params, OpData* data,
                            TfLiteTensor* input, TfLiteTensor* filter,
-                           TfLiteTensor* bias, TfLiteTensor* input_quantized,
-                           TfLiteTensor* output) {
+                           TfLiteTensor* bias, TfLiteTensor* output) {
   gemmlowp::GemmContext* gemm_context = gemm_support::GetFromContext(context);
 
   int32_t input_offset = -input->params.zero_point;
@@ -292,6 +291,8 @@ TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
   } else if (kernel_type == kPie) {
     if (input->type == kTfLiteFloat32) {
       // Pie currently only supports quantized models and float inputs/outputs.
+      TfLiteTensor* input_quantized =
+          &context->tensors[node->temporaries->data[0]];
       return EvalPieQuantized(context, node, params, data, input, filter, bias,
                               input_quantized, output);
     } else {
@@ -346,15 +347,13 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   TfLiteTensor* bias = GetOptionalInputTensor(context, node, kBiasTensor);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
-  TfLiteTensor* input_quantized = &context->tensors[node->temporaries->data[0]];
-
   switch (filter->type) {  // Already know in/out types are same.
     case kTfLiteFloat32:
       return EvalFloat<kernel_type>(context, node, params, data, input, filter,
                                     bias, output);
     case kTfLiteUInt8:
       return EvalQuantized<kernel_type>(context, node, params, data, input,
-                                        filter, bias, input_quantized, output);
+                                        filter, bias, output);
     default:
       context->ReportError(context, "Type not currently supported.");
       return kTfLiteError;
-- 
GitLab


From c8ae9e86f33053484b05e405dadd2c8a98b8b41b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 1 May 2018 17:59:59 -0700
Subject: [PATCH 131/395] Internal change

PiperOrigin-RevId: 195028221
---
 .../contrib/factorization/python/ops/factorization_ops.py     | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/contrib/factorization/python/ops/factorization_ops.py b/tensorflow/contrib/factorization/python/ops/factorization_ops.py
index 811fa89bc3..5cef4068ed 100644
--- a/tensorflow/contrib/factorization/python/ops/factorization_ops.py
+++ b/tensorflow/contrib/factorization/python/ops/factorization_ops.py
@@ -107,7 +107,7 @@ class WALSModel(object):
       # the prep_gramian_op for row(column) can be run.
       worker_init_op = model.worker_init
 
-      # To be run once per integration sweep before the row(column) update
+      # To be run once per iteration sweep before the row(column) update
       # initialize ops can be run. Note that in the distributed training
       # situations, this should only be run by the chief trainer. All other
       # trainers need to block until this is done.
@@ -436,7 +436,7 @@ class WALSModel(object):
       gramian: Variable storing the gramian calculated from the factors.
 
     Returns:
-      A op that updates the gramian with the calculated value from the factors.
+      An op that updates the gramian with the calculated value from the factors.
     """
     partial_gramians = []
     for f in factors:
-- 
GitLab


From 69b2c639f55b065a5dbf829351034441bebc8437 Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Tue, 1 May 2018 18:46:31 -0700
Subject: [PATCH 132/395] [XLA:CPU] Re-use the same llvm::GlobalVariable for
 identical literals

This isn't necessary today, but it will be after an optimization change I'm
about to make.

LLVM has a constant merging pass too, but one of the motivations here is to
avoid the LLVM compile time overhead of having many large arrays in the IR.

PiperOrigin-RevId: 195032900
---
 tensorflow/compiler/xla/layout_util.cc        |  22 +++
 tensorflow/compiler/xla/layout_util.h         |   3 +
 tensorflow/compiler/xla/literal_util.cc       |  22 +++
 tensorflow/compiler/xla/literal_util.h        |   4 +
 .../compiler/xla/service/cpu/cpu_compiler.cc  |   2 +
 .../compiler/xla/service/cpu/ir_emitter.cc    |  29 ++--
 .../compiler/xla/service/cpu/ir_emitter.h     |  16 +++
 .../compiler/xla/service/cpu/tests/BUILD      |  14 ++
 .../cpu/tests/cpu_literal_caching_test.cc     | 125 ++++++++++++++++++
 tensorflow/compiler/xla/shape_util.cc         |  23 ++++
 tensorflow/compiler/xla/shape_util.h          |   3 +
 .../xla/tests/llvm_irgen_test_base.cc         |   4 +-
 tensorflow/compiler/xla/xla_data.proto        |  18 ++-
 tensorflow/core/lib/hash/hash.h               |  19 ++-
 14 files changed, 289 insertions(+), 15 deletions(-)
 create mode 100644 tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc

diff --git a/tensorflow/compiler/xla/layout_util.cc b/tensorflow/compiler/xla/layout_util.cc
index fdc4bbdd8b..c6f8f6766e 100644
--- a/tensorflow/compiler/xla/layout_util.cc
+++ b/tensorflow/compiler/xla/layout_util.cc
@@ -29,6 +29,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/types.h"
 #include "tensorflow/compiler/xla/util.h"
 #include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/hash/hash.h"
 #include "tensorflow/core/lib/strings/numbers.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
@@ -465,4 +466,25 @@ std::ostream& operator<<(std::ostream& out, const Layout& layout) {
   return out;
 }
 
+/*static*/ size_t LayoutUtil::Hash(const Layout& layout) {
+  using tensorflow::hash;
+  using tensorflow::Hash64Combine;
+
+  size_t hash_value = hash<Format>()(layout.format());
+
+  for (int64 minor_to_major : layout.minor_to_major()) {
+    hash_value = Hash64Combine(hash_value, hash<int64>()(minor_to_major));
+  }
+
+  for (int64 padded_dim : layout.padded_dimensions()) {
+    hash_value = Hash64Combine(hash_value, hash<int64>()(padded_dim));
+  }
+
+  hash_value =
+      Hash64Combine(hash_value, hash<PaddingValue>()(layout.padding_value()));
+  hash_value = Hash64Combine(hash_value, layout.max_sparse_elements());
+
+  return hash_value;
+}
+
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/layout_util.h b/tensorflow/compiler/xla/layout_util.h
index 6c54eb2201..6cec750101 100644
--- a/tensorflow/compiler/xla/layout_util.h
+++ b/tensorflow/compiler/xla/layout_util.h
@@ -195,6 +195,9 @@ class LayoutUtil {
   static bool AreDimensionsConsecutive(const Layout& layout,
                                        tensorflow::gtl::ArraySlice<int64> dims);
 
+  // Compute a hash for `layout`.
+  static size_t Hash(const Layout& layout);
+
  private:
   TF_DISALLOW_COPY_AND_ASSIGN(LayoutUtil);
 };
diff --git a/tensorflow/compiler/xla/literal_util.cc b/tensorflow/compiler/xla/literal_util.cc
index bb6dd4f909..b3b5e34ba2 100644
--- a/tensorflow/compiler/xla/literal_util.cc
+++ b/tensorflow/compiler/xla/literal_util.cc
@@ -29,6 +29,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/util.h"
 #include "tensorflow/core/lib/core/casts.h"
 #include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/hash/hash.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/lib/strings/stringprintf.h"
@@ -2148,6 +2149,27 @@ string Literal::GetR1U8AsString() const {
   return LiteralView(literal, view_root);
 }
 
+size_t Literal::Hash() const {
+  using tensorflow::Hash64;
+  using tensorflow::Hash64Combine;
+
+  size_t hash_value = ShapeUtil::Hash(shape());
+
+  ShapeUtil::ForEachSubshape(
+      shape(), [&](const Shape& subshape, const ShapeIndex& index) {
+        if (ShapeUtil::IsTuple(subshape)) {
+          return;
+        }
+
+        CHECK(LayoutUtil::IsDense(subshape.layout()));
+        hash_value = Hash64Combine(
+            hash_value, Hash64(static_cast<const char*>(untyped_data(index)),
+                               size_bytes(index)));
+      });
+
+  return hash_value;
+}
+
 LiteralView::LiteralView(const Literal& literal, const ShapeIndex& view_root) {
   shape_ = ShapeUtil::GetSubshape(literal.shape(), view_root);
   pieces_ = ShapeTree<Piece>(shape_);
diff --git a/tensorflow/compiler/xla/literal_util.h b/tensorflow/compiler/xla/literal_util.h
index 956ff7d21c..290f388078 100644
--- a/tensorflow/compiler/xla/literal_util.h
+++ b/tensorflow/compiler/xla/literal_util.h
@@ -662,6 +662,10 @@ class Literal {
   // LayoutUtil::MaxSparseElements(SetSubshape(shape(), index).layout()).
   int64 sparse_element_count() const;
 
+  // Compute a hash for this literal.  This literal must not be a sparse tensor
+  // or a tuple containing a sparse tensor.
+  size_t Hash() const;
+
  protected:
   // Internal template helper for the Literal::CopySliceFrom(), matching its
   // arguments one by one.
diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
index d8ba289f29..e298d67e09 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
+++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
@@ -787,6 +787,8 @@ CpuCompiler::CompileAheadOfTime(std::vector<std::unique_ptr<HloModule>> modules,
       TF_RETURN_IF_ERROR(verify_status);
     }
 
+    XLA_VLOG_LINES(2, "LLVM IR:\n" + llvm_ir::DumpModuleToString(llvm_module));
+
     Disassembler disassembler(*target_machine);
     CompilerFunctor compiler_functor(
         target_machine.get(), &disassembler, opt_level,
diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
index d582b5aaae..e473389a29 100644
--- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
+++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
@@ -160,10 +160,8 @@ Status IrEmitter::HandleBitcast(HloInstruction* bitcast) {
   return Status::OK();
 }
 
-Status IrEmitter::HandleConstant(HloInstruction* constant) {
-  VLOG(2) << "HandleConstant: " << constant->ToString();
-  const Literal& literal = constant->literal();
-  llvm::GlobalVariable* global_for_const;
+llvm::GlobalVariable* IrEmitter::EmitGlobalForLiteral(const Literal& literal) {
+  llvm::GlobalVariable* result;
 
   // We avoid creating large constants in the LLVM IR since LLVM is not
   // efficient for large constant arrays.  We still emit "small enough" constant
@@ -174,27 +172,42 @@ Status IrEmitter::HandleConstant(HloInstruction* constant) {
       ByteSizeOf(literal.shape()) >= kMaxInternalConstantSizeInBytes) {
     string global_name = tensorflow::strings::StrCat(
         "constant_global_", external_global_constant_counter_++);
-    global_for_const = new llvm::GlobalVariable(
+    result = new llvm::GlobalVariable(
         /*Module=*/*module_,
         /*Type=*/IrShapeType(literal.shape()),
         /*isConstant=*/true,
         /*Linkage=*/llvm::GlobalValue::ExternalLinkage,
         /*Initializer=*/nullptr,
         /*Name=*/AsStringRef(global_name));
-    global_for_const->setAlignment(MinimumAlignmentForShape(literal.shape()));
+    result->setAlignment(MinimumAlignmentForShape(literal.shape()));
     external_constant_pool_->Insert(global_name, literal,
                                     MinimumAlignmentForShape(literal.shape()));
   } else {
     llvm::Constant* initializer =
         llvm_ir::ConvertLiteralToIrConstant(literal, module_);
-    global_for_const = new llvm::GlobalVariable(
+    result = new llvm::GlobalVariable(
         /*Module=*/*module_,
         /*Type=*/initializer->getType(),
         /*isConstant=*/true,
         /*Linkage=*/llvm::GlobalValue::PrivateLinkage,
         /*Initializer=*/initializer,
         /*Name=*/"");
-    global_for_const->setAlignment(MinimumAlignmentForShape(literal.shape()));
+    result->setAlignment(MinimumAlignmentForShape(literal.shape()));
+  }
+  return result;
+}
+
+Status IrEmitter::HandleConstant(HloInstruction* constant) {
+  VLOG(2) << "HandleConstant: " << constant->ToString();
+  const Literal& literal = constant->literal();
+  llvm::GlobalVariable* global_for_const;
+
+  auto it = emitted_literals_.find(&literal);
+  if (it != emitted_literals_.end()) {
+    global_for_const = it->second;
+  } else {
+    global_for_const = EmitGlobalForLiteral(literal);
+    emitted_literals_[&literal] = global_for_const;
   }
   emitted_value_[constant] = global_for_const;
   VLOG(2) << "  emitted value: " << llvm_ir::DumpToString(*global_for_const);
diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.h b/tensorflow/compiler/xla/service/cpu/ir_emitter.h
index 0f2f3d1817..5a04076080 100644
--- a/tensorflow/compiler/xla/service/cpu/ir_emitter.h
+++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.h
@@ -530,6 +530,8 @@ class IrEmitter : public DfsHloVisitorWithDefault {
   Status EmitXfeedTransfer(XfeedKind kind, const Shape& shape,
                            llvm::Value* program_buffer_address);
 
+  llvm::GlobalVariable* EmitGlobalForLiteral(const Literal& literal);
+
   const HloModuleConfig& hlo_module_config_;
 
   bool is_top_level_computation_;
@@ -539,6 +541,20 @@ class IrEmitter : public DfsHloVisitorWithDefault {
   int64 external_global_constant_counter_ = 0;
   ExternalConstantPool* external_constant_pool_;
 
+  struct LiteralPtrHashFunctor {
+    size_t operator()(const Literal* literal) const { return literal->Hash(); }
+  };
+
+  struct LiteralPtrEqualityFunctor {
+    bool operator()(const Literal* lhs, const Literal* rhs) const {
+      return *lhs == *rhs;
+    }
+  };
+
+  tensorflow::gtl::FlatMap<const Literal*, llvm::GlobalVariable*,
+                           LiteralPtrHashFunctor, LiteralPtrEqualityFunctor>
+      emitted_literals_;
+
   TF_DISALLOW_COPY_AND_ASSIGN(IrEmitter);
 };
 
diff --git a/tensorflow/compiler/xla/service/cpu/tests/BUILD b/tensorflow/compiler/xla/service/cpu/tests/BUILD
index bfd95c3fe0..4ddb7a85bc 100644
--- a/tensorflow/compiler/xla/service/cpu/tests/BUILD
+++ b/tensorflow/compiler/xla/service/cpu/tests/BUILD
@@ -147,3 +147,17 @@ tf_cc_test(
         "//tensorflow/core:test_main",
     ],
 )
+
+tf_cc_test(
+    name = "cpu_literal_caching_test",
+    srcs = ["cpu_literal_caching_test.cc"],
+    deps = [
+        "//tensorflow/compiler/xla/service:hlo",
+        "//tensorflow/compiler/xla/service/cpu:cpu_compiler",
+        "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test",
+        "//tensorflow/compiler/xla/tools/parser:hlo_parser",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc
new file mode 100644
index 0000000000..f0404d07d9
--- /dev/null
+++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc
@@ -0,0 +1,125 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/cpu/cpu_compiler.h"
+#include "tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h"
+#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h"
+
+namespace xla {
+namespace cpu {
+namespace {
+class CpuExternalConstantsTest : public CpuCodegenTest {};
+
+TEST_F(CpuExternalConstantsTest, RepeatedArrayConstants) {
+  // We use a while loop here to force the two constant HloInstructions to be in
+  // different computations.  Otherwise the HLO optimizer itself CSEs them.
+  const string hlo_text = R"(
+HloModule RepeatedConstants
+
+while_body {
+  arg_body = f32[2,3,2] parameter(0)
+  ROOT const = f32[2,3,2] constant(
+  f32[2,3,2]
+    {{{1, 2}, {1001, 1002}, {2001, 2002}},
+     {{2, 1}, {2001, 3002}, {2001, 2002}}})
+}
+
+while_cond {
+  arg_cond = f32[2,3,2] parameter(0)
+  ROOT unknown = pred[] infeed()
+}
+
+ENTRY main {
+  param = f32[2,3,2] parameter(0)
+  const_a = f32[2,3,2] constant(
+  f32[2,3,2]
+    {{{1, 2}, {1001, 1002}, {2001, 2002}},
+     {{2, 1}, {2001, 3002}, {2001, 2002}}})
+  const_b = f32[2,3,2] while(f32[2,3,2] const_a), condition=while_cond, body=while_body
+
+  out0 = () outfeed(f32[2,3,2] const_a)
+  out1 = () outfeed(f32[2,3,2] const_b)
+
+  ROOT root = f32[] constant(1)
+}
+)";
+
+  string filecheck_pattern = R"(
+CHECK: private constant [2 x [3 x [2 x float]]]
+CHECK-NOT: private constant [2 x [3 x [2 x float]]]
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_text));
+
+  CpuAotCompilationOptions options{
+      /*triple=*/"x86_64-pc-linux", /*cpu_name=*/"", /*features=*/"",
+      /*entry_point_name=*/"entry",
+      /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static};
+
+  CompileAheadOfTimeAndVerifyIr(std::move(module), options, filecheck_pattern,
+                                /*match_optimized_ir=*/false);
+}
+
+TEST_F(CpuExternalConstantsTest, RepeatedTupleConstants) {
+  // We use a while loop here to force the two constant HloInstructions to be in
+  // different computations.  Otherwise the HLO optimizer itself CSEs them.
+  const string hlo_text = R"(
+HloModule RepeatedConstants
+
+while_body {
+  arg_body = (f32[2,1]{1,0}, f32[2]{0}) parameter(0)
+  ROOT const = (f32[2,1]{1,0}, f32[2]{0}) constant((f32[2,1], f32[2]) ( f32[2,1] { { 1 }, { 2 } }, {2, 42} ))
+}
+
+while_cond {
+  arg_cond = (f32[2,1]{1,0}, f32[2]{0}) parameter(0)
+  ROOT unknown = pred[] infeed()
+}
+
+ENTRY main {
+  param = f32[2,3,2] parameter(0)
+  const_a = (f32[2,1]{1,0}, f32[2]{0}) constant((f32[2,1], f32[2]) ( f32[2,1] { { 1 }, { 2 } }, {2, 42} ))
+  const_b = (f32[2,1]{1,0}, f32[2]{0}) while((f32[2,1]{1,0}, f32[2]{0}) const_a), condition=while_cond, body=while_body
+
+  out0 = () outfeed((f32[2,1]{1,0}, f32[2]{0}) const_a)
+  out1 = () outfeed((f32[2,1]{1,0}, f32[2]{0}) const_b)
+
+  ROOT root = f32[] constant(1)
+}
+)";
+
+  string filecheck_pattern = R"(
+CHECK: private constant [2 x float]
+CHECK: private constant [2 x [1 x float]]
+CHECK-NOT: private constant [2 x float]
+CHECK-NOT: private constant [2 x [1 x float]]
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_text));
+
+  CpuAotCompilationOptions options{
+      /*triple=*/"x86_64-pc-linux", /*cpu_name=*/"", /*features=*/"",
+      /*entry_point_name=*/"entry",
+      /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static};
+
+  CompileAheadOfTimeAndVerifyIr(std::move(module), options, filecheck_pattern,
+                                /*match_optimized_ir=*/false);
+}
+
+}  // namespace
+}  // namespace cpu
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/shape_util.cc b/tensorflow/compiler/xla/shape_util.cc
index d58baa3220..c330473cda 100644
--- a/tensorflow/compiler/xla/shape_util.cc
+++ b/tensorflow/compiler/xla/shape_util.cc
@@ -32,6 +32,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/stringpiece.h"
 #include "tensorflow/core/lib/gtl/iterator_range.h"
 #include "tensorflow/core/lib/gtl/optional.h"
+#include "tensorflow/core/lib/hash/hash.h"
 #include "tensorflow/core/lib/strings/numbers.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
@@ -1472,4 +1473,26 @@ std::ostream& operator<<(std::ostream& out, const Shape& shape) {
   return out;
 }
 
+/*static*/ size_t ShapeUtil::Hash(const Shape& shape) {
+  using tensorflow::hash;
+  using tensorflow::Hash64Combine;
+
+  size_t hash_value = hash<PrimitiveType>()(shape.element_type());
+
+  if (shape.tuple_shapes().empty()) {
+    for (int64 dim : shape.dimensions()) {
+      hash_value = Hash64Combine(hash_value, hash<int64>()(dim));
+    }
+
+    hash_value = Hash64Combine(hash_value, LayoutUtil::Hash(shape.layout()));
+  } else {
+    hash_value = 0;
+    for (const Shape& subshape : shape.tuple_shapes()) {
+      hash_value = Hash64Combine(hash_value, ShapeUtil::Hash(subshape));
+    }
+  }
+
+  return hash_value;
+}
+
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/shape_util.h b/tensorflow/compiler/xla/shape_util.h
index 5fa728e7c2..cb8bf5a2b9 100644
--- a/tensorflow/compiler/xla/shape_util.h
+++ b/tensorflow/compiler/xla/shape_util.h
@@ -650,6 +650,9 @@ class ShapeUtil {
               .ok());
   }
 
+  // Compute a hash for `shape`.
+  static size_t Hash(const Shape& shape);
+
  private:
   // Validates all of the non-layout properties of the shape -- this is a helper
   // used by both the layout-optional and layout-required public method.
diff --git a/tensorflow/compiler/xla/tests/llvm_irgen_test_base.cc b/tensorflow/compiler/xla/tests/llvm_irgen_test_base.cc
index 3023df47cd..2c45f19c09 100644
--- a/tensorflow/compiler/xla/tests/llvm_irgen_test_base.cc
+++ b/tensorflow/compiler/xla/tests/llvm_irgen_test_base.cc
@@ -62,8 +62,8 @@ void LLVMIRGenTestBase::CompileAheadOfTimeAndVerifyIr(
     std::unique_ptr<HloModule> hlo_module, const AotCompilationOptions& options,
     const string& pattern, bool match_optimized_ir) {
   SetIrHook(match_optimized_ir);
-  ASSERT_TRUE(
-      CompileToAotCompilationResult(std::move(hlo_module), options).ok());
+  TF_ASSERT_OK(
+      CompileToAotCompilationResult(std::move(hlo_module), options).status());
   ResetIrHook();
 
   StatusOr<bool> filecheck_result = RunFileCheck(ir_, pattern);
diff --git a/tensorflow/compiler/xla/xla_data.proto b/tensorflow/compiler/xla/xla_data.proto
index d23f9e5918..750d72d797 100644
--- a/tensorflow/compiler/xla/xla_data.proto
+++ b/tensorflow/compiler/xla/xla_data.proto
@@ -134,6 +134,8 @@ enum Format {
 // example, Convert) are ignored.
 //
 // See the XLA documentation for more information on shapes and layouts.
+//
+// LINT.IfChange
 message Layout {
   // The method used to store the data in memory. The format determines which of
   // the other fields are used by the layout.
@@ -159,9 +161,12 @@ message Layout {
   // memory.  This field must be unset unless the format is SPARSE.
   int64 max_sparse_elements = 5;
 
-  // Important: if any field is added, be sure to modify ShapeUtil::Equal()
-  // appropriately to account for the new field.
+  // Important: if any field is added, be sure to modify ShapeUtil::Equal() and
+  // LayoutUtil::Hash appropriately to account for the new field.
 }
+// LINT.ThenChange( \
+//     https://www.tensorflow.org/code/tensorflow/compiler/xla/shape_util.cc,      \
+//     https://www.tensorflow.org/code/tensorflow/compiler/xla/layout_util.cc)
 
 // A shape describes the number of dimensions in the array, the size of each
 // dimension, and the primitive component type.
@@ -170,6 +175,8 @@ message Layout {
 // defined.
 //
 // See the XLA documentation for more information on shapes and layouts.
+//
+// LINT.IfChange
 message Shape {
   reserved 1;
   reserved "rank";
@@ -190,9 +197,12 @@ message Shape {
   // The layout used to back this shape.
   Layout layout = 5;
 
-  // Important: if any field is added, be sure to modify ShapeUtil::Equal() and
-  // ShapeUtil::Compatible() appropriately to account for the new field.
+  // Important: if any field is added, be sure to modify ShapeUtil::Equal(),
+  // ShapeUtil::Compatible() and ShapeUtil::Hash() appropriately to account for
+  // the new field.
 }
+// LINT.ThenChange( \
+//     https://www.tensorflow.org/code/tensorflow/compiler/xla/shape_util.cc)
 
 // Shape of the parameters and output of a computation (like a traditional
 // function signature).
diff --git a/tensorflow/core/lib/hash/hash.h b/tensorflow/core/lib/hash/hash.h
index ca05e6346e..3f85303c0f 100644
--- a/tensorflow/core/lib/hash/hash.h
+++ b/tensorflow/core/lib/hash/hash.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include <stddef.h>
 #include <stdint.h>
 
+#include <functional>
 #include <string>
 
 #include "tensorflow/core/lib/core/stringpiece.h"
@@ -49,11 +50,27 @@ inline uint64 Hash64Combine(uint64 a, uint64 b) {
 // In particular, tensorflow::hash is not the identity function for pointers.
 // This is important for power-of-two sized hashtables like FlatMap and FlatSet,
 // because otherwise they waste the majority of their hash buckets.
-template <typename T>
+//
+// The second type argument is only used for SFNIAE below.
+template <typename T, typename = void>
 struct hash {
   size_t operator()(const T& t) const { return std::hash<T>()(t); }
 };
 
+template <typename T>
+struct hash<T, typename std::enable_if<std::is_enum<T>::value>::type> {
+  size_t operator()(T value) const {
+    // This works around a defect in the std::hash C++ spec that isn't fixed in
+    // (at least) gcc 4.8.4:
+    // http://www.open-std.org/jtc1/sc22/wg21/docs/lwg-defects.html#2148
+    //
+    // We should be able to remove this and use the default
+    // tensorflow::hash<EnumTy>() once we stop building with GCC versions old
+    // enough to not have this defect fixed.
+    return std::hash<uint64>()(static_cast<uint64>(value));
+  }
+};
+
 template <typename T>
 struct hash<T*> {
   size_t operator()(const T* t) const {
-- 
GitLab


From c0f1080188c5c6955cfa3b3c086ac262b1e5ec02 Mon Sep 17 00:00:00 2001
From: Patrick Nguyen <drpng@google.com>
Date: Tue, 1 May 2018 19:02:10 -0700
Subject: [PATCH 133/395] Make the CRF work when sequence_lengths are int32.

PiperOrigin-RevId: 195034218
---
 tensorflow/contrib/crf/python/ops/crf.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensorflow/contrib/crf/python/ops/crf.py b/tensorflow/contrib/crf/python/ops/crf.py
index d2beff849e..2d2cbdc199 100644
--- a/tensorflow/contrib/crf/python/ops/crf.py
+++ b/tensorflow/contrib/crf/python/ops/crf.py
@@ -52,6 +52,7 @@ from __future__ import print_function
 
 import numpy as np
 
+from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.layers import utils
 from tensorflow.python.ops import array_ops
@@ -147,7 +148,9 @@ def crf_log_norm(inputs, sequence_lengths, transition_params):
     # partition function.
     forward_cell = CrfForwardRnnCell(transition_params)
     # Sequence length is not allowed to be less than zero.
-    sequence_lengths_less_one = math_ops.maximum(0, sequence_lengths - 1)
+    sequence_lengths_less_one = math_ops.maximum(
+        constant_op.constant(0, dtype=sequence_lengths.dtype),
+        sequence_lengths - 1)
     _, alphas = rnn.dynamic_rnn(
         cell=forward_cell,
         inputs=rest_of_input,
-- 
GitLab


From b50f6325143486eb82b5654f8794f0771b54dd4d Mon Sep 17 00:00:00 2001
From: Dan Moldovan <mdan@google.com>
Date: Tue, 1 May 2018 19:05:39 -0700
Subject: [PATCH 134/395] Minor refactor: establish some operator naming
 conventions and apply them, so that the interface is a bit more consistent.

PiperOrigin-RevId: 195034691
---
 .../autograph/converters/break_statements.py  |   4 +-
 .../autograph/converters/control_flow.py      |  24 ++--
 .../contrib/autograph/operators/__init__.py   |  16 ++-
 .../autograph/operators/control_flow.py       | 105 ++++++++++--------
 .../autograph/operators/control_flow_test.py  |  30 ++---
 5 files changed, 99 insertions(+), 80 deletions(-)

diff --git a/tensorflow/contrib/autograph/converters/break_statements.py b/tensorflow/contrib/autograph/converters/break_statements.py
index 91de82f0a7..1be1c96dd3 100644
--- a/tensorflow/contrib/autograph/converters/break_statements.py
+++ b/tensorflow/contrib/autograph/converters/break_statements.py
@@ -114,9 +114,9 @@ class BreakStatementTransformer(transformer.Base):
           template,
           var_name=break_var,
           for_stmt=node)
-      extra_cond = templates.replace_as_expression(
+      extra_test = templates.replace_as_expression(
           'not var_name', var_name=break_var)
-      anno.setanno(node[1], 'extra_cond', extra_cond)
+      anno.setanno(node[1], 'extra_test', extra_test)
 
     return node
 
diff --git a/tensorflow/contrib/autograph/converters/control_flow.py b/tensorflow/contrib/autograph/converters/control_flow.py
index 2e26cdb3d9..935a2786db 100644
--- a/tensorflow/contrib/autograph/converters/control_flow.py
+++ b/tensorflow/contrib/autograph/converters/control_flow.py
@@ -207,7 +207,7 @@ class ControlFlowTransformer(transformer.Base):
       def body_name(state_ssf):
         body
         return state_ssf,
-      state_ast_tuple = ag__.while_loop(
+      state_ast_tuple = ag__.while_stmt(
           test_name, body_name, (state,), (extra_deps,))
     """
     node = templates.replace(
@@ -252,31 +252,31 @@ class ControlFlowTransformer(transformer.Base):
       state_ast_tuple = gast.Tuple([n.ast() for n in state], None)
 
     node_body = ast_util.rename_symbols(node.body, ssf_map)
-    if anno.hasanno(node, 'extra_cond'):
-      extra_cond = anno.getanno(node, 'extra_cond')
-      extra_cond = ast_util.rename_symbols(extra_cond, ssf_map)
+    if anno.hasanno(node, 'extra_test'):
+      extra_test = anno.getanno(node, 'extra_test')
+      extra_test = ast_util.rename_symbols(extra_test, ssf_map)
     else:
-      extra_cond = parser.parse_expression('True')
+      extra_test = parser.parse_expression('True')
 
     template = """
-      def extra_cond_name(state_ssf):
-        return extra_cond_expr
+      def extra_test_name(state_ssf):
+        return extra_test_expr
       def body_name(iterate, state_ssf):
         body
         return state_ssf,
-      state_ast_tuple = ag__.for_loop(
-          iterated, extra_cond_name, body_name, (state,))
+      state_ast_tuple = ag__.for_stmt(
+          iter_, extra_test_name, body_name, (state,))
     """
     node = templates.replace(
         template,
         state=state,
         state_ssf=state_ssf,
         state_ast_tuple=state_ast_tuple,
-        iterated=node.iter,
+        iter_=node.iter,
         iterate=node.target,
-        extra_cond_name=self.context.namer.new_symbol('extra_cond',
+        extra_test_name=self.context.namer.new_symbol('extra_test',
                                                       all_referenced),
-        extra_cond_expr=extra_cond,
+        extra_test_expr=extra_test,
         body_name=self.context.namer.new_symbol('loop_body', all_referenced),
         body=node_body)
 
diff --git a/tensorflow/contrib/autograph/operators/__init__.py b/tensorflow/contrib/autograph/operators/__init__.py
index 04b4734551..38b761d97d 100644
--- a/tensorflow/contrib/autograph/operators/__init__.py
+++ b/tensorflow/contrib/autograph/operators/__init__.py
@@ -19,11 +19,19 @@ conditionals and loops, implemented in functional form, using for example
 closures for the body.
 """
 
+# Naming conventions:
+#  * operator names match the name usually used for the respective Python
+#    idiom; examples: for_stmt, list_append
+#  * operator arguments match either of:
+#    - the corresponding Python AST attribute (e.g. the condition of an if
+#      statement is called test) if the operator represents an AST construct
+#    - the names used in the Python docs, if the operator is a function (e.g.
+#      list_ and x for append, see
+#      https://docs.python.org/3.7/tutorial/datastructures.html)
+
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-# TODO(mdan): Add a container for implementation-specific toggles (throughout).
-
-from tensorflow.contrib.autograph.operators.control_flow import for_loop
-from tensorflow.contrib.autograph.operators.control_flow import while_loop
+from tensorflow.contrib.autograph.operators.control_flow import for_stmt
+from tensorflow.contrib.autograph.operators.control_flow import while_stmt
diff --git a/tensorflow/contrib/autograph/operators/control_flow.py b/tensorflow/contrib/autograph/operators/control_flow.py
index d9d8b0d593..9f7202821f 100644
--- a/tensorflow/contrib/autograph/operators/control_flow.py
+++ b/tensorflow/contrib/autograph/operators/control_flow.py
@@ -25,44 +25,55 @@ from tensorflow.python.framework import tensor_util
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import gen_math_ops
 
-# TODO(mdan): Rename _loop to _stmt to follow Python nomenclature.
-# TODO(mdan): Rename arguments to match the AST names.
 
-
-def for_loop(iterated, extra_cond, loop_body, init_state):
+def for_stmt(iter_, extra_test, body, init_state):
   """Functional form of a for statement.
 
-  The loop operates on a so-called state, which includes all symbols that are
-  variant across loop iterations, excluding the iterate. In what follows we
-  refer to state as either a tuple of entities that represent an actual state,
-  or a list of arguments of the corresponding types.
+  The loop operates on a state, which includes all symbols that are
+  variant across loop iterations, excluding the iterate as well as the
+  variables local to the loop.
+
+  For example, given the loop below that calculates the geometric and
+  arithmetic means or some numbers:
+
+    geo_mean = 1
+    arith_mean = 0
+    for i in range(n):
+      a = numbers[i]
+      geo_mean *= a
+      arith_mean += a
+
+  The state is represented by the variables geo_mean and arith_mean. The
+  argument for initial_state may contain the tuple (1, 0), the body will
+  include the arguments geo_mean and arith_mean and will return a tuple
+  representing the new values for geo_mean and respectively arith_mean.
 
   Args:
-    iterated: The entity being iterated over.
-    extra_cond: Callable with the state as arguments, and boolean return type.
+    iter_: The entity being iterated over.
+    extra_test: Callable with the state as arguments, and boolean return type.
         An additionnal loop condition.
-    loop_body: Callable with the iterate and the state as arguments, and
+    body: Callable with the iterate and the state as arguments, and
         state as return type. The actual loop body.
     init_state: Tuple containing the initial state.
 
   Returns:
     Tuple containing the final state.
   """
-  if tensor_util.is_tensor(iterated):
-    return _known_len_for_loop(iterated, extra_cond, loop_body, init_state)
-  elif isinstance(iterated, dataset_ops.Dataset):
-    return _dataset_for_loop(iterated, extra_cond, loop_body, init_state)
+  if tensor_util.is_tensor(iter_):
+    return _known_len_for_stmt(iter_, extra_test, body, init_state)
+  elif isinstance(iter_, dataset_ops.Dataset):
+    return _dataset_for_stmt(iter_, extra_test, body, init_state)
   else:
-    return _py_for_loop(iterated, extra_cond, loop_body, init_state)
+    return _py_for_stmt(iter_, extra_test, body, init_state)
 
 
-def _py_for_loop(iterated, extra_cond, loop_body, init_state):
-  """Overload of for_loop that executes a Python for loop."""
+def _py_for_stmt(iter_, extra_test, body, init_state):
+  """Overload of for_stmt that executes a Python for loop."""
   state = init_state
-  for iterate in iterated:
-    if not extra_cond(*state):
+  for target in iter_:
+    if not extra_test(*state):
       break
-    state = loop_body(iterate, *state)
+    state = body(target, *state)
 
   # TODO(mdan): Remove this special case.
   if len(state) == 1:
@@ -70,23 +81,23 @@ def _py_for_loop(iterated, extra_cond, loop_body, init_state):
   return state
 
 
-def _known_len_for_loop(iterated, extra_cond, loop_body, init_state):
-  """Overload of for_loop that iterates over objects that define a length."""
-  n = builtins.dynamic_len(iterated)
+def _known_len_for_stmt(iter_, extra_test, body, init_state):
+  """Overload of for_stmt that iterates over objects that define a length."""
+  n = builtins.dynamic_len(iter_)
 
   def while_body(iterate_index, *state):
-    iterate = iterated[iterate_index]
-    new_state = loop_body(iterate, *state)
+    iterate = iter_[iterate_index]
+    new_state = body(iterate, *state)
     return (iterate_index + 1,) + new_state
 
   def while_cond(iterate_index, *state):
-    return gen_math_ops.logical_and(iterate_index < n, extra_cond(*state))
+    return gen_math_ops.logical_and(iterate_index < n, extra_test(*state))
 
-  results = while_loop(
+  results = while_stmt(
       while_cond,
       while_body,
       init_state=(0,) + init_state,
-      extra_deps=(iterated,),
+      extra_deps=(iter_,),
       opts=dict(maximum_iterations=n))
   # Dropping the iteration index because it's not syntactically visible.
   results = results[1:]
@@ -97,8 +108,8 @@ def _known_len_for_loop(iterated, extra_cond, loop_body, init_state):
   return results
 
 
-def _dataset_for_loop(ds, extra_cond, loop_body, init_state):
-  """Overload of for_loop that iterates over TF Datasets."""
+def _dataset_for_stmt(ds, extra_test, body, init_state):
+  """Overload of for_stmt that iterates over TF Datasets."""
   # Because Datsets only expose get_next, in the style of Python iterators,
   # we are forced to unpack the loop as:
   #
@@ -117,15 +128,15 @@ def _dataset_for_loop(ds, extra_cond, loop_body, init_state):
     epoch_number, iterate = iterator.get_next()
 
     def while_body(epoch_number, iterate, *state):
-      new_state = loop_body(iterate, *state)
+      new_state = body(iterate, *state)
       epoch_number, iterate = iterator.get_next()
       return (epoch_number, iterate) + new_state
 
     def while_cond(epoch_number, iterate, *state):
       del iterate
-      return gen_math_ops.logical_and(epoch_number < 1, extra_cond(*state))
+      return gen_math_ops.logical_and(epoch_number < 1, extra_test(*state))
 
-    results = while_loop(
+    results = while_stmt(
         while_cond,
         while_body,
         init_state=(epoch_number, iterate) + init_state,
@@ -140,7 +151,7 @@ def _dataset_for_loop(ds, extra_cond, loop_body, init_state):
   return results
 
 
-def while_loop(loop_cond, loop_body, init_state, extra_deps, opts=None):
+def while_stmt(test, body, init_state, extra_deps, opts=None):
   """Functional form of a while statement.
 
   The loop operates on a so-called state, which includes all symbols that are
@@ -149,13 +160,13 @@ def while_loop(loop_cond, loop_body, init_state, extra_deps, opts=None):
   of the corresponding types.
 
   Args:
-    loop_cond: Callable with the state as arguments, and boolean return type.
+    test: Callable with the state as arguments, and boolean return type.
         The loop condition.
-    loop_body: Callable with the state as arguments, and state as return type.
+    body: Callable with the state as arguments, and state as return type.
         The actual loop body.
     init_state: Tuple containing the initial state.
     extra_deps: Tuple containing additional entities on which the loop may
-        depend, such as loop invariants referenced by loop_cond. Used
+        depend, such as loop invariants referenced by test. Used
         exclusively for dispatch control.
     opts: Optional dict of extra loop parameters.
 
@@ -166,24 +177,24 @@ def while_loop(loop_cond, loop_body, init_state, extra_deps, opts=None):
   # That could be somethins as simple as a collection of dispatch rules, with
   # some prioritization.
   if any(tensor_util.is_tensor(v) for v in init_state + extra_deps):
-    return _tf_while_loop(loop_cond, loop_body, init_state, opts)
+    return _tf_while_stmt(test, body, init_state, opts)
   else:
-    return _py_while_loop(loop_cond, loop_body, init_state, opts)
+    return _py_while_stmt(test, body, init_state, opts)
 
 
-def _tf_while_loop(loop_cond, loop_body, init_state, opts):
-  """Overload of while_loop that stages a TF while_loop."""
+def _tf_while_stmt(test, body, init_state, opts):
+  """Overload of while_stmt that stages a TF while_stmt."""
   if opts is None:
     opts = {}
-  return control_flow_ops.while_loop(loop_cond, loop_body, init_state, **opts)
+  return control_flow_ops.while_loop(test, body, init_state, **opts)
 
 
-def _py_while_loop(loop_cond, loop_body, init_state, opts):
-  """Overload of while_loop that executes a Python while loop."""
+def _py_while_stmt(test, body, init_state, opts):
+  """Overload of while_stmt that executes a Python while loop."""
   del opts
   state = init_state
-  while loop_cond(*state):
-    state = loop_body(*state)
+  while test(*state):
+    state = body(*state)
   return state
 
 
diff --git a/tensorflow/contrib/autograph/operators/control_flow_test.py b/tensorflow/contrib/autograph/operators/control_flow_test.py
index a0cd0bfa82..b14d7edba3 100644
--- a/tensorflow/contrib/autograph/operators/control_flow_test.py
+++ b/tensorflow/contrib/autograph/operators/control_flow_test.py
@@ -29,28 +29,28 @@ from tensorflow.python.platform import test
 class ForLoopTest(test.TestCase):
 
   def test_tensor(self):
-    s = control_flow.for_loop(
+    s = control_flow.for_stmt(
         constant_op.constant([1, 2, 3, 4]),
-        extra_cond=lambda s: True,
-        loop_body=lambda i, s: (s + i,),
+        extra_test=lambda s: True,
+        body=lambda i, s: (s + i,),
         init_state=(0,))
     with self.test_session() as sess:
       self.assertEqual((10,), sess.run(s))
 
   def test_python(self):
-    s = control_flow.for_loop(
+    s = control_flow.for_stmt(
         range(5),
-        extra_cond=lambda s: True,
-        loop_body=lambda i, s: (s + i,),
+        extra_test=lambda s: True,
+        body=lambda i, s: (s + i,),
         init_state=(0,))
     self.assertEqual(10, s)
 
   def test_dataset(self):
     to_int32 = lambda i: math_ops.cast(i, dtypes.int32)
-    s = control_flow.for_loop(
+    s = control_flow.for_stmt(
         dataset_ops.Dataset.range(5).map(to_int32),
-        extra_cond=lambda s: True,
-        loop_body=lambda i, s: (s + i,),
+        extra_test=lambda s: True,
+        body=lambda i, s: (s + i,),
         init_state=(0,))
     with self.test_session() as sess:
       self.assertEqual((10,), sess.run(s))
@@ -60,9 +60,9 @@ class WhileLoopTest(test.TestCase):
 
   def test_tensor(self):
     n = constant_op.constant(5)
-    results = control_flow.while_loop(
-        loop_cond=lambda i, s: i < n,
-        loop_body=lambda i, s: (i + 1, s + i,),
+    results = control_flow.while_stmt(
+        test=lambda i, s: i < n,
+        body=lambda i, s: (i + 1, s + i,),
         init_state=(0, 0),
         extra_deps=(n,))
     with self.test_session() as sess:
@@ -70,9 +70,9 @@ class WhileLoopTest(test.TestCase):
 
   def test_python(self):
     n = 5
-    results = control_flow.while_loop(
-        loop_cond=lambda i, s: i < n,
-        loop_body=lambda i, s: (i + 1, s + i),
+    results = control_flow.while_stmt(
+        test=lambda i, s: i < n,
+        body=lambda i, s: (i + 1, s + i),
         init_state=(0, 0),
         extra_deps=(n,))
     self.assertEqual((5, 10), results)
-- 
GitLab


From 7715b7b0650c2f20b47189a060580a45e510acd8 Mon Sep 17 00:00:00 2001
From: Mark Heffernan <meheff@google.com>
Date: Tue, 1 May 2018 20:14:29 -0700
Subject: [PATCH 135/395] Add missing colocated element in test in
 buffer_assignment_test. This was resulting in a flaky test because sometimes
 the live set would include this missing colocated element perhaps because the
 buffers in the allocation has some nondeterministic order (read from a map?).

PiperOrigin-RevId: 195039311
---
 tensorflow/compiler/xla/service/buffer_assignment_test.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/compiler/xla/service/buffer_assignment_test.cc b/tensorflow/compiler/xla/service/buffer_assignment_test.cc
index 40cf6483aa..f6d6b5c36a 100644
--- a/tensorflow/compiler/xla/service/buffer_assignment_test.cc
+++ b/tensorflow/compiler/xla/service/buffer_assignment_test.cc
@@ -1634,6 +1634,7 @@ TEST_F(BufferAssignmentTest, PeakBuffersWhile) {
   }
   EXPECT_EQ(bcast_buffer->instruction(), bcast);
   EXPECT_TRUE(
+      nonbcast_buffer->instruction() == copy ||
       nonbcast_buffer->instruction() == while_op ||
       nonbcast_buffer->instruction() == body->parameter_instruction(0) ||
       nonbcast_buffer->instruction() == body->root_instruction() ||
-- 
GitLab


From 5e1448f691afe6e9ba57bb67497311c45b855b82 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 01:36:18 -0700
Subject: [PATCH 136/395] BUGFIX: Convert inputs and list of gradients into
 tuple if they are not instance of tuple. Otherwise this causes "unhashable
 keys" error when we try to hash. Also fixed lint error.

PiperOrigin-RevId: 195061425
---
 tensorflow/contrib/kfac/python/ops/fisher_blocks.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensorflow/contrib/kfac/python/ops/fisher_blocks.py b/tensorflow/contrib/kfac/python/ops/fisher_blocks.py
index 32c776cb38..3a5c8eb5f9 100644
--- a/tensorflow/contrib/kfac/python/ops/fisher_blocks.py
+++ b/tensorflow/contrib/kfac/python/ops/fisher_blocks.py
@@ -673,9 +673,6 @@ class KroneckerProductFB(FisherBlock):
   output factors.
   """
 
-  def __init__(self, layer_collection):
-    super(KroneckerProductFB, self).__init__(layer_collection)
-
   def _setup_damping(self, damping, normalization=None):
     """Makes functions that compute the damping values for both factors."""
     def compute_damping():
@@ -1309,6 +1306,8 @@ class InputOutputMultiTowerMultiUse(InputOutputMultiTower):
       else:
         raise ValueError("Global config variable TOWER_STRATEGY must be one of "
                          "'concat' or 'separate'.")
+    else:
+      inputs = tuple(inputs)
 
     # Now we perform the analogous processing for grads_list
     if isinstance(grads_list[0][0], (list, tuple)):
@@ -1351,6 +1350,8 @@ class InputOutputMultiTowerMultiUse(InputOutputMultiTower):
       else:
         raise ValueError("Global config variable TOWER_STRATEGY must be one of "
                          "'concat' or 'separate'.")
+    else:
+      grads_list = tuple(tuple(grads) for grads in grads_list)
 
     if self._num_uses is None:
       raise ValueError("You must supply a value for the num_uses argument if "
-- 
GitLab


From af4c683798738000b067f60e5ab8abe0115b29c8 Mon Sep 17 00:00:00 2001
From: Sergii Khomenko <x-sam@brainscode.com>
Date: Wed, 2 May 2018 15:39:46 +0200
Subject: [PATCH 137/395] Fix string issue for temp_export_dir (#18951)

* Fix string issue for temp_export_dir
---
 .../learn/python/learn/utils/saved_model_export_utils.py     | 3 ++-
 tensorflow/python/saved_model/builder_impl.py                | 5 +++--
 tensorflow/python/training/saver.py                          | 2 +-
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils.py b/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils.py
index c7cdb41312..f8106d1e4a 100644
--- a/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils.py
+++ b/tensorflow/contrib/learn/python/learn/utils/saved_model_export_utils.py
@@ -343,7 +343,8 @@ def get_temp_export_dir(timestamped_export_dir):
   """
   (dirname, basename) = os.path.split(timestamped_export_dir)
   temp_export_dir = os.path.join(
-      compat.as_bytes(dirname), compat.as_bytes('temp-{}'.format(basename)))
+      compat.as_bytes(dirname),
+      compat.as_bytes('temp-{}'.format(compat.as_text(basename))))
   return temp_export_dir
 
 
diff --git a/tensorflow/python/saved_model/builder_impl.py b/tensorflow/python/saved_model/builder_impl.py
index 3447d917e9..01903ae596 100644
--- a/tensorflow/python/saved_model/builder_impl.py
+++ b/tensorflow/python/saved_model/builder_impl.py
@@ -132,7 +132,8 @@ class SavedModelBuilder(object):
       if not file_io.file_exists(asset_destination_filepath):
         file_io.copy(asset_source_filepath, asset_destination_filepath)
 
-    tf_logging.info("Assets written to: %s", assets_destination_dir)
+    tf_logging.info("Assets written to: %s",
+                    compat.as_text(assets_destination_dir))
 
   def _maybe_add_legacy_init_op(self, legacy_init_op=None):
     """Add legacy init op to the SavedModel.
@@ -441,7 +442,7 @@ class SavedModelBuilder(object):
           compat.as_bytes(self._export_dir),
           compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
       file_io.write_string_to_file(path, self._saved_model.SerializeToString())
-    tf_logging.info("SavedModel written to: %s", path)
+    tf_logging.info("SavedModel written to: %s", compat.as_text(path))
 
     return path
 
diff --git a/tensorflow/python/training/saver.py b/tensorflow/python/training/saver.py
index 53e821c995..8134fd74aa 100644
--- a/tensorflow/python/training/saver.py
+++ b/tensorflow/python/training/saver.py
@@ -1725,7 +1725,7 @@ class Saver(object):
       return
     if save_path is None:
       raise ValueError("Can't load save_path when it is None.")
-    logging.info("Restoring parameters from %s", save_path)
+    logging.info("Restoring parameters from %s", compat.as_text(save_path))
     try:
       if context.executing_eagerly():
         self._build_eager(save_path, build_save=False, build_restore=True)
-- 
GitLab


From 6a09fcdbf047f5ab3e154238ed142883dd89af44 Mon Sep 17 00:00:00 2001
From: Clayne Robison <clayne.b.robison@intel.com>
Date: Wed, 2 May 2018 06:47:58 -0700
Subject: [PATCH 138/395] Reverting changes from 495d511 that break
 install_pip_packages.sh in Ubuntu 16.04 containers, causing nightly mkl ci
 builds to fail. (#18888)

---
 .../tools/ci_build/install/install_pip_packages.sh    | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/tensorflow/tools/ci_build/install/install_pip_packages.sh b/tensorflow/tools/ci_build/install/install_pip_packages.sh
index 5aaf544afd..982161cefe 100755
--- a/tensorflow/tools/ci_build/install/install_pip_packages.sh
+++ b/tensorflow/tools/ci_build/install/install_pip_packages.sh
@@ -17,14 +17,9 @@
 set -e
 
 # We don't apt-get install so that we can install a newer version of pip.
-# Only needed for Ubuntu 14.04 ,and not needed for Ubuntu 16.04 / Debian 8,9
-if $(cat /etc/*-release | grep -q 14.04); then
-  easy_install -U pip==9.0.3
-  easy_install3 -U pip==9.0.3
-else
-  pip2 install --upgrade pip==9.0.3
-  pip3 install --upgrade pip==9.0.3
-fi
+# Only needed for Ubuntu 14.04 and 16.04; not needed for 18.04 and Debian 8,9?
+easy_install -U pip==9.0.3
+easy_install3 -U pip==9.0.3
 
 # Install pip packages from whl files to avoid the time-consuming process of
 # building from source.
-- 
GitLab


From e02d08f6d0b1637babf34b022e9326b25d8471e1 Mon Sep 17 00:00:00 2001
From: Paul Van Eck <pvaneck@us.ibm.com>
Date: Wed, 2 May 2018 07:08:30 -0700
Subject: [PATCH 139/395] Allow tfdbg mouse down scroll in curses UI (#18942)

* Allow tfdbg mouse down scroll in curses UI

This commit allows users to continuously scroll the screen when the mouse
is held down on the scroll bar when using the curses UI.

* Only allow click-hold scrolling on scroll bar arrows
---
 tensorflow/python/debug/cli/curses_ui.py | 36 +++++++++++++++++++++---
 1 file changed, 32 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/debug/cli/curses_ui.py b/tensorflow/python/debug/cli/curses_ui.py
index f66cefb427..7b87972d69 100644
--- a/tensorflow/python/debug/cli/curses_ui.py
+++ b/tensorflow/python/debug/cli/curses_ui.py
@@ -190,8 +190,6 @@ class ScrollBar(object):
     return layout
 
   def get_click_command(self, mouse_y):
-    # TODO(cais): Support continuous scrolling when the mouse button is held
-    # down.
     if self._output_num_rows <= 1:
       return None
     elif mouse_y == self._min_y:
@@ -271,6 +269,10 @@ class CursesUI(base_ui.BaseUI):
 
   _UI_WAIT_MESSAGE = "Processing..."
 
+  # The delay (in ms) between each update of the scroll bar when the mouse
+  # button is held down on the scroll bar. Controls how fast the screen scrolls.
+  _MOUSE_SCROLL_DELAY_MS = 100
+
   _single_instance_lock = threading.Lock()
 
   def __init__(self, on_ui_exit=None, config=None):
@@ -855,7 +857,30 @@ class CursesUI(base_ui.BaseUI):
       except curses.error:
         mouse_event_type = None
 
-      if mouse_event_type == curses.BUTTON1_RELEASED:
+      if mouse_event_type == curses.BUTTON1_PRESSED:
+        # Logic for held mouse-triggered scrolling.
+        if mouse_x >= self._max_x - 2:
+          # Disable blocking on checking for user input.
+          self._command_window.nodelay(True)
+
+          # Loop while mouse button is pressed.
+          while mouse_event_type == curses.BUTTON1_PRESSED:
+            # Sleep for a bit.
+            curses.napms(self._MOUSE_SCROLL_DELAY_MS)
+            scroll_command = self._scroll_bar.get_click_command(mouse_y)
+            if scroll_command in (_SCROLL_UP_A_LINE, _SCROLL_DOWN_A_LINE):
+              self._scroll_output(scroll_command)
+
+            # Check to see if different mouse event is in queue.
+            self._command_window.getch()
+            try:
+              _, _, _, _, mouse_event_type = self._screen_getmouse()
+            except curses.error:
+              pass
+
+          self._command_window.nodelay(False)
+          return x
+      elif mouse_event_type == curses.BUTTON1_RELEASED:
         # Logic for mouse-triggered scrolling.
         if mouse_x >= self._max_x - 2:
           scroll_command = self._scroll_bar.get_click_command(mouse_y)
@@ -1677,4 +1702,7 @@ class CursesUI(base_ui.BaseUI):
       self._redraw_output()
 
   def _screen_set_mousemask(self):
-    curses.mousemask(self._mouse_enabled)
+    if self._mouse_enabled:
+      curses.mousemask(curses.BUTTON1_RELEASED | curses.BUTTON1_PRESSED)
+    else:
+      curses.mousemask(0)
-- 
GitLab


From ba1c33faeb6df1ae363888e2e7330e219f0679ea Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 07:51:53 -0700
Subject: [PATCH 140/395] ArraysExtraInfo: Add name_regexp field and regexp
 name matching. PiperOrigin-RevId: 195091587

---
 tensorflow/contrib/lite/toco/BUILD            |  1 +
 .../contrib/lite/toco/model_flags.proto       |  3 +-
 tensorflow/contrib/lite/toco/tooling_util.cc  | 79 ++++++++++++-------
 3 files changed, 53 insertions(+), 30 deletions(-)

diff --git a/tensorflow/contrib/lite/toco/BUILD b/tensorflow/contrib/lite/toco/BUILD
index f16225fd66..a3eff8ac70 100644
--- a/tensorflow/contrib/lite/toco/BUILD
+++ b/tensorflow/contrib/lite/toco/BUILD
@@ -396,6 +396,7 @@ cc_library(
         ":toco_port",
         ":types_proto_cc",
         "//tensorflow/core:lib",
+        "//third_party/re2",
         "@com_google_absl//absl/strings",
         "@protobuf_archive//:protobuf_headers",
     ],
diff --git a/tensorflow/contrib/lite/toco/model_flags.proto b/tensorflow/contrib/lite/toco/model_flags.proto
index d23e80c464..6c1c53658c 100644
--- a/tensorflow/contrib/lite/toco/model_flags.proto
+++ b/tensorflow/contrib/lite/toco/model_flags.proto
@@ -96,8 +96,9 @@ message RnnState {
 // model that does not already contain such MinMax information.
 message ArraysExtraInfo {
   message Entry {
-    // Next ID to use: 7.
+    // Next ID to use: 8.
     optional string name = 1;
+    optional string name_regexp = 7;
     optional double min = 2;
     optional double max = 3;
     optional IODataType data_type = 4;
diff --git a/tensorflow/contrib/lite/toco/tooling_util.cc b/tensorflow/contrib/lite/toco/tooling_util.cc
index f334c51bbb..36f38ba8b0 100644
--- a/tensorflow/contrib/lite/toco/tooling_util.cc
+++ b/tensorflow/contrib/lite/toco/tooling_util.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "absl/strings/str_join.h"
 #include "absl/strings/str_replace.h"
 #include "absl/strings/str_split.h"
+#include "third_party/re2/re2.h"
 #include "tensorflow/contrib/lite/toco/dump_graphviz.h"
 #include "tensorflow/contrib/lite/toco/model_flags.pb.h"
 #include "tensorflow/contrib/lite/toco/toco_graphviz_dump_options.h"
@@ -1983,38 +1984,58 @@ void FinishBuildingRNNStates(Model* model) {
   }
 }
 
+// Returns the array names that match the ArraysExtraInfo's name and
+// name_regexp. The regexp match is for a full match.
+std::unordered_set<string> ScanArrayNames(
+    const Model& model, const toco::ArraysExtraInfo_Entry& entry) {
+  std::unordered_set<string> matches;
+  if (model.HasArray(entry.name())) {
+    matches.insert(entry.name());
+  }
+  if (!entry.name_regexp().empty()) {
+    const auto& arrays = model.GetArrayMap();
+    const RE2 name_regexp = {entry.name_regexp()};
+    for (auto it = arrays.begin(); it != arrays.end(); ++it) {
+      if (RE2::FullMatch(it->first, name_regexp)) {
+        matches.insert(it->first);
+      }
+    }
+  }
+  return matches;
+}
+
 void UseArraysExtraInfo(Model* model, bool quantize_output) {
   for (const auto& entry : model->flags.arrays_extra_info().entries()) {
-    if (!model->HasArray(entry.name())) {
-      continue;
-    }
-    auto& array = model->GetArray(entry.name());
-    if (entry.has_min() || entry.has_max()) {
-      CHECK_EQ(entry.has_min(), entry.has_max());
-      auto& minmax = array.GetOrCreateMinMax();
-      minmax.min = entry.min();
-      minmax.max = entry.max();
-    }
-    if (entry.has_data_type() && quantize_output) {
-      array.final_data_type =
-          ConvertIODataTypeToArrayDataType(entry.data_type());
-    }
-    if (entry.has_shape()) {
-      array.clear_shape();
-      // Make sure to create the shape even if there are no dims, to
-      // correctly record 0-D shapes.
-      array.mutable_shape();
-      for (int dim : entry.shape().dims()) {
-        array.mutable_shape()->mutable_dims()->push_back(dim);
+    const auto matches = ScanArrayNames(*model, entry);
+    for (const auto& matched_name : matches) {
+      auto& array = model->GetArray(matched_name);
+      if (entry.has_min() || entry.has_max()) {
+        CHECK_EQ(entry.has_min(), entry.has_max());
+        auto& minmax = array.GetOrCreateMinMax();
+        minmax.min = entry.min();
+        minmax.max = entry.max();
       }
-    }
-    if (entry.has_constant_float_value()) {
-      CHECK(array.has_shape());
-      if (array.data_type == ArrayDataType::kFloat) {
-        auto& data = array.GetMutableBuffer<ArrayDataType::kFloat>().data;
-        data.resize(RequiredBufferSizeForShape(array.shape()));
-        for (float& f : data) {
-          f = entry.constant_float_value();
+      if (entry.has_data_type() && quantize_output) {
+        array.final_data_type =
+            ConvertIODataTypeToArrayDataType(entry.data_type());
+      }
+      if (entry.has_shape()) {
+        array.clear_shape();
+        // Make sure to create the shape even if there are no dims, to
+        // correctly record 0-D shapes.
+        array.mutable_shape();
+        for (int dim : entry.shape().dims()) {
+          array.mutable_shape()->mutable_dims()->push_back(dim);
+        }
+      }
+      if (entry.has_constant_float_value()) {
+        CHECK(array.has_shape());
+        if (array.data_type == ArrayDataType::kFloat) {
+          auto& data = array.GetMutableBuffer<ArrayDataType::kFloat>().data;
+          data.resize(RequiredBufferSizeForShape(array.shape()));
+          for (float& f : data) {
+            f = entry.constant_float_value();
+          }
         }
       }
     }
-- 
GitLab


From 72fd2b8e97f301039ac0eb60435cbbddf36212f6 Mon Sep 17 00:00:00 2001
From: Priya Gupta <priyag@google.com>
Date: Wed, 2 May 2018 08:04:09 -0700
Subject: [PATCH 141/395] Use experimental auto_sharding in multi worker
 dataset.

PiperOrigin-RevId: 195092992
---
 tensorflow/contrib/distribute/python/BUILD     | 1 +
 tensorflow/contrib/distribute/python/values.py | 5 +++--
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensorflow/contrib/distribute/python/BUILD b/tensorflow/contrib/distribute/python/BUILD
index cdb3a8d65e..aaafc184bf 100644
--- a/tensorflow/contrib/distribute/python/BUILD
+++ b/tensorflow/contrib/distribute/python/BUILD
@@ -21,6 +21,7 @@ py_library(
     srcs = ["values.py"],
     visibility = ["//tensorflow:internal"],
     deps = [
+        ":input_ops",
         ":prefetching_ops_v2",
         "//tensorflow/contrib/data/python/ops:batching",
         "//tensorflow/contrib/eager/python:datasets",
diff --git a/tensorflow/contrib/distribute/python/values.py b/tensorflow/contrib/distribute/python/values.py
index 18afdaa7b0..aaf177d07e 100644
--- a/tensorflow/contrib/distribute/python/values.py
+++ b/tensorflow/contrib/distribute/python/values.py
@@ -27,6 +27,7 @@ import weakref
 import six
 
 from tensorflow.contrib.data.python.ops import batching
+from tensorflow.contrib.distribute.python import input_ops
 from tensorflow.contrib.distribute.python import prefetching_ops_v2
 from tensorflow.python.eager import context
 from tensorflow.python.framework import device as tf_device
@@ -651,8 +652,8 @@ class MultiWorkerDataset(object):
         six.iteritems(worker_device_map)):
       with ops.device(worker):
         worker_input = dataset_fn()
-        # TODO(yuefengz, priyag): support efficient sharding.
-        worker_input = worker_input.shard(len(worker_device_map), i)
+        worker_input = input_ops.auto_shard_dataset(
+            worker_input, len(worker_device_map), i)
         self._datasets[worker] = PerDeviceDataset(
             worker_input, worker_devices, prefetch_on_device=prefetch_on_device)
 
-- 
GitLab


From 22eed5405906de6df8846bd9ce4ee0a57917aa3c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 08:51:07 -0700
Subject: [PATCH 142/395] Automated g4 rollback of changelist 195091587

PiperOrigin-RevId: 195098224
---
 tensorflow/contrib/lite/toco/BUILD            |  1 -
 .../contrib/lite/toco/model_flags.proto       |  3 +-
 tensorflow/contrib/lite/toco/tooling_util.cc  | 79 +++++++------------
 3 files changed, 30 insertions(+), 53 deletions(-)

diff --git a/tensorflow/contrib/lite/toco/BUILD b/tensorflow/contrib/lite/toco/BUILD
index a3eff8ac70..f16225fd66 100644
--- a/tensorflow/contrib/lite/toco/BUILD
+++ b/tensorflow/contrib/lite/toco/BUILD
@@ -396,7 +396,6 @@ cc_library(
         ":toco_port",
         ":types_proto_cc",
         "//tensorflow/core:lib",
-        "//third_party/re2",
         "@com_google_absl//absl/strings",
         "@protobuf_archive//:protobuf_headers",
     ],
diff --git a/tensorflow/contrib/lite/toco/model_flags.proto b/tensorflow/contrib/lite/toco/model_flags.proto
index 6c1c53658c..d23e80c464 100644
--- a/tensorflow/contrib/lite/toco/model_flags.proto
+++ b/tensorflow/contrib/lite/toco/model_flags.proto
@@ -96,9 +96,8 @@ message RnnState {
 // model that does not already contain such MinMax information.
 message ArraysExtraInfo {
   message Entry {
-    // Next ID to use: 8.
+    // Next ID to use: 7.
     optional string name = 1;
-    optional string name_regexp = 7;
     optional double min = 2;
     optional double max = 3;
     optional IODataType data_type = 4;
diff --git a/tensorflow/contrib/lite/toco/tooling_util.cc b/tensorflow/contrib/lite/toco/tooling_util.cc
index 36f38ba8b0..f334c51bbb 100644
--- a/tensorflow/contrib/lite/toco/tooling_util.cc
+++ b/tensorflow/contrib/lite/toco/tooling_util.cc
@@ -26,7 +26,6 @@ limitations under the License.
 #include "absl/strings/str_join.h"
 #include "absl/strings/str_replace.h"
 #include "absl/strings/str_split.h"
-#include "third_party/re2/re2.h"
 #include "tensorflow/contrib/lite/toco/dump_graphviz.h"
 #include "tensorflow/contrib/lite/toco/model_flags.pb.h"
 #include "tensorflow/contrib/lite/toco/toco_graphviz_dump_options.h"
@@ -1984,58 +1983,38 @@ void FinishBuildingRNNStates(Model* model) {
   }
 }
 
-// Returns the array names that match the ArraysExtraInfo's name and
-// name_regexp. The regexp match is for a full match.
-std::unordered_set<string> ScanArrayNames(
-    const Model& model, const toco::ArraysExtraInfo_Entry& entry) {
-  std::unordered_set<string> matches;
-  if (model.HasArray(entry.name())) {
-    matches.insert(entry.name());
-  }
-  if (!entry.name_regexp().empty()) {
-    const auto& arrays = model.GetArrayMap();
-    const RE2 name_regexp = {entry.name_regexp()};
-    for (auto it = arrays.begin(); it != arrays.end(); ++it) {
-      if (RE2::FullMatch(it->first, name_regexp)) {
-        matches.insert(it->first);
-      }
-    }
-  }
-  return matches;
-}
-
 void UseArraysExtraInfo(Model* model, bool quantize_output) {
   for (const auto& entry : model->flags.arrays_extra_info().entries()) {
-    const auto matches = ScanArrayNames(*model, entry);
-    for (const auto& matched_name : matches) {
-      auto& array = model->GetArray(matched_name);
-      if (entry.has_min() || entry.has_max()) {
-        CHECK_EQ(entry.has_min(), entry.has_max());
-        auto& minmax = array.GetOrCreateMinMax();
-        minmax.min = entry.min();
-        minmax.max = entry.max();
-      }
-      if (entry.has_data_type() && quantize_output) {
-        array.final_data_type =
-            ConvertIODataTypeToArrayDataType(entry.data_type());
-      }
-      if (entry.has_shape()) {
-        array.clear_shape();
-        // Make sure to create the shape even if there are no dims, to
-        // correctly record 0-D shapes.
-        array.mutable_shape();
-        for (int dim : entry.shape().dims()) {
-          array.mutable_shape()->mutable_dims()->push_back(dim);
-        }
+    if (!model->HasArray(entry.name())) {
+      continue;
+    }
+    auto& array = model->GetArray(entry.name());
+    if (entry.has_min() || entry.has_max()) {
+      CHECK_EQ(entry.has_min(), entry.has_max());
+      auto& minmax = array.GetOrCreateMinMax();
+      minmax.min = entry.min();
+      minmax.max = entry.max();
+    }
+    if (entry.has_data_type() && quantize_output) {
+      array.final_data_type =
+          ConvertIODataTypeToArrayDataType(entry.data_type());
+    }
+    if (entry.has_shape()) {
+      array.clear_shape();
+      // Make sure to create the shape even if there are no dims, to
+      // correctly record 0-D shapes.
+      array.mutable_shape();
+      for (int dim : entry.shape().dims()) {
+        array.mutable_shape()->mutable_dims()->push_back(dim);
       }
-      if (entry.has_constant_float_value()) {
-        CHECK(array.has_shape());
-        if (array.data_type == ArrayDataType::kFloat) {
-          auto& data = array.GetMutableBuffer<ArrayDataType::kFloat>().data;
-          data.resize(RequiredBufferSizeForShape(array.shape()));
-          for (float& f : data) {
-            f = entry.constant_float_value();
-          }
+    }
+    if (entry.has_constant_float_value()) {
+      CHECK(array.has_shape());
+      if (array.data_type == ArrayDataType::kFloat) {
+        auto& data = array.GetMutableBuffer<ArrayDataType::kFloat>().data;
+        data.resize(RequiredBufferSizeForShape(array.shape()));
+        for (float& f : data) {
+          f = entry.constant_float_value();
         }
       }
     }
-- 
GitLab


From d6d4355a39a56a1b0d0abc7ce74d8307a1925459 Mon Sep 17 00:00:00 2001
From: Tony Wang <tonywy@google.com>
Date: Wed, 2 May 2018 09:52:10 -0700
Subject: [PATCH 143/395] Add Name String to GraphOptimizationPass and Log
 Registered Passes

Added a name string to GraphOptimization class and set to the class name through REGISTER_OPTIMIZATION macro.
Modified RunGrouping function to log the name and phase of optimization pass that's running.
Added two additional functions to log all registered optimization passes in the order of execution.

PiperOrigin-RevId: 195106355
---
 .../common_runtime/optimization_registry.cc   | 21 +++++++++++++++++++
 .../common_runtime/optimization_registry.h    | 18 ++++++++++++++--
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/common_runtime/optimization_registry.cc b/tensorflow/core/common_runtime/optimization_registry.cc
index 7f270b4d4e..bf49a758b2 100644
--- a/tensorflow/core/common_runtime/optimization_registry.cc
+++ b/tensorflow/core/common_runtime/optimization_registry.cc
@@ -36,6 +36,8 @@ Status OptimizationPassRegistry::RunGrouping(
     for (auto& phase : group->second) {
       VLOG(1) << "Running optimization phase " << phase.first;
       for (auto& pass : phase.second) {
+        VLOG(1) << "Running optimization pass: "
+                << pass->GetOptimizationPassName();
         Status s = pass->Run(options);
         if (!s.ok()) return s;
       }
@@ -44,4 +46,23 @@ Status OptimizationPassRegistry::RunGrouping(
   return Status::OK();
 }
 
+void OptimizationPassRegistry::LogGrouping(Grouping grouping, int vlog_level) {
+  auto group = groups_.find(grouping);
+  if (group != groups_.end()) {
+    for (auto& phase : group->second) {
+      for (auto& pass : phase.second) {
+        VLOG(vlog_level) << "Registered optimization pass grouping " << grouping
+                         << " phase " << phase.first << ": "
+                         << pass->GetOptimizationPassName();
+      }
+    }
+  }
+}
+
+void OptimizationPassRegistry::LogAllGroupings(int vlog_level) {
+  for (auto group = groups_.begin(); group != groups_.end(); ++group) {
+    LogGrouping(group->first, vlog_level);
+  }
+}
+
 }  // namespace tensorflow
diff --git a/tensorflow/core/common_runtime/optimization_registry.h b/tensorflow/core/common_runtime/optimization_registry.h
index a469c8aa4e..1b535faf19 100644
--- a/tensorflow/core/common_runtime/optimization_registry.h
+++ b/tensorflow/core/common_runtime/optimization_registry.h
@@ -65,6 +65,13 @@ class GraphOptimizationPass {
  public:
   virtual ~GraphOptimizationPass() {}
   virtual Status Run(const GraphOptimizationPassOptions& options) = 0;
+  void SetOptimizationPassName(string name) { _optimization_pass_name = name; }
+  string GetOptimizationPassName() { return _optimization_pass_name; }
+
+ private:
+  // The name of the opitmization pass, which is the same as the inherited class
+  // name.
+  string _optimization_pass_name;
 };
 
 // The key is a 'phase' number. Phases are executed in increasing
@@ -95,6 +102,10 @@ class OptimizationPassRegistry {
   // Returns the global registry of optimization passes.
   static OptimizationPassRegistry* Global();
 
+  // Prints registered optimization passes for debugging.
+  void LogGrouping(Grouping grouping, int vlog_level);
+  void LogAllGroupings(int vlog_level);
+
  private:
   std::map<Grouping, GraphOptimizationPasses> groups_;
 };
@@ -105,7 +116,9 @@ class OptimizationPassRegistration {
  public:
   OptimizationPassRegistration(OptimizationPassRegistry::Grouping grouping,
                                int phase,
-                               std::unique_ptr<GraphOptimizationPass> pass) {
+                               std::unique_ptr<GraphOptimizationPass> pass,
+                               string optimization_pass_name) {
+    pass->SetOptimizationPassName(optimization_pass_name);
     OptimizationPassRegistry::Global()->Register(grouping, phase,
                                                  std::move(pass));
   }
@@ -123,7 +136,8 @@ class OptimizationPassRegistration {
   static optimization_registration::OptimizationPassRegistration       \
       register_optimization_##ctr(                                     \
           grouping, phase,                                             \
-          std::unique_ptr<GraphOptimizationPass>(new optimization))
+          std::unique_ptr<GraphOptimizationPass>(new optimization()),  \
+          #optimization)
 
 }  // namespace tensorflow
 
-- 
GitLab


From c4394346027fa01f12261e6fea6a1b7f19ac21a9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 10:02:09 -0700
Subject: [PATCH 144/395] Instantiate SwapDimension1And2InTensor3 for
 Eigen::half

PiperOrigin-RevId: 195107839
---
 tensorflow/core/kernels/conv_ops_gpu_3.cu.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/core/kernels/conv_ops_gpu_3.cu.cc b/tensorflow/core/kernels/conv_ops_gpu_3.cu.cc
index 2503b475dc..180531b8c0 100644
--- a/tensorflow/core/kernels/conv_ops_gpu_3.cu.cc
+++ b/tensorflow/core/kernels/conv_ops_gpu_3.cu.cc
@@ -1027,6 +1027,7 @@ template struct functor::SwapDimension1And2InTensor3<GPUDevice, float2,
                                                      /*conjugate=*/true>;
 template struct functor::SwapDimension1And2InTensor3<GPUDevice, double2,
                                                      /*conjugate=*/true>;
+template struct functor::SwapDimension1And2InTensor3<GPUDevice, Eigen::half>;
 
 template struct functor::SwapDimension0And2InTensor3<GPUDevice, uint8>;
 template struct functor::SwapDimension0And2InTensor3<GPUDevice, uint16>;
-- 
GitLab


From 1f47bbd1e09a9ed4086d0484024e8989a65274a9 Mon Sep 17 00:00:00 2001
From: Benoit Steiner <bsteiner@google.com>
Date: Wed, 2 May 2018 10:07:13 -0700
Subject: [PATCH 145/395] Optimized the analysis of rank and size operations.

PiperOrigin-RevId: 195108832
---
 .../core/grappler/costs/graph_properties.cc   | 32 +++++++++++++++++++
 tensorflow/core/grappler/op_types.cc          |  4 +++
 tensorflow/core/grappler/op_types.h           |  2 ++
 3 files changed, 38 insertions(+)

diff --git a/tensorflow/core/grappler/costs/graph_properties.cc b/tensorflow/core/grappler/costs/graph_properties.cc
index 2c7b57971a..69b22561b2 100644
--- a/tensorflow/core/grappler/costs/graph_properties.cc
+++ b/tensorflow/core/grappler/costs/graph_properties.cc
@@ -475,6 +475,38 @@ class SymbolicShapeRefiner {
               }
             }
           }
+        } else if (IsRank(*input)) {
+          if (c->inference_context->RankKnown(c->inference_context->input(0))) {
+            int32 rank =
+                c->inference_context->Rank(c->inference_context->input(0));
+            Tensor t(DT_INT32, {});
+            t.flat<int32>()(0) = rank;
+            const_values[dst_input] = t;
+            input_tensors[dst_input] = &const_values[dst_input];
+          }
+        } else if (IsSize(*input)) {
+          DimensionHandle size =
+              c->inference_context->NumElements(c->inference_context->input(0));
+          if (c->inference_context->ValueKnown(size)) {
+            int64 sz = c->inference_context->Value(size);
+            bool valid = false;
+            if (input->attr().at("T").type() == DT_INT32) {
+              if (sz < std::numeric_limits<int32>::max()) {
+                Tensor t(DT_INT32, {});
+                t.flat<int32>()(0) = sz;
+                const_values[dst_input] = t;
+                valid = true;
+              }
+            } else {
+              Tensor t(DT_INT64, {});
+              t.flat<int64>()(0) = sz;
+              const_values[dst_input] = t;
+              valid = true;
+            }
+            if (valid) {
+              input_tensors[dst_input] = &const_values[dst_input];
+            }
+          }
         }
 
         if (c->output_tensors_as_shapes.size() > src_output) {
diff --git a/tensorflow/core/grappler/op_types.cc b/tensorflow/core/grappler/op_types.cc
index bf6d4c0921..7c936dfca1 100644
--- a/tensorflow/core/grappler/op_types.cc
+++ b/tensorflow/core/grappler/op_types.cc
@@ -262,6 +262,8 @@ bool IsRandomShuffle(const NodeDef& node) {
   return node.op() == "RandomShuffle";
 }
 
+bool IsRank(const NodeDef& node) { return node.op() == "Rank"; }
+
 bool IsReal(const NodeDef& node) { return node.op() == "Real"; }
 
 bool IsRealDiv(const NodeDef& node) { return node.op() == "RealDiv"; }
@@ -317,6 +319,8 @@ bool IsShuffle(const NodeDef& node) { return node.op() == "Shuffle"; }
 
 bool IsSigmoidGrad(const NodeDef& node) { return node.op() == "SigmoidGrad"; }
 
+bool IsSize(const NodeDef& node) { return node.op() == "Size"; }
+
 bool IsSlice(const NodeDef& node) { return node.op() == "Slice"; }
 
 bool IsSoftplusGrad(const NodeDef& node) { return node.op() == "SoftplusGrad"; }
diff --git a/tensorflow/core/grappler/op_types.h b/tensorflow/core/grappler/op_types.h
index 3dddf3f1ea..7a1b438768 100644
--- a/tensorflow/core/grappler/op_types.h
+++ b/tensorflow/core/grappler/op_types.h
@@ -100,6 +100,7 @@ bool IsProd(const NodeDef& node);
 bool IsPow(const NodeDef& node);
 bool IsQueue(const NodeDef& node);
 bool IsRandomShuffle(const NodeDef& node);
+bool IsRank(const NodeDef& node);
 bool IsReal(const NodeDef& node);
 bool IsRealDiv(const NodeDef& node);
 bool IsRelu6Grad(const NodeDef& node);
@@ -116,6 +117,7 @@ bool IsRsqrtGrad(const NodeDef& node);
 bool IsSelect(const NodeDef& node);
 bool IsSeluGrad(const NodeDef& node);
 bool IsSend(const NodeDef& node);
+bool IsSize(const NodeDef& node);
 bool IsSlice(const NodeDef& node);
 bool IsShape(const NodeDef& node);
 bool IsShapeN(const NodeDef& node);
-- 
GitLab


From bc86da090f2f2e850768bbdfd603c7217aecdb53 Mon Sep 17 00:00:00 2001
From: Andrew Selle <aselle@andyselle.com>
Date: Wed, 2 May 2018 10:33:07 -0700
Subject: [PATCH 146/395] Fix Makefile to not use benchmark anymore (switch to
 minimal) (#19019)

Minimal uses nothing and does almost nothing, but it does nothing
requiring protos or rest of tensorflow runtime.

Benchmark_model originally was more like this, but it became
useful for actually benchmarking, making it less useful as a minimal
example.
---
 tensorflow/contrib/lite/Makefile              | 19 +++--
 .../contrib/lite/examples/minimal/minimal.cc  | 71 +++++++++++++++++++
 2 files changed, 80 insertions(+), 10 deletions(-)
 create mode 100644 tensorflow/contrib/lite/examples/minimal/minimal.cc

diff --git a/tensorflow/contrib/lite/Makefile b/tensorflow/contrib/lite/Makefile
index 65fba52d46..e4f86e258a 100644
--- a/tensorflow/contrib/lite/Makefile
+++ b/tensorflow/contrib/lite/Makefile
@@ -1,4 +1,3 @@
-
 # Find where we're running from, so we can store generated files here.
 ifeq ($(origin MAKEFILE_DIR), undefined)
 	MAKEFILE_DIR := $(shell dirname $(realpath $(lastword $(MAKEFILE_LIST))))
@@ -69,12 +68,12 @@ LIB_NAME := libtensorflow-lite.a
 LIB_PATH := $(LIBDIR)$(LIB_NAME)
 
 # A small example program that shows how to link against the library.
-BENCHMARK_PATH := $(BINDIR)benchmark_model
+MINIMAL_PATH := $(BINDIR)minimal
 
-BENCHMARK_SRCS := \
-tensorflow/contrib/lite/tools/benchmark_model.cc
-BENCHMARK_OBJS := $(addprefix $(OBJDIR), \
-$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(BENCHMARK_SRCS))))
+MINIMAL_SRCS := \
+tensorflow/contrib/lite/examples/minimal/minimal.cc
+MINIMAL_OBJS := $(addprefix $(OBJDIR), \
+$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(MINIMAL_SRCS))))
 
 # What sources we want to compile, must be kept in sync with the main Bazel
 # build files.
@@ -100,7 +99,7 @@ $(wildcard tensorflow/contrib/lite/*/*test.cc) \
 $(wildcard tensorflow/contrib/lite/*/*/*test.cc) \
 $(wildcard tensorflow/contrib/lite/*/*/*/*test.cc) \
 $(wildcard tensorflow/contrib/lite/kernels/test_util.cc) \
-$(BENCHMARK_SRCS)
+$(MINIMAL_SRCS)
 # Filter out all the excluded files.
 TF_LITE_CC_SRCS := $(filter-out $(CORE_CC_EXCLUDE_SRCS), $(CORE_CC_ALL_SRCS))
 # File names of the intermediate files target compilation generates.
@@ -119,17 +118,17 @@ $(OBJDIR)%.o: %.c
 	$(CC) $(CCFLAGS) $(INCLUDES) -c $< -o $@
 
 # The target that's compiled if there's no command-line arguments.
-all: $(LIB_PATH) $(BENCHMARK_PATH)
+all: $(LIB_PATH)  $(MINIMAL_PATH)
 
 # Gathers together all the objects we've compiled into a single '.a' archive.
 $(LIB_PATH): $(LIB_OBJS)
 	@mkdir -p $(dir $@)
 	$(AR) $(ARFLAGS) $(LIB_PATH) $(LIB_OBJS)
 
-$(BENCHMARK_PATH): $(BENCHMARK_OBJS) $(LIB_PATH)
+$(MINIMAL_PATH): $(MINIMAL_OBJS) $(LIB_PATH)
 	@mkdir -p $(dir $@)
 	$(CXX) $(CXXFLAGS) $(INCLUDES) \
-	-o $(BENCHMARK_PATH) $(BENCHMARK_OBJS) \
+	-o $(MINIMAL_PATH) $(MINIMAL_OBJS) \
 	$(LIBFLAGS) $(LIB_PATH) $(LDFLAGS) $(LIBS)
 
 # Gets rid of all generated files.
diff --git a/tensorflow/contrib/lite/examples/minimal/minimal.cc b/tensorflow/contrib/lite/examples/minimal/minimal.cc
new file mode 100644
index 0000000000..106e3b0270
--- /dev/null
+++ b/tensorflow/contrib/lite/examples/minimal/minimal.cc
@@ -0,0 +1,71 @@
+/* Copyright 2018 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.
+==============================================================================*/
+#include "tensorflow/contrib/lite/model.h"
+#include "tensorflow/contrib/lite/interpreter.h"
+#include "tensorflow/contrib/lite/kernels/register.h"
+#include <cstdio>
+
+// This is an example that is minimal to read a model
+// from disk and perform inference. There is no data being loaded
+// that is up to you to add as a user.
+//
+// NOTE: Do not add any dependencies to this that cannot be built with
+// the minimal makefile. This example must remain trivial to build with
+// the minimal build tool.
+//
+// Usage: minimal <tflite model>
+
+using namespace tflite;
+
+#define TFLITE_MINIMAL_CHECK(x) \
+  if(!(x)) {                                                    \
+    fprintf(stderr, "Error at %s:%d\n",  __FILE__, __LINE__); \
+    exit(1); \
+  }
+
+
+int main(int argc, char *argv[]) {
+  if(argc != 2) {
+    fprintf(stderr, "Usage: %s <model>\n");
+    return 1;
+  }
+  const char* filename = argv[1];
+
+  // Load model
+  std::unique_ptr<tflite::FlatBufferModel> model
+      = tflite::FlatBufferModel::BuildFromFile(filename);
+  TFLITE_MINIMAL_CHECK(model != nullptr);
+
+  // Build the interpreter
+  tflite::ops::builtin::BuiltinOpResolver resolver;
+  InterpreterBuilder builder(*model.get(), resolver);
+  std::unique_ptr<Interpreter> interpreter;
+  builder(&interpreter);
+  TFLITE_MINIMAL_CHECK(interpreter != nullptr);
+
+  // Allocate tensor buffers.
+  TFLITE_MINIMAL_CHECK(interpreter->AllocateTensors() == kTfLiteOk);
+
+  // Fill input buffers
+  // TODO(user): Insert code to fill input tensors
+
+  // Run inference
+  TFLITE_MINIMAL_CHECK(interpreter->Invoke() == kTfLiteOk);
+
+  // Read output buffers
+  // TODO(user): Insert getting data out code.
+
+  return 0;
+}
-- 
GitLab


From e408e8171540386d3dfed0a7c4fa1d0e1cc812cd Mon Sep 17 00:00:00 2001
From: Justin Lebar <jlebar@google.com>
Date: Wed, 2 May 2018 10:36:26 -0700
Subject: [PATCH 147/395] Internal-only change.

PiperOrigin-RevId: 195113702
---
 tensorflow/tensorflow.bzl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl
index e5cc886b32..b2cec7655f 100644
--- a/tensorflow/tensorflow.bzl
+++ b/tensorflow/tensorflow.bzl
@@ -1507,6 +1507,7 @@ def tf_py_wrap_cc(name,
 # 2. When --define=no_tensorflow_py_deps=false (by default), it's a normal py_test.
 def py_test(deps=[], data=[], **kwargs):
   native.py_test(
+      # TODO(jlebar): Ideally we'd use tcmalloc here.,
       deps=select({
           "//conditions:default": deps,
           clean_dep("//tensorflow:no_tensorflow_py_deps"): [],
-- 
GitLab


From 489640a0d00ea7b5826937781cd1bf3a520b0b5d Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Wed, 2 May 2018 10:43:25 -0700
Subject: [PATCH 148/395] Fix some nits in cpu_literal_caching_test that I
 noticed after submission

PiperOrigin-RevId: 195114829
---
 .../service/cpu/tests/cpu_literal_caching_test.cc  | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc
index f0404d07d9..b10eb74635 100644
--- a/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc
+++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc
@@ -20,9 +20,9 @@ limitations under the License.
 namespace xla {
 namespace cpu {
 namespace {
-class CpuExternalConstantsTest : public CpuCodegenTest {};
+class CpuDuplicateConstantsTest : public CpuCodegenTest {};
 
-TEST_F(CpuExternalConstantsTest, RepeatedArrayConstants) {
+TEST_F(CpuDuplicateConstantsTest, RepeatedArrayConstants) {
   // We use a while loop here to force the two constant HloInstructions to be in
   // different computations.  Otherwise the HLO optimizer itself CSEs them.
   const string hlo_text = R"(
@@ -56,6 +56,10 @@ ENTRY main {
 }
 )";
 
+  // TODO(b/78879738): The fake "f32[] constant(1)" root is only needed to work
+  // around b/78879738.  Once b/78879738 is fixed, we can set one of the
+  // outfeeds as the root.
+
   string filecheck_pattern = R"(
 CHECK: private constant [2 x [3 x [2 x float]]]
 CHECK-NOT: private constant [2 x [3 x [2 x float]]]
@@ -73,7 +77,7 @@ CHECK-NOT: private constant [2 x [3 x [2 x float]]]
                                 /*match_optimized_ir=*/false);
 }
 
-TEST_F(CpuExternalConstantsTest, RepeatedTupleConstants) {
+TEST_F(CpuDuplicateConstantsTest, RepeatedTupleConstants) {
   // We use a while loop here to force the two constant HloInstructions to be in
   // different computations.  Otherwise the HLO optimizer itself CSEs them.
   const string hlo_text = R"(
@@ -101,6 +105,10 @@ ENTRY main {
 }
 )";
 
+  // TODO(b/78879738): The fake "f32[] constant(1)" root is only needed to work
+  // around b/78879738.  Once b/78879738 is fixed, we can set one of the
+  // outfeeds as the root.
+
   string filecheck_pattern = R"(
 CHECK: private constant [2 x float]
 CHECK: private constant [2 x [1 x float]]
-- 
GitLab


From 55972370f3243ca061b882120383545d70642cf8 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 10:57:13 -0700
Subject: [PATCH 149/395] Initialize all members of CollectiveParams at
 construction time to avoid warnings of uninit memory access by dynamic
 analysis tools.

PiperOrigin-RevId: 195117321
---
 tensorflow/core/framework/collective.h | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/tensorflow/core/framework/collective.h b/tensorflow/core/framework/collective.h
index f6fe12e7ef..f8d27d3868 100644
--- a/tensorflow/core/framework/collective.h
+++ b/tensorflow/core/framework/collective.h
@@ -52,7 +52,8 @@ struct CollGroupParams {
   DeviceType device_type;
   int32 num_tasks;  // number of distinct tasks in group
   string ToString() const;
-  CollGroupParams() : device_type(DEVICE_CPU) {}
+  CollGroupParams()
+      : group_key(0), group_size(0), device_type(DEVICE_CPU), num_tasks(0) {}
 };
 
 // The best implementation of a collective op depends on many factors
@@ -71,10 +72,11 @@ struct CollImplDetails {
 
 // Data common to all members of a collective instance.
 struct CollInstanceParams {
-  int32 instance_key;  // Identifies all participating graph nodes.
-  CollectiveType type;
-  DataType data_type;
-  TensorShape shape;
+  // Identifies all participating graph nodes.
+  int32 instance_key = -1;
+  CollectiveType type = UNDEFINED_COLLECTIVE;
+  DataType data_type = DT_FLOAT;
+  TensorShape shape = {0};
   // Fully qualified name of device for each member, in default rank order.
   std::vector<string> device_names;
   // Task name prefix of corresponding device name.
@@ -99,8 +101,8 @@ struct CollectiveParams {
   CollInstanceParams instance;
   CollTaskParams task;
 
-  string name;             // node name used only for log or error messages
-  int default_rank;        // index of this op within device_names
+  string name = "";        // node name used only for log or error messages
+  int default_rank = -1;   // index of this op within device_names
   bool is_source = false;  // broadcast only
   // Rank of this device in each subdivision permutation.
   std::vector<int> subdiv_rank;
-- 
GitLab


From c08bf79144b3acc731018147e92fd389bcb60b2d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 10:57:49 -0700
Subject: [PATCH 150/395] Renames _regression_head_with_mean_squared_error_loss
 to _regression_head.

PiperOrigin-RevId: 195117425
---
 .../estimator/python/estimator/head.py        |   7 +-
 .../python/estimator/canned/baseline.py       |   2 +-
 .../python/estimator/canned/boosted_trees.py  |   2 +-
 tensorflow/python/estimator/canned/dnn.py     |   3 +-
 .../estimator/canned/dnn_linear_combined.py   |   3 +-
 tensorflow/python/estimator/canned/head.py    |   2 +-
 .../python/estimator/canned/head_test.py      | 105 ++++++++----------
 tensorflow/python/estimator/canned/linear.py  |   2 +-
 8 files changed, 53 insertions(+), 73 deletions(-)

diff --git a/tensorflow/contrib/estimator/python/estimator/head.py b/tensorflow/contrib/estimator/python/estimator/head.py
index 2a6d17e81b..5d19bf4714 100644
--- a/tensorflow/contrib/estimator/python/estimator/head.py
+++ b/tensorflow/contrib/estimator/python/estimator/head.py
@@ -235,7 +235,7 @@ def regression_head(weight_column=None,
   Raises:
     ValueError: If `label_dimension` or `loss_reduction` is invalid.
   """
-  return head_lib._regression_head_with_mean_squared_error_loss(  # pylint:disable=protected-access
+  return head_lib._regression_head(  # pylint:disable=protected-access
       weight_column=weight_column,
       label_dimension=label_dimension,
       loss_reduction=loss_reduction,
@@ -297,7 +297,7 @@ def poisson_regression_head(
   def _poisson_loss(labels, logits):
     return nn.log_poisson_loss(
         targets=labels, log_input=logits, compute_full_loss=compute_full_loss)
-  return head_lib._regression_head_with_mean_squared_error_loss(  # pylint:disable=protected-access
+  return head_lib._regression_head(  # pylint:disable=protected-access
       weight_column=weight_column,
       label_dimension=label_dimension,
       loss_reduction=loss_reduction,
@@ -360,8 +360,7 @@ def logistic_regression_head(
         labels, n_classes=2, message='Labels must be in range [0, 1]')
     return nn.sigmoid_cross_entropy_with_logits(
         labels=labels, logits=logits)
-  # TODO(roumposg): Rename to _regression_head, since it supports loss_fn arg.
-  return head_lib._regression_head_with_mean_squared_error_loss(  # pylint:disable=protected-access
+  return head_lib._regression_head(  # pylint:disable=protected-access
       weight_column=weight_column,
       label_dimension=1,
       loss_reduction=loss_reduction,
diff --git a/tensorflow/python/estimator/canned/baseline.py b/tensorflow/python/estimator/canned/baseline.py
index 3e92a77543..980c057372 100644
--- a/tensorflow/python/estimator/canned/baseline.py
+++ b/tensorflow/python/estimator/canned/baseline.py
@@ -344,7 +344,7 @@ class BaselineRegressor(estimator.Estimator):
       A `BaselineRegressor` estimator.
     """
 
-    head = head_lib._regression_head_with_mean_squared_error_loss(  # pylint: disable=protected-access
+    head = head_lib._regression_head(  # pylint: disable=protected-access
         label_dimension=label_dimension,
         weight_column=weight_column,
         loss_reduction=loss_reduction)
diff --git a/tensorflow/python/estimator/canned/boosted_trees.py b/tensorflow/python/estimator/canned/boosted_trees.py
index d281fd90ea..6d7a3299f7 100644
--- a/tensorflow/python/estimator/canned/boosted_trees.py
+++ b/tensorflow/python/estimator/canned/boosted_trees.py
@@ -690,7 +690,7 @@ def _create_regression_head(label_dimension, weight_column=None):
     raise ValueError('For now only 1 dimension regression is supported.'
                      'label_dimension given as {}'.format(label_dimension))
   # pylint: disable=protected-access
-  return head_lib._regression_head_with_mean_squared_error_loss(
+  return head_lib._regression_head(
       label_dimension=label_dimension,
       weight_column=weight_column,
       loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE)
diff --git a/tensorflow/python/estimator/canned/dnn.py b/tensorflow/python/estimator/canned/dnn.py
index 6382622e0b..973a6ec747 100644
--- a/tensorflow/python/estimator/canned/dnn.py
+++ b/tensorflow/python/estimator/canned/dnn.py
@@ -481,8 +481,7 @@ class DNNRegressor(estimator.Estimator):
           features=features,
           labels=labels,
           mode=mode,
-          head=head_lib.  # pylint: disable=protected-access
-          _regression_head_with_mean_squared_error_loss(
+          head=head_lib._regression_head(  # pylint: disable=protected-access
               label_dimension=label_dimension, weight_column=weight_column,
               loss_reduction=loss_reduction),
           hidden_units=hidden_units,
diff --git a/tensorflow/python/estimator/canned/dnn_linear_combined.py b/tensorflow/python/estimator/canned/dnn_linear_combined.py
index f47706db2f..95efc0a028 100644
--- a/tensorflow/python/estimator/canned/dnn_linear_combined.py
+++ b/tensorflow/python/estimator/canned/dnn_linear_combined.py
@@ -553,8 +553,7 @@ class DNNLinearCombinedRegressor(estimator.Estimator):
           features=features,
           labels=labels,
           mode=mode,
-          head=head_lib.  # pylint: disable=protected-access
-          _regression_head_with_mean_squared_error_loss(
+          head=head_lib._regression_head(  # pylint: disable=protected-access
               label_dimension=label_dimension, weight_column=weight_column,
               loss_reduction=loss_reduction),
           linear_feature_columns=linear_feature_columns,
diff --git a/tensorflow/python/estimator/canned/head.py b/tensorflow/python/estimator/canned/head.py
index efa4bdf598..48f448d7f5 100644
--- a/tensorflow/python/estimator/canned/head.py
+++ b/tensorflow/python/estimator/canned/head.py
@@ -1197,7 +1197,7 @@ class _BinaryLogisticHeadWithSigmoidCrossEntropyLoss(_Head):
         train_op=train_op)
 
 
-def _regression_head_with_mean_squared_error_loss(
+def _regression_head(
     weight_column=None,
     label_dimension=1,
     loss_reduction=losses.Reduction.SUM,
diff --git a/tensorflow/python/estimator/canned/head_test.py b/tensorflow/python/estimator/canned/head_test.py
index 7da3df01dc..32a6339936 100644
--- a/tensorflow/python/estimator/canned/head_test.py
+++ b/tensorflow/python/estimator/canned/head_test.py
@@ -2607,26 +2607,24 @@ class BinaryLogisticHeadWithSigmoidCrossEntropyLossTest(test.TestCase):
           rtol=tol, atol=tol)
 
 
-class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
+class RegressionHead(test.TestCase):
 
   def setUp(self):
     ops.reset_default_graph()
 
   def test_invalid_label_dimension(self):
     with self.assertRaisesRegexp(ValueError, r'Invalid label_dimension'):
-      head_lib._regression_head_with_mean_squared_error_loss(label_dimension=-1)
+      head_lib._regression_head(label_dimension=-1)
     with self.assertRaisesRegexp(ValueError, r'Invalid label_dimension'):
-      head_lib._regression_head_with_mean_squared_error_loss(label_dimension=0)
+      head_lib._regression_head(label_dimension=0)
 
   def test_invalid_loss_reduction(self):
     with self.assertRaisesRegexp(
         ValueError, r'Invalid loss_reduction: invalid_loss_reduction'):
-      head_lib._regression_head_with_mean_squared_error_loss(
-          loss_reduction='invalid_loss_reduction')
+      head_lib._regression_head(loss_reduction='invalid_loss_reduction')
     with self.assertRaisesRegexp(
         ValueError, r'Invalid loss_reduction: none'):
-      head_lib._regression_head_with_mean_squared_error_loss(
-          loss_reduction=losses.Reduction.NONE)
+      head_lib._regression_head(loss_reduction=losses.Reduction.NONE)
 
   def test_loss_fn_arg_labels_missing(self):
     def _loss_fn(logits):
@@ -2635,7 +2633,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
         ValueError,
         r'loss_fn must contain argument: labels\. '
         r'Given arguments: \(\'logits\',\)'):
-      head_lib._regression_head_with_mean_squared_error_loss(loss_fn=_loss_fn)
+      head_lib._regression_head(loss_fn=_loss_fn)
 
   def test_loss_fn_arg_logits_missing(self):
     def _loss_fn(labels):
@@ -2644,12 +2642,12 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
         ValueError,
         r'loss_fn must contain argument: logits\. '
         r'Given arguments: \(\'labels\',\)'):
-      head_lib._regression_head_with_mean_squared_error_loss(loss_fn=_loss_fn)
+      head_lib._regression_head(loss_fn=_loss_fn)
 
   def test_loss_fn_arg_features_ok(self):
     def _loss_fn(labels, logits, features):
       del labels, logits, features  # Unused
-      head_lib._regression_head_with_mean_squared_error_loss(loss_fn=_loss_fn)
+      head_lib._regression_head(loss_fn=_loss_fn)
 
   def test_loss_fn_arg_invalid(self):
     def _loss_fn(labels, logits, name=None):
@@ -2657,11 +2655,10 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
     with self.assertRaisesRegexp(
         ValueError,
         r'loss_fn has unexpected args: \[\'name\'\]'):
-      head_lib._regression_head_with_mean_squared_error_loss(loss_fn=_loss_fn)
+      head_lib._regression_head(loss_fn=_loss_fn)
 
   def test_invalid_logits(self):
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        label_dimension=3)
+    head = head_lib._regression_head(label_dimension=3)
     self.assertEqual(3, head.logits_dimension)
     logits_1d = np.array(((45.,), (41.,),))
 
@@ -2685,8 +2682,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
         })
 
   def test_incompatible_labels_eval(self):
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        label_dimension=3)
+    head = head_lib._regression_head(label_dimension=3)
     self.assertEqual(3, head.logits_dimension)
     values_3d = np.array(((45., 46., 47.), (41., 42., 43.),))
     values_1d = np.array(((43.,), (44.,),))
@@ -2732,8 +2728,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
         })
 
   def test_incompatible_labels_train(self):
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        label_dimension=3)
+    head = head_lib._regression_head(label_dimension=3)
     self.assertEqual(3, head.logits_dimension)
     values_3d = np.array(((45., 46., 47.), (41., 42., 43.),))
     values_1d = np.array(((43.,), (44.,),))
@@ -2784,12 +2779,11 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
         })
 
   def test_name(self):
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        name='foo')
+    head = head_lib._regression_head(name='foo')
     self.assertEqual('foo', head.name)
 
   def test_predict(self):
-    head = head_lib._regression_head_with_mean_squared_error_loss()
+    head = head_lib._regression_head()
     self.assertEqual(1, head.logits_dimension)
 
     # Create estimator spec.
@@ -2826,8 +2820,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
   def test_predict_with_inverse_link_fn(self):
     def _inverse_link_fn(logits):
       return logits - 10.
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        inverse_link_fn=_inverse_link_fn)
+    head = head_lib._regression_head(inverse_link_fn=_inverse_link_fn)
 
     # Create estimator spec.
     logits = np.array(((45,), (41,),), dtype=np.int32)
@@ -2866,7 +2859,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
           logits, spec.export_outputs['predict'].outputs['logits'].eval())
 
   def test_eval_create_loss(self):
-    head = head_lib._regression_head_with_mean_squared_error_loss()
+    head = head_lib._regression_head()
     logits = np.array(((45,), (41,),), dtype=np.float32)
     labels = np.array(((43,), (44,),), dtype=np.int32)
     features = {'x': np.array(((42,),), dtype=np.float32)}
@@ -2895,8 +2888,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
           data=[logits])
       with ops.control_dependencies([check_labels, check_logits]):
         return constant_op.constant(loss)
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        label_dimension=2, loss_fn=_loss_fn)
+    head = head_lib._regression_head(label_dimension=2, loss_fn=_loss_fn)
 
     actual_training_loss = head.create_loss(
         features={'x': np.array(((42,),), dtype=np.int32)},
@@ -2913,8 +2905,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
     def _loss_fn(labels, logits):
       del labels, logits  # Unused
       return constant_op.constant(loss)
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        label_dimension=2, loss_fn=_loss_fn)
+    head = head_lib._regression_head(label_dimension=2, loss_fn=_loss_fn)
 
     logits = np.array([[-1., 1.], [-2., 2.]], dtype=np.float32)
     labels = np.array([[1., 0.], [2., -1.]], dtype=np.float32)
@@ -2933,7 +2924,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_eval_labels_none(self):
     """Tests that error is raised when labels is None."""
-    head = head_lib._regression_head_with_mean_squared_error_loss()
+    head = head_lib._regression_head()
 
     with self.assertRaisesRegexp(
         ValueError, r'You must provide a labels Tensor\. Given: None\.'):
@@ -2944,7 +2935,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
           labels=None)
 
   def test_eval(self):
-    head = head_lib._regression_head_with_mean_squared_error_loss()
+    head = head_lib._regression_head()
     self.assertEqual(1, head.logits_dimension)
 
     logits = np.array(((45,), (41,),), dtype=np.float32)
@@ -2986,8 +2977,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
       self.assertAllClose(expected_loss_mean, loss_mean_value_op.eval())
 
   def test_eval_metric_ops_with_head_name_for_regression(self):
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        name='some_regression_head')
+    head = head_lib._regression_head(name='some_regression_head')
     logits = np.array(((1,), (9,)), dtype=np.float32)
     labels = np.array(((1,), (1,)), dtype=np.int64)
     features = {'x': np.array(((42,),), dtype=np.int32)}
@@ -3004,7 +2994,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
     self.assertItemsEqual(expected_metric_keys, spec.eval_metric_ops.keys())
 
   def test_eval_with_regularization_losses(self):
-    head = head_lib._regression_head_with_mean_squared_error_loss(
+    head = head_lib._regression_head(
         loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE)
     self.assertEqual(1, head.logits_dimension)
 
@@ -3049,7 +3039,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
           expected_metrics, {k: value_ops[k].eval() for k in value_ops})
 
   def test_train_create_loss(self):
-    head = head_lib._regression_head_with_mean_squared_error_loss()
+    head = head_lib._regression_head()
     logits = np.array(((45,), (41,),), dtype=np.float32)
     labels = np.array(((43,), (44,),), dtype=np.int32)
     features = {'x': np.array(((42,),), dtype=np.float32)}
@@ -3073,7 +3063,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_train_create_loss_loss_reduction(self):
     """Tests create_loss with loss_reduction."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
+    head = head_lib._regression_head(
         loss_reduction=losses.Reduction.SUM_BY_NONZERO_WEIGHTS)
     logits = np.array(((45,), (41,),), dtype=np.float32)
     labels = np.array(((43,), (44,),), dtype=np.int32)
@@ -3098,7 +3088,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_train_labels_none(self):
     """Tests that error is raised when labels is None."""
-    head = head_lib._regression_head_with_mean_squared_error_loss()
+    head = head_lib._regression_head()
     def _no_op_train_fn(loss):
       del loss
       return control_flow_ops.no_op()
@@ -3113,7 +3103,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
           train_op_fn=_no_op_train_fn)
 
   def test_train(self):
-    head = head_lib._regression_head_with_mean_squared_error_loss()
+    head = head_lib._regression_head()
     self.assertEqual(1, head.logits_dimension)
 
     # Create estimator spec.
@@ -3163,7 +3153,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
       }, summary_str)
 
   def test_train_with_optimizer(self):
-    head = head_lib._regression_head_with_mean_squared_error_loss()
+    head = head_lib._regression_head()
     self.assertEqual(1, head.logits_dimension)
 
     # Create estimator spec.
@@ -3197,8 +3187,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
       self.assertEqual(expected_train_result, train_result)
 
   def test_train_summaries_with_head_name(self):
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        name='some_regression_head')
+    head = head_lib._regression_head(name='some_regression_head')
     self.assertEqual(1, head.logits_dimension)
 
     # Create estimator spec.
@@ -3237,7 +3226,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
           summary_str)
 
   def test_train_with_regularization_losses(self):
-    head = head_lib._regression_head_with_mean_squared_error_loss(
+    head = head_lib._regression_head(
         loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE)
     self.assertEqual(1, head.logits_dimension)
 
@@ -3285,8 +3274,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_weighted_multi_example_eval(self):
     """1d label, 3 examples, 1 batch."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        weight_column='label_weights')
+    head = head_lib._regression_head(weight_column='label_weights')
     self.assertEqual(1, head.logits_dimension)
 
     # Create estimator spec.
@@ -3330,7 +3318,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_weight_with_numeric_column(self):
     """1d label, 3 examples, 1 batch."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
+    head = head_lib._regression_head(
         weight_column=feature_column_lib.numeric_column(
             'label_weights', normalizer_fn=lambda x: x + 1.))
 
@@ -3356,8 +3344,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_weighted_multi_example_train(self):
     """1d label, 3 examples, 1 batch."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        weight_column='label_weights')
+    head = head_lib._regression_head(weight_column='label_weights')
     self.assertEqual(1, head.logits_dimension)
 
     # Create estimator spec.
@@ -3408,8 +3395,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_train_one_dim_create_loss(self):
     """Tests create_loss with 1D labels and weights (shape [batch_size])."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        weight_column='label_weights')
+    head = head_lib._regression_head(weight_column='label_weights')
     logits = np.array(((45,), (41,), (44,)), dtype=np.float32)
     x_feature_rank_1 = np.array((42., 43., 44.,), dtype=np.float32)
     weight_rank_1 = np.array((1., .1, 1.5,), dtype=np.float64)
@@ -3435,8 +3421,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_train_one_dim(self):
     """Tests train with 1D labels and weights (shape [batch_size])."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        weight_column='label_weights')
+    head = head_lib._regression_head(weight_column='label_weights')
     self.assertEqual(1, head.logits_dimension)
 
     # Create estimator spec.
@@ -3493,7 +3478,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_weighted_multi_value_eval_create_loss(self):
     """3d label, 1 example, 1 batch."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
+    head = head_lib._regression_head(
         weight_column='label_weights', label_dimension=3)
     logits = np.array(((45., 41., 44.),))
     labels = np.array(((35., 42., 45.),))
@@ -3515,7 +3500,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_weighted_multi_value_eval(self):
     """3d label, 1 example, 1 batch."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
+    head = head_lib._regression_head(
         weight_column='label_weights', label_dimension=3)
     self.assertEqual(3, head.logits_dimension)
 
@@ -3562,7 +3547,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_weighted_multi_value_train_create_loss(self):
     """3d label, 1 example, 1 batch."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
+    head = head_lib._regression_head(
         weight_column='label_weights', label_dimension=3)
     logits = np.array(((45., 41., 44.),))
     labels = np.array(((35., 42., 45.),))
@@ -3584,7 +3569,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_weighted_multi_value_train(self):
     """3d label, 1 example, 1 batch."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
+    head = head_lib._regression_head(
         weight_column='label_weights', label_dimension=3)
     self.assertEqual(3, head.logits_dimension)
 
@@ -3639,8 +3624,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_weighted_multi_batch_eval(self):
     """1d label, 1 example, 3 batches."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        weight_column='label_weights')
+    head = head_lib._regression_head(weight_column='label_weights')
     self.assertEqual(1, head.logits_dimension)
 
     # Create estimator spec.
@@ -3705,8 +3689,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_weighted_multi_batch_train(self):
     """1d label, 1 example, 3 batches."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
-        weight_column='label_weights')
+    head = head_lib._regression_head(weight_column='label_weights')
     self.assertEqual(1, head.logits_dimension)
 
     # Create estimator spec.
@@ -3755,7 +3738,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
   def test_multi_dim_weighted_train_create_loss(self):
     """Logits, labels of shape [2, 2, 3], weight shape [2, 2]."""
     label_dimension = 3
-    head = head_lib._regression_head_with_mean_squared_error_loss(
+    head = head_lib._regression_head(
         weight_column='label_weights', label_dimension=label_dimension)
     logits = np.array([[[00., 01., 02.], [10., 11., 12.]],
                        [[20., 21., 22.], [30., 31., 32.]]])
@@ -3785,7 +3768,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_multi_dim_weighted_train(self):
     """Logits, labels of shape [2, 2, 3], weight shape [2, 2]."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
+    head = head_lib._regression_head(
         weight_column='label_weights', label_dimension=3)
     logits = np.array([[[00., 01., 02.], [10., 11., 12.]],
                        [[20., 21., 22.], [30., 31., 32.]]])
@@ -3816,7 +3799,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_multi_dim_train_weights_wrong_inner_dim(self):
     """Logits, labels of shape [2, 2, 3], weight shape [2, 1]."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
+    head = head_lib._regression_head(
         weight_column='label_weights', label_dimension=3)
     logits = np.array([[[00., 01., 02.], [10., 11., 12.]],
                        [[20., 21., 22.], [30., 31., 32.]]])
@@ -3844,7 +3827,7 @@ class RegressionHeadWithMeanSquaredErrorLossTest(test.TestCase):
 
   def test_multi_dim_train_weights_wrong_outer_dim(self):
     """Logits, labels of shape [2, 2, 3], weight shape [2, 2, 2]."""
-    head = head_lib._regression_head_with_mean_squared_error_loss(
+    head = head_lib._regression_head(
         weight_column='label_weights', label_dimension=3)
     logits = np.array([[[00., 01., 02.], [10., 11., 12.]],
                        [[20., 21., 22.], [30., 31., 32.]]])
diff --git a/tensorflow/python/estimator/canned/linear.py b/tensorflow/python/estimator/canned/linear.py
index e7ec417991..81657f0c01 100644
--- a/tensorflow/python/estimator/canned/linear.py
+++ b/tensorflow/python/estimator/canned/linear.py
@@ -415,7 +415,7 @@ class LinearRegressor(estimator.Estimator):
       loss_reduction: One of `tf.losses.Reduction` except `NONE`. Describes how
         to reduce training loss over batch. Defaults to `SUM`.
     """
-    head = head_lib._regression_head_with_mean_squared_error_loss(  # pylint: disable=protected-access
+    head = head_lib._regression_head(  # pylint: disable=protected-access
         label_dimension=label_dimension, weight_column=weight_column,
         loss_reduction=loss_reduction)
 
-- 
GitLab


From 1cc225858f9d7fb4d8772a7f0e962b71f780ad54 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Wed, 2 May 2018 11:12:58 -0700
Subject: [PATCH 151/395] Automated g4 rollback of changelist 194981511

PiperOrigin-RevId: 195120627
---
 tensorflow/core/common_runtime/device.h               | 11 -----------
 tensorflow/core/common_runtime/device_mgr.cc          |  3 ---
 .../process_function_library_runtime.cc               |  3 +--
 3 files changed, 1 insertion(+), 16 deletions(-)

diff --git a/tensorflow/core/common_runtime/device.h b/tensorflow/core/common_runtime/device.h
index b537666492..5918cd9bbf 100644
--- a/tensorflow/core/common_runtime/device.h
+++ b/tensorflow/core/common_runtime/device.h
@@ -51,8 +51,6 @@ limitations under the License.
 
 namespace tensorflow {
 
-class DeviceMgr;
-
 class Device : public DeviceBase {
  public:
   Device(Env* env, const DeviceAttributes& device_attributes);
@@ -135,10 +133,6 @@ class Device : public DeviceBase {
   // Returns the resource manager associated w/ this device.
   virtual ResourceMgr* resource_manager() { return rmgr_; }
 
-  // Returns the device manager that owns this device, or nullptr if this Device
-  // is not owned by a device manager.
-  DeviceMgr* device_mgr() const { return device_mgr_; }
-
   // Summarizes the status of this Device, for debugging.
   string DebugString() const { return ProtoDebugString(device_attributes_); }
 
@@ -164,11 +158,6 @@ class Device : public DeviceBase {
   }
 
  private:
-  friend class DeviceMgr;
-
-  // Pointer to the device manager that owns this device. Not owned.
-  DeviceMgr* device_mgr_ = nullptr;
-
   const DeviceAttributes device_attributes_;
   DeviceNameUtils::ParsedName parsed_name_;
 
diff --git a/tensorflow/core/common_runtime/device_mgr.cc b/tensorflow/core/common_runtime/device_mgr.cc
index 470abc1431..a77601ba79 100644
--- a/tensorflow/core/common_runtime/device_mgr.cc
+++ b/tensorflow/core/common_runtime/device_mgr.cc
@@ -27,9 +27,6 @@ namespace tensorflow {
 DeviceMgr::DeviceMgr(const std::vector<Device*>& devices)
     : name_backing_store_(128) {
   for (Device* d : devices) {
-    CHECK(d->device_mgr_ == nullptr);
-    d->device_mgr_ = this;
-
     devices_.push_back(d);
 
     // Register under the (1) full name and (2) canonical name.
diff --git a/tensorflow/core/common_runtime/process_function_library_runtime.cc b/tensorflow/core/common_runtime/process_function_library_runtime.cc
index 668ce87749..e61ed8c479 100644
--- a/tensorflow/core/common_runtime/process_function_library_runtime.cc
+++ b/tensorflow/core/common_runtime/process_function_library_runtime.cc
@@ -144,8 +144,7 @@ Status ProcessFunctionLibraryRuntime::GetDeviceContext(
   }
   Device* device = flr->device();
   string device_type = device->parsed_name().type;
-  if (device_type == "CPU" || device_type == "TPU_SYSTEM" ||
-      device_type == "TPU") {
+  if (device_type == "CPU" || device_type == "TPU_SYSTEM") {
     // "TPU_SYSTEM" indicates that `device` is a CPU.
     return Status::OK();
   }
-- 
GitLab


From 156483f1dc6b2e706482976f09f866d226a4dfee Mon Sep 17 00:00:00 2001
From: Benjamin Kramer <kramerb@google.com>
Date: Wed, 2 May 2018 11:17:47 -0700
Subject: [PATCH 152/395] [XLA:GPU] Unroll unfused elementwise op kernels.

So far we only unrolled loop fusions, elementwise ops is a logical extension.
We don't spend a lot of time in unfused elementwise ops in benchmarks, so this
is only worth a small speedup on V100.

PiperOrigin-RevId: 195121530
---
 .../xla/service/gpu/ir_emitter_unnested.cc    | 46 +++++++++++++------
 .../compiler/xla/tests/hlo_test_base.cc       |  5 +-
 2 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
index 26e497762f..9f37235d32 100644
--- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
+++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
@@ -257,8 +257,36 @@ llvm::Function* IrEmitterUnnested::BuildKernelPrototype(
   return kernel;
 }
 
+namespace {
+// Computes the maximum valid unroll factor for a given instruction.
+int ComputeMaxUnrollFactor(const HloInstruction* hlo) {
+  int max_unroll_factor = hlo->GetModule()
+                              ->config()
+                              .debug_options()
+                              .xla_gpu_max_kernel_unroll_factor();
+
+  // Find the largest possible power of two to unroll by.
+  // TODO(kramerb): Make this smarter.
+  int64 num_elements = ShapeUtil::ElementsIn(hlo->shape());
+  for (int i = max_unroll_factor; i > 1; i /= 2) {
+    if (num_elements % i == 0) {
+      return i;
+    }
+  }
+
+  // Cannot unroll.
+  return 1;
+}
+}  // namespace
+
 Status IrEmitterUnnested::DefaultAction(HloInstruction* hlo) {
-  thunk_sequence_->emplace_back(BuildKernelThunk(hlo));
+  int unroll_factor = 1;
+  // Unfused elementwise operations are usually memory bound, unroll them.
+  if (hlo->IsElementwise()) {
+    unroll_factor = ComputeMaxUnrollFactor(hlo);
+  }
+
+  thunk_sequence_->emplace_back(BuildKernelThunk(hlo, unroll_factor));
   return IrEmitter::DefaultAction(hlo);
 }
 
@@ -537,23 +565,11 @@ Status IrEmitterUnnested::HandleFusion(HloInstruction* fusion) {
     return Status::OK();
   }
 
-  int max_unroll_factor = fusion->GetModule()
-                              ->config()
-                              .debug_options()
-                              .xla_gpu_max_kernel_unroll_factor();
-
-  // Find the largest possible power of two to unroll by.
-  // TODO(kramerb): Make this smarter.
   int unroll_factor = 1;
+  // TODO(kramerb): Unrolling multi-output loop fusions too.
   if (!fusion->IsMultiOutputFusion()) {
     CHECK(fusion->fusion_kind() == HloInstruction::FusionKind::kLoop);
-    int64 num_elements = ShapeUtil::ElementsIn(fusion->shape());
-    for (int i = max_unroll_factor; i > 1; i /= 2) {
-      if (num_elements % i == 0) {
-        unroll_factor = i;
-        break;
-      }
-    }
+    unroll_factor = ComputeMaxUnrollFactor(fusion);
   }
 
   thunk_sequence_->emplace_back(BuildKernelThunk(fusion, unroll_factor));
diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.cc b/tensorflow/compiler/xla/tests/hlo_test_base.cc
index 8b64f2e631..12598579c7 100644
--- a/tensorflow/compiler/xla/tests/hlo_test_base.cc
+++ b/tensorflow/compiler/xla/tests/hlo_test_base.cc
@@ -95,7 +95,10 @@ HloTestBase::HloTestBase(se::Platform* test_platform,
 /* static */
 std::unique_ptr<HloModule> HloTestBase::CreateNewModule(const string& name) {
   HloModuleConfig config;
-  config.set_debug_options(GetDebugOptionsForTest());
+  auto debug_options = HloTestBase::GetDebugOptionsForTest();
+  debug_options.set_xla_gpu_max_kernel_unroll_factor(1);
+  config.set_debug_options(debug_options);
+
   return MakeUnique<HloModule>(name, VersionedComputationHandle(), config);
 }
 
-- 
GitLab


From df9bb02c647e395dbb2da393f7de085320e7c5c9 Mon Sep 17 00:00:00 2001
From: gracehoney <31743510+aaroey@users.noreply.github.com>
Date: Wed, 2 May 2018 11:21:02 -0700
Subject: [PATCH 153/395] Fix formatting and linter issues

---
 tensorflow/contrib/tensorrt/BUILD             |  5 ++-
 .../contrib/tensorrt/convert/convert_graph.cc | 20 +++------
 .../contrib/tensorrt/convert/convert_nodes.cc | 20 ++++-----
 .../tensorrt/convert/trt_optimization_pass.cc |  5 +--
 .../tensorrt/convert/trt_optimization_pass.h  | 15 +++++--
 .../contrib/tensorrt/kernels/trt_engine_op.cc | 14 +++---
 .../contrib/tensorrt/kernels/trt_engine_op.h  |  8 ++--
 .../tensorrt/resources/trt_allocator.cc       |  3 ++
 .../tensorrt/resources/trt_allocator.h        | 23 +++++-----
 .../tensorrt/resources/trt_resources.h        | 45 ++++++++++---------
 .../contrib/tensorrt/segment/segment.cc       | 31 +++++++------
 tensorflow/contrib/tensorrt/segment/segment.h |  2 +-
 12 files changed, 98 insertions(+), 93 deletions(-)

diff --git a/tensorflow/contrib/tensorrt/BUILD b/tensorflow/contrib/tensorrt/BUILD
index 1792fa310a..675f0b1fd6 100644
--- a/tensorflow/contrib/tensorrt/BUILD
+++ b/tensorflow/contrib/tensorrt/BUILD
@@ -88,7 +88,6 @@ cc_library(
         ":trt_logging",
         ":trt_resources",
         "//tensorflow/core:gpu_headers_lib",
-        "//tensorflow/core:gpu_runtime",
         "//tensorflow/core:lib_proto_parsing",
         "//tensorflow/core:stream_executor_headers_lib",
     ] + if_tensorrt([
@@ -212,7 +211,6 @@ tf_cuda_library(
         ":trt_logging",
         "//tensorflow/core:framework_headers_lib",
         "//tensorflow/core:framework_lite",
-        "//tensorflow/core:core_cpu_lib",
         "//tensorflow/core:lib_proto_parsing",
     ] + if_tensorrt([
         "@local_config_tensorrt//:nv_infer",
@@ -236,6 +234,9 @@ tf_cuda_library(
         ":segment",
         ":trt_logging",
         ":trt_resources",
+        "//tensorflow/core/grappler/clusters:cluster",
+        "//tensorflow/core/grappler/optimizers:custom_graph_optimizer",
+        "//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry",
         "//tensorflow/core/grappler:grappler_item",
         "//tensorflow/core/grappler:utils",
         "//tensorflow/core:framework",
diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.cc b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
index 8459ad4a61..a8c07df4a0 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_graph.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
@@ -30,13 +30,10 @@ limitations under the License.
 #include "tensorflow/core/graph/algorithm.h"
 #include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/graph/graph_constructor.h"
-#include "tensorflow/core/grappler/clusters/utils.h"
 #include "tensorflow/core/grappler/clusters/virtual_cluster.h"
 #include "tensorflow/core/grappler/costs/graph_properties.h"
 #include "tensorflow/core/grappler/devices.h"
 #include "tensorflow/core/grappler/grappler_item.h"
-#include "tensorflow/core/grappler/optimizers/constant_folding.h"
-#include "tensorflow/core/grappler/optimizers/layout_optimizer.h"
 #include "tensorflow/core/grappler/optimizers/meta_optimizer.h"
 #include "tensorflow/core/grappler/utils.h"
 #include "tensorflow/core/lib/core/errors.h"
@@ -206,7 +203,7 @@ static tensorflow::Status FillSubGraphEdgeSets(ConvertGraphParams* p) {
                              subgraph_outputs_set.begin(),
                              subgraph_outputs_set.end());
   return tensorflow::Status::OK();
-};
+}
 
 tensorflow::Status GetCalibNode(ConvertGraphParams* params) {
   TF_RETURN_IF_ERROR(FillSubGraphEdgeSets(params));
@@ -345,18 +342,11 @@ tensorflow::Status ConvertGraphDefToTensorRT(
   // optimization pass
   tensorflow::grappler::GrapplerItem item;
   item.fetch = output_names;
-  tensorflow::GraphDef gdef;
 
-  // Layout optimization
-  item.graph = graph_def;
-  tensorflow::grappler::Cluster* cluster;
-
-  // virtual cluster
   tensorflow::DeviceProperties device_properties;
-
   device_properties.set_type("GPU");
   device_properties.mutable_environment()->insert({"architecture", "6"});
-  cluster =
+  tensorflow::grappler::Cluster* cluster =
       new tensorflow::grappler::VirtualCluster({{"/GPU:0", device_properties}});
 
   // single machine
@@ -366,6 +356,7 @@ tensorflow::Status ConvertGraphDefToTensorRT(
   VLOG(2) << "gpus: " << num_gpus;
   tensorflow::RewriterConfig rw_cfg;
   tensorflow::grappler::MetaOptimizer meta_opt(nullptr, rw_cfg);
+  tensorflow::GraphDef gdef;
   TF_RETURN_IF_ERROR(meta_opt.Optimize(cluster, item, &gdef));
   item.graph = gdef;
 
@@ -416,9 +407,12 @@ tensorflow::Status ConvertAfterShapes(
   for (auto s : segments) {
     total_num_nodes_in_segments += s.first.size();
   }
-  // We are creating the map here since cluster may not be available in all cases
+  // We create the map here since cluster may not be available in all cases.
   std::map<string, tensorflow::Device*> name_to_device_map;
   if (cluster) {
+    // TODO(aaroey): consider using DeviceSet::FindDeviceByName(), as in a
+    // distributed environment, devices from different workers can have same
+    // short name.
     for (const auto dm : cluster->GetDeviceSet()->devices()) {
       name_to_device_map[dm->name()] = dm;
     }
diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
index ae0e861be5..4d3710a514 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
@@ -482,7 +482,7 @@ class Converter {
     weights.SetValues(weight_store_->store_.back().data());
     return weights;
   }
-  bool isFP16() { return fp16_; };
+  bool isFP16() { return fp16_; }
   TRT_ShapedWeights get_temp_weights_like(const TRT_ShapedWeights& weights) {
     return this->get_temp_weights(weights.type_, weights.shape_);
   }
@@ -673,7 +673,7 @@ std::function<Eigen::half(Eigen::half)> LambdaFactory::unary<Eigen::half>() {
     case OP_CATEGORY::RSQRT: {
       VLOG(2) << "RSQRT GETS DONE";
       return [](Eigen::half t) -> Eigen::half {
-        return Eigen::half(1.0 / sqrt(float(t)));
+        return Eigen::half(1.0 / sqrt(static_cast<float>(t)));
       };
     }
     case OP_CATEGORY::NEG:
@@ -2328,8 +2328,8 @@ tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) {
             << ", at node: " << node_name
             << "with output entry from shape_map: " << op_info_vec.size();
     // TODO(ben,jie): update TRT input format/dimension
-    nvinfer1::DimsCHW input_dim_psuedo_chw;
-    for (int i = 0; i < 3; i++) input_dim_psuedo_chw.d[i] = 1;
+    nvinfer1::DimsCHW input_dim_pseudo_chw;
+    for (int i = 0; i < 3; i++) input_dim_pseudo_chw.d[i] = 1;
 
     // TODO(jie): TRT 3.x only support 4 dimensional input tensor.
     //            update the code once TRT 4.0 comes out.
@@ -2343,7 +2343,7 @@ tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) {
     for (int i = 1; i < op_info.shape().dim_size(); i++) {
       VLOG(2) << "dimension: " << i
               << " , size: " << op_info.shape().dim(i).size();
-      input_dim_psuedo_chw.d[i - 1] = op_info.shape().dim(i).size();
+      input_dim_pseudo_chw.d[i - 1] = op_info.shape().dim(i).size();
     }
 
     // TODO(ben,jie): proper way to restore input tensor name?
@@ -2354,7 +2354,7 @@ tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) {
 
     input_names.push_back(input_tensor_name);
     nvinfer1::ITensor* input_tensor = converter.network()->addInput(
-        input_tensor_name.c_str(), dtype, input_dim_psuedo_chw);
+        input_tensor_name.c_str(), dtype, input_dim_pseudo_chw);
 
     if (!input_tensor)
       return tensorflow::errors::InvalidArgument(
@@ -2572,8 +2572,8 @@ tensorflow::Status ConvertSubGraphToTensorRTNodeDef(
             << ", at node: " << node_name
             << " with output entry from shape_map: " << op_info_vec.size();
     // TODO(ben,jie): update TRT input format/dimension
-    nvinfer1::DimsCHW input_dim_psuedo_chw;
-    for (int i = 0; i < 3; i++) input_dim_psuedo_chw.d[i] = 1;
+    nvinfer1::DimsCHW input_dim_pseudo_chw;
+    for (int i = 0; i < 3; i++) input_dim_pseudo_chw.d[i] = 1;
 
     // TODO(jie): TRT 3.x only support 4 dimensional input tensor.
     //            update the code once TRT 4.0 comes out.
@@ -2587,7 +2587,7 @@ tensorflow::Status ConvertSubGraphToTensorRTNodeDef(
     for (int i = 1; i < op_info.shape().dim_size(); i++) {
       VLOG(2) << "dimension: " << i
               << " , size: " << op_info.shape().dim(i).size();
-      input_dim_psuedo_chw.d[i - 1] = op_info.shape().dim(i).size();
+      input_dim_pseudo_chw.d[i - 1] = op_info.shape().dim(i).size();
     }
 
     // TODO(ben,jie): proper way to restore input tensor name?
@@ -2598,7 +2598,7 @@ tensorflow::Status ConvertSubGraphToTensorRTNodeDef(
 
     input_names.push_back(input_tensor_name);
     nvinfer1::ITensor* input_tensor = converter.network()->addInput(
-        input_tensor_name.c_str(), dtype, input_dim_psuedo_chw);
+        input_tensor_name.c_str(), dtype, input_dim_pseudo_chw);
 
     if (!input_tensor)
       return tensorflow::errors::InvalidArgument(
diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
index 21013fbf9e..8f634b1f74 100644
--- a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
+++ b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc
@@ -20,7 +20,6 @@ limitations under the License.
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/logging.h"
-#include "tensorflow/core/public/session_options.h"
 
 #if GOOGLE_CUDA
 #if GOOGLE_TENSORRT
@@ -64,7 +63,7 @@ tensorflow::Status TRTOptimizationPass::Init(
     }
   }
   return tensorflow::Status::OK();
-};
+}
 
 void TRTOptimizationPass::PrintDebugInfo(
     tensorflow::grappler::Cluster* cluster,
@@ -218,8 +217,6 @@ void TRTOptimizationPass::Feedback(
     const tensorflow::grappler::GrapplerItem& item,
     const GraphDef& optimized_graph, double result) {}
 
-
-
 }  // namespace convert
 }  // namespace tensorrt
 }  // namespace tensorflow
diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
index c554a5d784..d8ecead23e 100644
--- a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
+++ b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h
@@ -28,6 +28,7 @@ limitations under the License.
 namespace tensorflow {
 namespace tensorrt {
 namespace convert {
+
 class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer {
  public:
   TRTOptimizationPass(const string& name = "TRTOptimizationPass")
@@ -37,17 +38,21 @@ class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer {
         maximum_batch_size_(-1),
         maximum_workspace_size_(-1) {
     VLOG(1) << "Constructing " << name_;
-  };
+  }
+
   string name() const override { return name_; };
+
   tensorflow::Status Init(const tensorflow::RewriterConfig_CustomGraphOptimizer*
                               config = nullptr) override;
 
   tensorflow::Status Optimize(tensorflow::grappler::Cluster* cluster,
                               const tensorflow::grappler::GrapplerItem& item,
                               GraphDef* optimized_graph) override;
+
   void Feedback(tensorflow::grappler::Cluster* cluster,
                 const tensorflow::grappler::GrapplerItem& item,
                 const GraphDef& optimized_graph, double result) override;
+
   void PrintDebugInfo(tensorflow::grappler::Cluster* cluster,
                       const tensorflow::grappler::GrapplerItem& item);
 
@@ -58,9 +63,11 @@ class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer {
   int maximum_batch_size_;
   int64_t maximum_workspace_size_;
 };
+
 }  // namespace convert
 }  // namespace tensorrt
 }  // namespace tensorflow
-#endif
-#endif
-#endif
+
+#endif  // GOOGLE_CUDA
+#endif  // GOOGLE_TENSORRT
+#endif  // TENSORFLOW_CONTRIB_TENSORRT_CONVERT_TRT_OPTIMIZATION_PASS_H_
diff --git a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
index f10b10edec..5c5b2e3c07 100644
--- a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
+++ b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc
@@ -38,7 +38,6 @@ TRTEngineOp::TRTEngineOp(OpKernelConstruction* context) : OpKernel(context) {
   // register input output node name in trt_sub_graph
   OP_REQUIRES_OK(context, context->GetAttr("input_nodes", &input_nodes_));
   OP_REQUIRES_OK(context, context->GetAttr("output_nodes", &output_nodes_));
-
 }
 
 void TRTEngineOp::Compute(OpKernelContext* context) {
@@ -49,15 +48,12 @@ void TRTEngineOp::Compute(OpKernelContext* context) {
   if (!trt_execution_context_ptr_) {
     IRuntime* infer = nvinfer1::createInferRuntime(logger);
 #if NV_TENSORRT_MAJOR > 3
-    auto device=context->device();
-    auto dev_allocator=device->getAllocator(tensorflow::AllocatorAttributes())
-    // tensorflow::TfGpuId tf_gpu_id(
-    //     context->device()->tensorflow_gpu_device_info()->gpu_id);
-    // tensorflow::GPUOptions gpuoptions;
-    // auto pm = tensorflow::ProcessState::singleton();
-    // auto dev_allocator = pm->GetGPUAllocator(gpuoptions, tf_gpu_id, 1);
+    auto device = context->device();
+    auto dev_allocator =
+        device->GetAllocator(tensorflow::AllocatorAttributes());
     if (!dev_allocator) {
-      LOG(FATAL) << "Can't find device allocator for gpu device" << tf_gpu_id;
+      LOG(FATAL) << "Can't find device allocator for gpu device "
+                 << device->name();
     }
     allocator_ = std::make_shared<TRTDeviceAllocator>(dev_allocator);
     infer->setGpuAllocator(allocator_.get());
diff --git a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h
index fec4bd728b..e613a71422 100644
--- a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h
+++ b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h
@@ -17,15 +17,15 @@ limitations under the License.
 #define TENSORFLOW_CONTRIB_TENSORRT_KERNELS_TRT_ENGINE_OP_H_
 
 #include <memory>
-#include <string>
 #include <vector>
 
-#if GOOGLE_CUDA
-#if GOOGLE_TENSORRT
-#include "cuda/include/cuda_runtime_api.h"
 #include "tensorflow/contrib/tensorrt/resources/trt_allocator.h"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/op_kernel.h"
+
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
+#include "cuda/include/cuda_runtime_api.h"
 #include "tensorrt/include/NvInfer.h"
 
 namespace tensorflow {
diff --git a/tensorflow/contrib/tensorrt/resources/trt_allocator.cc b/tensorflow/contrib/tensorrt/resources/trt_allocator.cc
index b94f8a2da7..0f0508331c 100644
--- a/tensorflow/contrib/tensorrt/resources/trt_allocator.cc
+++ b/tensorflow/contrib/tensorrt/resources/trt_allocator.cc
@@ -16,8 +16,10 @@ limitations under the License.
 #include "tensorflow/contrib/tensorrt/resources/trt_allocator.h"
 
 #include "tensorflow/core/platform/logging.h"
+
 #if GOOGLE_CUDA
 #if GOOGLE_TENSORRT
+
 #if NV_TENSORRT_MAJOR > 2
 #include "cuda/include/cuda_runtime_api.h"
 
@@ -54,6 +56,7 @@ void TRTDeviceAllocator::free(void* memory) {
 
 }  // namespace tensorrt
 }  // namespace tensorflow
+
 #endif
 #endif
 #endif
diff --git a/tensorflow/contrib/tensorrt/resources/trt_allocator.h b/tensorflow/contrib/tensorrt/resources/trt_allocator.h
index dd4f8c7943..a0c2540a76 100644
--- a/tensorflow/contrib/tensorrt/resources/trt_allocator.h
+++ b/tensorflow/contrib/tensorrt/resources/trt_allocator.h
@@ -16,35 +16,34 @@ limitations under the License.
 #ifndef TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_ALLOCATOR_H_
 #define TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_ALLOCATOR_H_
 
-#include <list>
-#include <sstream>
-#include <string>
-#include <thread>
-#include <vector>
 
 #include "tensorflow/contrib/tensorrt/log/trt_logger.h"
 #include "tensorflow/core/framework/allocator.h"
-#include "tensorflow/core/framework/resource_mgr.h"
+
 #if GOOGLE_CUDA
 #if GOOGLE_TENSORRT
 #include "tensorrt/include/NvInfer.h"
+
 #if NV_TENSORRT_MAJOR == 3
-// define interface here temporarily until TRT 4.0 is released
+// Define interface here temporarily until TRT 4.0 is released
 namespace nvinfer1 {
 class IGpuAllocator {
+ public:
   virtual void* allocate(uint64_t size, uint64_t alignment, uint32_t flags) = 0;
   virtual void free(void* memory) = 0;
 };
 }  // namespace nvinfer1
 #endif
+
 namespace tensorflow {
 namespace tensorrt {
+
 class TRTCudaAllocator : public nvinfer1::IGpuAllocator {
   // Allocator implementation that is using cuda allocator instead of device
   // allocator in case we can't get device allocator from TF.
  public:
   TRTCudaAllocator() {}
-  virtual ~TRTCudaAllocator(){};
+  virtual ~TRTCudaAllocator() {}
   void* allocate(uint64_t size, uint64_t alignment, uint32_t flags) override;
   void free(void* memory) override;
 };
@@ -53,7 +52,7 @@ class TRTDeviceAllocator : public nvinfer1::IGpuAllocator {
   // Allocator implementation wrapping TF device allocators.
  public:
   TRTDeviceAllocator(tensorflow::Allocator* allocator);
-  virtual ~TRTDeviceAllocator(){};
+  virtual ~TRTDeviceAllocator() {}
   void* allocate(uint64_t size, uint64_t alignment, uint32_t flags) override;
   void free(void* memory) override;
 
@@ -64,6 +63,6 @@ class TRTDeviceAllocator : public nvinfer1::IGpuAllocator {
 }  // namespace tensorrt
 }  // namespace tensorflow
 
-#endif
-#endif
-#endif
+#endif  // GOOGLE_TENSORRT
+#endif  // GOOGLE_CUDA
+#endif  // TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_ALLOCATOR_H_
diff --git a/tensorflow/contrib/tensorrt/resources/trt_resources.h b/tensorflow/contrib/tensorrt/resources/trt_resources.h
index 166ca9c3de..e3469124ac 100644
--- a/tensorflow/contrib/tensorrt/resources/trt_resources.h
+++ b/tensorflow/contrib/tensorrt/resources/trt_resources.h
@@ -13,22 +13,24 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#ifndef TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRTRESOURCES_H_
-#define TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRTRESOURCES_H_
+#ifndef TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_RESOURCES_H_
+#define TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_RESOURCES_H_
 
 #include <list>
 #include <sstream>
 #include <string>
 #include <thread>
 #include <vector>
+
 #include "tensorflow/contrib/tensorrt/log/trt_logger.h"
+#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h"
+#include "tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h"
 #include "tensorflow/core/framework/resource_mgr.h"
 
 #if GOOGLE_CUDA
 #if GOOGLE_TENSORRT
-#include "tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h"
+
 #include "tensorrt/include/NvInfer.h"
-#include "tensorflow/contrib/tensorrt/resources/trt_allocator.h"
 
 namespace tensorflow {
 namespace tensorrt {
@@ -41,6 +43,11 @@ class TRTCalibrationResource : public tensorflow::ResourceBase {
         engine_(nullptr),
         logger_(nullptr),
         thr_(nullptr) {}
+
+  ~TRTCalibrationResource() {
+    VLOG(0) << "Destroying Calibration Resource " << std::endl << DebugString();
+  }
+
   string DebugString() override {
     std::stringstream oss;
     oss << " Calibrator = " << std::hex << calibrator_ << std::dec << std::endl
@@ -48,13 +55,12 @@ class TRTCalibrationResource : public tensorflow::ResourceBase {
         << " Network    = " << std::hex << network_ << std::dec << std::endl
         << " Engine     = " << std::hex << engine_ << std::dec << std::endl
         << " Logger     = " << std::hex << logger_ << std::dec << std::endl
-        << " Allocator  = " << std::hex << allocator_.get()<< std::dec << std::endl
+        << " Allocator  = " << std::hex << allocator_.get() << std::dec
+        << std::endl
         << " Thread     = " << std::hex << thr_ << std::dec << std::endl;
     return oss.str();
   }
-  ~TRTCalibrationResource() {
-    VLOG(0) << "Destroying Calibration Resource " << std::endl << DebugString();
-  }
+
   TRTInt8Calibrator* calibrator_;
   nvinfer1::IBuilder* builder_;
   nvinfer1::INetworkDefinition* network_;
@@ -68,31 +74,28 @@ class TRTCalibrationResource : public tensorflow::ResourceBase {
 class TRTWeightStore : public tensorflow::ResourceBase {
  public:
   TRTWeightStore() {}
-  std::list<std::vector<uint8_t>> store_;
+
+  virtual ~TRTWeightStore() { VLOG(1) << "Destroying store" << DebugString(); }
+
   string DebugString() override {
     std::stringstream oss;
-    size_t lenBytes = 0;
+    size_t len_bytes = 0;
     for (const auto& v : store_) {
-      lenBytes += v.size() * sizeof(uint8_t);
+      len_bytes += v.size() * sizeof(uint8_t);
     }
     oss << " Number of entries     = " << store_.size() << std::endl
         << " Total number of bytes = "
-        << store_.size() * sizeof(std::vector<uint8_t>) + lenBytes << std::endl;
+        << store_.size() * sizeof(std::vector<uint8_t>) + len_bytes
+        << std::endl;
     return oss.str();
   }
-  virtual ~TRTWeightStore() { VLOG(1) << "Destroying store" << DebugString(); }
-};
 
-class TRTEngineResource : public tensorflow::ResourceBase {
- public:
-  TRTEngineResource() : runtime_(nullptr), ctx_(nullptr){};
-  string DebugString() override { return string(""); }
-  nvinfer1::IRuntime* runtime_;
-  nvinfer1::IExecutionContext* ctx_;
+  std::list<std::vector<uint8_t>> store_;
 };
 
 }  // namespace tensorrt
 }  // namespace tensorflow
-#endif  // TENSORFLOW_CONTRIB_TENSORRT_RESOURCEMGR_TRTRESOURCES_H_
+
 #endif
 #endif
+#endif  // TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_RESOURCES_H_
diff --git a/tensorflow/contrib/tensorrt/segment/segment.cc b/tensorflow/contrib/tensorrt/segment/segment.cc
index 4901e30a87..cc42913eca 100644
--- a/tensorflow/contrib/tensorrt/segment/segment.cc
+++ b/tensorflow/contrib/tensorrt/segment/segment.cc
@@ -47,14 +47,15 @@ class SimpleEdge {
         src_port_(src_port),
         dst_(dst),
         dst_port_(dst_port),
-        control_(is_control){};
+        control_(is_control) {}
+  ~SimpleEdge() {}
+
   SimpleNode* src() const { return src_; }
   SimpleNode* dst() const { return dst_; }
   int src_output() const { return src_port_; }
   int dst_input() const { return dst_port_; }
   int id() const { return id_; }
   bool IsControlEdge() const { return control_; }
-  ~SimpleEdge() {}
 
  private:
   int id_;
@@ -64,11 +65,13 @@ class SimpleEdge {
   int dst_port_;
   bool control_;
 };
+
 class SimpleNode {
  public:
   SimpleNode(const tensorflow::Node* node, const int id);
-  const std::vector<SimpleEdge*>& in_edges() const { return in_edges_; };
-  const std::vector<SimpleEdge*>& out_edges() const { return out_edges_; };
+
+  const std::vector<SimpleEdge*>& in_edges() const { return in_edges_; }
+  const std::vector<SimpleEdge*>& out_edges() const { return out_edges_; }
   std::vector<SimpleNode*> in_nodes() const {
     std::vector<SimpleNode*> res;
     res.reserve(in_edges_.size());
@@ -92,15 +95,18 @@ class SimpleNode {
 
 class SimpleGraph {
  public:
-  SimpleGraph(const tensorflow::Graph* g);
+  explicit SimpleGraph(const tensorflow::Graph* g);
+  ~SimpleGraph();
+
   void AddControlEdge(SimpleNode* src, SimpleNode* dst);
   void AddEdge(SimpleNode* src, int out_port, SimpleNode* dst, int in_port);
   void RemoveEdge(const SimpleEdge*);
   SimpleNode* FindNodeId(int node_id) {
-    if (node_id < 0 || node_id > (int)nodes_.size()) return nullptr;
+    if (node_id < 0 || node_id > static_cast<int>(nodes_.size())) {
+      return nullptr;
+    }
     return nodes_[node_id];
   }
-  ~SimpleGraph();
   int num_node_ids() const { return nodes_.size(); }
   const SimpleNode* source_node() const {
     return nodes_[tensorflow::Graph::kSourceId];
@@ -163,7 +169,7 @@ SimpleGraph::SimpleGraph(const tensorflow::Graph* g) : g_(g) {
 void SimpleGraph::AddEdge(SimpleNode* src, int out_port, SimpleNode* dst,
                           int in_port) {
   int i = edges_.size();
-  if (free_edge_ids_.size()) {
+  if (!free_edge_ids_.empty()) {
     auto it = free_edge_ids_.begin();
     i = *it;
     free_edge_ids_.erase(it);
@@ -275,7 +281,7 @@ bool CanContractEdge(const SimpleEdge* edge,
 }
 }  // namespace
 
-void ContractEdge(SimpleEdge* edge, std::unique_ptr<SimpleGraph>& graph,
+void ContractEdge(SimpleEdge* edge, SimpleGraph* graph,
                   std::vector<const SimpleEdge*>* remove_edges) {
   // Transfer all inputs and outputs of 'dst' to 'src' except edges
   // connecting the two.
@@ -352,7 +358,6 @@ tensorflow::Status SegmentGraph(
     tensorflow::Graph* tf_graph,
     const std::function<bool(const tensorflow::Node*)>& candidate_fn,
     const SegmentOptions& options, SegmentNodesVector* segments) {
-
   auto graph = std::unique_ptr<SimpleGraph>(new SimpleGraph(tf_graph));
   // Use a union-find to collect the nodes that belong to the same
   // segment. A node value of nullptr indicates that the node is not a candidate
@@ -440,7 +445,7 @@ tensorflow::Status SegmentGraph(
         // don't visit them again.
         SimpleEdge* e = const_cast<SimpleEdge*>(contract_edge);
         std::vector<const SimpleEdge*> remove_edges;
-        ContractEdge(e, graph, &remove_edges);
+        ContractEdge(e, graph.get(), &remove_edges);
 
         for (const SimpleEdge* r : remove_edges) {
           contract_edges.erase(r);
@@ -466,7 +471,7 @@ tensorflow::Status SegmentGraph(
       if (tf_node->has_assigned_device_name()) {
         device_maps[u.ParentValue()->name()].insert(
             tf_node->assigned_device_name());
-      } else if (tf_node->requested_device().size() > 0) {
+      } else if (!tf_node->requested_device().empty()) {
         device_maps[u.ParentValue()->name()].insert(
             tf_node->requested_device());
       } else {
@@ -497,7 +502,7 @@ tensorflow::Status SegmentGraph(
     }
     // TODO(sami): Make segmenter placement aware once trtscopes are in place
     const auto& dev_itr = device_maps.find(itr.first);
-    if (dev_itr == device_maps.end() || dev_itr->second.size() == 0) {
+    if (dev_itr == device_maps.end() || dev_itr->second.empty()) {
       VLOG(1) << "No device assigned to segment " << segments->size();
       segments->emplace_back(std::make_pair(segment_node_names, string()));
     } else if (dev_itr->second.size() > 1) {
diff --git a/tensorflow/contrib/tensorrt/segment/segment.h b/tensorflow/contrib/tensorrt/segment/segment.h
index c5aca4bf04..1568dd9153 100644
--- a/tensorflow/contrib/tensorrt/segment/segment.h
+++ b/tensorflow/contrib/tensorrt/segment/segment.h
@@ -63,7 +63,7 @@ tensorflow::Status SegmentGraph(
 // all the NodeDefs in that subgraph.
 // @return the status.
 tensorflow::Status SegmentGraph(
-    tensorflow::Graph* graph,
+    tensorflow::Graph* tf_graph,
     const std::function<bool(const tensorflow::Node*)>& candidate_fn,
     const SegmentOptions& options, SegmentNodesVector* segments);
 
-- 
GitLab


From fb820662625d30e7e137580dae142a3eaf933335 Mon Sep 17 00:00:00 2001
From: Anna R <annarev@google.com>
Date: Wed, 2 May 2018 11:38:44 -0700
Subject: [PATCH 154/395] Copy module list before iterating over it. Also,
 import python module for clarity

---
 tensorflow/tools/api/generator/create_python_api.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensorflow/tools/api/generator/create_python_api.py b/tensorflow/tools/api/generator/create_python_api.py
index c06a39bfbd..d1e7f23fbc 100644
--- a/tensorflow/tools/api/generator/create_python_api.py
+++ b/tensorflow/tools/api/generator/create_python_api.py
@@ -23,6 +23,7 @@ import collections
 import os
 import sys
 
+from tensorflow import python  # pylint: disable=unused-import
 from tensorflow.python.util import tf_decorator
 
 
@@ -158,7 +159,8 @@ def get_api_init_text():
 
   # Traverse over everything imported above. Specifically,
   # we want to traverse over TensorFlow Python modules.
-  for module in sys.modules.values():
+  module_list = list(sys.modules.values())
+  for module in module_list:
     # Only look at tensorflow modules.
     if (not module or not hasattr(module, '__name__') or
         'tensorflow.' not in module.__name__):
-- 
GitLab


From 9b6cba1d739e36ec2da59a593afb09bf17307650 Mon Sep 17 00:00:00 2001
From: Justin Lebar <jlebar@google.com>
Date: Wed, 2 May 2018 11:40:09 -0700
Subject: [PATCH 155/395] Internal-only change.

PiperOrigin-RevId: 195125476
---
 tensorflow/python/kernel_tests/linalg/BUILD | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/kernel_tests/linalg/BUILD b/tensorflow/python/kernel_tests/linalg/BUILD
index 6573cb9a1a..052f11f92e 100644
--- a/tensorflow/python/kernel_tests/linalg/BUILD
+++ b/tensorflow/python/kernel_tests/linalg/BUILD
@@ -62,7 +62,10 @@ cuda_py_test(
         "//tensorflow/python:platform_test",
     ],
     shard_count = 5,
-    tags = ["noasan"],  # times out b/63678675
+    tags = [
+        "noasan",  # times out, b/63678675
+        "optonly",  # times out
+    ],
 )
 
 cuda_py_test(
-- 
GitLab


From 1ea4a77c6ccd2c783aedb2ccaf76f46b018c12c5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 11:45:15 -0700
Subject: [PATCH 156/395] Replaced calls to tensorflow::StringPiece::ToString
 with std::string conversions. That is, instances of sp.ToString() are
 replaced with std::string(sp).

This will allow tensorflow::StringPiece::ToString to be removed, which is necessary before it can be replaced with absl::string_view.

PiperOrigin-RevId: 195126422
---
 tensorflow/core/kernels/gpu_utils.h           |  3 +-
 .../kernels/merge_v2_checkpoints_op_test.cc   |  2 +-
 .../remote_fused_graph_execute_utils.cc       | 32 ++++----
 .../core/kernels/save_restore_v2_ops.cc       |  4 +-
 tensorflow/core/kernels/string_strip_op.cc    |  2 +-
 tensorflow/core/kernels/tensor_array_ops.cc   |  2 +-
 .../core/kernels/whole_file_read_ops.cc       |  2 +-
 tensorflow/core/lib/strings/numbers.h         |  4 +-
 tensorflow/core/lib/strings/scanner_test.cc   | 82 +++++++++----------
 tensorflow/core/lib/strings/str_util.cc       |  4 +-
 tensorflow/core/lib/strings/str_util.h        |  2 +-
 tensorflow/core/platform/env.cc               |  4 +-
 tensorflow/core/platform/env_test.cc          |  2 +-
 tensorflow/core/platform/file_system.cc       |  2 +-
 .../core/platform/file_system_helper.cc       |  2 +-
 tensorflow/core/platform/file_system_test.cc  |  2 +-
 tensorflow/core/util/command_line_flags.cc    |  2 +-
 tensorflow/core/util/env_var.cc               |  8 +-
 .../core/util/example_proto_fast_parsing.cc   |  2 +-
 19 files changed, 82 insertions(+), 81 deletions(-)

diff --git a/tensorflow/core/kernels/gpu_utils.h b/tensorflow/core/kernels/gpu_utils.h
index 2f64619afc..c7dbefa0b4 100644
--- a/tensorflow/core/kernels/gpu_utils.h
+++ b/tensorflow/core/kernels/gpu_utils.h
@@ -123,7 +123,8 @@ class AutoTuneMap {
   string GetActionSummary(StringPiece action, const Parameters& params,
                           const Config& config) {
     return strings::Printf("autotune_map %s %s: %s -> (%s)", name_.c_str(),
-                           action.ToString().c_str(), params.ToString().c_str(),
+                           std::string(action).c_str(),
+                           params.ToString().c_str(),
                            config.ToString().c_str());
   }
 
diff --git a/tensorflow/core/kernels/merge_v2_checkpoints_op_test.cc b/tensorflow/core/kernels/merge_v2_checkpoints_op_test.cc
index 3b9e9e9b75..10e468ce46 100644
--- a/tensorflow/core/kernels/merge_v2_checkpoints_op_test.cc
+++ b/tensorflow/core/kernels/merge_v2_checkpoints_op_test.cc
@@ -115,7 +115,7 @@ class MergeV2CheckpointsOpTest : public OpsTestBase {
     for (int i = 0; i < 2; ++i) {
       int directory_found =
           Env::Default()
-              ->IsDirectory(io::Dirname(prefixes[i]).ToString())
+              ->IsDirectory(std::string(io::Dirname(prefixes[i])))
               .code();
       if (delete_old_dirs) {
         EXPECT_EQ(error::NOT_FOUND, directory_found);
diff --git a/tensorflow/core/kernels/remote_fused_graph_execute_utils.cc b/tensorflow/core/kernels/remote_fused_graph_execute_utils.cc
index cc4d9a49a0..194a711d98 100644
--- a/tensorflow/core/kernels/remote_fused_graph_execute_utils.cc
+++ b/tensorflow/core/kernels/remote_fused_graph_execute_utils.cc
@@ -47,7 +47,7 @@ std::unordered_set<string> BuildNodeSetFromNodeNamesAndPorts(
   std::unordered_set<string> retval;
   for (const string& node_name_and_port : node_names_and_ports) {
     const TensorId tid = ParseTensorName(node_name_and_port);
-    retval.emplace(tid.first.ToString());
+    retval.emplace(std::string(tid.first));
   }
   return retval;
 }
@@ -64,7 +64,7 @@ Node* FindMutableNodeByName(const string& name, Graph* graph) {
 const NodeDef* FindNodeDefByName(const string& input,
                                  const GraphDef& graph_def) {
   const TensorId tid = ParseTensorName(input);
-  const string name = tid.first.ToString();
+  const string name = std::string(tid.first);
   for (const NodeDef& node_def : graph_def.node()) {
     if (node_def.name() == name) {
       return &node_def;
@@ -77,7 +77,7 @@ bool IsSameNodeName(const NodeDef& node_def, const string& node_name_and_port,
                     TensorId* tid) {
   CHECK_NOTNULL(tid);
   *tid = ParseTensorName(node_name_and_port);
-  if (node_def.name() == tid->first.ToString()) {
+  if (node_def.name() == tid->first) {
     return true;
   }
   return false;
@@ -326,7 +326,7 @@ RemoteFusedGraphExecuteUtils::GetExecutorBuildRegistry() {
     const string& node_name) {
   for (const std::pair<string, Tensor>& pair : input_tensor_vector) {
     const TensorId tid = ParseTensorName(pair.first);
-    if (node_name == tid.first.ToString()) {
+    if (node_name == tid.first) {
       return true;
     }
   }
@@ -423,7 +423,7 @@ RemoteFusedGraphExecuteUtils::AddOutputTensorShapeTypeByTensorShapeMap(
   std::vector<DataType> data_types;
   std::vector<TensorShape> shapes;
   const TensorId tid = ParseTensorName(name_and_port);
-  const string node_name = tid.first.ToString();
+  const string node_name = std::string(tid.first);
   const int port = tid.second;
   const NodeDef* node_def = FindNodeDefByName(node_name, graph_def);
   CHECK_NOTNULL(node_def);
@@ -522,7 +522,7 @@ RemoteFusedGraphExecuteUtils::GetTensorShapeType(
     const TensorShapeMap& tensor_shape_map, const string& node_name) {
   if (node_name.find(':') != string::npos) {
     const TensorId tid = ParseTensorName(node_name);
-    return GetTensorShapeType(tensor_shape_map, tid.first.ToString(),
+    return GetTensorShapeType(tensor_shape_map, std::string(tid.first),
                               tid.second);
   } else {
     return GetTensorShapeType(tensor_shape_map, node_name, 0);
@@ -570,7 +570,7 @@ RemoteFusedGraphExecuteUtils::BuildRemoteGraphInputsAndOutputsFromProto(
   const TensorId tid = ParseTensorName(name);
   CHECK_EQ(tensor_shape_map->count(name), 0);
   tensor_shape_map->emplace(
-      tid.first.ToString(),
+      std::string(tid.first),
       std::make_pair(tid.second,
                      std::make_pair(tensor.dtype(), tensor.shape())));
 }
@@ -692,7 +692,7 @@ RemoteFusedGraphExecuteUtils::BuildRemoteFusedGraphExecuteOpNode(
   std::vector<NodeBuilder::NodeOut> node_out_list;
   for (const string& input : inputs) {
     const TensorId tid = ParseTensorName(input);
-    Node* node = FindMutableNodeByName(tid.first.ToString(), graph);
+    Node* node = FindMutableNodeByName(std::string(tid.first), graph);
     CHECK_NOTNULL(node);
     node_out_list.emplace_back(node, tid.second);
   }
@@ -848,7 +848,7 @@ RemoteFusedGraphExecuteUtils::BuildRemoteFusedGraphExecuteOpNode(
 
   for (const string& subgraph_input : std::get<1>(cluster)) {
     const TensorId tid = ParseTensorName(subgraph_input);
-    const string subgraph_input_name = tid.first.ToString();
+    const string subgraph_input_name = std::string(tid.first);
     const int subgraph_input_port = tid.second;
     const NodeDef* node_def = FindNodeDefByName(subgraph_input_name, graph_def);
     CHECK_NOTNULL(node_def);
@@ -895,7 +895,7 @@ RemoteFusedGraphExecuteUtils::BuildRemoteFusedGraphExecuteOpNode(
   std::deque<const Node*> queue;
   for (const string& output : border_outputs) {
     const TensorId tid = ParseTensorName(output);
-    const string& output_node_name = tid.first.ToString();
+    const string& output_node_name = std::string(tid.first);
     for (const Node* node : graph.nodes()) {
       if (output_node_name == node->name()) {
         queue.push_back(node);
@@ -916,8 +916,7 @@ RemoteFusedGraphExecuteUtils::BuildRemoteFusedGraphExecuteOpNode(
       bool input_found = false;
       for (const string& input : border_inputs) {
         const TensorId tid = ParseTensorName(input);
-        if (tid.first.ToString() == src_node->name() &&
-            tid.second == src_port) {
+        if (tid.first == src_node->name() && tid.second == src_port) {
           input_found = true;
           border_input_nodes.insert(src_node);
         }
@@ -976,7 +975,7 @@ RemoteFusedGraphExecuteUtils::BuildRemoteFusedGraphExecuteOpNode(
       for (int j = 0; j < border_outputs.size(); ++j) {
         const string& output = border_outputs.at(j);
         const TensorId tid = ParseTensorName(output);
-        const string output_name = tid.first.ToString();
+        const string output_name = std::string(tid.first);
         Node* src_node = edge->src();
         if (src_node != nullptr && src_node->name() == output_name &&
             edge->src_output() == tid.second) {
@@ -996,11 +995,12 @@ RemoteFusedGraphExecuteUtils::BuildRemoteFusedGraphExecuteOpNode(
   // RemoteFusedGraphExecuteOpNode
   for (const string& output : outputs) {
     const TensorId output_tid = ParseTensorName(output);
-    const string output_name = output_tid.first.ToString();
+    const string output_name = std::string(output_tid.first);
     for (size_t i = 0; i < border_outputs.size(); ++i) {
       const TensorId subgraph_output_tid =
           ParseTensorName(border_outputs.at(i));
-      const string& subgraph_output_name = subgraph_output_tid.first.ToString();
+      const string& subgraph_output_name =
+          std::string(subgraph_output_tid.first);
       if (output_name == subgraph_output_name) {
         LOG(INFO) << "As graph output and subgraph output are same, "
                   << "the graph output node is replaced by identity node";
@@ -1435,7 +1435,7 @@ RemoteFusedGraphExecuteUtils::BuildNodeMapFromOpsDefinitions(
     GraphDef* graph_def) {
   const TensorId tid = ParseTensorName(input);
   CHECK_EQ(0, tid.second);
-  const string node_name = tid.first.ToString();
+  const string node_name = std::string(tid.first);
   for (NodeDef& node : *graph_def->mutable_node()) {
     if (node.name() != node_name) {
       continue;
diff --git a/tensorflow/core/kernels/save_restore_v2_ops.cc b/tensorflow/core/kernels/save_restore_v2_ops.cc
index 3acf290ea2..ab4de6c815 100644
--- a/tensorflow/core/kernels/save_restore_v2_ops.cc
+++ b/tensorflow/core/kernels/save_restore_v2_ops.cc
@@ -220,9 +220,9 @@ class MergeV2Checkpoints : public OpKernel {
         context, tensorflow::MergeBundles(env, input_prefixes, merged_prefix));
 
     if (delete_old_dirs_) {
-      const string& merged_dir = io::Dirname(merged_prefix).ToString();
+      const string& merged_dir = std::string(io::Dirname(merged_prefix));
       for (const string& input_prefix : input_prefixes) {
-        const string& dirname = io::Dirname(input_prefix).ToString();
+        const string& dirname = std::string(io::Dirname(input_prefix));
         if (dirname == merged_dir) continue;
         Status status = env->DeleteDir(dirname);
         // For sharded save, only the first delete will go through and all
diff --git a/tensorflow/core/kernels/string_strip_op.cc b/tensorflow/core/kernels/string_strip_op.cc
index ae700f4294..2aeafa28c4 100644
--- a/tensorflow/core/kernels/string_strip_op.cc
+++ b/tensorflow/core/kernels/string_strip_op.cc
@@ -43,7 +43,7 @@ class StringStripOp : public OpKernel {
     for (int64 i = 0; i < input.size(); ++i) {
       StringPiece entry(input(i));
       str_util::RemoveWhitespaceContext(&entry);
-      output(i) = entry.ToString();
+      output(i) = std::string(entry);
     }
   }
 };
diff --git a/tensorflow/core/kernels/tensor_array_ops.cc b/tensorflow/core/kernels/tensor_array_ops.cc
index 7ec26d95e6..ef9748b1aa 100644
--- a/tensorflow/core/kernels/tensor_array_ops.cc
+++ b/tensorflow/core/kernels/tensor_array_ops.cc
@@ -293,7 +293,7 @@ class TensorArrayGradOp : public TensorArrayCreationOp {
                                        resource.name());
       }
       tensor_array_name =
-          StringPiece(resource.name()).substr(container.size()).ToString();
+          std::string(StringPiece(resource.name()).substr(container.size()));
     }
 
     auto output_handle = tensor_array_output_handle->flat<string>();
diff --git a/tensorflow/core/kernels/whole_file_read_ops.cc b/tensorflow/core/kernels/whole_file_read_ops.cc
index 17a39ce29b..ed2bf3e8e2 100644
--- a/tensorflow/core/kernels/whole_file_read_ops.cc
+++ b/tensorflow/core/kernels/whole_file_read_ops.cc
@@ -134,7 +134,7 @@ class WriteFileOp : public OpKernel {
                     "Contents tensor must be scalar, but had shape: ",
                     contents_input->shape().DebugString()));
     const string& filename = filename_input->scalar<string>()();
-    const string dir = io::Dirname(filename).ToString();
+    const string dir = std::string(io::Dirname(filename));
     if (!context->env()->FileExists(dir).ok()) {
       OP_REQUIRES_OK(context, context->env()->RecursivelyCreateDir(dir));
     }
diff --git a/tensorflow/core/lib/strings/numbers.h b/tensorflow/core/lib/strings/numbers.h
index e9add42849..9cb56415cb 100644
--- a/tensorflow/core/lib/strings/numbers.h
+++ b/tensorflow/core/lib/strings/numbers.h
@@ -140,11 +140,11 @@ inline bool ProtoParseNumeric(StringPiece s, uint64* value) {
 }
 
 inline bool ProtoParseNumeric(StringPiece s, float* value) {
-  return safe_strtof(s.ToString().c_str(), value);
+  return safe_strtof(std::string(s).c_str(), value);
 }
 
 inline bool ProtoParseNumeric(StringPiece s, double* value) {
-  return safe_strtod(s.ToString().c_str(), value);
+  return safe_strtod(std::string(s).c_str(), value);
 }
 
 // Convert strings to number of type T.
diff --git a/tensorflow/core/lib/strings/scanner_test.cc b/tensorflow/core/lib/strings/scanner_test.cc
index 55ff3405c3..b0f568a03e 100644
--- a/tensorflow/core/lib/strings/scanner_test.cc
+++ b/tensorflow/core/lib/strings/scanner_test.cc
@@ -42,24 +42,24 @@ TEST_F(ScannerTest, Any) {
                   .Any(Scanner::DIGIT)
                   .Any(Scanner::LETTER)
                   .GetResult(&remaining, &match));
-  EXPECT_EQ("   horse", match.ToString());
-  EXPECT_EQ("0123", remaining.ToString());
+  EXPECT_EQ("   horse", match);
+  EXPECT_EQ("0123", remaining);
 
   EXPECT_TRUE(Scanner("")
                   .Any(Scanner::SPACE)
                   .Any(Scanner::DIGIT)
                   .Any(Scanner::LETTER)
                   .GetResult(&remaining, &match));
-  EXPECT_EQ("", remaining.ToString());
-  EXPECT_EQ("", match.ToString());
+  EXPECT_EQ("", remaining);
+  EXPECT_EQ("", match);
 
   EXPECT_TRUE(Scanner("----")
                   .Any(Scanner::SPACE)
                   .Any(Scanner::DIGIT)
                   .Any(Scanner::LETTER)
                   .GetResult(&remaining, &match));
-  EXPECT_EQ("----", remaining.ToString());
-  EXPECT_EQ("", match.ToString());
+  EXPECT_EQ("----", remaining);
+  EXPECT_EQ("", match);
 }
 
 TEST_F(ScannerTest, AnySpace) {
@@ -69,8 +69,8 @@ TEST_F(ScannerTest, AnySpace) {
                   .One(Scanner::LETTER)
                   .AnySpace()
                   .GetResult(&remaining, &match));
-  EXPECT_EQ("  a ", match.ToString());
-  EXPECT_EQ("b ", remaining.ToString());
+  EXPECT_EQ("  a ", match);
+  EXPECT_EQ("b ", remaining);
 }
 
 TEST_F(ScannerTest, AnyEscapedNewline) {
@@ -143,8 +143,8 @@ TEST_F(ScannerTest, ScanUntil) {
                   .ScanUntil('\'')
                   .OneLiteral("'")
                   .GetResult(&remaining, &match));
-  EXPECT_EQ(R"( \\'rest)", remaining.ToString());
-  EXPECT_EQ(R"(' \1 \2 \3 \')", match.ToString());
+  EXPECT_EQ(R"( \\'rest)", remaining);
+  EXPECT_EQ(R"(' \1 \2 \3 \')", match);
 
   // The "scan until" character is not present.
   remaining = match = "unset";
@@ -152,15 +152,15 @@ TEST_F(ScannerTest, ScanUntil) {
                    .OneLiteral("'")
                    .ScanUntil('\'')
                    .GetResult(&remaining, &match));
-  EXPECT_EQ("unset", remaining.ToString());
-  EXPECT_EQ("unset", match.ToString());
+  EXPECT_EQ("unset", remaining);
+  EXPECT_EQ("unset", match);
 
   // Scan until an escape character.
   remaining = match = "";
   EXPECT_TRUE(
       Scanner(R"(123\456)").ScanUntil('\\').GetResult(&remaining, &match));
-  EXPECT_EQ(R"(\456)", remaining.ToString());
-  EXPECT_EQ("123", match.ToString());
+  EXPECT_EQ(R"(\456)", remaining);
+  EXPECT_EQ("123", match);
 }
 
 TEST_F(ScannerTest, ScanEscapedUntil) {
@@ -170,8 +170,8 @@ TEST_F(ScannerTest, ScanEscapedUntil) {
                   .ScanEscapedUntil('\'')
                   .OneLiteral("'")
                   .GetResult(&remaining, &match));
-  EXPECT_EQ("rest", remaining.ToString());
-  EXPECT_EQ(R"(' \1 \2 \3 \' \\')", match.ToString());
+  EXPECT_EQ("rest", remaining);
+  EXPECT_EQ(R"(' \1 \2 \3 \' \\')", match);
 
   // The "scan until" character is not present.
   remaining = match = "unset";
@@ -179,27 +179,27 @@ TEST_F(ScannerTest, ScanEscapedUntil) {
                    .OneLiteral("'")
                    .ScanEscapedUntil('\'')
                    .GetResult(&remaining, &match));
-  EXPECT_EQ("unset", remaining.ToString());
-  EXPECT_EQ("unset", match.ToString());
+  EXPECT_EQ("unset", remaining);
+  EXPECT_EQ("unset", match);
 }
 
 TEST_F(ScannerTest, ZeroOrOneLiteral) {
   StringPiece remaining, match;
   EXPECT_TRUE(
       Scanner("abc").ZeroOrOneLiteral("abC").GetResult(&remaining, &match));
-  EXPECT_EQ("abc", remaining.ToString());
-  EXPECT_EQ("", match.ToString());
+  EXPECT_EQ("abc", remaining);
+  EXPECT_EQ("", match);
 
   EXPECT_TRUE(
       Scanner("abcd").ZeroOrOneLiteral("ab").ZeroOrOneLiteral("c").GetResult(
           &remaining, &match));
-  EXPECT_EQ("d", remaining.ToString());
-  EXPECT_EQ("abc", match.ToString());
+  EXPECT_EQ("d", remaining);
+  EXPECT_EQ("abc", match);
 
   EXPECT_TRUE(
       Scanner("").ZeroOrOneLiteral("abc").GetResult(&remaining, &match));
-  EXPECT_EQ("", remaining.ToString());
-  EXPECT_EQ("", match.ToString());
+  EXPECT_EQ("", remaining);
+  EXPECT_EQ("", match);
 }
 
 // Test output of GetResult (including the forms with optional params),
@@ -215,24 +215,24 @@ TEST_F(ScannerTest, CaptureAndGetResult) {
                   .StopCapture()
                   .Any(Scanner::SPACE)
                   .GetResult(&remaining, &match));
-  EXPECT_EQ("second", remaining.ToString());
-  EXPECT_EQ("first", match.ToString());
+  EXPECT_EQ("second", remaining);
+  EXPECT_EQ("first", match);
   EXPECT_TRUE(scan.GetResult());
   remaining = "";
   EXPECT_TRUE(scan.GetResult(&remaining));
-  EXPECT_EQ("second", remaining.ToString());
+  EXPECT_EQ("second", remaining);
   remaining = "";
   match = "";
   EXPECT_TRUE(scan.GetResult(&remaining, &match));
-  EXPECT_EQ("second", remaining.ToString());
-  EXPECT_EQ("first", match.ToString());
+  EXPECT_EQ("second", remaining);
+  EXPECT_EQ("first", match);
 
   scan.RestartCapture().One(Scanner::LETTER).One(Scanner::LETTER);
   remaining = "";
   match = "";
   EXPECT_TRUE(scan.GetResult(&remaining, &match));
-  EXPECT_EQ("cond", remaining.ToString());
-  EXPECT_EQ("se", match.ToString());
+  EXPECT_EQ("cond", remaining);
+  EXPECT_EQ("se", match);
 }
 
 // Tests that if StopCapture is not called, then calling GetResult, then
@@ -242,14 +242,14 @@ TEST_F(ScannerTest, MultipleGetResultExtendsCapture) {
 
   Scanner scan("one2three");
   EXPECT_TRUE(scan.Many(Scanner::LETTER).GetResult(&remaining, &match));
-  EXPECT_EQ("2three", remaining.ToString());
-  EXPECT_EQ("one", match.ToString());
+  EXPECT_EQ("2three", remaining);
+  EXPECT_EQ("one", match);
   EXPECT_TRUE(scan.Many(Scanner::DIGIT).GetResult(&remaining, &match));
-  EXPECT_EQ("three", remaining.ToString());
-  EXPECT_EQ("one2", match.ToString());
+  EXPECT_EQ("three", remaining);
+  EXPECT_EQ("one2", match);
   EXPECT_TRUE(scan.Many(Scanner::LETTER).GetResult(&remaining, &match));
-  EXPECT_EQ("", remaining.ToString());
-  EXPECT_EQ("one2three", match.ToString());
+  EXPECT_EQ("", remaining);
+  EXPECT_EQ("one2three", match);
 }
 
 TEST_F(ScannerTest, FailedMatchDoesntChangeResult) {
@@ -258,8 +258,8 @@ TEST_F(ScannerTest, FailedMatchDoesntChangeResult) {
   StringPiece remaining = "rem";
   StringPiece match = "match";
   EXPECT_FALSE(scan.One(Scanner::SPACE).GetResult(&remaining, &match));
-  EXPECT_EQ("rem", remaining.ToString());
-  EXPECT_EQ("match", match.ToString());
+  EXPECT_EQ("rem", remaining);
+  EXPECT_EQ("match", match);
 }
 
 TEST_F(ScannerTest, DefaultCapturesAll) {
@@ -271,8 +271,8 @@ TEST_F(ScannerTest, DefaultCapturesAll) {
                   .AnySpace()
                   .Any(Scanner::LETTER)
                   .GetResult(&remaining, &match));
-  EXPECT_EQ("", remaining.ToString());
-  EXPECT_EQ("a b", match.ToString());
+  EXPECT_EQ("", remaining);
+  EXPECT_EQ("a b", match);
 }
 
 TEST_F(ScannerTest, AllCharClasses) {
diff --git a/tensorflow/core/lib/strings/str_util.cc b/tensorflow/core/lib/strings/str_util.cc
index 4598b8ccc7..cab8f81585 100644
--- a/tensorflow/core/lib/strings/str_util.cc
+++ b/tensorflow/core/lib/strings/str_util.cc
@@ -332,7 +332,7 @@ string StringReplace(StringPiece s, StringPiece oldsub, StringPiece newsub,
                      bool replace_all) {
   // TODO(jlebar): We could avoid having to shift data around in the string if
   // we had a StringPiece::find() overload that searched for a StringPiece.
-  string res = s.ToString();
+  string res = std::string(s);
   size_t pos = 0;
   while ((pos = res.find(oldsub.data(), pos, oldsub.size())) != string::npos) {
     res.replace(pos, oldsub.size(), newsub.data(), newsub.size());
@@ -449,7 +449,7 @@ bool SplitAndParseAsFloats(StringPiece text, char delim,
   return SplitAndParseAsInts<float>(text, delim,
                                     [](StringPiece str, float* value) {
                                       return strings::safe_strtof(
-                                          str.ToString().c_str(), value);
+                                          std::string(str).c_str(), value);
                                     },
                                     result);
 }
diff --git a/tensorflow/core/lib/strings/str_util.h b/tensorflow/core/lib/strings/str_util.h
index e97d00b975..c887db7eff 100644
--- a/tensorflow/core/lib/strings/str_util.h
+++ b/tensorflow/core/lib/strings/str_util.h
@@ -205,7 +205,7 @@ std::vector<string> Split(StringPiece text, StringPiece delims, Predicate p) {
       if ((i == text.size()) || (delims.find(text[i]) != StringPiece::npos)) {
         StringPiece token(text.data() + token_start, i - token_start);
         if (p(token)) {
-          result.push_back(token.ToString());
+          result.push_back(std::string(token));
         }
         token_start = i + 1;
       }
diff --git a/tensorflow/core/platform/env.cc b/tensorflow/core/platform/env.cc
index b9a9ef85eb..fe7d0aa7d1 100644
--- a/tensorflow/core/platform/env.cc
+++ b/tensorflow/core/platform/env.cc
@@ -92,7 +92,7 @@ Env::Env() : file_system_registry_(new FileSystemRegistryImpl) {}
 Status Env::GetFileSystemForFile(const string& fname, FileSystem** result) {
   StringPiece scheme, host, path;
   io::ParseURI(fname, &scheme, &host, &path);
-  FileSystem* file_system = file_system_registry_->Lookup(scheme.ToString());
+  FileSystem* file_system = file_system_registry_->Lookup(std::string(scheme));
   if (!file_system) {
     if (scheme.empty()) {
       scheme = "[local]";
@@ -166,7 +166,7 @@ bool Env::FilesExist(const std::vector<string>& files,
   for (const auto& file : files) {
     StringPiece scheme, host, path;
     io::ParseURI(file, &scheme, &host, &path);
-    files_per_fs[scheme.ToString()].push_back(file);
+    files_per_fs[std::string(scheme)].push_back(file);
   }
 
   std::unordered_map<string, Status> per_file_status;
diff --git a/tensorflow/core/platform/env_test.cc b/tensorflow/core/platform/env_test.cc
index a70a417e6a..c461a40086 100644
--- a/tensorflow/core/platform/env_test.cc
+++ b/tensorflow/core/platform/env_test.cc
@@ -357,7 +357,7 @@ TEST_F(DefaultEnvTest, LocalTempFilename) {
   CHECK_EQ(error::OUT_OF_RANGE,
            file_to_read->Read(0 /* offset */, 1024 /* n */, &content, scratch)
                .code());
-  EXPECT_EQ("Null", content.ToString());
+  EXPECT_EQ("Null", content);
 
   // Delete the temporary file.
   TF_CHECK_OK(env->DeleteFile(filename));
diff --git a/tensorflow/core/platform/file_system.cc b/tensorflow/core/platform/file_system.cc
index b55e94d552..922773684b 100644
--- a/tensorflow/core/platform/file_system.cc
+++ b/tensorflow/core/platform/file_system.cc
@@ -158,7 +158,7 @@ Status FileSystem::RecursivelyCreateDir(const string& dirname) {
   std::reverse(sub_dirs.begin(), sub_dirs.end());
 
   // Now create the directories.
-  string built_path = remaining_dir.ToString();
+  string built_path = std::string(remaining_dir);
   for (const StringPiece sub_dir : sub_dirs) {
     built_path = io::JoinPath(built_path, sub_dir);
     Status status = CreateDir(io::CreateURI(scheme, host, built_path));
diff --git a/tensorflow/core/platform/file_system_helper.cc b/tensorflow/core/platform/file_system_helper.cc
index 22c5057281..0ba0e6304f 100644
--- a/tensorflow/core/platform/file_system_helper.cc
+++ b/tensorflow/core/platform/file_system_helper.cc
@@ -59,7 +59,7 @@ Status GetMatchingPaths(FileSystem* fs, Env* env, const string& pattern,
   string fixed_prefix = pattern.substr(0, pattern.find_first_of("*?[\\"));
   string eval_pattern = pattern;
   std::vector<string> all_files;
-  string dir = io::Dirname(fixed_prefix).ToString();
+  string dir = std::string(io::Dirname(fixed_prefix));
   // If dir is empty then we need to fix up fixed_prefix and eval_pattern to
   // include . as the top level directory.
   if (dir.empty()) {
diff --git a/tensorflow/core/platform/file_system_test.cc b/tensorflow/core/platform/file_system_test.cc
index f261b8f576..c0a16c95f9 100644
--- a/tensorflow/core/platform/file_system_test.cc
+++ b/tensorflow/core/platform/file_system_test.cc
@@ -125,7 +125,7 @@ class InterPlanetaryFileSystem : public NullFileSystem {
     ASSERT_EQ(scheme, "ipfs");
     ASSERT_EQ(host, "solarsystem");
     str_util::ConsumePrefix(&path, "/");
-    *parsed_path = path.ToString();
+    *parsed_path = std::string(path);
   }
 
   std::map<string, std::set<string>> celestial_bodies_ = {
diff --git a/tensorflow/core/util/command_line_flags.cc b/tensorflow/core/util/command_line_flags.cc
index 480ce94fca..8c27d01917 100644
--- a/tensorflow/core/util/command_line_flags.cc
+++ b/tensorflow/core/util/command_line_flags.cc
@@ -32,7 +32,7 @@ bool ParseStringFlag(tensorflow::StringPiece arg, tensorflow::StringPiece flag,
   if (str_util::ConsumePrefix(&arg, "--") &&
       str_util::ConsumePrefix(&arg, flag) &&
       str_util::ConsumePrefix(&arg, "=")) {
-    *value_parsing_ok = hook(arg.ToString());
+    *value_parsing_ok = hook(std::string(arg));
     return true;
   }
 
diff --git a/tensorflow/core/util/env_var.cc b/tensorflow/core/util/env_var.cc
index c844850179..8d43bcc927 100644
--- a/tensorflow/core/util/env_var.cc
+++ b/tensorflow/core/util/env_var.cc
@@ -28,7 +28,7 @@ namespace tensorflow {
 Status ReadBoolFromEnvVar(StringPiece env_var_name, bool default_val,
                           bool* value) {
   *value = default_val;
-  const char* tf_env_var_val = getenv(env_var_name.ToString().c_str());
+  const char* tf_env_var_val = getenv(std::string(env_var_name).c_str());
   if (tf_env_var_val == nullptr) {
     return Status::OK();
   }
@@ -48,7 +48,7 @@ Status ReadBoolFromEnvVar(StringPiece env_var_name, bool default_val,
 Status ReadInt64FromEnvVar(StringPiece env_var_name, int64 default_val,
                            int64* value) {
   *value = default_val;
-  const char* tf_env_var_val = getenv(env_var_name.ToString().c_str());
+  const char* tf_env_var_val = getenv(std::string(env_var_name).c_str());
   if (tf_env_var_val == nullptr) {
     return Status::OK();
   }
@@ -62,11 +62,11 @@ Status ReadInt64FromEnvVar(StringPiece env_var_name, int64 default_val,
 
 Status ReadStringFromEnvVar(StringPiece env_var_name, StringPiece default_val,
                             string* value) {
-  const char* tf_env_var_val = getenv(env_var_name.ToString().c_str());
+  const char* tf_env_var_val = getenv(std::string(env_var_name).c_str());
   if (tf_env_var_val != nullptr) {
     *value = tf_env_var_val;
   } else {
-    *value = default_val.ToString();
+    *value = std::string(default_val);
   }
   return Status::OK();
 }
diff --git a/tensorflow/core/util/example_proto_fast_parsing.cc b/tensorflow/core/util/example_proto_fast_parsing.cc
index 7946fa1782..3ce7988057 100644
--- a/tensorflow/core/util/example_proto_fast_parsing.cc
+++ b/tensorflow/core/util/example_proto_fast_parsing.cc
@@ -353,7 +353,7 @@ bool TestFastParse(const string& serialized, Example* example) {
     // I.e. last entry in the map overwrites all the previous ones.
     parsed::FeatureMapEntry& name_and_feature =
         parsed_example[parsed_example_size - i - 1];
-    string name = name_and_feature.first.ToString();
+    string name = std::string(name_and_feature.first);
     if ((*features.mutable_feature()).count(name) > 0) continue;
 
     auto& value = (*features.mutable_feature())[name];
-- 
GitLab


From 3c0afb1cf6679097c2316fda8803b3679b37871f Mon Sep 17 00:00:00 2001
From: Bixia Zheng <bixia@google.com>
Date: Wed, 2 May 2018 11:57:24 -0700
Subject: [PATCH 157/395] Turn on two half precision tests for GPU.

PiperOrigin-RevId: 195128326
---
 tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc | 8 --------
 1 file changed, 8 deletions(-)

diff --git a/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc b/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc
index 7fa61eb33c..6cb470caf8 100644
--- a/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc
+++ b/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc
@@ -52,12 +52,7 @@ class MatOpsSimpleTest : public ClientLibraryTestBase {};
 template <typename T>
 class MatOpsSimpleTest_F16F32 : public MatOpsSimpleTest {};
 
-// TODO(bixia): This test for F16 failed on GPU 02-25-2018.
-#ifdef XLA_TEST_BACKEND_GPU
-TYPED_TEST_CASE(MatOpsSimpleTest_F16F32, ::testing::Types<float>);
-#else
 TYPED_TEST_CASE(MatOpsSimpleTest_F16F32, TypesF16F32);
-#endif
 
 XLA_TYPED_TEST(MatOpsSimpleTest_F16F32, ExpTwoByTwoValues) {
   using T = TypeParam;
@@ -171,11 +166,8 @@ string PrintTestLinspaceMaxParam(
 }
 
 #ifndef XLA_BACKEND_DOES_NOT_SUPPORT_FLOAT16
-// TODO(bixia): This test failed on GPU 02-25-2018
-#ifdef XLA_TEST_BACKEND_CPU
 XLA_TEST_P(TestLinspaceMaxParametric, TestF16) { TestImpl<Eigen::half>(); }
 #endif
-#endif
 XLA_TEST_P(TestLinspaceMaxParametric, TestF32) { TestImpl<float>(); }
 
 INSTANTIATE_TEST_CASE_P(
-- 
GitLab


From a08db2f231e303017efb1378bec191c87a0faed7 Mon Sep 17 00:00:00 2001
From: Smit Shilu <shilu_smit@yahoo.com>
Date: Wed, 2 May 2018 15:23:31 -0400
Subject: [PATCH 158/395] command Typo

---
 tensorflow/contrib/lite/g3doc/rpi.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/lite/g3doc/rpi.md b/tensorflow/contrib/lite/g3doc/rpi.md
index 7a3a231626..ab50789307 100644
--- a/tensorflow/contrib/lite/g3doc/rpi.md
+++ b/tensorflow/contrib/lite/g3doc/rpi.md
@@ -32,7 +32,7 @@ This has been tested on Raspberry Pi 3b, Raspbian GNU/Linux 9.1 (stretch), gcc v
 
 Log in to you RPI, install the toolchain.
 ```bash
-sudo apt-get instal build-essential
+sudo apt-get install build-essential
 ```
 
 First, clone this TensorFlow repository. Run this at the root of the repository:
-- 
GitLab


From ce0ef2275bda40a6edcd738ccede61ccd3dd824b Mon Sep 17 00:00:00 2001
From: Asim Shankar <ashankar@google.com>
Date: Wed, 2 May 2018 12:32:28 -0700
Subject: [PATCH 159/395] docs: Link to the appropriately branched version of
 the live colab notebooks.

And update that link on release changes.

PiperOrigin-RevId: 195133689
---
 tensorflow/docs_src/get_started/eager.md    |  2 +-
 tensorflow/tools/ci_build/update_version.py | 10 ++++++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/tensorflow/docs_src/get_started/eager.md b/tensorflow/docs_src/get_started/eager.md
index ad89f0154c..f08ac74425 100644
--- a/tensorflow/docs_src/get_started/eager.md
+++ b/tensorflow/docs_src/get_started/eager.md
@@ -1,3 +1,3 @@
 # Get Started with Eager Execution
 
-[Colab notebook](https://colab.research.google.com/github/tensorflow/models/blob/master/samples/core/get_started/eager.ipynb)
+[Colab notebook](https://colab.research.google.com/github/tensorflow/models/blob/r1.8.0/samples/core/get_started/eager.ipynb)
diff --git a/tensorflow/tools/ci_build/update_version.py b/tensorflow/tools/ci_build/update_version.py
index 52a0da9a14..9ddb219048 100755
--- a/tensorflow/tools/ci_build/update_version.py
+++ b/tensorflow/tools/ci_build/update_version.py
@@ -248,6 +248,16 @@ def update_md_files(old_version, new_version):
     replace_string_in_line(r"<version>%s<\/version>" % old_version,
                            "<version>%s</version>" % new_version, filepath)
 
+  # Update any links to colab notebooks.
+  def colab_url(version):
+    version_string = "%d.%d.%d" % (version.major, version.minor, version.patch)
+    prefix = "https://colab.research.google.com/github/tensorflow/models/blob/r"
+    return prefix + version_string + "/"
+
+  replace_string_in_line(
+      colab_url(old_version), colab_url(new_version),
+      "%s/docs_src/get_started/eager.md" % TF_SRC_DIR)
+
 
 def major_minor_change(old_version, new_version):
   """Check if a major or minor change occurred."""
-- 
GitLab


From 262b176e27a3bcd01d518bea6d57683625df42b6 Mon Sep 17 00:00:00 2001
From: Benoit Steiner <bsteiner@google.com>
Date: Wed, 2 May 2018 12:58:06 -0700
Subject: [PATCH 160/395] Added support for packing of symbolic shapes

PiperOrigin-RevId: 195137239
---
 .../core/grappler/costs/graph_properties.cc   | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/tensorflow/core/grappler/costs/graph_properties.cc b/tensorflow/core/grappler/costs/graph_properties.cc
index 69b22561b2..23d25cba8d 100644
--- a/tensorflow/core/grappler/costs/graph_properties.cc
+++ b/tensorflow/core/grappler/costs/graph_properties.cc
@@ -770,6 +770,29 @@ class SymbolicShapeRefiner {
           c->output_tensors_as_shapes.resize(1);
           c->output_tensors_as_shapes[0] = result;
         }
+      } else if (IsPack(node)) {
+        // A Pack node concatenating scalars is often used to generate a shape.
+        std::vector<DimensionHandle> dims;
+        bool valid = true;
+        for (int i = 0; i < ic->num_inputs(); ++i) {
+          const Tensor* t = ic->input_tensor(i);
+          if (t) {
+            if (t->dims() != 0 ||
+                (t->dtype() != DT_INT32 && t->dtype() != DT_INT64)) {
+              valid = false;
+              break;
+            }
+            int64 size = t->dtype() == DT_INT32 ? t->scalar<int32>()()
+                                                : t->scalar<int64>()();
+            dims.push_back(size < 0 ? ic->UnknownDim() : ic->MakeDim(size));
+          } else {
+            dims.push_back(ic->UnknownDim());
+          }
+        }
+        if (valid) {
+          c->output_tensors_as_shapes.resize(1);
+          c->output_tensors_as_shapes[0] = ic->MakeShape(dims);
+        }
       } else if (IsSlice(node)) {
         ShapeHandle input = ic->input_tensors_as_shapes()[0];
         bool valid = ic->RankKnown(input);
-- 
GitLab


From ad491ad2c258fdb71cc0cea5bffe7931622e749f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 13:05:15 -0700
Subject: [PATCH 161/395] [XLA] Redesign: Dump HloSnapshot in local service as
 well. And support replaying HloSnapshot.

PiperOrigin-RevId: 195138472
---
 .../xla/service/compile_only_service.cc       | 16 +++++++
 tensorflow/compiler/xla/tools/BUILD           |  1 +
 .../compiler/xla/tools/replay_computation.cc  | 44 +++++++++++++++++--
 3 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/xla/service/compile_only_service.cc b/tensorflow/compiler/xla/service/compile_only_service.cc
index c9f78a0f9f..d39fd7307a 100644
--- a/tensorflow/compiler/xla/service/compile_only_service.cc
+++ b/tensorflow/compiler/xla/service/compile_only_service.cc
@@ -70,6 +70,22 @@ CompileOnlyService::CompileAheadOfTime(
     TF_RET_CHECK(instance.computation.has_program_shape());
 
     const DebugOptions& debug_options = options.debug_options();
+
+    // Dump computation proto if flag is set.
+    const string& directory_path = debug_options.xla_dump_computations_to();
+    if (!directory_path.empty()) {
+      HloSnapshot hlo_snapshot;
+      *hlo_snapshot.mutable_hlo()->mutable_hlo_module() = instance.computation;
+      string filename = tensorflow::strings::StrCat(
+          "computation_", instance.computation.id(), "__",
+          instance.computation.entry_computation_name());
+      const string& per_host_path = tensorflow::io::JoinPath(
+          directory_path, tensorflow::port::Hostname());
+
+      TF_RETURN_IF_ERROR(
+          Executable::DumpToDirectory(per_host_path, filename, hlo_snapshot));
+    }
+
     const auto& program_shape = instance.computation.program_shape();
     ExecutionOptions execution_options;
     *execution_options.mutable_debug_options() = debug_options;
diff --git a/tensorflow/compiler/xla/tools/BUILD b/tensorflow/compiler/xla/tools/BUILD
index 0bc4045a54..78ab2dccaf 100644
--- a/tensorflow/compiler/xla/tools/BUILD
+++ b/tensorflow/compiler/xla/tools/BUILD
@@ -88,6 +88,7 @@ cc_library(
         "//tensorflow/compiler/xla/client:global_data",
         "//tensorflow/compiler/xla/client:local_client",
         "//tensorflow/compiler/xla/client/lib:testing",
+        "//tensorflow/compiler/xla/service:hlo_proto",
         "//tensorflow/compiler/xla/service:session_proto",
         "//tensorflow/compiler/xla/tests:test_utils",
         "//tensorflow/core:framework_internal",
diff --git a/tensorflow/compiler/xla/tools/replay_computation.cc b/tensorflow/compiler/xla/tools/replay_computation.cc
index 62a353ad09..d8cedad65e 100644
--- a/tensorflow/compiler/xla/tools/replay_computation.cc
+++ b/tensorflow/compiler/xla/tools/replay_computation.cc
@@ -42,6 +42,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/client/local_client.h"
 #include "tensorflow/compiler/xla/execution_options_util.h"
 #include "tensorflow/compiler/xla/literal_util.h"
+#include "tensorflow/compiler/xla/service/hlo.pb.h"
 #include "tensorflow/compiler/xla/service/session.pb.h"
 #include "tensorflow/compiler/xla/shape_util.h"
 #include "tensorflow/compiler/xla/status_macros.h"
@@ -75,9 +76,14 @@ struct Options {
 //
 // Similarly, infeeds fake data of shape fake_infeed_shape if it is provided;
 // otherwise, no infeed is performed.
-StatusOr<std::unique_ptr<Literal>> ReplayComputation(
-    const SessionModule& module, Client* client, const Options& opts) {
-  TF_ASSIGN_OR_RETURN(Computation computation, client->LoadSnapshot(module));
+template <typename ModuleT>
+StatusOr<std::unique_ptr<Literal>> ReplayComputation(const ModuleT& module,
+                                                     Client* client,
+                                                     const Options& opts) {
+  static_assert(std::is_same<ModuleT, HloSnapshot>::value ||
+                    std::is_same<ModuleT, SessionModule>::value,
+                "Proto must be in HloSnapshot or SessionModule format");
+  TF_ASSIGN_OR_RETURN(auto computation, client->LoadSnapshot(module));
 
   std::vector<std::unique_ptr<GlobalData>> arguments;
   if (opts.use_fake_data) {
@@ -153,6 +159,38 @@ int RealMain(tensorflow::gtl::ArraySlice<char*> args, const Options& opts) {
   tensorflow::Env* env = tensorflow::Env::Default();
   int exit_status = EXIT_SUCCESS;
   for (char* arg : args) {
+    HloSnapshot snapshot;
+    auto status = tensorflow::ReadBinaryProto(env, arg, &snapshot);
+    if (status.ok()) {
+      StatusOr<std::unique_ptr<Literal>> result_status =
+          ReplayComputation(snapshot, client, opts);
+      if (!result_status.ok()) {
+        fprintf(stderr, "%s: error: %s\n", arg,
+                result_status.status().ToString().c_str());
+        exit_status = EXIT_FAILURE;
+        continue;
+      }
+
+      std::unique_ptr<Literal> result = result_status.ConsumeValueOrDie();
+      if (result != nullptr) {
+        fprintf(stdout, "%s: %s :: %s:%s\n", arg,
+                snapshot.hlo().hlo_module().name().c_str(),
+                ShapeUtil::HumanString(result->shape()).c_str(),
+                result->ToString().c_str());
+        if (snapshot.has_result()) {
+          std::unique_ptr<Literal> literal =
+              Literal::CreateFromProto(snapshot.result()).ConsumeValueOrDie();
+          fprintf(stdout, "was %s:%s\n",
+                  ShapeUtil::HumanString(snapshot.result().shape()).c_str(),
+                  literal->ToString().c_str());
+        }
+      }
+
+      continue;
+    }
+    fprintf(stderr, "%s: is not HloSnapshot: %s. Trying as SessionModule...\n",
+            arg, status.ToString().c_str());
+
     SessionModule module;
     TF_CHECK_OK(tensorflow::ReadBinaryProto(env, arg, &module));
     StatusOr<std::unique_ptr<Literal>> result_status =
-- 
GitLab


From b182fd88e10b1d36f30a349e312c3a7ae5f3cc95 Mon Sep 17 00:00:00 2001
From: Jiri Simsa <jsimsa@google.com>
Date: Wed, 2 May 2018 13:05:51 -0700
Subject: [PATCH 162/395] Increasing test size to reflect recent additions and
 prevent test timeouts.

PiperOrigin-RevId: 195138565
---
 tensorflow/contrib/data/python/kernel_tests/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/data/python/kernel_tests/BUILD b/tensorflow/contrib/data/python/kernel_tests/BUILD
index d59dd17aea..7643c2a9fc 100644
--- a/tensorflow/contrib/data/python/kernel_tests/BUILD
+++ b/tensorflow/contrib/data/python/kernel_tests/BUILD
@@ -32,7 +32,7 @@ py_test(
 
 py_test(
     name = "bucketing_test",
-    size = "small",
+    size = "medium",
     srcs = ["bucketing_test.py"],
     srcs_version = "PY2AND3",
     deps = [
-- 
GitLab


From 79f6d50d784cf27c6e1fb5200ca5022a334198fe Mon Sep 17 00:00:00 2001
From: Saurabh Saxena <srbs@google.com>
Date: Wed, 2 May 2018 13:23:56 -0700
Subject: [PATCH 163/395] Fix tsan failure in batch_dataset_op_test. The error
 was being caused because we were trying to save invocation_results_` while
 the function call was in progress. Now we wait for all invocations to finish
 before saving both `invocation_results_` and `batch_results_`.

Did local A/B testing for tsan.
Before: 12/100 failed
After: All passed
PiperOrigin-RevId: 195141349
---
 .../kernels/data/map_and_batch_dataset_op.cc  | 22 +++++++++++--------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc b/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc
index 7bc43e2072..c9551fbf16 100644
--- a/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc
@@ -272,6 +272,15 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
         TF_RETURN_IF_ERROR(writer->WriteScalar(full_name("current_batch_index"),
                                                current_batch_index_));
         TF_RETURN_IF_ERROR(SaveParent(writer, input_impl_));
+        // Wait for the map_fn dispatches made in `InvokeFunctionLocked` to
+        // finish. This may delay saving a checkpoint by a bit but keeps the
+        // code clean and also saves us from checkpointing the state of the
+        // `BlockingCounter`.
+        std::vector<int64> num_elements(batch_results_.size());
+        for (size_t i = 0; i < batch_results_.size(); i++) {
+          WaitForBatch(i, &num_elements[i]).IgnoreError();
+        }
+
         TF_RETURN_IF_ERROR(writer->WriteScalar(
             full_name("invocation_results_size"), invocation_results_.size()));
         for (size_t i = 0; i < invocation_results_.size(); ++i) {
@@ -280,7 +289,8 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
         TF_RETURN_IF_ERROR(writer->WriteScalar(full_name("batch_results_size"),
                                                batch_results_.size()));
         for (size_t i = 0; i < batch_results_.size(); ++i) {
-          TF_RETURN_IF_ERROR(WriteBatchResultLocked(writer, i));
+          TF_RETURN_IF_ERROR(
+              WriteBatchResultLocked(writer, i, num_elements[i]));
         }
         return Status::OK();
       }
@@ -567,15 +577,9 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
         return Status::OK();
       }
 
-      Status WriteBatchResultLocked(IteratorStateWriter* writer, size_t index)
+      Status WriteBatchResultLocked(IteratorStateWriter* writer, size_t index,
+                                    int64 num_elements)
           EXCLUSIVE_LOCKS_REQUIRED(mu_) {
-        // Wait for the map_fn dispatches made in `InvokeFunctionLocked` to
-        // finish. This may delay saving a checkpoint by a bit but keeps the
-        // code clean and also saves us from checkpointing the state of the
-        // `BlockingCounter`.
-        int64 num_elements = 0;
-        WaitForBatch(index, &num_elements).IgnoreError();
-
         const BatchResult& result = batch_results_[index];
         string prefix = strings::StrCat("batch_results_", index);
         {
-- 
GitLab


From 2706eeb1fbd4a2cf0e1af8efa3c7f3539944079e Mon Sep 17 00:00:00 2001
From: Jiri Simsa <jsimsa@google.com>
Date: Wed, 2 May 2018 13:29:01 -0700
Subject: [PATCH 164/395] Re-enabling a test.

PiperOrigin-RevId: 195142105
---
 .../contrib/data/python/kernel_tests/batch_dataset_op_test.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py
index a4a0ce79b6..6588fd04ac 100644
--- a/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py
@@ -630,9 +630,7 @@ class BatchDatasetSerializationTest(
         lambda x: array_ops.fill([x], x)).apply(
             batching.dense_to_sparse_batch(4, [12]))
 
-  # TODO(b/70988345): Re-enable when sparse tensors are properly supported by
-  # the DatasetSerializationTestBase.
-  def _testDenseToSparseBatchDatasetCore(self):
+  def testDenseToSparseBatchDatasetCore(self):
     components = np.random.randint(5, size=(40,)).astype(np.int32)
     diff_comp = np.random.randint(2, size=(100,)).astype(np.int32)
 
-- 
GitLab


From f9e8a75036154a73f256783eccf53bca6612d606 Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Wed, 2 May 2018 13:36:31 -0700
Subject: [PATCH 165/395] [XLA] Add new optimization that sinks constants into
 while loop bodies

Example transformation:

   state = (..., const, ...)
   while (pred(state)) {
     (..., v, ...) = state
     use(v)
     state = (..., v, ...)
   }

 =>

   state = (..., const, ...)
   while (pred(state)) {
     (..., v, ...) = state
     use(const)
     state = (..., v, ...)
   }

PiperOrigin-RevId: 195143323
---
 tensorflow/compiler/xla/service/BUILD         |  27 +++
 tensorflow/compiler/xla/service/cpu/BUILD     |   1 +
 .../compiler/xla/service/cpu/cpu_compiler.cc  |   2 +
 tensorflow/compiler/xla/service/hlo_module.cc |   8 +
 tensorflow/compiler/xla/service/hlo_module.h  |   5 +
 .../service/while_loop_constant_sinking.cc    | 128 +++++++++++
 .../xla/service/while_loop_constant_sinking.h |  68 ++++++
 .../while_loop_constant_sinking_test.cc       | 200 ++++++++++++++++++
 .../while_loop_invariant_code_motion.cc       |  27 +--
 tensorflow/compiler/xla/service/while_util.cc |  17 ++
 tensorflow/compiler/xla/service/while_util.h  |   6 +
 .../compiler/xla/service/while_util_test.cc   |  37 ++++
 12 files changed, 506 insertions(+), 20 deletions(-)
 create mode 100644 tensorflow/compiler/xla/service/while_loop_constant_sinking.cc
 create mode 100644 tensorflow/compiler/xla/service/while_loop_constant_sinking.h
 create mode 100644 tensorflow/compiler/xla/service/while_loop_constant_sinking_test.cc

diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD
index 6e2510aa10..17964cdd59 100644
--- a/tensorflow/compiler/xla/service/BUILD
+++ b/tensorflow/compiler/xla/service/BUILD
@@ -2687,6 +2687,33 @@ tf_cc_test(
     ],
 )
 
+cc_library(
+    name = "while_loop_constant_sinking",
+    srcs = ["while_loop_constant_sinking.cc"],
+    hdrs = ["while_loop_constant_sinking.h"],
+    deps = [
+        ":hlo",
+        ":hlo_pass",
+        ":while_util",
+        "//tensorflow/compiler/xla:statusor",
+        "//tensorflow/compiler/xla:util",
+        "//tensorflow/core:lib",
+    ],
+)
+
+tf_cc_test(
+    name = "while_loop_constant_sinking_test",
+    srcs = ["while_loop_constant_sinking_test.cc"],
+    deps = [
+        ":hlo_matchers",
+        ":while_loop_constant_sinking",
+        "//tensorflow/compiler/xla:test",
+        "//tensorflow/compiler/xla/tests:hlo_verified_test_base",
+        "//tensorflow/compiler/xla/tools/parser:hlo_parser",
+        "//tensorflow/core:test",
+    ],
+)
+
 cc_library(
     name = "despecializer",
     srcs = ["despecializer.cc"],
diff --git a/tensorflow/compiler/xla/service/cpu/BUILD b/tensorflow/compiler/xla/service/cpu/BUILD
index 2fc6c6bd55..cb81e413a3 100644
--- a/tensorflow/compiler/xla/service/cpu/BUILD
+++ b/tensorflow/compiler/xla/service/cpu/BUILD
@@ -131,6 +131,7 @@ cc_library(
         "//tensorflow/compiler/xla/service:reshape_mover",
         "//tensorflow/compiler/xla/service:transpose_folding",
         "//tensorflow/compiler/xla/service:tuple_simplifier",
+        "//tensorflow/compiler/xla/service:while_loop_constant_sinking",
         "//tensorflow/compiler/xla/service:while_loop_invariant_code_motion",
         "//tensorflow/compiler/xla/service:while_loop_simplifier",
         "//tensorflow/compiler/xla/service:zero_sized_hlo_elimination",
diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
index e298d67e09..91ed6e427a 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
+++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
@@ -87,6 +87,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/service/reshape_mover.h"
 #include "tensorflow/compiler/xla/service/transpose_folding.h"
 #include "tensorflow/compiler/xla/service/tuple_simplifier.h"
+#include "tensorflow/compiler/xla/service/while_loop_constant_sinking.h"
 #include "tensorflow/compiler/xla/service/while_loop_invariant_code_motion.h"
 #include "tensorflow/compiler/xla/service/while_loop_simplifier.h"
 #include "tensorflow/compiler/xla/service/zero_sized_hlo_elimination.h"
@@ -270,6 +271,7 @@ Status CpuCompiler::RunHloPasses(HloModule* module, bool is_aot_compile) {
 
     pass.AddPass<WhileLoopInvariantCodeMotion>();
     pass.AddPass<TupleSimplifier>();
+    pass.AddPass<WhileLoopConstantSinking>();
     pass.AddPass<WhileLoopSimplifier>();
     pass.AddPass<HloDCE>();
     pass.AddPass<ReshapeMover>();
diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc
index 987c4b2719..c7a7192867 100644
--- a/tensorflow/compiler/xla/service/hlo_module.cc
+++ b/tensorflow/compiler/xla/service/hlo_module.cc
@@ -540,6 +540,14 @@ uint64 HloModule::RandomNew64() const {
   return rng_();
 }
 
+HloComputation* HloModule::GetComputationWithName(
+    tensorflow::StringPiece name) {
+  auto it = c_find_if(computations(), [&](HloComputation* computation) {
+    return computation->name() == name;
+  });
+  return it == computations().end() ? nullptr : *it;
+}
+
 /* static */ std::atomic<int> HloModule::next_unique_module_id_(0);
 
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_module.h b/tensorflow/compiler/xla/service/hlo_module.h
index 82d790ec3b..f9674df812 100644
--- a/tensorflow/compiler/xla/service/hlo_module.h
+++ b/tensorflow/compiler/xla/service/hlo_module.h
@@ -32,6 +32,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/service/name_uniquer.h"
 #include "tensorflow/compiler/xla/service/versioned_computation_handle.h"
 #include "tensorflow/compiler/xla/types.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
 #include "tensorflow/core/lib/gtl/array_slice.h"
 #include "tensorflow/core/lib/gtl/iterator_range.h"
 #include "tensorflow/core/platform/logging.h"
@@ -138,6 +139,10 @@ class HloModule {
             MakeUnwrappingIterator(computations_.end())};
   }
 
+  // Returns the computation in this module that has the name `name`.  Returns
+  // null if there is no such computation.
+  HloComputation* GetComputationWithName(tensorflow::StringPiece name);
+
   // Gets the number of computations in this module.
   int64 computation_count() const { return computations_.size(); }
 
diff --git a/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc b/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc
new file mode 100644
index 0000000000..10fc4958fa
--- /dev/null
+++ b/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc
@@ -0,0 +1,128 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/while_loop_constant_sinking.h"
+#include "tensorflow/compiler/xla/service/while_util.h"
+#include "tensorflow/compiler/xla/util.h"
+#include "tensorflow/core/lib/gtl/flatmap.h"
+#include "tensorflow/core/lib/gtl/inlined_vector.h"
+
+namespace xla {
+
+// Replaces all uses of old_instr with new_instr except the use at
+// `while_body_root` (which must be a tuple instruction) at index `tuple_index`.
+// This utility helps us replace an instruction in the while body with a
+// constant while still keeping it trivially loop invariant.
+static Status ReplaceUsesWhileKeepingLoopInvariance(
+    HloInstruction* old_instr, HloInstruction* new_instr,
+    HloInstruction* while_body_root, int64 tuple_index) {
+  CHECK_EQ(while_body_root->opcode(), HloOpcode::kTuple);
+
+  std::vector<HloInstruction*> users;
+  users.reserve(old_instr->user_count());
+  c_copy(old_instr->users(), std::back_inserter(users));
+
+  for (auto* user : users) {
+    for (int64 i = 0, e = user->operand_count(); i < e; i++) {
+      if (user->operand(i) == old_instr &&
+          !(user == while_body_root && i == tuple_index)) {
+        TF_RETURN_IF_ERROR(user->ReplaceOperandWith(i, new_instr));
+      }
+    }
+  }
+
+  return Status::OK();
+}
+
+StatusOr<bool> WhileLoopConstantSinking::TrySinkingConstantsIntoWhileBody(
+    HloInstruction* while_instr) {
+  HloComputation* while_body = while_instr->while_body();
+
+  const HloInstruction& init_value = *while_instr->operand(0);
+  if (init_value.opcode() != HloOpcode::kTuple) {
+    return false;
+  }
+
+  bool changed = false;
+
+  for (HloInstruction* invariant_gte :
+       WhileUtil::GetInvariantGTEsForWhileBody(*while_body)) {
+    int64 index = invariant_gte->tuple_index();
+    const HloInstruction& invariant_value = *init_value.operand(index);
+    if (invariant_value.opcode() == HloOpcode::kConstant) {
+      auto* constant_instr =
+          while_body->AddInstruction(invariant_value.Clone(/*suffix=*/".sunk"));
+      TF_RETURN_IF_ERROR(ReplaceUsesWhileKeepingLoopInvariance(
+          invariant_gte, constant_instr, while_body->root_instruction(),
+          index));
+      changed = true;
+    }
+  }
+
+  return changed;
+}
+
+StatusOr<bool> WhileLoopConstantSinking::Run(HloModule* module) {
+  VLOG(2) << "HLO module before WhileLoopConstantSinking:";
+  XLA_VLOG_LINES(2, module->ToString());
+
+  bool changed = false;
+  std::vector<HloInstruction*> while_instrs;
+  for (auto* comp : module->MakeNonfusionComputations()) {
+    // Right now we don't particulary care about optimizing while-of-while
+    // patterns.  If/When we do, we'll want to visit the outer while (while_0)
+    // before we visit the inner while (while_1):
+    //
+    // while_1_body(state) {
+    //   val = gte(state, 0) // Loop invariant
+    //   use(val)
+    // }
+    //
+    // while_0_body(state) {
+    //   val = gte(state, 0) // Loop invariant
+    //   while_1 = while(init=tuple(val, ...), body=while_1_body, ...)
+    //   ...
+    // }
+    //
+    // main {
+    //   while_0 = while(init=(constant, ...), body=while_0_body, ...)
+    // }
+    //
+    // This will let us sink the constant into the outer while first and then
+    // into the inner while in a single run of this pass.
+    c_copy_if(comp->instructions(), std::back_inserter(while_instrs),
+              [](const HloInstruction* instr) {
+                return instr->opcode() == HloOpcode::kWhile;
+              });
+  }
+
+  for (HloInstruction* while_instr : while_instrs) {
+    // We only sink into while loop bodies, but this can be extended to
+    // transform conditions as well.
+    TF_ASSIGN_OR_RETURN(bool result,
+                        TrySinkingConstantsIntoWhileBody(while_instr));
+    changed |= result;
+  }
+
+  if (changed) {
+    VLOG(2) << "HLO module after WhileLoopConstantSinking:";
+    XLA_VLOG_LINES(2, module->ToString());
+  } else {
+    VLOG(2) << "HLO module unchanged after WhileLoopConstantSinking";
+  }
+
+  return changed;
+}
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/while_loop_constant_sinking.h b/tensorflow/compiler/xla/service/while_loop_constant_sinking.h
new file mode 100644
index 0000000000..21fb8568a8
--- /dev/null
+++ b/tensorflow/compiler/xla/service/while_loop_constant_sinking.h
@@ -0,0 +1,68 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_CONSTANT_SINKING_H_
+#define TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_CONSTANT_SINKING_H_
+
+#include "tensorflow/compiler/xla/service/hlo_module.h"
+#include "tensorflow/compiler/xla/service/hlo_pass_interface.h"
+#include "tensorflow/compiler/xla/statusor.h"
+
+namespace xla {
+
+// Sinks while loop invariant values that happen to be constants into the while
+// loop body.  This is probably not a win in isolation but may unlock further
+// optimizations like constant folding.
+//
+//   state = (..., const, ...)
+//   while (pred(state)) {
+//     (..., v, ...) = state
+//     use(v)
+//     state = (..., v, ...)
+//   }
+//
+// =>
+//
+//   state = (..., const, ...)
+//   while (pred(state)) {
+//     (..., v, ...) = state
+//     use(const)
+//     state = (..., v, ...)
+//   }
+//
+// Note that it leaves the `v` in place to keep that component of the state
+// tuple trivially loop invariant.  WhileLoopSimplifier will later get rid of
+// `v`.
+//
+// We only sink into while loop bodies, but this can be extended to transform
+// conditions as well.
+//
+// TODO(b/79121449):  We should also sink broadcasts of constants.
+class WhileLoopConstantSinking : public HloPassInterface {
+ public:
+  ~WhileLoopConstantSinking() override = default;
+
+  tensorflow::StringPiece name() const override {
+    return "while-loop-invariant-code-motion";
+  }
+
+  StatusOr<bool> Run(HloModule* module) override;
+
+ private:
+  StatusOr<bool> TrySinkingConstantsIntoWhileBody(HloInstruction* while_instr);
+};
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_CONSTANT_SINKING_H_
diff --git a/tensorflow/compiler/xla/service/while_loop_constant_sinking_test.cc b/tensorflow/compiler/xla/service/while_loop_constant_sinking_test.cc
new file mode 100644
index 0000000000..0d2288d8ea
--- /dev/null
+++ b/tensorflow/compiler/xla/service/while_loop_constant_sinking_test.cc
@@ -0,0 +1,200 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/while_loop_constant_sinking.h"
+
+#include "tensorflow/compiler/xla/service/hlo_matchers.h"
+#include "tensorflow/compiler/xla/test.h"
+#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+
+namespace xla {
+namespace {
+
+namespace op = xla::testing::opcode_matchers;
+using ::testing::_;
+
+class WhileLoopConstantSinkingTest : public ::testing::Test {};
+
+TEST_F(WhileLoopConstantSinkingTest, SinkOneConstant) {
+  const char* const hlo_string = R"(
+HloModule ModuleWithWhile
+
+body {
+  p_body = (f32[2],f32[2]) parameter(0)
+  p_body.0 = f32[2] get-tuple-element((f32[2],f32[2]) p_body), index=0
+  p_body.1 = f32[2] get-tuple-element((f32[2],f32[2]) p_body), index=1
+
+  add.0 = f32[2] add(p_body.0, p_body.1)
+  ROOT root = (f32[2],f32[2]) tuple(add.0, p_body.1)
+}
+
+condition {
+  p_cond = (f32[2],f32[2]) parameter(0)
+  ROOT result = pred[] constant(true)
+}
+
+ENTRY entry {
+  const_0 = f32[2] constant({1, 2})
+  const_1 = f32[2] constant({2, 1})
+  while_init = (f32[2],f32[2]) tuple(const_0, const_1)
+  ROOT while = (f32[2],f32[2]) while(while_init), condition=condition, body=body
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_string));
+
+  TF_ASSERT_OK_AND_ASSIGN(bool changed,
+                          WhileLoopConstantSinking{}.Run(module.get()));
+  ASSERT_TRUE(changed);
+
+  auto* while_body = module->GetComputationWithName("body");
+  EXPECT_THAT(while_body->root_instruction(),
+              op::Tuple(op::Add(_, op::Constant()), _));
+}
+
+TEST_F(WhileLoopConstantSinkingTest, KeepConstantsLoopInvariant) {
+  const char* const hlo_string = R"(
+HloModule ModuleWithWhile
+
+body {
+  p_body = (f32[2],f32[2],f32[2]) parameter(0)
+  p_body.0 = f32[2] get-tuple-element((f32[2],f32[2],f32[2]) p_body), index=0
+  p_body.1 = f32[2] get-tuple-element((f32[2],f32[2],f32[2]) p_body), index=1
+  p_body.2 = f32[2] get-tuple-element((f32[2],f32[2],f32[2]) p_body), index=2
+
+  add.0 = f32[2] add(p_body.1, p_body.2)
+  ROOT root = (f32[2],f32[2],f32[2]) tuple(add.0, p_body.1, p_body.2)
+}
+
+condition {
+  p_cond = (f32[2],f32[2],f32[2]) parameter(0)
+  ROOT result = pred[] constant(true)
+}
+
+ENTRY entry {
+  const_0 = f32[2] constant({1, 2})
+  const_1 = f32[2] constant({2, 1})
+  const_2 = f32[2] constant({3, 1})
+  while_init = (f32[2],f32[2],f32[2]) tuple(const_0, const_1, const_2)
+  ROOT while = (f32[2],f32[2],f32[2]) while(while_init), condition=condition, body=body
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_string));
+
+  TF_ASSERT_OK_AND_ASSIGN(bool changed,
+                          WhileLoopConstantSinking{}.Run(module.get()));
+  ASSERT_TRUE(changed);
+
+  auto* while_body = module->GetComputationWithName("body");
+  EXPECT_THAT(while_body->root_instruction(),
+              op::Tuple(op::Add(op::Constant(), op::Constant()),
+                        op::GetTupleElement(op::Parameter(0)),
+                        op::GetTupleElement(op::Parameter(0))));
+}
+
+TEST_F(WhileLoopConstantSinkingTest, TupleShapedConstants) {
+  const char* const hlo_string = R"(
+HloModule ModuleWithWhile
+
+body {
+  p_b = (f32[2],(f32[2],f32[2])) parameter(0)
+  p_b.0 = f32[2] get-tuple-element((f32[2],f32[2],f32[2]) p_b), index=0
+  p_b.1 = (f32[2],f32[2]) get-tuple-element((f32[2],(f32[2],f32[2])) p_b), index=1
+
+  p_b.1.1 = f32[2] get-tuple-element(p_b.1), index=0
+
+  ROOT root = (f32[2],f32[2],f32[2]) tuple(p_b.1.1, p_b.1)
+}
+
+condition {
+  p_cond = (f32[2],(f32[2],f32[2])) parameter(0)
+  ROOT result = pred[] constant(true)
+}
+
+ENTRY entry {
+  const_0 = f32[2] constant({1, 2})
+  const_1 = (f32[2], f32[2]) constant((f32[2], f32[2]) ({2, 1},{3,1}))
+  while_init = (f32[2],(f32[2],f32[2])) tuple(const_0, const_1)
+  ROOT while = (f32[2],(f32[2],f32[2])) while(while_init), condition=condition, body=body
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_string));
+
+  TF_ASSERT_OK_AND_ASSIGN(bool changed,
+                          WhileLoopConstantSinking{}.Run(module.get()));
+  ASSERT_TRUE(changed);
+
+  auto* while_body = module->GetComputationWithName("body");
+  EXPECT_THAT(while_body->root_instruction(),
+              op::Tuple(op::GetTupleElement(op::Constant(), 0),
+                        op::GetTupleElement(op::Parameter(0))));
+}
+
+TEST_F(WhileLoopConstantSinkingTest, DuplicateGTEs) {
+  // This test shows that the pass fails to optimize non-canonical IR.
+  //
+  // Even though the input IR has a constant value for p_b.2.dup,
+  // WhileLoopConstantSinking doesn't try to detect this.  Instead, it relies on
+  // prior runs of HLO CSE to have commoned these identical GTE instructions.
+
+  const char* const hlo_string = R"(
+HloModule ModuleWithWhile
+
+body {
+  p_b = (f32[2],f32[2],f32[2]) parameter(0)
+
+  p_b.1     = f32[2] get-tuple-element((f32[2],f32[2],f32[2]) p_b), index=1
+  p_b.2     = f32[2] get-tuple-element((f32[2],f32[2],f32[2]) p_b), index=2
+  p_b.2.dup = f32[2] get-tuple-element((f32[2],f32[2],f32[2]) p_b), index=2
+
+  add.0 = f32[2] add(p_b.1, p_b.2.dup)
+  ROOT root = (f32[2],f32[2],f32[2]) tuple(add.0, p_b.1, p_b.2)
+}
+
+condition {
+  p_cond = (f32[2],f32[2],f32[2]) parameter(0)
+  ROOT result = pred[] constant(true)
+}
+
+ENTRY entry {
+  const_0 = f32[2] constant({1, 2})
+  const_1 = f32[2] constant({2, 1})
+  const_2 = f32[2] constant({3, 1})
+  while_init = (f32[2],f32[2],f32[2]) tuple(const_0, const_1, const_2)
+  ROOT while = (f32[2],f32[2],f32[2]) while(while_init), condition=condition, body=body
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_string));
+
+  TF_ASSERT_OK_AND_ASSIGN(bool changed,
+                          WhileLoopConstantSinking{}.Run(module.get()));
+  ASSERT_TRUE(changed);
+
+  auto* while_body = module->GetComputationWithName("body");
+  EXPECT_THAT(while_body->root_instruction(),
+              op::Tuple(op::Add(op::Constant(), ::testing::Not(op::Constant())),
+                        op::GetTupleElement(op::Parameter(0)),
+                        op::GetTupleElement(op::Parameter(0))));
+}
+}  // namespace
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc
index 3ef0cdff67..321fdeb1ea 100644
--- a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc
+++ b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc
@@ -115,25 +115,6 @@ static bool NotWorthHoistingIndividually(const HloInstruction& instruction) {
   }
 }
 
-// Populates `gte_set` with the GetTupleElement instructions in `while_body`
-// that access elements in the parameter tuple that don't change across
-// iterations.  Assumes `while_body` is the body computation of the while loop
-// in question.
-static void GatherInvariantGTEs(HloComputation* while_body,
-                                FlatSet<HloInstruction*>* gte_set) {
-  const HloInstruction::InstructionVector root_operands =
-      while_body->root_instruction()->operands();
-  for (int i = 0; i < root_operands.size(); i++) {
-    HloInstruction* instr = root_operands[i];
-    if (instr->opcode() == HloOpcode::kGetTupleElement &&
-        instr->tuple_index() == i &&
-        instr->operand(0) == while_body->parameter_instruction(0) &&
-        ShapeUtil::IsArray(instr->shape())) {
-      InsertOrDie(gte_set, instr);
-    }
-  }
-}
-
 static StatusOr<bool> TryHoistingInvariantInstructionsFromWhileBody(
     HloInstruction* while_instr) {
   auto print_no_metadata = HloPrintOptions{}.set_print_metadata(false);
@@ -172,7 +153,13 @@ static StatusOr<bool> TryHoistingInvariantInstructionsFromWhileBody(
   // unhoisted_invariant_instructions -- they can be legally hoisted, but there
   // is no benefit to hoisting them unless something that uses it is also
   // hoisted.
-  GatherInvariantGTEs(while_body, &unhoisted_invariant_instructions);
+  for (auto* instr : WhileUtil::GetInvariantGTEsForWhileBody(*while_body)) {
+    if (ShapeUtil::IsArray(instr->shape())) {
+      // TODO(b/79147885): We should try to generalize this to tuples for
+      // uniformity's sake, if nothing else.
+      InsertOrDie(&unhoisted_invariant_instructions, instr);
+    }
+  }
 
   if (unhoisted_invariant_instructions.empty()) {
     // There are no obviously loop invariant elements in the state being
diff --git a/tensorflow/compiler/xla/service/while_util.cc b/tensorflow/compiler/xla/service/while_util.cc
index bd07941843..ed20b36292 100644
--- a/tensorflow/compiler/xla/service/while_util.cc
+++ b/tensorflow/compiler/xla/service/while_util.cc
@@ -244,4 +244,21 @@ static Shape MakeLoopStateShape(const WhileUtil::LoopStateTy& init_values) {
   }
   return result;
 }
+
+/*static*/ std::vector<HloInstruction*> WhileUtil::GetInvariantGTEsForWhileBody(
+    const HloComputation& while_body) {
+  std::vector<HloInstruction*> result;
+  const HloInstruction::InstructionVector root_operands =
+      while_body.root_instruction()->operands();
+  for (int i = 0; i < root_operands.size(); i++) {
+    HloInstruction* instr = root_operands[i];
+    if (instr->opcode() == HloOpcode::kGetTupleElement &&
+        instr->tuple_index() == i &&
+        instr->operand(0) == while_body.parameter_instruction(0)) {
+      result.push_back(instr);
+    }
+  }
+  return result;
+}
+
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/service/while_util.h b/tensorflow/compiler/xla/service/while_util.h
index 1688d46742..322d27b88c 100644
--- a/tensorflow/compiler/xla/service/while_util.h
+++ b/tensorflow/compiler/xla/service/while_util.h
@@ -74,6 +74,12 @@ class WhileUtil {
       HloComputation* computation, int32 trip_count,
       const LoopStateTy& init_values,
       const LoopBodyGeneratorTy& loop_body_generator);
+
+  // Returns the GetTupleElement instructions in `while_body` that access
+  // elements in the parameter tuple that don't change across iterations.
+  // Assumes `while_body` is the body computation of the while loop in question.
+  static std::vector<HloInstruction*> GetInvariantGTEsForWhileBody(
+      const HloComputation& while_body);
 };
 }  // namespace xla
 
diff --git a/tensorflow/compiler/xla/service/while_util_test.cc b/tensorflow/compiler/xla/service/while_util_test.cc
index cf0d0db99b..974bc542a3 100644
--- a/tensorflow/compiler/xla/service/while_util_test.cc
+++ b/tensorflow/compiler/xla/service/while_util_test.cc
@@ -126,5 +126,42 @@ TEST(WhileUtilTest, MakeTwoInstructionsLive) {
                         op::GetTupleElement(op::Parameter(0), 3)));
 }
 
+TEST(WhileUtilTest, GetInvariantGTEsForWhileBody) {
+  const char* const hlo_string = R"(
+HloModule ModuleWithWhile
+
+body {
+  param.b = (s32[], s32[]) parameter(0)
+  gte.0 = s32[] get-tuple-element(param.b), index=0
+  gte.1 = s32[] get-tuple-element(param.b), index=1
+  add = s32[] add(gte.0, gte.1)
+  ROOT tuple = (s32[], s32[]) tuple(gte.0, add)
+}
+
+cond {
+  param.c = (s32[], s32[]) parameter(0)
+  ROOT constant = pred[] constant(true)
+}
+
+ENTRY main {
+  init = (s32[], s32[]) parameter(0)
+  ROOT while = (s32[], s32[]) while(init), condition=cond, body=body
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_string));
+
+  HloComputation* while_body = module->GetComputationWithName("body");
+
+  ASSERT_NE(while_body, nullptr)
+      << "Expected exactly one while_body computation";
+
+  std::vector<HloInstruction*> gte_list =
+      WhileUtil::GetInvariantGTEsForWhileBody(*while_body);
+
+  ASSERT_EQ(gte_list.size(), 1);
+  EXPECT_EQ((*gte_list.begin())->name(), "gte.0");
+}
 }  // namespace
 }  // namespace xla
-- 
GitLab


From d750a1310d5312b934e5f1689d0abd467847b7d1 Mon Sep 17 00:00:00 2001
From: Anna R <annarev@google.com>
Date: Wed, 2 May 2018 13:44:32 -0700
Subject: [PATCH 166/395] Copy module list file in the while statement instead
 of creating a new variable

---
 tensorflow/tools/api/generator/create_python_api.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensorflow/tools/api/generator/create_python_api.py b/tensorflow/tools/api/generator/create_python_api.py
index d1e7f23fbc..788f6d3573 100644
--- a/tensorflow/tools/api/generator/create_python_api.py
+++ b/tensorflow/tools/api/generator/create_python_api.py
@@ -159,8 +159,7 @@ def get_api_init_text():
 
   # Traverse over everything imported above. Specifically,
   # we want to traverse over TensorFlow Python modules.
-  module_list = list(sys.modules.values())
-  for module in module_list:
+  for module in list(sys.modules.values()):
     # Only look at tensorflow modules.
     if (not module or not hasattr(module, '__name__') or
         'tensorflow.' not in module.__name__):
-- 
GitLab


From bd6c00aabe9a34715a5b2026eeccac4bc2a8d0de Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 13:42:40 -0700
Subject: [PATCH 167/395] Fix a bug in create_python_api.py

I got an error complaining about "RuntimeError: dictionary changed size during iteration", this change fixes it.

PiperOrigin-RevId: 195144333
---
 tensorflow/tools/api/generator/create_python_api.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/tools/api/generator/create_python_api.py b/tensorflow/tools/api/generator/create_python_api.py
index c06a39bfbd..65baa6e4b4 100644
--- a/tensorflow/tools/api/generator/create_python_api.py
+++ b/tensorflow/tools/api/generator/create_python_api.py
@@ -158,7 +158,7 @@ def get_api_init_text():
 
   # Traverse over everything imported above. Specifically,
   # we want to traverse over TensorFlow Python modules.
-  for module in sys.modules.values():
+  for module in list(sys.modules.values()):
     # Only look at tensorflow modules.
     if (not module or not hasattr(module, '__name__') or
         'tensorflow.' not in module.__name__):
-- 
GitLab


From 8f610384b61f7b1b62302b9a861c1d4b19b36b33 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 13:44:30 -0700
Subject: [PATCH 168/395] Updated ABSL to latest version in workspace.bzl.

PiperOrigin-RevId: 195144612
---
 tensorflow/workspace.bzl | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 16da59c5cf..f4f7bc4615 100644
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -96,11 +96,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""):
   tf_http_archive(
       name = "com_google_absl",
       urls = [
-          "https://mirror.bazel.build/github.com/abseil/abseil-cpp/archive/720c017e30339fd1786ce4aac68bc8559736e53f.tar.gz",
-          "https://github.com/abseil/abseil-cpp/archive/720c017e30339fd1786ce4aac68bc8559736e53f.tar.gz",
+          "https://mirror.bazel.build/github.com/abseil/abseil-cpp/archive/9613678332c976568272c8f4a78631a29159271d.tar.gz",
+          "https://github.com/abseil/abseil-cpp/archive/9613678332c976568272c8f4a78631a29159271d.tar.gz",
       ],
-     sha256 = "5996380e3e8b981f55d1c8d58e709c00dbb4806ba367be75d0925a68cc2f6478",
-     strip_prefix = "abseil-cpp-720c017e30339fd1786ce4aac68bc8559736e53f",
+     sha256 = "1273a1434ced93bc3e703a48c5dced058c95e995c8c009e9bdcb24a69e2180e9",
+     strip_prefix = "abseil-cpp-9613678332c976568272c8f4a78631a29159271d",
      build_file = clean_dep("//third_party:com_google_absl.BUILD"),
   )
 
@@ -299,11 +299,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""):
   tf_http_archive(
       name = "absl_py",
       urls = [
-          "https://mirror.bazel.build/github.com/abseil/abseil-py/archive/acec853355ef987eae48a8d87a79351c15dff593.tar.gz",
-          "https://github.com/abseil/abseil-py/archive/acec853355ef987eae48a8d87a79351c15dff593.tar.gz",
+          "https://mirror.bazel.build/github.com/abseil/abseil-py/archive/ea8c4d2ddbf3fba610c4d613260561699b776db8.tar.gz",
+          "https://github.com/abseil/abseil-py/archive/ea8c4d2ddbf3fba610c4d613260561699b776db8.tar.gz",
       ],
-      sha256 = "29e4584e778bee13aa4093824133d131d927cc160561892880118d9ff7b95a6a",
-      strip_prefix = "abseil-py-acec853355ef987eae48a8d87a79351c15dff593",
+      sha256 = "c30b48e0d2580ef1412e55c5c0e1dab8db2ee4ab56e2075eccff29c90c7c7059",
+      strip_prefix = "abseil-py-ea8c4d2ddbf3fba610c4d613260561699b776db8",
   )
 
   tf_http_archive(
-- 
GitLab


From bf70368d36df3ee9a16f5285940d73fb54d911c0 Mon Sep 17 00:00:00 2001
From: gracehoney <31743510+aaroey@users.noreply.github.com>
Date: Wed, 2 May 2018 14:46:12 -0700
Subject: [PATCH 169/395] Fix breaking tests

---
 tensorflow/contrib/tensorrt/convert/convert_graph.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.cc b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
index a8c07df4a0..4df54a749f 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_graph.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_graph.cc
@@ -342,6 +342,7 @@ tensorflow::Status ConvertGraphDefToTensorRT(
   // optimization pass
   tensorflow::grappler::GrapplerItem item;
   item.fetch = output_names;
+  item.graph = graph_def;
 
   tensorflow::DeviceProperties device_properties;
   device_properties.set_type("GPU");
-- 
GitLab


From 08fec96547a673084589e1be45f4bde0246f6fdf Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 14:58:57 -0700
Subject: [PATCH 170/395] Fix support for batch_normalization with mixed
 precision

When the type of the input tensor `x` is not the same as the type of
the parameters `mean`, `variance`, `offset`, and `scale`, a cast is
required.

This mixed precision case occurs when using the BatchNormalization
layer with a data type of float16 or bfloat16.

PiperOrigin-RevId: 195157279
---
 tensorflow/python/ops/nn_batchnorm_test.py | 20 ++++++++++++++------
 tensorflow/python/ops/nn_impl.py           |  6 ++++--
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/ops/nn_batchnorm_test.py b/tensorflow/python/ops/nn_batchnorm_test.py
index 3ac2c8eb17..1508ff44ce 100644
--- a/tensorflow/python/ops/nn_batchnorm_test.py
+++ b/tensorflow/python/ops/nn_batchnorm_test.py
@@ -292,12 +292,16 @@ class BatchNormalizationTest(test.TestCase):
             self.assertAllClose(
                 tf_batch_norm, keep_dims_tf_batch_norm, atol=0.000001)
 
-  def _testBatchNormArbitraryShapes(self, x_shape, param_shape, atol=0.0001):
-    x_val = np.random.random_sample(x_shape).astype(np.float32)
-    m_val = np.random.random_sample(param_shape).astype(np.float32)
-    v_val = np.random.random_sample(param_shape).astype(np.float32)
-    beta_val = np.random.random_sample(param_shape).astype(np.float32)
-    gamma_val = np.random.random_sample(param_shape).astype(np.float32)
+  def _testBatchNormArbitraryShapes(self, x_shape, param_shape, atol=0.0001,
+                                    dtype=dtypes.float32,
+                                    param_dtype=dtypes.float32):
+    numpy_dtype = dtype.as_numpy_dtype
+    numpy_param_dtype = param_dtype.as_numpy_dtype
+    x_val = np.random.random_sample(x_shape).astype(numpy_dtype)
+    m_val = np.random.random_sample(param_shape).astype(numpy_param_dtype)
+    v_val = np.random.random_sample(param_shape).astype(numpy_param_dtype)
+    beta_val = np.random.random_sample(param_shape).astype(numpy_param_dtype)
+    gamma_val = np.random.random_sample(param_shape).astype(numpy_param_dtype)
     for use_gpu in [True, False]:
       with self.test_session(use_gpu=use_gpu) as sess:
         x = constant_op.constant(x_val, name="x")
@@ -332,6 +336,10 @@ class BatchNormalizationTest(test.TestCase):
     self._testBatchNormArbitraryShapes(
         (2, 3, 2, 4, 5), (1, 1, 1, 4, 5), atol=0.005)
 
+  def testBatchNormMixedPrecision(self):
+    self._testBatchNormArbitraryShapes((3, 3), (1, 3), dtype=dtypes.float16,
+                                       param_dtype=dtypes.float32, atol=0.001)
+
 
 @test_util.with_c_api
 class SufficientStatisticsTest(test.TestCase):
diff --git a/tensorflow/python/ops/nn_impl.py b/tensorflow/python/ops/nn_impl.py
index 576627e78e..783d485892 100644
--- a/tensorflow/python/ops/nn_impl.py
+++ b/tensorflow/python/ops/nn_impl.py
@@ -830,8 +830,10 @@ def batch_normalization(x,
     inv = math_ops.rsqrt(variance + variance_epsilon)
     if scale is not None:
       inv *= scale
-    return x * inv + (
-        offset - mean * inv if offset is not None else -mean * inv)
+    # Note: tensorflow/contrib/quantize/python/fold_batch_norms.py depends on
+    # the precise order of ops that are generated by the expression below.
+    return x * math_ops.cast(inv, x.dtype) + math_ops.cast(
+        offset - mean * inv if offset is not None else -mean * inv, x.dtype)
 
 
 @tf_export("nn.fused_batch_norm")
-- 
GitLab


From d030ea951a477e2e141c13fed42681bcc5e97b4a Mon Sep 17 00:00:00 2001
From: Chris Ying <chrisying@google.com>
Date: Wed, 2 May 2018 15:03:33 -0700
Subject: [PATCH 171/395] Add steps_per_run to LoggingTensorHook and
 StepCounterHook and other logging bug fixes.

PiperOrigin-RevId: 195158238
---
 .../contrib/tpu/python/tpu/tpu_estimator.py   | 61 ++++++++----
 tensorflow/python/estimator/estimator.py      | 21 +++--
 .../training/basic_session_run_hooks.py       | 16 ++--
 .../training/basic_session_run_hooks_test.py  | 93 +++++++++++++++++--
 ...sorflow.train.-checkpoint-saver-hook.pbtxt |  2 +-
 5 files changed, 151 insertions(+), 42 deletions(-)

diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
index eb537b7b6a..534042b42c 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
@@ -459,11 +459,9 @@ class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook):
     session.run(self._init_ops,
                 options=config_pb2.RunOptions(timeout_in_ms=5 * 60 * 1000))
 
-    logging.info('Start infeed thread controller')
     self._infeed_controller = self._create_infeed_controller(
         name='InfeedController', target=self._run_infeed, args=(session,))
 
-    logging.info('Start outfeed thread controller')
     self._outfeed_controller = _OpQueueContext(
         name='OutfeedController', target=self._run_outfeed, args=(session,))
 
@@ -1553,7 +1551,7 @@ class _OutfeedHostCallHook(session_run_hook.SessionRunHook):
 
 
 class ExamplesPerSecondHook(basic_session_run_hooks.StepCounterHook):
-  """"Calculate and report the number of examples/sec during training."""
+  """Calculate and report global_step/sec and examples/sec during runtime."""
 
   def __init__(self,
                batch_size,
@@ -1569,12 +1567,18 @@ class ExamplesPerSecondHook(basic_session_run_hooks.StepCounterHook):
         summary_writer=summary_writer)
 
   def _log_and_record(self, elapsed_steps, elapsed_time, global_step):
-    examples_per_sec = self._batch_size * elapsed_steps / elapsed_time
+    global_step_per_sec = elapsed_steps / elapsed_time
+    examples_per_sec = self._batch_size * global_step_per_sec
     if self._summary_writer is not None:
+      global_step_summary = Summary(value=[
+          Summary.Value(tag='global_step/sec', simple_value=global_step_per_sec)
+      ])
       example_summary = Summary(value=[
-          Summary.Value(tag='examples_sec', simple_value=examples_per_sec)
+          Summary.Value(tag='examples/sec', simple_value=examples_per_sec)
       ])
+      self._summary_writer.add_summary(global_step_summary, global_step)
       self._summary_writer.add_summary(example_summary, global_step)
+    logging.info('global_step/sec: %g', global_step_per_sec)
     logging.info('examples/sec: %g', examples_per_sec)
 
 
@@ -1844,6 +1848,12 @@ class TPUEstimator(estimator_lib.Estimator):
     # config.model_dir.
     model_function = self._augment_model_fn(model_fn, batch_axis)
 
+    # Overwrite log_step_count_steps to disable TensorLoggingHook and
+    # StepCounterHook from being created in Estimator. TPUEstimator already
+    # added equivalent hooks in _augment_model_fn above.
+    self._log_every_n_steps = config.log_step_count_steps
+    config = config.replace(log_step_count_steps=None)
+
     # Passing non-None params as wrapped model_fn has it.
     params = params or {}
     super(TPUEstimator, self).__init__(
@@ -2039,39 +2049,50 @@ class TPUEstimator(estimator_lib.Estimator):
           host_ops = host_call.create_tpu_hostcall()
           if host_ops is None:
             host_ops = []
-
           shutdown_hooks = []
           if os.environ.get('TF_TPU_GRACEFUL_SHUTDOWN', '0') != '0':
             shutdown_hooks.append(session_support.GracefulShutdownHook())
-
-          hooks = [
+          with ops.control_dependencies([loss]):
+            global_step = array_ops.identity(training.get_global_step())
+          hooks = input_hooks + shutdown_hooks
+          logging_hook_frequency = (    # Divide and round up
+              (self._log_every_n_steps +
+               self._config.tpu_config.iterations_per_loop - 1) //
+              self._config.tpu_config.iterations_per_loop)
+          hooks.extend([
               TPUInfeedOutfeedSessionHook(
                   ctx,
                   enqueue_ops,
                   host_ops,
                   run_infeed_loop_on_coordinator=(
                       run_infeed_loop_on_coordinator)),
-              ExamplesPerSecondHook(
-                  ctx.global_batch_size, output_dir=self.model_dir),
               InstallSignalHandlerHook(),
               training.LoggingTensorHook(
                   {
                       'loss': array_ops.identity(loss),
-                      'step': training.get_global_step()
+                      'step': global_step,
                   },
-                  every_n_secs=30)
-          ] + input_hooks + shutdown_hooks
+                  every_n_iter=logging_hook_frequency)
+          ])
+          examples_hook = ExamplesPerSecondHook(
+              ctx.global_batch_size,
+              output_dir=self.model_dir,
+              every_n_steps=self._log_every_n_steps)
+          examples_hook._set_steps_per_run(   # pylint: disable=protected-access
+              self._config.tpu_config.iterations_per_loop)
+          hooks.append(examples_hook)
 
           chief_hooks = []
           if (self._config.save_checkpoints_secs or
               self._config.save_checkpoints_steps):
-            chief_hooks.append(
-                training.CheckpointSaverHook(
-                    self.model_dir,
-                    save_secs=self._config.save_checkpoints_secs,
-                    save_steps=self._config.save_checkpoints_steps,
-                    steps_per_run=self._config.tpu_config.iterations_per_loop,
-                    scaffold=scaffold))
+            checkpoint_hook = training.CheckpointSaverHook(
+                self.model_dir,
+                save_secs=self._config.save_checkpoints_secs,
+                save_steps=self._config.save_checkpoints_steps,
+                scaffold=scaffold)
+            checkpoint_hook._set_steps_per_run(   # pylint: disable=protected-access
+                self._config.tpu_config.iterations_per_loop)
+            chief_hooks.append(checkpoint_hook)
           summary.scalar(model_fn_lib.LOSS_METRIC_KEY, loss)
           with ops.control_dependencies([loss]):
             update_ops = _sync_variables_ops()
diff --git a/tensorflow/python/estimator/estimator.py b/tensorflow/python/estimator/estimator.py
index 3691c99dda..946f093ba7 100644
--- a/tensorflow/python/estimator/estimator.py
+++ b/tensorflow/python/estimator/estimator.py
@@ -994,15 +994,18 @@ class Estimator(object):
       summary.scalar('loss', estimator_spec.loss)
     ops.add_to_collection(ops.GraphKeys.LOSSES, estimator_spec.loss)
     worker_hooks.extend(hooks)
-    worker_hooks.extend([
-        training.NanTensorHook(estimator_spec.loss),
-        training.LoggingTensorHook(
-            {
-                'loss': estimator_spec.loss,
-                'step': global_step_tensor
-            },
-            every_n_iter=self._config.log_step_count_steps)
-    ])
+    worker_hooks.append(
+        training.NanTensorHook(estimator_spec.loss)
+    )
+    if self._config.log_step_count_steps is not None:
+      worker_hooks.append(
+          training.LoggingTensorHook(
+              {
+                  'loss': estimator_spec.loss,
+                  'step': global_step_tensor
+              },
+              every_n_iter=self._config.log_step_count_steps)
+      )
     worker_hooks.extend(estimator_spec.training_hooks)
 
     if not (estimator_spec.scaffold.saver or
diff --git a/tensorflow/python/training/basic_session_run_hooks.py b/tensorflow/python/training/basic_session_run_hooks.py
index d1cc7d8ce3..abcf76a220 100644
--- a/tensorflow/python/training/basic_session_run_hooks.py
+++ b/tensorflow/python/training/basic_session_run_hooks.py
@@ -380,8 +380,7 @@ class CheckpointSaverHook(session_run_hook.SessionRunHook):
                saver=None,
                checkpoint_basename="model.ckpt",
                scaffold=None,
-               listeners=None,
-               steps_per_run=1):
+               listeners=None):
     """Initializes a `CheckpointSaverHook`.
 
     Args:
@@ -394,9 +393,6 @@ class CheckpointSaverHook(session_run_hook.SessionRunHook):
       listeners: List of `CheckpointSaverListener` subclass instances.
         Used for callbacks that run immediately before or after this hook saves
         the checkpoint.
-      steps_per_run: `int`, number of steps that occur between each invocation
-        of the hook. Primarily used for TPU workloads which run multiple steps
-        in a while loop in a single Session.run.
 
     Raises:
       ValueError: One of `save_steps` or `save_secs` should be set.
@@ -412,6 +408,9 @@ class CheckpointSaverHook(session_run_hook.SessionRunHook):
     self._timer = SecondOrStepTimer(every_secs=save_secs,
                                     every_steps=save_steps)
     self._listeners = listeners or []
+    self._steps_per_run = 1
+
+  def _set_steps_per_run(self, steps_per_run):
     self._steps_per_run = steps_per_run
 
   def begin(self):
@@ -522,6 +521,10 @@ class StepCounterHook(session_run_hook.SessionRunHook):
     self._output_dir = output_dir
     self._last_global_step = None
     self._global_step_check_count = 0
+    self._steps_per_run = 1
+
+  def _set_steps_per_run(self, steps_per_run):
+    self._steps_per_run = steps_per_run
 
   def begin(self):
     if self._summary_writer is None and self._output_dir:
@@ -547,7 +550,8 @@ class StepCounterHook(session_run_hook.SessionRunHook):
     _ = run_context
 
     stale_global_step = run_values.results
-    if self._timer.should_trigger_for_step(stale_global_step+1):
+    if self._timer.should_trigger_for_step(
+        stale_global_step + self._steps_per_run):
       # get the real value after train op.
       global_step = run_context.session.run(self._global_step_tensor)
       if self._timer.should_trigger_for_step(global_step):
diff --git a/tensorflow/python/training/basic_session_run_hooks_test.py b/tensorflow/python/training/basic_session_run_hooks_test.py
index 31898562f8..7344ce2758 100644
--- a/tensorflow/python/training/basic_session_run_hooks_test.py
+++ b/tensorflow/python/training/basic_session_run_hooks_test.py
@@ -764,8 +764,8 @@ class CheckpointSaverHookMultiStepTest(test.TestCase):
       hook = basic_session_run_hooks.CheckpointSaverHook(
           self.model_dir,
           save_steps=2*self.steps_per_run,
-          scaffold=self.scaffold,
-          steps_per_run=self.steps_per_run)
+          scaffold=self.scaffold)
+      hook._set_steps_per_run(self.steps_per_run)
       hook.begin()
       self.scaffold.finalize()
       with session_lib.Session() as sess:
@@ -781,8 +781,8 @@ class CheckpointSaverHookMultiStepTest(test.TestCase):
       hook = basic_session_run_hooks.CheckpointSaverHook(
           self.model_dir,
           save_steps=2*self.steps_per_run,
-          scaffold=self.scaffold,
-          steps_per_run=self.steps_per_run)
+          scaffold=self.scaffold)
+      hook._set_steps_per_run(self.steps_per_run)
       hook.begin()
       self.scaffold.finalize()
       with session_lib.Session() as sess:
@@ -823,8 +823,8 @@ class CheckpointSaverHookMultiStepTest(test.TestCase):
       hook = basic_session_run_hooks.CheckpointSaverHook(
           self.model_dir,
           save_steps=2*self.steps_per_run,
-          scaffold=self.scaffold,
-          steps_per_run=self.steps_per_run)
+          scaffold=self.scaffold)
+      hook._set_steps_per_run(self.steps_per_run)
       hook.begin()
       self.scaffold.finalize()
       with session_lib.Session() as sess:
@@ -997,6 +997,87 @@ class StepCounterHookTest(test.TestCase):
             'global step.*has not been increased')
       hook.end(sess)
 
+  def _setup_steps_per_run_test(self,
+                                every_n_steps,
+                                steps_per_run,
+                                graph,
+                                sess):
+    variables.get_or_create_global_step()
+    self.train_op = training_util._increment_global_step(steps_per_run)
+    self.summary_writer = fake_summary_writer.FakeSummaryWriter(
+        self.log_dir, graph)
+    self.hook = basic_session_run_hooks.StepCounterHook(
+        summary_writer=self.summary_writer, every_n_steps=every_n_steps)
+    self.hook._set_steps_per_run(steps_per_run)
+    self.hook.begin()
+    sess.run(variables_lib.global_variables_initializer())
+    self.mon_sess = monitored_session._HookedSession(sess, [self.hook])
+
+  def test_steps_per_run_less_than_every_n_steps(self):
+    with ops.Graph().as_default() as g, session_lib.Session() as sess:
+      self._setup_steps_per_run_test(10, 5, g, sess)
+
+      # Logs at 15, 25
+      for _ in range(5):
+        time.sleep(0.01)
+        self.mon_sess.run(self.train_op)
+
+      self.hook.end(sess)
+      self.summary_writer.assert_summaries(
+          test_case=self,
+          expected_logdir=self.log_dir,
+          expected_graph=g,
+          expected_summaries={})
+      self.assertItemsEqual([15, 25], self.summary_writer.summaries.keys())
+      for step in [15, 25]:
+        summary_value = self.summary_writer.summaries[step][0].value[0]
+        self.assertEqual('global_step/sec', summary_value.tag)
+        self.assertGreater(summary_value.simple_value, 0)
+
+  def test_steps_per_run_equal_every_n_steps(self):
+    with ops.Graph().as_default() as g, session_lib.Session() as sess:
+      self._setup_steps_per_run_test(5, 5, g, sess)
+
+      # Logs at 10, 15, 20, 25
+      for _ in range(5):
+        time.sleep(0.01)
+        self.mon_sess.run(self.train_op)
+
+      self.hook.end(sess)
+      self.summary_writer.assert_summaries(
+          test_case=self,
+          expected_logdir=self.log_dir,
+          expected_graph=g,
+          expected_summaries={})
+      self.assertItemsEqual([10, 15, 20, 25],
+                            self.summary_writer.summaries.keys())
+      for step in [10, 15, 20, 25]:
+        summary_value = self.summary_writer.summaries[step][0].value[0]
+        self.assertEqual('global_step/sec', summary_value.tag)
+        self.assertGreater(summary_value.simple_value, 0)
+
+  def test_steps_per_run_greater_than_every_n_steps(self):
+    with ops.Graph().as_default() as g, session_lib.Session() as sess:
+      self._setup_steps_per_run_test(5, 10, g, sess)
+
+      # Logs at 20, 30, 40, 50
+      for _ in range(5):
+        time.sleep(0.01)
+        self.mon_sess.run(self.train_op)
+
+      self.hook.end(sess)
+      self.summary_writer.assert_summaries(
+          test_case=self,
+          expected_logdir=self.log_dir,
+          expected_graph=g,
+          expected_summaries={})
+      self.assertItemsEqual([20, 30, 40, 50],
+                            self.summary_writer.summaries.keys())
+      for step in [20, 30, 40, 50]:
+        summary_value = self.summary_writer.summaries[step][0].value[0]
+        self.assertEqual('global_step/sec', summary_value.tag)
+        self.assertGreater(summary_value.simple_value, 0)
+
 
 class SummarySaverHookTest(test.TestCase):
 
diff --git a/tensorflow/tools/api/golden/tensorflow.train.-checkpoint-saver-hook.pbtxt b/tensorflow/tools/api/golden/tensorflow.train.-checkpoint-saver-hook.pbtxt
index 327799729c..c3037baa8c 100644
--- a/tensorflow/tools/api/golden/tensorflow.train.-checkpoint-saver-hook.pbtxt
+++ b/tensorflow/tools/api/golden/tensorflow.train.-checkpoint-saver-hook.pbtxt
@@ -5,7 +5,7 @@ tf_class {
   is_instance: "<type \'object\'>"
   member_method {
     name: "__init__"
-    argspec: "args=[\'self\', \'checkpoint_dir\', \'save_secs\', \'save_steps\', \'saver\', \'checkpoint_basename\', \'scaffold\', \'listeners\', \'steps_per_run\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'model.ckpt\', \'None\', \'None\', \'1\'], "
+    argspec: "args=[\'self\', \'checkpoint_dir\', \'save_secs\', \'save_steps\', \'saver\', \'checkpoint_basename\', \'scaffold\', \'listeners\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'model.ckpt\', \'None\', \'None\'], "
   }
   member_method {
     name: "after_create_session"
-- 
GitLab


From 1d92d5037e1cec1a5099234a5568b68c7e675576 Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Wed, 2 May 2018 15:05:58 -0700
Subject: [PATCH 172/395] [TF:XLA] Bump open source llvm revision to r331338

PiperOrigin-RevId: 195158710
---
 tensorflow/workspace.bzl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index f4f7bc4615..94cac4f8fa 100644
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -452,11 +452,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""):
   tf_http_archive(
       name = "llvm",
       urls = [
-          "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/068c967842b83d22007eee4515b57e8d9aaccb82.tar.gz",
-          "https://github.com/llvm-mirror/llvm/archive/068c967842b83d22007eee4515b57e8d9aaccb82.tar.gz",
+          "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/a5108a08ceab35886a7df07c86f96aedd3d94bb7.tar.gz",
+          "https://github.com/llvm-mirror/llvm/archive/a5108a08ceab35886a7df07c86f96aedd3d94bb7.tar.gz",
       ],
-      sha256 = "4950432fb5cc68e5bf1f87a30b17dfdc69a5b93dac1e89d5274242d3ce7dae7c",
-      strip_prefix = "llvm-068c967842b83d22007eee4515b57e8d9aaccb82",
+      sha256 = "79cae03ebbdfd812bb69c460e1325ca069b5c576f7c7071f8216cf2b0975e36f",
+      strip_prefix = "llvm-a5108a08ceab35886a7df07c86f96aedd3d94bb7",
       build_file = clean_dep("//third_party/llvm:llvm.BUILD"),
   )
 
-- 
GitLab


From 9180cc254dff42368af126aa68eb82823ef67736 Mon Sep 17 00:00:00 2001
From: Yuanzhong Xu <yuanzx@google.com>
Date: Wed, 2 May 2018 15:14:08 -0700
Subject: [PATCH 173/395] [XLA] BF16 propagation: do not change if propagation
 is confined inside a fusion.

We now use a set to track all the potential changes, and do the actual changes
on the HLOs at the end. This also makes the boolean return value (whether
anything is changed) correct.

PiperOrigin-RevId: 195160025
---
 .../xla/service/bfloat16_propagation.cc       | 421 ++++++++++++------
 .../xla/service/bfloat16_propagation.h        |  93 +++-
 .../xla/service/bfloat16_propagation_test.cc  |  31 ++
 3 files changed, 387 insertions(+), 158 deletions(-)

diff --git a/tensorflow/compiler/xla/service/bfloat16_propagation.cc b/tensorflow/compiler/xla/service/bfloat16_propagation.cc
index 43ebe92c5e..ed0746980f 100644
--- a/tensorflow/compiler/xla/service/bfloat16_propagation.cc
+++ b/tensorflow/compiler/xla/service/bfloat16_propagation.cc
@@ -33,7 +33,7 @@ BFloat16Propagation::BFloat16Propagation(
     const BFloat16Support* bfloat16_support)
     : bfloat16_support_(bfloat16_support) {}
 
-void BFloat16Propagation::DetermineAndMutateFusionComputationPrecision(
+void BFloat16Propagation::DetermineFusionComputationPrecision(
     HloInstruction* fusion) {
   CHECK_EQ(fusion->opcode(), HloOpcode::kFusion);
   if (!bfloat16_support_->SupportsMixedPrecisions(*fusion)) {
@@ -48,15 +48,13 @@ void BFloat16Propagation::DetermineAndMutateFusionComputationPrecision(
   auto root = fusion->fused_instructions_computation()->root_instruction();
 
   // Adjust root's element types according to the fusion's output shape.
-  ShapeUtil::ForEachMutableSubshape(
-      root->mutable_shape(), [&](Shape* subshape, const ShapeIndex& index) {
-        if (subshape->element_type() != F32) {
+  ShapeUtil::ForEachSubshape(
+      root->shape(), [&](const Shape& subshape, const ShapeIndex& index) {
+        if (subshape.element_type() != F32) {
           return;
         }
-        if (ShapeUtil::GetSubshape(fusion->shape(), index).element_type() ==
-            BF16) {
-          subshape->set_element_type(BF16);
-          changed_ = true;
+        if (OutputTypeAfterChange(fusion, index) == BF16) {
+          AddToOrRemoveFromBF16ChangeSet(root, index, BF16);
           VLOG(2) << "Fused root " << root->ToString() << " at shape index "
                   << index << " changed to BF16 precision for fusion "
                   << fusion->ToString();
@@ -67,13 +65,101 @@ void BFloat16Propagation::DetermineAndMutateFusionComputationPrecision(
   auto insts =
       fusion->fused_instructions_computation()->MakeInstructionPostOrder();
   for (auto inst_it = insts.rbegin(); inst_it != insts.rend(); ++inst_it) {
-    DetermineAndMutateInstructionPrecision(*inst_it, /*skip_parameters=*/false);
+    DetermineInstructionPrecision(*inst_it, /*skip_parameters=*/false);
   }
-  computations_visited_in_mutation_pass_.insert(
+  computations_visited_in_backward_pass_.insert(
       fusion->fused_instructions_computation());
+
+  RevertIfFusionInternalBF16Changes(fusion);
+}
+
+void BFloat16Propagation::RevertIfFusionInternalBF16Changes(
+    HloInstruction* fusion) {
+  auto has_changes = [this](HloInstruction* inst) {
+    auto it = changes_to_bf16_.find(inst);
+    return it != changes_to_bf16_.end() && !it->second.empty();
+  };
+
+  auto root = fusion->fused_instructions_computation()->root_instruction();
+  tensorflow::gtl::FlatSet<const HloValue*> changed_root_buffers;
+
+  auto root_changes_it = changes_to_bf16_.find(root);
+  if (root_changes_it != changes_to_bf16_.end()) {
+    for (const auto& index : root_changes_it->second) {
+      for (const HloValue* value :
+           dataflow_->GetValueSet(root, index).values()) {
+        changed_root_buffers.insert(value);
+      }
+    }
+  }
+
+  auto aliases_changed_root_buffer =
+      [this, &changed_root_buffers](const HloInstruction* inst) {
+        bool aliasing = false;
+        ShapeUtil::ForEachSubshape(
+            inst->shape(), [&](const Shape& subshape, const ShapeIndex& index) {
+              if (aliasing) {
+                // Skip if aliasing is already found.
+                return;
+              }
+              // Only F32 buffers are considered for changing to BF16 in this
+              // pass.
+              if (subshape.element_type() != F32) {
+                return;
+              }
+              for (const HloValue* value :
+                   dataflow_->GetValueSet(inst, index).values()) {
+                if (ContainsKey(changed_root_buffers, value)) {
+                  aliasing = true;
+                  break;
+                }
+              }
+            });
+        return aliasing;
+      };
+
+  for (auto inst :
+       fusion->fused_instructions_computation()->MakeInstructionPostOrder()) {
+    if (inst->opcode() == HloOpcode::kParameter) {
+      continue;
+    }
+    if (aliases_changed_root_buffer(inst)) {
+      continue;
+    }
+    if (inst->opcode() == HloOpcode::kFusion) {
+      bool parameter_reverted = false;
+      for (int64 i = 0; i < inst->operand_count(); ++i) {
+        if (has_changes(inst->mutable_operand(i))) {
+          // Changes on the operand have not been reverted.
+          continue;
+        }
+        auto* fused_parameter = inst->fused_parameter(i);
+        if (has_changes(fused_parameter)) {
+          changes_to_bf16_.erase(fused_parameter);
+          parameter_reverted = true;
+        }
+      }
+      if (parameter_reverted) {
+        RevertIfFusionInternalBF16Changes(inst);
+      }
+    }
+    if (!has_changes(inst)) {
+      continue;
+    }
+    bool revert_changes = true;
+    for (auto operand : inst->operands()) {
+      if (has_changes(operand)) {
+        revert_changes = false;
+        break;
+      }
+    }
+    if (revert_changes) {
+      changes_to_bf16_.erase(inst);
+    }
+  }
 }
 
-void BFloat16Propagation::DetermineAndMutateWhileComputationsPrecision(
+void BFloat16Propagation::DetermineWhileComputationsPrecision(
     HloInstruction* while_hlo) {
   CHECK_EQ(while_hlo->opcode(), HloOpcode::kWhile);
 
@@ -86,16 +172,14 @@ void BFloat16Propagation::DetermineAndMutateWhileComputationsPrecision(
   auto body_root = body->root_instruction();
   HloComputation* condition = while_hlo->while_condition();
 
-  ShapeUtil::ForEachMutableSubshape(
-      body_root->mutable_shape(),
-      [this, while_hlo, body_root](Shape* subshape, const ShapeIndex& index) {
-        if (subshape->element_type() != F32) {
+  ShapeUtil::ForEachSubshape(
+      body_root->shape(), [this, while_hlo, body_root](
+                              const Shape& subshape, const ShapeIndex& index) {
+        if (subshape.element_type() != F32) {
           return;
         }
-        if (ShapeUtil::GetSubshape(while_hlo->shape(), index).element_type() ==
-            BF16) {
-          subshape->set_element_type(BF16);
-          changed_ = true;
+        if (OutputTypeAfterChange(while_hlo, index) == BF16) {
+          AddToOrRemoveFromBF16ChangeSet(body_root, index, BF16);
           VLOG(2) << "While body root " << body_root->ToString()
                   << " at shape index " << index
                   << " changed to BF16 precision for while "
@@ -106,30 +190,30 @@ void BFloat16Propagation::DetermineAndMutateWhileComputationsPrecision(
   auto body_insts = body->MakeInstructionPostOrder();
   for (auto inst_it = body_insts.rbegin(); inst_it != body_insts.rend();
        ++inst_it) {
-    DetermineAndMutateInstructionPrecision(*inst_it, /*skip_parameters=*/false);
+    DetermineInstructionPrecision(*inst_it, /*skip_parameters=*/false);
   }
-  computations_visited_in_mutation_pass_.insert(body);
+  computations_visited_in_backward_pass_.insert(body);
 
   auto condition_insts = condition->MakeInstructionPostOrder();
   for (auto inst_it = condition_insts.rbegin();
        inst_it != condition_insts.rend(); ++inst_it) {
-    DetermineAndMutateInstructionPrecision(*inst_it, /*skip_parameters=*/false);
+    DetermineInstructionPrecision(*inst_it, /*skip_parameters=*/false);
   }
-  computations_visited_in_mutation_pass_.insert(condition);
+  computations_visited_in_backward_pass_.insert(condition);
 }
 
 bool BFloat16Propagation::AllUsersConsumeBF16(const HloInstruction& hlo,
                                               const ShapeIndex& index) const {
-  auto value_set = dataflow_->GetValueSet(&hlo, index);
+  auto& value_set = dataflow_->GetValueSet(&hlo, index);
   for (const HloValue* value : value_set.values()) {
     if (ContainsKey(values_that_must_be_kept_as_f32_, value)) {
       return false;
     }
-    if (value->shape().element_type() == BF16) {
+    if (ValueTypeAfterChange(value) == BF16) {
       continue;
     }
     for (const HloUse& use : value->uses()) {
-      if (!ContainsKey(instructions_visited_in_mutation_pass_,
+      if (!ContainsKey(instructions_visited_in_backward_pass_,
                        use.instruction)) {
         // We don't know yet whether use.instruction will consume BF16 since it
         // hasn't been visited. Although we visit instructions in reverse
@@ -145,26 +229,23 @@ bool BFloat16Propagation::AllUsersConsumeBF16(const HloInstruction& hlo,
       // precision, or a called computation's parameters have been changed to
       // BF16 for fusions or whiles.
       if (use.instruction->opcode() == HloOpcode::kFusion) {
-        const auto* fused_parameter =
+        auto* fused_parameter =
             use.instruction->fused_parameter(use.operand_number);
-        if (ShapeUtil::GetSubshape(fused_parameter->shape(), use.operand_index)
-                .element_type() != BF16) {
+        if (OutputTypeAfterChange(fused_parameter, use.operand_index) != BF16) {
           return false;
         }
         continue;
       } else if (use.instruction->opcode() == HloOpcode::kWhile) {
-        const auto* cond_parameter =
+        auto* cond_parameter =
             use.instruction->while_condition()->parameter_instruction(
                 use.operand_number);
-        if (ShapeUtil::GetSubshape(cond_parameter->shape(), use.operand_index)
-                .element_type() != BF16) {
+        if (OutputTypeAfterChange(cond_parameter, use.operand_index) != BF16) {
           return false;
         }
-        const auto* body_parameter =
+        auto* body_parameter =
             use.instruction->while_body()->parameter_instruction(
                 use.operand_number);
-        if (ShapeUtil::GetSubshape(body_parameter->shape(), use.operand_index)
-                .element_type() != BF16) {
+        if (OutputTypeAfterChange(body_parameter, use.operand_index) != BF16) {
           return false;
         }
         continue;
@@ -174,19 +255,20 @@ bool BFloat16Propagation::AllUsersConsumeBF16(const HloInstruction& hlo,
         continue;
       }
       // If the op propagates precision and it outputs a BF16, then it's OK to
-      // supply BF16 also as the input. In the backward mutation pass, the users
-      // shapes should have already been processed.
+      // supply BF16 also as the input. In the backward pass, the users shapes
+      // should have already been processed.
       PrimitiveType user_output_type = PRIMITIVE_TYPE_INVALID;
       if (use.instruction->opcode() == HloOpcode::kTuple ||
           (use.instruction->opcode() == HloOpcode::kCrossReplicaSum &&
            ShapeUtil::IsTuple(use.instruction->shape()))) {
-        user_output_type = ShapeUtil::GetSubshape(
-                               ShapeUtil::GetSubshape(use.instruction->shape(),
-                                                      {use.operand_number}),
-                               use.operand_index)
-                               .element_type();
+        ShapeIndex use_output_index{use.operand_number};
+        for (int64 i : use.operand_index) {
+          use_output_index.push_back(i);
+        }
+        user_output_type =
+            OutputTypeAfterChange(use.instruction, use_output_index);
       } else {
-        user_output_type = use.instruction->shape().element_type();
+        user_output_type = OutputTypeAfterChange(use.instruction, {});
       }
       if (bfloat16_support_->EffectiveOperandPrecisionIsOutputPrecision(
               *use.instruction, use.operand_number) &&
@@ -199,8 +281,8 @@ bool BFloat16Propagation::AllUsersConsumeBF16(const HloInstruction& hlo,
   return true;
 }
 
-void BFloat16Propagation::DetermineAndMutateInstructionPrecision(
-    HloInstruction* hlo, bool skip_parameters) {
+void BFloat16Propagation::DetermineInstructionPrecision(HloInstruction* hlo,
+                                                        bool skip_parameters) {
   // We handle any fusion computation or while body/condition after the
   // instruction is handled, because we need to know the output shape of a
   // fusion or while before propagating inside its  computations.
@@ -209,12 +291,12 @@ void BFloat16Propagation::DetermineAndMutateInstructionPrecision(
       [this, hlo, &postpone_processing_called_computations] {
         if (!postpone_processing_called_computations) {
           if (hlo->opcode() == HloOpcode::kFusion) {
-            DetermineAndMutateFusionComputationPrecision(hlo);
+            DetermineFusionComputationPrecision(hlo);
           } else if (hlo->opcode() == HloOpcode::kWhile) {
-            DetermineAndMutateWhileComputationsPrecision(hlo);
+            DetermineWhileComputationsPrecision(hlo);
           }
         }
-        instructions_visited_in_mutation_pass_.insert(hlo);
+        instructions_visited_in_backward_pass_.insert(hlo);
       });
 
   if (hlo->opcode() == HloOpcode::kWhile &&
@@ -245,9 +327,9 @@ void BFloat16Propagation::DetermineAndMutateInstructionPrecision(
   CHECK(hlo->parent() != nullptr);
   if (hlo == hlo->parent()->root_instruction()) {
     if (!hlo->parent()->IsFusionComputation()) {
-      ShapeUtil::ForEachSubshape(hlo->shape(), [&](const Shape& subshape,
+      ShapeUtil::ForEachSubshape(hlo->shape(), [&](const Shape& /* subshape */,
                                                    const ShapeIndex& index) {
-        if (subshape.element_type() != F32) {
+        if (OutputTypeAfterChange(hlo, index) != F32) {
           return;
         }
         for (const auto* value : dataflow_->GetValueSet(hlo, index).values()) {
@@ -269,13 +351,12 @@ void BFloat16Propagation::DetermineAndMutateInstructionPrecision(
     return;
   }
 
-  ShapeUtil::ForEachMutableSubshape(
-      hlo->mutable_shape(),
-      [hlo, this](Shape* subshape, const ShapeIndex& index) {
-        if (subshape->element_type() == F32 &&
+  ShapeUtil::ForEachSubshape(
+      hlo->shape(),
+      [hlo, this](const Shape& /* subshape */, const ShapeIndex& index) {
+        if (OutputTypeAfterChange(hlo, index) == F32 &&
             AllUsersConsumeBF16(*hlo, index)) {
-          subshape->set_element_type(BF16);
-          changed_ = true;
+          AddToOrRemoveFromBF16ChangeSet(hlo, index, BF16);
           VLOG(2) << "HloInstruction output at shape index " << index
                   << " changed to BF16 precision: " << hlo->ToString();
         }
@@ -308,26 +389,24 @@ void BFloat16Propagation::AdjustCalledComputationParameters(
         CHECK_EQ(operands.size(), computation->num_parameters());
         for (int64 i = 0; i < operands.size(); ++i) {
           auto parameter = computation->parameter_instruction(i);
-          ShapeUtil::ForEachMutableSubshape(
-              parameter->mutable_shape(),
-              [this, i, hlo, &operands, parameter](Shape* subshape,
+          ShapeUtil::ForEachSubshape(
+              parameter->shape(),
+              [this, i, hlo, &operands, parameter](const Shape& /* subshape */,
                                                    const ShapeIndex& index) {
                 if (!ShapeUtil::IsLeafIndex(parameter->shape(), index)) {
                   return;
                 }
                 PrimitiveType operand_type =
-                    ShapeUtil::GetSubshape(operands[i]->shape(), index)
-                        .element_type();
-                if (subshape->element_type() == operand_type) {
+                    OutputTypeAfterChange(operands[i], index);
+                if (OutputTypeAfterChange(parameter, index) == operand_type) {
                   return;
                 }
-                CHECK(operand_type == F32 || operand_type == BF16);
-                subshape->set_element_type(operand_type);
-                changed_ = true;
+                AddToOrRemoveFromBF16ChangeSet(parameter, index, operand_type);
                 VLOG(2) << "Called computation parameter "
                         << parameter->ToString() << " at shape index " << index
-                        << " adjusted to match operand in HLO "
-                        << hlo->ToString();
+                        << " adjusted to "
+                        << (operand_type == BF16 ? "BF16" : "F32")
+                        << " to match operand in HLO " << hlo->ToString();
               });
         }
       };
@@ -348,51 +427,48 @@ void BFloat16Propagation::AdjustCalledComputationParameters(
 
 void BFloat16Propagation::AdjustCalledComputationRoot(HloInstruction* hlo) {
   auto adjust_computation = [this, hlo](HloComputation* computation,
-                                        const Shape& output_shape) {
+                                        HloInstruction* output) {
     // Adjust root.
     HloInstruction* root = computation->root_instruction();
-    ShapeUtil::ForEachMutableSubshape(
-        root->mutable_shape(), [this, hlo, root, &output_shape](
-                                   Shape* subshape, const ShapeIndex& index) {
-          if (!ShapeUtil::IsLeafIndex(hlo->shape(), index)) {
-            return;
-          }
-          const PrimitiveType output_type =
-              ShapeUtil::GetSubshape(output_shape, index).element_type();
-          if (subshape->element_type() == output_type) {
-            return;
-          }
-          CHECK(output_type == F32 || output_type == BF16);
-          subshape->set_element_type(output_type);
-          // It's possible that output_type is F32, but the root instruction's
-          // type is BF16; e.g., a fusion node's output was changed to BF16
-          // initially but then adjusted back to F32, and the fusion computation
-          // is now being adjusted after the fusion node.
-          if (output_type == F32) {
-            for (const auto* value :
-                 dataflow_->GetValueSet(root, index).values()) {
-              // We rely on the fact that this adjustment works in reverse
-              // topological order so that called computation will be
-              // processed later. Adding the value to
-              // values_that_must_be_kept_as_f32_ will ensure the
-              // correctness of the adjustment for HLOs that will be
-              // processed later.
-              values_that_must_be_kept_as_f32_.insert(value);
-            }
-          }
-          changed_ = true;
-          VLOG(2) << "Called computation root " << root->ToString()
-                  << " at shape index " << index
-                  << " adjusted to match output shape of " << hlo->ToString();
-        });
+    ShapeUtil::ForEachSubshape(root->shape(), [this, hlo, root, output](
+                                                  const Shape& /* subshape */,
+                                                  const ShapeIndex& index) {
+      if (!ShapeUtil::IsLeafIndex(hlo->shape(), index)) {
+        return;
+      }
+      const PrimitiveType output_type = OutputTypeAfterChange(output, index);
+      if (OutputTypeAfterChange(root, index) == output_type) {
+        return;
+      }
+      AddToOrRemoveFromBF16ChangeSet(root, index, output_type);
+      // It's possible that output_type is F32, but the root instruction's
+      // type is BF16; e.g., a fusion node's output was changed to BF16
+      // initially but then adjusted back to F32, and the fusion computation
+      // is now being adjusted after the fusion node.
+      if (output_type == F32) {
+        for (const auto* value : dataflow_->GetValueSet(root, index).values()) {
+          // We rely on the fact that this adjustment works in reverse
+          // topological order so that called computation will be
+          // processed later. Adding the value to
+          // values_that_must_be_kept_as_f32_ will ensure the
+          // correctness of the adjustment for HLOs that will be
+          // processed later.
+          values_that_must_be_kept_as_f32_.insert(value);
+        }
+      }
+      VLOG(2) << "Called computation root " << root->ToString()
+              << " at shape index " << index << " adjusted to "
+              << (output_type == BF16 ? "BF16" : "F32")
+              << " to match output shape of " << hlo->ToString();
+    });
   };
 
   switch (hlo->opcode()) {
     case HloOpcode::kFusion:
-      adjust_computation(hlo->fused_instructions_computation(), hlo->shape());
+      adjust_computation(hlo->fused_instructions_computation(), hlo);
       break;
     case HloOpcode::kWhile:
-      adjust_computation(hlo->while_body(), hlo->shape());
+      adjust_computation(hlo->while_body(), hlo);
       break;
     default:
       break;
@@ -409,16 +485,19 @@ bool BFloat16Propagation::ResolveInconsistencyOfAliasingBuffersHelper(
   for (auto inst_it = insts.rbegin(); inst_it != insts.rend(); ++inst_it) {
     auto hlo = *inst_it;
     auto adjust_hlo_output = [this, hlo, &parameter_changed](
-                                 Shape* subshape, const ShapeIndex& index) {
-      if (subshape->element_type() != F32 && subshape->element_type() != BF16) {
+                                 const Shape& /* subshape */,
+                                 const ShapeIndex& index) {
+      auto output_type = OutputTypeAfterChange(hlo, index);
+      if (output_type != F32 && output_type != BF16) {
         return;
       }
       PrimitiveType type = BF16;
       for (const auto* value : dataflow_->GetValueSet(hlo, index).values()) {
-        if (value->shape().element_type() == BF16) {
+        auto value_type = ValueTypeAfterChange(value);
+        if (value_type == BF16) {
           continue;
         }
-        CHECK_EQ(value->shape().element_type(), F32);
+        CHECK_EQ(value_type, F32);
         type = F32;
         break;
       }
@@ -437,16 +516,17 @@ bool BFloat16Propagation::ResolveInconsistencyOfAliasingBuffersHelper(
           values_that_must_be_kept_as_f32_.insert(value);
         }
       }
-      if (type != subshape->element_type()) {
-        subshape->set_element_type(type);
+      if (type != output_type) {
+        AddToOrRemoveFromBF16ChangeSet(hlo, index, type);
         VLOG(2) << "HloInstruction output at shape index " << index
-                << " adjusted to " << *subshape << ": " << hlo->ToString();
+                << " adjusted to " << (type == BF16 ? "BF16" : "F32") << ": "
+                << hlo->ToString();
         if (hlo->opcode() == HloOpcode::kParameter) {
           parameter_changed = true;
         }
       }
     };
-    ShapeUtil::ForEachMutableSubshape(hlo->mutable_shape(), adjust_hlo_output);
+    ShapeUtil::ForEachSubshape(hlo->shape(), adjust_hlo_output);
     AdjustCalledComputationRoot(hlo);
     if (hlo->opcode() == HloOpcode::kWhile) {
       // We need to run on the while body and condition repeatedly until a fixed
@@ -463,8 +543,7 @@ bool BFloat16Propagation::ResolveInconsistencyOfAliasingBuffersHelper(
              ResolveInconsistencyOfAliasingBuffersHelper(hlo->while_body(),
                                                          &visited_in_while)) {
         visited_in_while.clear();
-        ShapeUtil::ForEachMutableSubshape(hlo->mutable_shape(),
-                                          adjust_hlo_output);
+        ShapeUtil::ForEachSubshape(hlo->shape(), adjust_hlo_output);
         AdjustCalledComputationRoot(hlo);
       }
       visited_computations->insert(visited_in_while.begin(),
@@ -478,7 +557,7 @@ bool BFloat16Propagation::ResolveInconsistencyOfAliasingBuffersHelper(
   return parameter_changed;
 }
 
-Status BFloat16Propagation::ResolveInconsistencyOfAliasingBuffers(
+void BFloat16Propagation::ResolveInconsistencyOfAliasingBuffers(
     HloModule* module) {
   std::list<HloComputation*> computations_topological_order =
       module->MakeComputationPostOrder();
@@ -490,7 +569,9 @@ Status BFloat16Propagation::ResolveInconsistencyOfAliasingBuffers(
     }
     ResolveInconsistencyOfAliasingBuffersHelper(*comp_it, &resolved);
   }
+}
 
+Status BFloat16Propagation::ResolveInconsistentFusions(HloModule* module) {
   // We could have changed a fusion computation's root shape to have a different
   // precision than the fusion node's output, if the fusion root does not
   // define a buffer (e.g., a tuple). Now we add conversions after such fusion
@@ -517,7 +598,7 @@ Status BFloat16Propagation::ResolveInconsistencyOfAliasingBuffers(
   // (2) after adding conversion
   // (3) after tuple simplifier and DCE.
   bool needs_tuple_simplifier = false;
-  for (auto computation : computations_topological_order) {
+  for (auto computation : module->MakeComputationPostOrder()) {
     auto insts = computation->MakeInstructionPostOrder();
     for (auto inst_it = insts.rbegin(); inst_it != insts.rend(); ++inst_it) {
       auto hlo = *inst_it;
@@ -587,7 +668,14 @@ Status BFloat16Propagation::ResolveInconsistencyOfAliasingBuffers(
       needs_tuple_simplifier |= ShapeUtil::IsTuple(hlo->shape());
     }
   }
+  if (needs_tuple_simplifier) {
+    TupleSimplifier tuple_simplifier;
+    TF_RETURN_IF_ERROR(tuple_simplifier.Run(module).status());
+  }
+  return Status::OK();
+}
 
+Status BFloat16Propagation::ResolveConvertedConstants(HloModule* module) {
   // We may have converted some constants from F32 to BF16, so adjust the
   // constant literals in such cases. We do this here instead of when the
   // constant node's is changed because 1) the HloInstruction interface does not
@@ -598,8 +686,7 @@ Status BFloat16Propagation::ResolveInconsistencyOfAliasingBuffers(
   // can avoid repeated conversions.
   //
   // TODO(b/73833576): Consider resetting literal in HloInstruction.
-  bool needs_dce = needs_tuple_simplifier;
-  for (auto computation : computations_topological_order) {
+  for (auto computation : module->MakeComputationPostOrder()) {
     for (auto hlo : computation->MakeInstructionPostOrder()) {
       if (hlo->opcode() != HloOpcode::kConstant) {
         continue;
@@ -612,23 +699,13 @@ Status BFloat16Propagation::ResolveInconsistencyOfAliasingBuffers(
         auto new_constant = computation->AddInstruction(
             HloInstruction::CreateConstant(std::move(converted_literal)));
         TF_RETURN_IF_ERROR(hlo->ReplaceAllUsesWith(new_constant));
-        needs_dce = true;
       }
     }
   }
-
-  if (needs_tuple_simplifier) {
-    TupleSimplifier tuple_simplifier;
-    TF_RETURN_IF_ERROR(tuple_simplifier.Run(module).status());
-  }
-  if (needs_dce) {
-    HloDCE dce;
-    TF_RETURN_IF_ERROR(dce.Run(module).status());
-  }
   return Status::OK();
 }
 
-Status BFloat16Propagation::RemoveNoopConversions(HloModule* module) {
+Status BFloat16Propagation::SkipNoopConversions(HloModule* module) {
   for (auto computation : module->computations()) {
     for (auto hlo : computation->MakeInstructionPostOrder()) {
       if (hlo->opcode() != HloOpcode::kConvert) {
@@ -643,7 +720,6 @@ Status BFloat16Propagation::RemoveNoopConversions(HloModule* module) {
       if (is_root) {
         computation->set_root_instruction(source);
       }
-      TF_RETURN_IF_ERROR(computation->RemoveInstructionAndUnusedOperands(hlo));
     }
   }
   return Status::OK();
@@ -652,8 +728,18 @@ Status BFloat16Propagation::RemoveNoopConversions(HloModule* module) {
 // The algorithm first does a forward pass (parameters to root) to determine a
 // set of instructions to consider using bfloat16, then does a backward pass to
 // determine the precisions of those instructions according to the need of
-// their users.
+// their users. During the backward pass, the potential changes are stored in
+// changes_to_bf16_ which are subject to further adjustments then applied to the
+// HLOs.
 StatusOr<bool> BFloat16Propagation::Run(HloModule* module) {
+  consider_using_bfloat16_.clear();
+  instructions_visited_in_backward_pass_.clear();
+  computations_visited_in_backward_pass_.clear();
+  values_that_must_be_kept_as_f32_.clear();
+  caller_counts_.clear();
+  changes_to_bf16_.clear();
+  changed_ = false;
+
   TF_ASSIGN_OR_RETURN(dataflow_, HloDataflowAnalysis::Run(*module));
 
   std::list<HloComputation*> computations_topological_order =
@@ -686,8 +772,24 @@ StatusOr<bool> BFloat16Propagation::Run(HloModule* module) {
     }
     auto insts = (*comp_it)->MakeInstructionPostOrder();
     for (auto inst_it = insts.rbegin(); inst_it != insts.rend(); ++inst_it) {
-      DetermineAndMutateInstructionPrecision(*inst_it,
-                                             /*skip_parameters=*/true);
+      DetermineInstructionPrecision(*inst_it,
+                                    /*skip_parameters=*/true);
+    }
+  }
+
+  // It's possible that an instruction does not define a buffer, but the
+  // defining instruction's shape has changed. So we need to adjust the output
+  // shapes of instructions according to the HLO values they refer to.
+  ResolveInconsistencyOfAliasingBuffers(module);
+
+  // Apply the changes in changes_to_bf16_.
+  for (auto& change : changes_to_bf16_) {
+    auto shape = change.first->mutable_shape();
+    for (const auto& index : change.second) {
+      auto subshape = ShapeUtil::GetMutableSubshape(shape, index);
+      CHECK_EQ(subshape->element_type(), F32);
+      subshape->set_element_type(BF16);
+      changed_ = true;
     }
   }
 
@@ -695,15 +797,56 @@ StatusOr<bool> BFloat16Propagation::Run(HloModule* module) {
     return false;
   }
 
-  // It's possible that an instruction does not define a buffer, but the
-  // defining instruction's shape has changed. So we need to adjust the output
-  // shapes of instructions according to the HLO values they refer to.
-  TF_RETURN_IF_ERROR(ResolveInconsistencyOfAliasingBuffers(module));
+  TF_RETURN_IF_ERROR(ResolveInconsistentFusions(module));
+  TF_RETURN_IF_ERROR(ResolveConvertedConstants(module));
 
   // This pass could have turned an F32 -> BF16 conversion to a no-op (BF16 ->
-  // BF16), so we remove them now.
-  TF_RETURN_IF_ERROR(RemoveNoopConversions(module));
+  // BF16), so we skip them now.
+  TF_RETURN_IF_ERROR(SkipNoopConversions(module));
+
+  {
+    // We may have dead HLOs after ResolveInconsistentFusions,
+    // ResolveConvertedConstants and SkipNoopConversions.
+    HloDCE dce;
+    TF_RETURN_IF_ERROR(dce.Run(module).status());
+  }
   return true;
 }
 
+PrimitiveType BFloat16Propagation::OutputTypeAfterChange(
+    HloInstruction* hlo, const ShapeIndex& index) const {
+  PrimitiveType type_on_hlo =
+      ShapeUtil::GetSubshape(hlo->shape(), index).element_type();
+  if (type_on_hlo != F32) {
+    return type_on_hlo;
+  }
+  auto it = changes_to_bf16_.find(hlo);
+  if (it == changes_to_bf16_.end()) {
+    return type_on_hlo;
+  }
+  return ContainsKey(it->second, index) ? BF16 : F32;
+}
+
+PrimitiveType BFloat16Propagation::ValueTypeAfterChange(
+    const HloValue* value) const {
+  auto hlo = value->defining_instruction();
+  const auto& position = value->defining_position();
+  return OutputTypeAfterChange(hlo, position.index);
+}
+
+void BFloat16Propagation::AddToOrRemoveFromBF16ChangeSet(
+    HloInstruction* hlo, const ShapeIndex& index, PrimitiveType target_type) {
+  if (target_type == BF16) {
+    auto& entry = changes_to_bf16_[hlo];
+    entry.insert(index);
+  } else {
+    CHECK_EQ(target_type, F32);
+    auto it = changes_to_bf16_.find(hlo);
+    if (it == changes_to_bf16_.end()) {
+      return;
+    }
+    it->second.erase(index);
+  }
+}
+
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/service/bfloat16_propagation.h b/tensorflow/compiler/xla/service/bfloat16_propagation.h
index 1744e9db90..de0355ddfc 100644
--- a/tensorflow/compiler/xla/service/bfloat16_propagation.h
+++ b/tensorflow/compiler/xla/service/bfloat16_propagation.h
@@ -26,6 +26,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
 #include "tensorflow/compiler/xla/service/hlo_module.h"
 #include "tensorflow/compiler/xla/service/hlo_pass_interface.h"
+#include "tensorflow/core/lib/hash/hash.h"
 
 namespace xla {
 
@@ -85,30 +86,39 @@ class BFloat16Propagation : public HloPassInterface {
   tensorflow::gtl::FlatSet<const HloInstruction*> consider_using_bfloat16_;
 
   // ***************************
-  // Functions called and state produced by the backward mutation pass (from
-  // root to parameters).
+  // Functions called and state produced by the backward pass (from root to
+  // parameters) that finds opportunities to use BF16.
 
-  // Determines the precision for the given instruction in the mutation pass.
-  void DetermineAndMutateInstructionPrecision(HloInstruction* hlo,
-                                              bool skip_parameters);
+  // Determines the precision for the given instruction in the
+  // opportunity-finding pass.
+  void DetermineInstructionPrecision(HloInstruction* hlo, bool skip_parameters);
 
-  // Special handling in the mutation pass for fusion computations.
+  // Special handling in the opportunity-finding pass for fusion computations.
   //
   // Precondition: hlo->opcode() == kFusion
-  void DetermineAndMutateFusionComputationPrecision(HloInstruction* fusion);
+  void DetermineFusionComputationPrecision(HloInstruction* fusion);
 
-  // Special handling in the mutation pass for while computations.
+  // Reverts changes to BF16 that will not propagate outside a fusion
+  // computation. This avoids BF16 casts overhead inside a fusion which won't
+  // save memory bandwidth.
+  //
+  // Precondition: hlo->opcode() == kFusion
+  void RevertIfFusionInternalBF16Changes(HloInstruction* fusion);
+
+  // Special handling in the opportunity-finding pass for while computations.
   //
   // Precondition: hlo->opcode() == kWhile
-  void DetermineAndMutateWhileComputationsPrecision(HloInstruction* while_hlo);
+  void DetermineWhileComputationsPrecision(HloInstruction* while_hlo);
 
-  // The set of HloInstructions that have been visited in the mutation pass.
+  // The set of HloInstructions that have been visited in the
+  // opportunity-finding pass.
   tensorflow::gtl::FlatSet<const HloInstruction*>
-      instructions_visited_in_mutation_pass_;
+      instructions_visited_in_backward_pass_;
 
-  // The set of HloComputations that have been visited in the mutation pass.
+  // The set of HloComputations that have been visited in the
+  // opportunity-finding pass.
   tensorflow::gtl::FlatSet<const HloComputation*>
-      computations_visited_in_mutation_pass_;
+      computations_visited_in_backward_pass_;
 
   // ***************************
   // Functions called by the final inconsistency resolving pass.
@@ -116,7 +126,7 @@ class BFloat16Propagation : public HloPassInterface {
   // Adjusts the output shapes of HloInstructions such that if two
   // HloInstructions have aliasing buffers in their outputs, they must have the
   // same precision.
-  Status ResolveInconsistencyOfAliasingBuffers(HloModule* module);
+  void ResolveInconsistencyOfAliasingBuffers(HloModule* module);
 
   // Resolves inconsistency of aliasing buffers for the given computation, and
   // recursively runs on a while instruction's condition and body until a fixed
@@ -134,9 +144,19 @@ class BFloat16Propagation : public HloPassInterface {
   void AdjustCalledComputationRoot(HloInstruction* hlo);
 
   // ***************************
-  // Removes no-op conversions (same source and target shapes) that can be
-  // produced this pass.
-  Status RemoveNoopConversions(HloModule* module);
+  // Functions called after changes in changes_to_bf16_ are applied.
+
+  // Resolves inconsistencies introduced by this pass for fusions with
+  // tuple-type output.
+  Status ResolveInconsistentFusions(HloModule* module);
+
+  // Converts the literals in kConstant HLOs which have their types changed to
+  // BF16 by this pass.
+  Status ResolveConvertedConstants(HloModule* module);
+
+  // Skips no-op conversions (same source and target shapes) that can be
+  // produced this pass, i.e., replaces them in their uses with their operands.
+  Status SkipNoopConversions(HloModule* module);
 
   // ***************************
   // Functions called and state used by two or more passes.
@@ -146,6 +166,23 @@ class BFloat16Propagation : public HloPassInterface {
   bool AllUsersConsumeBF16(const HloInstruction& hlo,
                            const ShapeIndex& index) const;
 
+  // The output element type of the HLO at the given shape index after changes
+  // in changes_to_bf16_ are applied.
+  PrimitiveType OutputTypeAfterChange(HloInstruction* hlo,
+                                      const ShapeIndex& index) const;
+
+  // The element type of the HLO value after changes in changes_to_bf16_ are
+  // applied.
+  PrimitiveType ValueTypeAfterChange(const HloValue* value) const;
+
+  // If target_type == BF16, adds the HLO at the given index to
+  // changes_to_bf16_; otherwise, target_type must be F32 and this function
+  // removes the HLO at the given index from changes_to_bf16_ if it was earlier
+  // added.
+  void AddToOrRemoveFromBF16ChangeSet(HloInstruction* hlo,
+                                      const ShapeIndex& index,
+                                      PrimitiveType target_type);
+
   // The set of F32 HLO values that must be kept in F32.
   tensorflow::gtl::FlatSet<const HloValue*> values_that_must_be_kept_as_f32_;
 
@@ -153,10 +190,28 @@ class BFloat16Propagation : public HloPassInterface {
   // module. Populated at the beginning of this pass.
   tensorflow::gtl::FlatMap<const HloComputation*, int64> caller_counts_;
 
+  // We first store the potential F32-to-BF16 changes to changes_to_bf16_, which
+  // are subject to further adjustment, then finally applied to the HLOs. This
+  // avoids setting changed_ to true but all changes are reverted during
+  // adjustment.
+  struct IndexHasher {
+    int64 operator()(const ShapeIndex& index) const {
+      int64 hash = 0;
+      for (int64 i : index) {
+        hash = tensorflow::Hash64Combine(hash, std::hash<int64>()(i));
+      }
+      return hash;
+    }
+  };
+  tensorflow::gtl::FlatMap<HloInstruction*,
+                           tensorflow::gtl::FlatSet<ShapeIndex, IndexHasher>>
+      changes_to_bf16_;
+
+  // Whether the last processed HLO module has been changed by this pass.
+  bool changed_ = false;
+
   const BFloat16Support* bfloat16_support_;
   std::unique_ptr<HloDataflowAnalysis> dataflow_;
-
-  bool changed_ = false;
 };
 
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc b/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc
index 183db1652e..313910a861 100644
--- a/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc
+++ b/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc
@@ -323,6 +323,37 @@ TEST_F(BFloat16PropagationTest, PropagateThroughFusion) {
   EXPECT_TRUE(OutputsBF16(b_f1));
 }
 
+// Tests that changes to BF16 that cannot be propagated outside a fusion are
+// discarded.
+TEST_F(BFloat16PropagationTest, DiscardFusionInternalBF16Changes) {
+  auto module = CreateNewModule();
+  auto builder = HloComputation::Builder(TestName());
+  Shape shape = ShapeUtil::MakeShape(F32, {2, 4});
+
+  HloInstruction* param = builder.AddInstruction(
+      HloInstruction::CreateParameter(0, shape, "param"));
+  HloInstruction* add = builder.AddInstruction(
+      HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param, param));
+
+  auto builder_f = HloComputation::Builder("fusion");
+  HloInstruction* a_f =
+      builder_f.AddInstruction(HloInstruction::CreateParameter(0, shape, "a"));
+  HloInstruction* b_f =
+      builder_f.AddInstruction(HloInstruction::CreateParameter(1, shape, "b"));
+  HloInstruction* add_f = builder_f.AddInstruction(
+      HloInstruction::CreateBinary(shape, HloOpcode::kAdd, a_f, b_f));
+  HloInstruction* dot_f = builder_f.AddInstruction(HloInstruction::CreateBinary(
+      ShapeUtil::MakeShape(F32, {4, 4}), HloOpcode::kDot, add_f, add_f));
+  auto comp_f = module->AddEmbeddedComputation(builder_f.Build());
+  auto fusion = builder.AddInstruction(HloInstruction::CreateFusion(
+      dot_f->shape(), HloInstruction::FusionKind::kCustom, {add, add}, comp_f));
+
+  auto computation = module->AddEntryComputation(builder.Build());
+
+  EXPECT_FALSE(PropagatePrecision(module.get()));
+  EXPECT_EQ(computation->root_instruction(), fusion);
+}
+
 // Tests that if 1) the root instruction of a fusion is a tuple, 2) the fusion
 // outputs are only used by a dot, and 3) one element of the tuple is used by
 // an add in the fusion computation, then the propagation pass should create a
-- 
GitLab


From 4704ae7af1918755d72f159f49d98d35da6eb6fa Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 15:21:17 -0700
Subject: [PATCH 174/395] Optimize LogicalOr and LogicalAnd with all true or
 false inputs:

LogicalOr(x, true) = true
LogicalOr(x, false) = x
LogicalAnd(x, true) = x
LogicalAnd(x, false) = false

and similar if the first argument is constant.

PiperOrigin-RevId: 195161140
---
 .../grappler/optimizers/constant_folding.cc   | 113 +++++++++++-------
 .../grappler/optimizers/constant_folding.h    |  13 +-
 .../optimizers/constant_folding_test.cc       |  50 ++++++--
 .../feature_column/feature_column_test.py     |  20 +++-
 4 files changed, 132 insertions(+), 64 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/constant_folding.cc b/tensorflow/core/grappler/optimizers/constant_folding.cc
index 4801f18619..47d8827686 100644
--- a/tensorflow/core/grappler/optimizers/constant_folding.cc
+++ b/tensorflow/core/grappler/optimizers/constant_folding.cc
@@ -866,6 +866,25 @@ Status CreateConstantTensorAttrValue(DataType type, double value,
 }
 
 #undef SET_TENSOR_CAL_CASE
+
+DataType GetDataTypeFromNodeOrProps(const NodeDef& node,
+                                    const GraphProperties& properties) {
+  DataType dtype = DT_INVALID;
+  if (node.attr().count("T") == 1) {
+    dtype = node.attr().at("T").type();
+  } else if (node.attr().count("dtype") == 1) {
+    dtype = node.attr().at("dtype").type();
+  } else if (IsLogicalOr(node) || IsLogicalAnd(node)) {
+    dtype = DT_BOOL;
+  } else {
+    auto output_props = properties.GetOutputProperties(node.name());
+    if (!output_props.empty()) {
+      dtype = output_props[0].dtype();
+    }
+  }
+  return dtype;
+}
+
 }  // namespace
 
 // static
@@ -1412,6 +1431,7 @@ bool ConstantFolding::IsOnes(const NodeDef& node) const {
   }
   const auto dtype = node.attr().at("dtype").type();
   switch (dtype) {
+    IS_ONES_CASE(DT_BOOL);
     IS_ONES_CASE(DT_HALF);
     IS_ONES_CASE(DT_BFLOAT16);
     IS_ONES_CASE(DT_FLOAT);
@@ -1447,6 +1467,7 @@ bool ConstantFolding::IsZeros(const NodeDef& node) const {
   }
   const auto dtype = node.attr().at("dtype").type();
   switch (dtype) {
+    IS_ZEROS_CASE(DT_BOOL);
     IS_ZEROS_CASE(DT_HALF);
     IS_ZEROS_CASE(DT_BFLOAT16);
     IS_ZEROS_CASE(DT_FLOAT);
@@ -1466,14 +1487,15 @@ bool ConstantFolding::IsZeros(const NodeDef& node) const {
   return false;
 }
 
-void ConstantFolding::ReplaceOperationWithIdentity(int input_to_forward,
-                                                   NodeDef* node,
-                                                   GraphDef* graph) {
+void ConstantFolding::ReplaceOperationWithIdentity(
+    int input_to_forward, const GraphProperties& properties, NodeDef* node,
+    GraphDef* graph) {
+  const DataType dtype = GetDataTypeFromNodeOrProps(*node, properties);
+  if (dtype == DT_INVALID) return;
+
   node->set_op("Identity");
-  DataType dtype = node->attr().at("T").type();
   node->clear_attr();
   (*node->mutable_attr())["T"].set_type(dtype);
-
   // Propagate the designated input through the identity.
   node->mutable_input()->SwapElements(0, input_to_forward);
   // Add all other inputs as control dependencies.
@@ -1489,14 +1511,15 @@ void ConstantFolding::ReplaceOperationWithIdentity(int input_to_forward,
   graph_modified_ = true;
 }
 
-void ConstantFolding::ReplaceOperationWithSnapshot(int input_to_forward,
-                                                   NodeDef* node,
-                                                   GraphDef* graph) {
+void ConstantFolding::ReplaceOperationWithSnapshot(
+    int input_to_forward, const GraphProperties& properties, NodeDef* node,
+    GraphDef* graph) {
+  const DataType dtype = GetDataTypeFromNodeOrProps(*node, properties);
+  if (dtype == DT_INVALID) return;
+
   node->set_op("Snapshot");
-  DataType dtype = node->attr().at("T").type();
   node->clear_attr();
   (*node->mutable_attr())["T"].set_type(dtype);
-
   // Propagate the designated input through the Snapshot.
   node->mutable_input()->SwapElements(0, input_to_forward);
   // Add all other inputs as control dependencies.
@@ -1535,15 +1558,18 @@ void ConstantFolding::ReplaceSubtractionFromZeroByNegation(NodeDef* node,
 }
 
 Status ConstantFolding::ReplaceOperationWithConstant(
-    double value, const AttrValue& dtype_attr, const TensorShapeProto& shape,
-    NodeDef* node, GraphDef* graph) {
+    double value, const GraphProperties& properties,
+    const TensorShapeProto& shape, NodeDef* node, GraphDef* graph) {
+  const DataType dtype = GetDataTypeFromNodeOrProps(*node, properties);
+  if (dtype == DT_INVALID) return Status::OK();
+
   AttrValue tensor_attr;
-  TF_RETURN_IF_ERROR(CreateConstantTensorAttrValue(dtype_attr.type(), value,
-                                                   shape, &tensor_attr));
+  TF_RETURN_IF_ERROR(
+      CreateConstantTensorAttrValue(dtype, value, shape, &tensor_attr));
+  node->set_op("Const");
   node->clear_attr();
-  node->mutable_attr()->insert({"dtype", dtype_attr});
+  (*node->mutable_attr())["dtype"].set_type(dtype);
   node->mutable_attr()->insert({"value", tensor_attr});
-  node->set_op("Const");
   // Convert all inputs to control dependencies.
   for (int i = 0; i < node->input_size(); ++i) {
     if (IsControlInput(node->input(i))) {
@@ -1566,12 +1592,12 @@ Status ConstantFolding::SimplifyGraph(GraphDef* optimized_graph,
     NodeDef* node = optimized_graph->mutable_node(i);
 
     if (IsSplit(*node) && node->attr().at("num_split").i() == 1) {
-      ReplaceOperationWithIdentity(1, node, optimized_graph);
+      ReplaceOperationWithIdentity(1, *properties, node, optimized_graph);
       continue;
     }
 
     if (IsSplitV(*node) && node->attr().at("num_split").i() == 1) {
-      ReplaceOperationWithIdentity(0, node, optimized_graph);
+      ReplaceOperationWithIdentity(0, *properties, node, optimized_graph);
       continue;
     }
 
@@ -1611,7 +1637,7 @@ Status ConstantFolding::SimplifyGraph(GraphDef* optimized_graph,
           replaceable &= shape.dim(j).size() == 1 || j == permutation[j];
         }
         if (replaceable) {
-          ReplaceOperationWithIdentity(0, node, optimized_graph);
+          ReplaceOperationWithIdentity(0, *properties, node, optimized_graph);
           continue;
         }
       }
@@ -1626,7 +1652,7 @@ Status ConstantFolding::SimplifyGraph(GraphDef* optimized_graph,
       // unknown_rank == false && (dim_size == 0 || first dim is of size 1)
       if (!shape.unknown_rank() &&
           (shape.dim_size() == 0 || shape.dim(0).size() == 1)) {
-        ReplaceOperationWithIdentity(0, node, optimized_graph);
+        ReplaceOperationWithIdentity(0, *properties, node, optimized_graph);
         continue;
       }
     }
@@ -1651,11 +1677,11 @@ Status ConstantFolding::SimplifyGraph(GraphDef* optimized_graph,
         for (int j = 0; j < axis.NumElements(); ++j) {
           // value of axis can be negative.
           if (axis.dtype() == DT_INT64) {
-            target_axes.insert(
-                (axis.vec<int64>()(j) + shape.dim_size()) % shape.dim_size());
+            target_axes.insert((axis.vec<int64>()(j) + shape.dim_size()) %
+                               shape.dim_size());
           } else {
-            target_axes.insert(
-                (axis.vec<int>()(j) + shape.dim_size()) % shape.dim_size());
+            target_axes.insert((axis.vec<int>()(j) + shape.dim_size()) %
+                               shape.dim_size());
           }
         }
 
@@ -1669,7 +1695,7 @@ Status ConstantFolding::SimplifyGraph(GraphDef* optimized_graph,
                          target_axes.find(j) == target_axes.end();
         }
         if (replaceable) {
-          ReplaceOperationWithIdentity(0, node, optimized_graph);
+          ReplaceOperationWithIdentity(0, *properties, node, optimized_graph);
           continue;
         }
       }
@@ -1711,7 +1737,7 @@ Status ConstantFolding::SimplifyGraph(GraphDef* optimized_graph,
           }
         }
         if (replaceable) {
-          ReplaceOperationWithIdentity(0, node, optimized_graph);
+          ReplaceOperationWithIdentity(0, *properties, node, optimized_graph);
           continue;
         }
       }
@@ -1740,7 +1766,7 @@ Status ConstantFolding::SimplifyGraph(GraphDef* optimized_graph,
           }
         }
         if (replaceable) {
-          ReplaceOperationWithIdentity(0, node, optimized_graph);
+          ReplaceOperationWithIdentity(0, *properties, node, optimized_graph);
           continue;
         }
       }
@@ -1764,7 +1790,7 @@ Status ConstantFolding::SimplifyGraph(GraphDef* optimized_graph,
           replaceable &= flatten(j) == 0;
         }
         if (replaceable) {
-          ReplaceOperationWithIdentity(0, node, optimized_graph);
+          ReplaceOperationWithIdentity(0, *properties, node, optimized_graph);
           continue;
         }
       }
@@ -1784,7 +1810,7 @@ Status ConstantFolding::SimplifyGraph(GraphDef* optimized_graph,
         replaceable &= shape.dim(j).size() > 1;
       }
       if (replaceable) {
-        ReplaceOperationWithIdentity(0, node, optimized_graph);
+        ReplaceOperationWithIdentity(0, *properties, node, optimized_graph);
         continue;
       }
     }
@@ -1996,9 +2022,9 @@ Status ConstantFolding::SimplifyGraph(GraphDef* optimized_graph,
       continue;
     }
 
-    const bool is_mul = IsMul(*node);
+    const bool is_mul = IsMul(*node) || IsLogicalAnd(*node);
     const bool is_matmul = IsMatMul(*node);
-    const bool is_add = IsAdd(*node) || IsBiasAdd(*node);
+    const bool is_add = IsAdd(*node) || IsBiasAdd(*node) || IsLogicalOr(*node);
     const bool is_sub = IsSub(*node);
     const bool is_any_div = IsAnyDiv(*node);
     // Simplify arithmetic operations with ones or zeros.
@@ -2025,7 +2051,7 @@ Status ConstantFolding::SimplifyGraph(GraphDef* optimized_graph,
       if (y_matches_output_shape &&
           ((is_mul && x_is_one) || (is_add && x_is_zero))) {
         // 1 * y = y or 0 + y = y.
-        ReplaceOperationWithSnapshot(1, node, optimized_graph);
+        ReplaceOperationWithSnapshot(1, *properties, node, optimized_graph);
         continue;
       }
 
@@ -2052,10 +2078,18 @@ Status ConstantFolding::SimplifyGraph(GraphDef* optimized_graph,
       if (x_matches_output_shape && (((is_mul || is_any_div) && y_is_one) ||
                                      ((is_add || is_sub) && y_is_zero))) {
         // x * 1 = x or x / 1 = x or x +/- 0 = x
-        ReplaceOperationWithSnapshot(0, node, optimized_graph);
+        ReplaceOperationWithSnapshot(0, *properties, node, optimized_graph);
         continue;
       }
 
+      // x OR true = true OR y = true.
+      const PartialTensorShape shp(output_shape);
+      if (shp.IsFullyDefined() && IsLogicalOr(*node) &&
+          (y_is_one || x_is_one)) {
+        TF_RETURN_IF_ERROR(ReplaceOperationWithConstant(
+            1, *properties, output_shape, node, optimized_graph));
+      }
+
       // Simplify multiplication and matmul by zeros.
       // Also optimize zeros divided by a tensor, but only if we are in
       // aggressive mode, since we might get rid of divisions by zero.
@@ -2063,26 +2097,19 @@ Status ConstantFolding::SimplifyGraph(GraphDef* optimized_graph,
           is_any_div && x_is_zero && is_aggressive;
       if ((x_is_zero || y_is_zero) &&
           (is_mul || is_matmul || optimize_zeros_divided_by_y)) {
-        const PartialTensorShape shp(output_shape);
         if (shp.IsFullyDefined()) {
-          AttrValue dtype_attr;
-          if (node->op() == "SparseMatMul") {
-            dtype_attr.set_type(DT_FLOAT);
-          } else {
-            dtype_attr = node->attr().at("T");
-          }
           TF_RETURN_IF_ERROR(ReplaceOperationWithConstant(
-              0, dtype_attr, output_shape, node, optimized_graph));
+              0, *properties, output_shape, node, optimized_graph));
           continue;
         }
         // Even if an input shape is only partially known, we may known that it
         // matches the output shape and thus forward the corresponding zero
         // input.
         if ((is_mul || is_any_div) && x_is_zero && x_matches_output_shape) {
-          ReplaceOperationWithIdentity(0, node, optimized_graph);
+          ReplaceOperationWithIdentity(0, *properties, node, optimized_graph);
           continue;
         } else if (is_mul && y_is_zero && y_matches_output_shape) {
-          ReplaceOperationWithIdentity(1, node, optimized_graph);
+          ReplaceOperationWithIdentity(1, *properties, node, optimized_graph);
           continue;
         }
       }
diff --git a/tensorflow/core/grappler/optimizers/constant_folding.h b/tensorflow/core/grappler/optimizers/constant_folding.h
index eb06cd081f..a694f1721a 100644
--- a/tensorflow/core/grappler/optimizers/constant_folding.h
+++ b/tensorflow/core/grappler/optimizers/constant_folding.h
@@ -78,12 +78,15 @@ class ConstantFolding : public GraphOptimizer {
 
   bool IsOnes(const NodeDef& node) const;
   bool IsZeros(const NodeDef& node) const;
-  void ReplaceOperationWithIdentity(int input_to_forward, NodeDef* node,
-                                    GraphDef* graph);
-  void ReplaceOperationWithSnapshot(int input_to_forward, NodeDef* node,
-                                    GraphDef* graph);
+  void ReplaceOperationWithIdentity(int input_to_forward,
+                                    const GraphProperties& properties,
+                                    NodeDef* node, GraphDef* graph);
+  void ReplaceOperationWithSnapshot(int input_to_forward,
+                                    const GraphProperties& properties,
+                                    NodeDef* node, GraphDef* graph);
   void ReplaceSubtractionFromZeroByNegation(NodeDef* node, GraphDef* graph);
-  Status ReplaceOperationWithConstant(double value, const AttrValue& dtype_attr,
+  Status ReplaceOperationWithConstant(double value,
+                                      const GraphProperties& properties,
                                       const TensorShapeProto& shape,
                                       NodeDef* node, GraphDef* graph);
   void ReplaceDivisionOfOnesByReciprocal(NodeDef* node, GraphDef* graph);
diff --git a/tensorflow/core/grappler/optimizers/constant_folding_test.cc b/tensorflow/core/grappler/optimizers/constant_folding_test.cc
index 306ddd22d7..f018b217e6 100644
--- a/tensorflow/core/grappler/optimizers/constant_folding_test.cc
+++ b/tensorflow/core/grappler/optimizers/constant_folding_test.cc
@@ -47,18 +47,30 @@ class ConstantFoldingTest : public GrapplerTest {
     }
     Output zeros = ops::Const(s.WithOpName("zeros"), zeros_t);
     Output ones = ops::Const(s.WithOpName("ones"), ones_t);
-    Output mul1 = ops::Mul(s.WithOpName("mul1"), x, zeros);
-    Output mul2 = ops::Mul(s.WithOpName("mul2"), x, ones);
-
+    Output mul1;
+    Output mul2;
+    Output add1;
+    Output add2;
+    if (DTYPE == DT_BOOL) {
+      mul1 = ops::LogicalAnd(s.WithOpName("mul1"), x, zeros);
+      mul2 = ops::LogicalAnd(s.WithOpName("mul2"), x, ones);
+      add1 = ops::LogicalOr(s.WithOpName("add1"), x, zeros);
+      add2 = ops::LogicalOr(s.WithOpName("add2"), x, ones);
+    } else {
+      mul1 = ops::Mul(s.WithOpName("mul1"), x, zeros);
+      mul2 = ops::Mul(s.WithOpName("mul2"), x, ones);
+      add1 = ops::Add(s.WithOpName("add1"), x, zeros);
+      add1 = ops::Add(s.WithOpName("add2"), x, ones);
+    }
     GrapplerItem item;
     TF_CHECK_OK(s.ToGraphDef(&item.graph));
-    item.fetch = {"mul1", "mul2"};
+    item.fetch = {"mul1", "mul2", "add1", "add2"};
     ConstantFolding optimizer(nullptr /* cpu_device */);
     GraphDef output;
     Status status = optimizer.Optimize(nullptr, item, &output);
     TF_EXPECT_OK(status);
-    LOG(INFO) << output.DebugString();
-    EXPECT_EQ(5, output.node_size());
+
+    EXPECT_EQ(7, output.node_size());
     for (int i = 0; i < output.node_size(); ++i) {
       const NodeDef& node = output.node(i);
       const string& name = node.name();
@@ -70,14 +82,27 @@ class ConstantFoldingTest : public GrapplerTest {
         EXPECT_EQ("Snapshot", node.op());
         EXPECT_EQ("x", node.input(0));
         EXPECT_EQ("^ones", node.input(1));
+      } else if (name == "add1") {
+        EXPECT_EQ("Snapshot", node.op());
+        EXPECT_EQ("x", node.input(0));
+        EXPECT_EQ("^zeros", node.input(1));
+      } else if (name == "add2") {
+        if (DTYPE == DT_BOOL) {
+          EXPECT_EQ("Const", node.op());
+          EXPECT_EQ("^x", node.input(0));
+          EXPECT_EQ("^ones", node.input(1));
+        } else {
+          EXPECT_EQ("Add", node.op());
+          EXPECT_EQ("x", node.input(0));
+          EXPECT_EQ("ones", node.input(1));
+        }
       }
     }
-    auto tensors_expected =
-        EvaluateNodes(item.graph, {"mul1", "mul2"}, {{"x", x_t}});
-    auto tensors = EvaluateNodes(output, {"mul1", "mul2"}, {{"x", x_t}});
-    EXPECT_EQ(2, tensors_expected.size());
-    EXPECT_EQ(2, tensors.size());
-    for (int i = 0; i < 2; ++i) {
+    auto tensors_expected = EvaluateNodes(item.graph, item.fetch, {{"x", x_t}});
+    auto tensors = EvaluateNodes(output, item.fetch, {{"x", x_t}});
+    EXPECT_EQ(4, tensors_expected.size());
+    EXPECT_EQ(4, tensors.size());
+    for (int i = 0; i < item.fetch.size(); ++i) {
       test::ExpectTensorEqual<T>(tensors_expected[i], tensors[i]);
     }
   }
@@ -393,6 +418,7 @@ TEST_F(ConstantFoldingTest, NeutralElement) {
 }
 
 TEST_F(ConstantFoldingTest, NeutralElement_ShortFloats) {
+  SimpleNeutralElementTest<DT_BOOL>();
   SimpleNeutralElementTest<DT_HALF>();
   SimpleNeutralElementTest<DT_BFLOAT16>();
 }
diff --git a/tensorflow/python/feature_column/feature_column_test.py b/tensorflow/python/feature_column/feature_column_test.py
index d963dd9b55..b06540489f 100644
--- a/tensorflow/python/feature_column/feature_column_test.py
+++ b/tensorflow/python/feature_column/feature_column_test.py
@@ -25,6 +25,8 @@ import numpy as np
 
 from tensorflow.core.example import example_pb2
 from tensorflow.core.example import feature_pb2
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import rewriter_config_pb2
 from tensorflow.python.client import session
 from tensorflow.python.eager import backprop
 from tensorflow.python.eager import context
@@ -54,8 +56,8 @@ from tensorflow.python.training import coordinator
 from tensorflow.python.training import queue_runner_impl
 
 
-def _initialized_session():
-  sess = session.Session()
+def _initialized_session(config=None):
+  sess = session.Session(config=config)
   sess.run(variables_lib.global_variables_initializer())
   sess.run(lookup_ops.tables_initializer())
   return sess
@@ -6191,7 +6193,12 @@ class WeightedCategoricalColumnTest(test.TestCase):
               'values': ((.5,), (1.,))
           }, (column,),
           sparse_combiner='mean')
-      with _initialized_session():
+      # Disabling the constant folding optimizer here since it changes the
+      # error message differently on CPU and GPU.
+      config = config_pb2.ConfigProto()
+      config.graph_options.rewrite_options.constant_folding = (
+          rewriter_config_pb2.RewriterConfig.OFF)
+      with _initialized_session(config):
         with self.assertRaisesRegexp(errors.OpError, 'Incompatible shapes'):
           predictions.eval()
 
@@ -6284,7 +6291,12 @@ class WeightedCategoricalColumnTest(test.TestCase):
               'values': ((.5,), (1.,))
           }, (column,),
           sparse_combiner='mean')
-      with _initialized_session():
+      # Disabling the constant folding optimizer here since it changes the
+      # error message differently on CPU and GPU.
+      config = config_pb2.ConfigProto()
+      config.graph_options.rewrite_options.constant_folding = (
+          rewriter_config_pb2.RewriterConfig.OFF)
+      with _initialized_session(config):
         with self.assertRaisesRegexp(errors.OpError, 'Incompatible shapes'):
           predictions.eval()
 
-- 
GitLab


From 5e9e6967b47989aa9084fb328f5ef0c40fc146ef Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Wed, 2 May 2018 16:22:41 -0700
Subject: [PATCH 175/395] Replaced calls to tensorflow::StringPiece::ToString
 with std::string conversions. That is, instances of sp.ToString() are
 replaced with std::string(sp).

This will allow tensorflow::StringPiece::ToString to be removed, which is necessary before it can be replaced with absl::string_view.

PiperOrigin-RevId: 195162126
---
 tensorflow/c/c_api.cc                            |  2 +-
 tensorflow/c/c_api_test.cc                       |  4 ++--
 tensorflow/c/checkpoint_reader.cc                |  6 +++---
 tensorflow/compiler/xla/shape_util.cc            |  8 ++++----
 tensorflow/compiler/xla/text_literal_reader.cc   | 10 +++++-----
 tensorflow/compiler/xla/text_literal_writer.cc   |  2 +-
 tensorflow/core/common_runtime/bfc_allocator.cc  |  2 +-
 tensorflow/core/common_runtime/graph_runner.cc   |  4 ++--
 tensorflow/core/common_runtime/session_state.cc  |  2 +-
 .../core/common_runtime/step_stats_collector.cc  |  6 +++---
 .../core/lib/monitoring/collection_registry.cc   |  8 ++++----
 .../core/lib/monitoring/collection_registry.h    |  4 ++--
 tensorflow/core/lib/monitoring/metric_def.h      |  4 ++--
 .../core/platform/cloud/curl_http_request.cc     |  4 ++--
 .../core/platform/cloud/gcs_file_system.cc       | 14 +++++++-------
 tensorflow/core/platform/cloud/oauth_client.cc   |  4 ++--
 .../core/platform/cloud/oauth_client_test.cc     |  8 ++++----
 tensorflow/stream_executor/lib/env.h             |  2 +-
 tensorflow/stream_executor/lib/path.cc           |  2 +-
 tensorflow/stream_executor/lib/str_util.h        |  2 +-
 .../freeze_requantization_ranges.cc              |  5 ++---
 .../graph_transforms/sparsify_gather_test.cc     |  4 ++--
 .../tools/graph_transforms/transform_graph.cc    | 16 ++++++++--------
 .../tools/graph_transforms/transform_utils.cc    |  2 +-
 24 files changed, 62 insertions(+), 63 deletions(-)

diff --git a/tensorflow/c/c_api.cc b/tensorflow/c/c_api.cc
index 18eeb28168..b86b277ac3 100644
--- a/tensorflow/c/c_api.cc
+++ b/tensorflow/c/c_api.cc
@@ -2097,7 +2097,7 @@ static void GraphImportGraphDefLocked(TF_Graph* graph, const GraphDef& def,
 
   for (int i = 0; i < size; ++i) {
     TensorId id = results.missing_unused_input_map_keys[i];
-    tf_results->missing_unused_key_names_data.push_back(id.first.ToString());
+    tf_results->missing_unused_key_names_data.push_back(std::string(id.first));
     tf_results->missing_unused_key_names[i] =
         tf_results->missing_unused_key_names_data.back().c_str();
     tf_results->missing_unused_key_indexes[i] = id.second;
diff --git a/tensorflow/c/c_api_test.cc b/tensorflow/c/c_api_test.cc
index 9b86425aa5..577f10c5e6 100644
--- a/tensorflow/c/c_api_test.cc
+++ b/tensorflow/c/c_api_test.cc
@@ -1368,7 +1368,7 @@ TEST(CAPI, SavedModel) {
   }
 
   const tensorflow::string input_op_name =
-      tensorflow::ParseTensorName(input_name).first.ToString();
+      std::string(tensorflow::ParseTensorName(input_name).first);
   TF_Operation* input_op =
       TF_GraphOperationByName(graph, input_op_name.c_str());
   ASSERT_TRUE(input_op != nullptr);
@@ -1376,7 +1376,7 @@ TEST(CAPI, SavedModel) {
   ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
 
   const tensorflow::string output_op_name =
-      tensorflow::ParseTensorName(output_name).first.ToString();
+      std::string(tensorflow::ParseTensorName(output_name).first);
   TF_Operation* output_op =
       TF_GraphOperationByName(graph, output_op_name.c_str());
   ASSERT_TRUE(output_op != nullptr);
diff --git a/tensorflow/c/checkpoint_reader.cc b/tensorflow/c/checkpoint_reader.cc
index b1f7bdaa54..74bc25a491 100644
--- a/tensorflow/c/checkpoint_reader.cc
+++ b/tensorflow/c/checkpoint_reader.cc
@@ -125,7 +125,7 @@ CheckpointReader::BuildV2VarMaps() {
       const auto& slice_proto = entry.slices(i);
       CHECK(filtered_keys
                 .insert(EncodeTensorNameSlice(
-                    v2_reader_->key().ToString() /* full var's name */,
+                    std::string(v2_reader_->key()) /* full var's name */,
                     TensorSlice(slice_proto)))
                 .second);
     }
@@ -138,11 +138,11 @@ CheckpointReader::BuildV2VarMaps() {
       new TensorSliceReader::VarToDataTypeMap);
   v2_reader_->Seek(kHeaderEntryKey);
   for (v2_reader_->Next(); v2_reader_->Valid(); v2_reader_->Next()) {
-    if (filtered_keys.count(v2_reader_->key().ToString()) > 0) continue;
+    if (filtered_keys.count(std::string(v2_reader_->key())) > 0) continue;
     CHECK(entry.ParseFromArray(v2_reader_->value().data(),
                                v2_reader_->value().size()))
         << entry.InitializationErrorString();
-    string key = v2_reader_->key().ToString();
+    string key = std::string(v2_reader_->key());
     (*var_to_shape_map)[key] = TensorShape(entry.shape());
     (*var_to_data_type_map)[key] = DataType(entry.dtype());
   }
diff --git a/tensorflow/compiler/xla/shape_util.cc b/tensorflow/compiler/xla/shape_util.cc
index c330473cda..7a897f6f8f 100644
--- a/tensorflow/compiler/xla/shape_util.cc
+++ b/tensorflow/compiler/xla/shape_util.cc
@@ -511,7 +511,7 @@ StatusOr<Shape> ParseShapeStringInternal(tensorflow::StringPiece* s) {
         break;
       } else if (must_end) {
         return InvalidArgument("Expected end of tuple; got: \"%s\"",
-                               s->ToString().c_str());
+                               std::string(*s).c_str());
       }
       shapes.emplace_back();
       TF_ASSIGN_OR_RETURN(shapes.back(), ParseShapeStringInternal(s));
@@ -541,7 +541,7 @@ StatusOr<Shape> ParseShapeStringInternal(tensorflow::StringPiece* s) {
       if (!tensorflow::strings::safe_strto64(input.c_str(), &element)) {
         return InvalidArgument(
             "Invalid s64 value in parsed shape string: \"%s\" in \"%s\"",
-            input.c_str(), s->ToString().c_str());
+            input.c_str(), std::string(*s).c_str());
       }
       return element;
     };
@@ -594,7 +594,7 @@ StatusOr<Shape> ParseShapeStringInternal(tensorflow::StringPiece* s) {
   }
 
   return InvalidArgument("Invalid shape string to parse: \"%s\"",
-                         s->ToString().c_str());
+                         std::string(*s).c_str());
 }
 }  // namespace
 
@@ -603,7 +603,7 @@ StatusOr<Shape> ParseShapeStringInternal(tensorflow::StringPiece* s) {
   TF_ASSIGN_OR_RETURN(Shape shape, ParseShapeStringInternal(&s));
   if (!s.empty()) {
     return InvalidArgument("Invalid shape string to parse: \"%s\"",
-                           s.ToString().c_str());
+                           std::string(s).c_str());
   }
   return shape;
 }
diff --git a/tensorflow/compiler/xla/text_literal_reader.cc b/tensorflow/compiler/xla/text_literal_reader.cc
index 44f874cd2a..56702feab9 100644
--- a/tensorflow/compiler/xla/text_literal_reader.cc
+++ b/tensorflow/compiler/xla/text_literal_reader.cc
@@ -42,7 +42,7 @@ StatusOr<std::unique_ptr<Literal>> TextLiteralReader::ReadPath(
       << "TextLiteralReader no longer supports reading .gz files";
   std::unique_ptr<tensorflow::RandomAccessFile> file;
   Status s =
-      tensorflow::Env::Default()->NewRandomAccessFile(path.ToString(), &file);
+      tensorflow::Env::Default()->NewRandomAccessFile(std::string(path), &file);
   if (!s.ok()) {
     return s;
   }
@@ -92,7 +92,7 @@ StatusOr<std::unique_ptr<Literal>> TextLiteralReader::ReadAllLines() {
 
   tensorflow::StringPiece sp(shape_string);
   if (tensorflow::str_util::RemoveWhitespaceContext(&sp) > 0) {
-    string tmp = sp.ToString();
+    string tmp = std::string(sp);
     shape_string = tmp;
   }
   TF_ASSIGN_OR_RETURN(Shape shape, ShapeUtil::ParseShapeString(shape_string));
@@ -124,10 +124,10 @@ StatusOr<std::unique_ptr<Literal>> TextLiteralReader::ReadAllLines() {
                              line.c_str());
     }
     float value;
-    if (!tensorflow::strings::safe_strtof(value_string.ToString().c_str(),
+    if (!tensorflow::strings::safe_strtof(std::string(value_string).c_str(),
                                           &value)) {
       return InvalidArgument("could not parse value as float: \"%s\"",
-                             value_string.ToString().c_str());
+                             std::string(value_string).c_str());
     }
     SplitByDelimToStringPieces(coordinates_string, ',', &coordinates);
     coordinate_values.clear();
@@ -136,7 +136,7 @@ StatusOr<std::unique_ptr<Literal>> TextLiteralReader::ReadAllLines() {
       if (!tensorflow::strings::safe_strto64(piece, &coordinate_value)) {
         return InvalidArgument(
             "could not parse coordinate member as int64: \"%s\"",
-            piece.ToString().c_str());
+            std::string(piece).c_str());
       }
       coordinate_values.push_back(coordinate_value);
     }
diff --git a/tensorflow/compiler/xla/text_literal_writer.cc b/tensorflow/compiler/xla/text_literal_writer.cc
index 3fee467594..6e3061b78a 100644
--- a/tensorflow/compiler/xla/text_literal_writer.cc
+++ b/tensorflow/compiler/xla/text_literal_writer.cc
@@ -33,7 +33,7 @@ namespace xla {
 /* static */ tensorflow::Status TextLiteralWriter::WriteToPath(
     const Literal& literal, tensorflow::StringPiece path) {
   std::unique_ptr<tensorflow::WritableFile> f;
-  auto s = tensorflow::Env::Default()->NewWritableFile(path.ToString(), &f);
+  auto s = tensorflow::Env::Default()->NewWritableFile(std::string(path), &f);
   if (!s.ok()) {
     return s;
   }
diff --git a/tensorflow/core/common_runtime/bfc_allocator.cc b/tensorflow/core/common_runtime/bfc_allocator.cc
index e9f839289a..8f2a419756 100644
--- a/tensorflow/core/common_runtime/bfc_allocator.cc
+++ b/tensorflow/core/common_runtime/bfc_allocator.cc
@@ -616,7 +616,7 @@ string BFCAllocator::RenderOccupancy() {
     region_offset += region.memory_size();
   }
 
-  return StringPiece(rendered, resolution).ToString();
+  return std::string(rendered, resolution);
 }
 
 void BFCAllocator::DumpMemoryLog(size_t num_bytes) {
diff --git a/tensorflow/core/common_runtime/graph_runner.cc b/tensorflow/core/common_runtime/graph_runner.cc
index 790f2eaa1e..adf2ef6f44 100644
--- a/tensorflow/core/common_runtime/graph_runner.cc
+++ b/tensorflow/core/common_runtime/graph_runner.cc
@@ -56,7 +56,7 @@ class SimpleRendezvous : public Rendezvous {
     }
 
     mutex_lock l(mu_);
-    string edge_name = parsed.edge_name.ToString();
+    string edge_name = std::string(parsed.edge_name);
     if (table_.count(edge_name) > 0) {
       return errors::Internal("Send of an already sent tensor");
     }
@@ -69,7 +69,7 @@ class SimpleRendezvous : public Rendezvous {
     Tensor tensor;
     Status status = Status::OK();
     {
-      string key = parsed.edge_name.ToString();
+      string key = std::string(parsed.edge_name);
       mutex_lock l(mu_);
       if (table_.count(key) <= 0) {
         status = errors::Internal("Did not find key ", key);
diff --git a/tensorflow/core/common_runtime/session_state.cc b/tensorflow/core/common_runtime/session_state.cc
index 6befa53dff..65ff356e73 100644
--- a/tensorflow/core/common_runtime/session_state.cc
+++ b/tensorflow/core/common_runtime/session_state.cc
@@ -70,7 +70,7 @@ Status TensorStore::SaveTensors(const std::vector<string>& output_names,
     // Save only the tensors in output_names in the session.
     for (const string& name : output_names) {
       TensorId id(ParseTensorName(name));
-      const string& op_name = id.first.ToString();
+      const string& op_name = std::string(id.first);
       auto it = tensors_.find(op_name);
       if (it != tensors_.end()) {
         // Save the tensor to the session state.
diff --git a/tensorflow/core/common_runtime/step_stats_collector.cc b/tensorflow/core/common_runtime/step_stats_collector.cc
index f21536d586..af6880c6b3 100644
--- a/tensorflow/core/common_runtime/step_stats_collector.cc
+++ b/tensorflow/core/common_runtime/step_stats_collector.cc
@@ -94,7 +94,7 @@ static int ExtractGpuWithStreamAll(string device_name) {
   } else {
     // Convert the captured string into an integer. But first we need to put
     // the digits back in order
-    string ordered_capture = capture.ToString();
+    string ordered_capture = std::string(capture);
     std::reverse(ordered_capture.begin(), ordered_capture.end());
     int gpu_id;
     CHECK(strings::safe_strto32(ordered_capture, &gpu_id));
@@ -123,7 +123,7 @@ static int ExtractGpuWithoutStream(string device_name) {
   } else {
     // Convert the captured string into an integer. But first we need to put
     // the digits back in order
-    string ordered_capture = capture.ToString();
+    string ordered_capture = std::string(capture);
     std::reverse(ordered_capture.begin(), ordered_capture.end());
     int gpu_id;
     CHECK(strings::safe_strto32(ordered_capture, &gpu_id));
@@ -170,7 +170,7 @@ void StepStatsCollector::BuildCostModel(
 
   for (auto& itr : per_device_stats) {
     const StringPiece device_name = itr.first;
-    const int gpu_id = ExtractGpuWithoutStream(device_name.ToString());
+    const int gpu_id = ExtractGpuWithoutStream(std::string(device_name));
     if (gpu_id >= 0) {
       // Reference the gpu hardware stats in addition to the regular stats
       // for this gpu device if they're available.
diff --git a/tensorflow/core/lib/monitoring/collection_registry.cc b/tensorflow/core/lib/monitoring/collection_registry.cc
index d3fd7132de..8c28620ff9 100644
--- a/tensorflow/core/lib/monitoring/collection_registry.cc
+++ b/tensorflow/core/lib/monitoring/collection_registry.cc
@@ -38,15 +38,15 @@ void Collector::CollectMetricDescriptor(
     mutex_lock l(mu_);
     return collected_metrics_->metric_descriptor_map
         .insert(std::make_pair(
-            metric_def->name().ToString(),
+            std::string(metric_def->name()),
             std::unique_ptr<MetricDescriptor>(new MetricDescriptor())))
         .first->second.get();
   }();
-  metric_descriptor->name = metric_def->name().ToString();
-  metric_descriptor->description = metric_def->description().ToString();
+  metric_descriptor->name = std::string(metric_def->name());
+  metric_descriptor->description = std::string(metric_def->description());
 
   for (const StringPiece label_name : metric_def->label_descriptions()) {
-    metric_descriptor->label_names.push_back(label_name.ToString());
+    metric_descriptor->label_names.push_back(std::string(label_name));
   }
 
   metric_descriptor->metric_kind = metric_def->kind();
diff --git a/tensorflow/core/lib/monitoring/collection_registry.h b/tensorflow/core/lib/monitoring/collection_registry.h
index 63cc0f550d..20f0444f8b 100644
--- a/tensorflow/core/lib/monitoring/collection_registry.h
+++ b/tensorflow/core/lib/monitoring/collection_registry.h
@@ -72,7 +72,7 @@ class MetricCollector {
         registration_time_millis_(registration_time_millis),
         collector_(collector),
         point_set_(point_set) {
-    point_set_->metric_name = metric_def->name().ToString();
+    point_set_->metric_name = std::string(metric_def->name());
   }
 
   const MetricDef<metric_kind, Value, NumLabels>* const metric_def_;
@@ -261,7 +261,7 @@ class Collector {
     auto* const point_set = [&]() {
       mutex_lock l(mu_);
       return collected_metrics_->point_set_map
-          .insert(std::make_pair(metric_def->name().ToString(),
+          .insert(std::make_pair(std::string(metric_def->name()),
                                  std::unique_ptr<PointSet>(new PointSet())))
           .first->second.get();
     }();
diff --git a/tensorflow/core/lib/monitoring/metric_def.h b/tensorflow/core/lib/monitoring/metric_def.h
index 5ecadcc427..6f94685665 100644
--- a/tensorflow/core/lib/monitoring/metric_def.h
+++ b/tensorflow/core/lib/monitoring/metric_def.h
@@ -98,8 +98,8 @@ class AbstractMetricDef {
                     const std::vector<string>& label_descriptions)
       : kind_(kind),
         value_type_(value_type),
-        name_(name.ToString()),
-        description_(description.ToString()),
+        name_(std::string(name)),
+        description_(std::string(description)),
         label_descriptions_(std::vector<string>(label_descriptions.begin(),
                                                 label_descriptions.end())) {}
 
diff --git a/tensorflow/core/platform/cloud/curl_http_request.cc b/tensorflow/core/platform/cloud/curl_http_request.cc
index 1ac6a7531b..081d4cf043 100644
--- a/tensorflow/core/platform/cloud/curl_http_request.cc
+++ b/tensorflow/core/platform/cloud/curl_http_request.cc
@@ -407,9 +407,9 @@ size_t CurlHttpRequest::HeaderCallback(const void* ptr, size_t size,
           .StopCapture()
           .OneLiteral(": ")
           .GetResult(&value, &name)) {
-    string str_value = value.ToString();
+    string str_value = std::string(value);
     str_util::StripTrailingWhitespace(&str_value);
-    that->response_headers_[name.ToString()] = str_value;
+    that->response_headers_[std::string(name)] = str_value;
   }
   return size * nmemb;
 }
diff --git a/tensorflow/core/platform/cloud/gcs_file_system.cc b/tensorflow/core/platform/cloud/gcs_file_system.cc
index 2d9c99c124..f1e18403ec 100644
--- a/tensorflow/core/platform/cloud/gcs_file_system.cc
+++ b/tensorflow/core/platform/cloud/gcs_file_system.cc
@@ -167,13 +167,13 @@ Status ParseGcsPath(StringPiece fname, bool empty_object_ok, string* bucket,
     return errors::InvalidArgument("GCS path doesn't start with 'gs://': ",
                                    fname);
   }
-  *bucket = bucketp.ToString();
+  *bucket = std::string(bucketp);
   if (bucket->empty() || *bucket == ".") {
     return errors::InvalidArgument("GCS path doesn't contain a bucket name: ",
                                    fname);
   }
   str_util::ConsumePrefix(&objectp, "/");
-  *object = objectp.ToString();
+  *object = std::string(objectp);
   if (!empty_object_ok && object->empty()) {
     return errors::InvalidArgument("GCS path doesn't contain an object name: ",
                                    fname);
@@ -212,7 +212,7 @@ std::set<string> AddAllSubpaths(const std::vector<string>& paths) {
   for (const string& path : paths) {
     StringPiece subpath = io::Dirname(path);
     while (!subpath.empty()) {
-      result.emplace(subpath.ToString());
+      result.emplace(std::string(subpath));
       subpath = io::Dirname(subpath);
     }
   }
@@ -704,7 +704,7 @@ GcsFileSystem::GcsFileSystem()
 
       if (!header_name.empty() && !header_value.empty()) {
         additional_header_.reset(new std::pair<const string, const string>(
-            header_name.ToString(), header_value.ToString()));
+            std::string(header_name), std::string(header_value)));
 
         VLOG(1) << "GCS additional header ENABLED. "
                 << "Name: " << additional_header_->first << ", "
@@ -1095,7 +1095,7 @@ Status GcsFileSystem::GetMatchingPaths(const string& pattern,
         // Find the fixed prefix by looking for the first wildcard.
         const string& fixed_prefix =
             pattern.substr(0, pattern.find_first_of("*?[\\"));
-        const string& dir = io::Dirname(fixed_prefix).ToString();
+        const string& dir = std::string(io::Dirname(fixed_prefix));
         if (dir.empty()) {
           return errors::InvalidArgument(
               "A GCS pattern doesn't have a bucket name: ", pattern);
@@ -1192,7 +1192,7 @@ Status GcsFileSystem::GetChildrenBounded(const string& dirname,
               " doesn't match the prefix ", object_prefix));
         }
         if (!relative_path.empty() || include_self_directory_marker) {
-          result->emplace_back(relative_path.ToString());
+          result->emplace_back(std::string(relative_path));
         }
         if (++retrieved_results >= max_results) {
           return Status::OK();
@@ -1220,7 +1220,7 @@ Status GcsFileSystem::GetChildrenBounded(const string& dirname,
               "Unexpected response: the returned folder name ", prefix_str,
               " doesn't match the prefix ", object_prefix);
         }
-        result->emplace_back(relative_path.ToString());
+        result->emplace_back(std::string(relative_path));
         if (++retrieved_results >= max_results) {
           return Status::OK();
         }
diff --git a/tensorflow/core/platform/cloud/oauth_client.cc b/tensorflow/core/platform/cloud/oauth_client.cc
index 06849f9093..59ad3cbcc2 100644
--- a/tensorflow/core/platform/cloud/oauth_client.cc
+++ b/tensorflow/core/platform/cloud/oauth_client.cc
@@ -216,7 +216,7 @@ Status OAuthClient::GetTokenFromServiceAccountJson(
   // Send the request to the Google OAuth 2.0 server to get the token.
   std::unique_ptr<HttpRequest> request(http_request_factory_->Create());
   std::vector<char> response_buffer;
-  request->SetUri(oauth_server_uri.ToString());
+  request->SetUri(std::string(oauth_server_uri));
   request->SetPostFromBuffer(request_body.c_str(), request_body.size());
   request->SetResultBuffer(&response_buffer);
   TF_RETURN_IF_ERROR(request->Send());
@@ -248,7 +248,7 @@ Status OAuthClient::GetTokenFromRefreshTokenJson(
 
   std::unique_ptr<HttpRequest> request(http_request_factory_->Create());
   std::vector<char> response_buffer;
-  request->SetUri(oauth_server_uri.ToString());
+  request->SetUri(std::string(oauth_server_uri));
   request->SetPostFromBuffer(request_body.c_str(), request_body.size());
   request->SetResultBuffer(&response_buffer);
   TF_RETURN_IF_ERROR(request->Send());
diff --git a/tensorflow/core/platform/cloud/oauth_client_test.cc b/tensorflow/core/platform/cloud/oauth_client_test.cc
index ad569758cc..4ffa72288b 100644
--- a/tensorflow/core/platform/cloud/oauth_client_test.cc
+++ b/tensorflow/core/platform/cloud/oauth_client_test.cc
@@ -124,11 +124,11 @@ TEST(OAuthClientTest, GetTokenFromServiceAccountJson) {
                   .OneLiteral("&assertion=")
                   .GetResult(&assertion, &grant_type));
   EXPECT_EQ("urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer",
-            grant_type.ToString());
+            grant_type);
 
-  int last_dot = assertion.ToString().find_last_of(".");
-  string header_dot_claim = assertion.ToString().substr(0, last_dot);
-  string signature_encoded = assertion.ToString().substr(last_dot + 1);
+  int last_dot = std::string(assertion).find_last_of(".");
+  string header_dot_claim = std::string(assertion.substr(0, last_dot));
+  string signature_encoded = std::string(assertion.substr(last_dot + 1));
 
   // Check that 'signature' signs 'header_dot_claim'.
 
diff --git a/tensorflow/stream_executor/lib/env.h b/tensorflow/stream_executor/lib/env.h
index 776eba0408..3ef8deb72e 100644
--- a/tensorflow/stream_executor/lib/env.h
+++ b/tensorflow/stream_executor/lib/env.h
@@ -32,7 +32,7 @@ inline Status FileExists(const string& filename) {
 }
 
 inline Status FileExists(const port::StringPiece& filename) {
-  return Env::Default()->FileExists(filename.ToString());
+  return Env::Default()->FileExists(std::string(filename));
 }
 
 }  // namespace port
diff --git a/tensorflow/stream_executor/lib/path.cc b/tensorflow/stream_executor/lib/path.cc
index 56e08c316f..58a862206c 100644
--- a/tensorflow/stream_executor/lib/path.cc
+++ b/tensorflow/stream_executor/lib/path.cc
@@ -33,7 +33,7 @@ string JoinPathImpl(std::initializer_list<port::StringPiece> paths) {
     if (path.empty()) continue;
 
     if (result.empty()) {
-      result = path.ToString();
+      result = std::string(path);
       continue;
     }
 
diff --git a/tensorflow/stream_executor/lib/str_util.h b/tensorflow/stream_executor/lib/str_util.h
index a81c666818..b02fe4f56f 100644
--- a/tensorflow/stream_executor/lib/str_util.h
+++ b/tensorflow/stream_executor/lib/str_util.h
@@ -31,7 +31,7 @@ inline string StripSuffixString(port::StringPiece str, port::StringPiece suffix)
   if (tensorflow::str_util::EndsWith(str, suffix)) {
     str.remove_suffix(suffix.size());
   }
-  return str.ToString();
+  return std::string(str);
 }
 
 using tensorflow::str_util::Lowercase;
diff --git a/tensorflow/tools/graph_transforms/freeze_requantization_ranges.cc b/tensorflow/tools/graph_transforms/freeze_requantization_ranges.cc
index f401723808..c8dc2a7c4d 100644
--- a/tensorflow/tools/graph_transforms/freeze_requantization_ranges.cc
+++ b/tensorflow/tools/graph_transforms/freeze_requantization_ranges.cc
@@ -92,9 +92,8 @@ Status ExtractMinMaxRecords(const string& log_file_name,
     if (!str_util::EndsWith(name_string, print_suffix)) {
       continue;
     }
-    string name =
-        name_string.substr(0, name_string.size() - print_suffix.size())
-            .ToString();
+    string name = std::string(
+        name_string.substr(0, name_string.size() - print_suffix.size()));
     records->push_back({name, min, max});
   }
   return Status::OK();
diff --git a/tensorflow/tools/graph_transforms/sparsify_gather_test.cc b/tensorflow/tools/graph_transforms/sparsify_gather_test.cc
index d41321c9a6..dd95779a1f 100644
--- a/tensorflow/tools/graph_transforms/sparsify_gather_test.cc
+++ b/tensorflow/tools/graph_transforms/sparsify_gather_test.cc
@@ -42,8 +42,8 @@ class SparsifyGatherTest : public ::testing::Test {
                       const std::vector<NodeDef*>& inputs, GraphDef* graph_def,
                       bool control_dep = false) {
     NodeDef* node_def = graph_def->add_node();
-    node_def->set_name(name.ToString());
-    node_def->set_op(op.ToString());
+    node_def->set_name(std::string(name));
+    node_def->set_op(std::string(op));
     if (!control_dep) {
       std::for_each(inputs.begin(), inputs.end(), [&node_def](NodeDef* input) {
         node_def->add_input(input->name());
diff --git a/tensorflow/tools/graph_transforms/transform_graph.cc b/tensorflow/tools/graph_transforms/transform_graph.cc
index 8ce8f5e24b..3b9dd3dd2d 100644
--- a/tensorflow/tools/graph_transforms/transform_graph.cc
+++ b/tensorflow/tools/graph_transforms/transform_graph.cc
@@ -65,19 +65,19 @@ Status ParseTransformParameters(const string& transforms_string,
               .GetResult(&remaining, &transform_name);
       if (!found_transform_name) {
         return errors::InvalidArgument("Looking for transform name, but found ",
-                                       remaining.ToString().c_str());
+                                       std::string(remaining).c_str());
       }
       if (Scanner(remaining).OneLiteral("(").GetResult(&remaining, &match)) {
         state = TRANSFORM_PARAM_NAME;
       } else {
         // Add a transform with no parameters.
-        params_list->push_back({transform_name.ToString(), func_parameters});
+        params_list->push_back({std::string(transform_name), func_parameters});
         transform_name = "";
         state = TRANSFORM_NAME;
       }
     } else if (state == TRANSFORM_PARAM_NAME) {
       if (Scanner(remaining).OneLiteral(")").GetResult(&remaining, &match)) {
-        params_list->push_back({transform_name.ToString(), func_parameters});
+        params_list->push_back({std::string(transform_name), func_parameters});
         transform_name = "";
         state = TRANSFORM_NAME;
       } else {
@@ -92,13 +92,13 @@ Status ParseTransformParameters(const string& transforms_string,
         if (!found_parameter_name) {
           return errors::InvalidArgument(
               "Looking for parameter name, but found ",
-              remaining.ToString().c_str());
+              std::string(remaining).c_str());
         }
         if (Scanner(remaining).OneLiteral("=").GetResult(&remaining, &match)) {
           state = TRANSFORM_PARAM_VALUE;
         } else {
           return errors::InvalidArgument("Looking for =, but found ",
-                                         remaining.ToString().c_str());
+                                         std::string(remaining).c_str());
         }
       }
     } else if (state == TRANSFORM_PARAM_VALUE) {
@@ -120,10 +120,10 @@ Status ParseTransformParameters(const string& transforms_string,
       }
       if (!found_parameter_value) {
         return errors::InvalidArgument("Looking for parameter name, but found ",
-                                       remaining.ToString().c_str());
+                                       std::string(remaining).c_str());
       }
-      func_parameters[parameter_name.ToString()].push_back(
-          parameter_value.ToString());
+      func_parameters[std::string(parameter_name)].push_back(
+          std::string(parameter_value));
       // Eat up any trailing quotes.
       Scanner(remaining).ZeroOrOneLiteral("\"").GetResult(&remaining, &match);
       Scanner(remaining).ZeroOrOneLiteral("'").GetResult(&remaining, &match);
diff --git a/tensorflow/tools/graph_transforms/transform_utils.cc b/tensorflow/tools/graph_transforms/transform_utils.cc
index 367048965d..af17fd75bc 100644
--- a/tensorflow/tools/graph_transforms/transform_utils.cc
+++ b/tensorflow/tools/graph_transforms/transform_utils.cc
@@ -93,7 +93,7 @@ void NodeNamePartsFromInput(const string& input_name, string* prefix,
   } else {
     *prefix = "";
   }
-  *node_name = node_name_piece.ToString();
+  *node_name = std::string(node_name_piece);
 }
 
 string NodeNameFromInput(const string& input_name) {
-- 
GitLab


From 4c256cda4f29ce5f634be44628e2c4c639974dc3 Mon Sep 17 00:00:00 2001
From: Priya Gupta <priyag@google.com>
Date: Wed, 2 May 2018 15:30:30 -0700
Subject: [PATCH 176/395] Add prefetching to one device distribution strategy.

PiperOrigin-RevId: 195162570
---
 .../contrib/distribute/python/one_device_strategy.py       | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensorflow/contrib/distribute/python/one_device_strategy.py b/tensorflow/contrib/distribute/python/one_device_strategy.py
index 646d2a5c3b..64aa369201 100644
--- a/tensorflow/contrib/distribute/python/one_device_strategy.py
+++ b/tensorflow/contrib/distribute/python/one_device_strategy.py
@@ -36,9 +36,10 @@ class OneDeviceStrategy(distribute_lib.DistributionStrategy):
   # doing something that won't work with other DistributionStrategy
   # implementations?
 
-  def __init__(self, device):
+  def __init__(self, device, prefetch_on_device=None):
     super(OneDeviceStrategy, self).__init__()
     self._device = device
+    self._prefetch_on_device = prefetch_on_device
 
   def _create_variable(self, next_creator, *args, **kwargs):
     # No need to distinguish tower-local variables when not mirroring,
@@ -61,7 +62,9 @@ class OneDeviceStrategy(distribute_lib.DistributionStrategy):
       return next_creator(*args, **kwargs)
 
   def distribute_dataset(self, dataset_fn):
-    return self._call_dataset_fn(dataset_fn)
+    return values.PerDeviceDataset(
+        self._call_dataset_fn(dataset_fn), [self._device],
+        self._prefetch_on_device)
 
   def _broadcast(self, tensor, destinations):
     return tensor
-- 
GitLab


From 85566b2420833a4ba59241330eeceedea4f98e3c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 15:35:11 -0700
Subject: [PATCH 177/395] Adding a version of rolled triangular solver code for
 the right-multiply case, which is used in Cholesky decomposition.  Replacing
 the unrolled version with a While loop drastically reduces XLA compilation
 times which allows much larger models to be run on TPU.

PiperOrigin-RevId: 195163298
---
 .../compiler/tf2xla/lib/triangular_solve.cc   | 179 +++++++++++++++---
 .../compiler/tf2xla/lib/triangular_solve.h    |   6 +
 tensorflow/compiler/tf2xla/lib/util.cc        |   7 +
 tensorflow/compiler/tf2xla/lib/util.h         |   5 +
 4 files changed, 173 insertions(+), 24 deletions(-)

diff --git a/tensorflow/compiler/tf2xla/lib/triangular_solve.cc b/tensorflow/compiler/tf2xla/lib/triangular_solve.cc
index d0279d4412..b4503601f9 100644
--- a/tensorflow/compiler/tf2xla/lib/triangular_solve.cc
+++ b/tensorflow/compiler/tf2xla/lib/triangular_solve.cc
@@ -1,4 +1,4 @@
-/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2018 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.
@@ -82,13 +82,6 @@ xla::StatusOr<xla::XlaOp> TriangularSolve(xla::XlaBuilder* builder,
         block_size);
   }
 
-  // Applies a complex conjugation operation if `a` is complex and `conjugate_a`
-  // is true, otherwise returns its argument.
-  auto maybe_conj = [&](xla::XlaBuilder* builder, xla::XlaOp x) {
-    auto perform_conj = a_shape.element_type() == xla::C64 && conjugate_a;
-    return perform_conj ? builder->Conj(x) : x;
-  };
-
   std::map<int, xla::XlaComputation> base_computations;
   auto get_base_triangular_solve =
       [&](int k) -> xla::StatusOr<xla::XlaComputation*> {
@@ -117,16 +110,21 @@ xla::StatusOr<xla::XlaOp> TriangularSolve(xla::XlaBuilder* builder,
               PrependMajorDims(sub.get(), batch_dimensions, b_lastd)),
           "b");
 
-      // We use a left-looking subroutine on the block diagonal in some common
-      // cases, while falling back to a recursive call in unsupported cases. The
-      // left-looking subroutine is written with a While loop and so yields much
-      // faster compile times. Moreover, the left-looking variant can give
-      // higher performance on smaller (sub)problems.
+      // We use a left-looking or right-looking subroutine on the block diagonal
+      // in the lower=true cases, while falling back to a recursive call in
+      // others. The left-looking and right-looking subroutines are written with
+      // a While loop and so yields much faster compile times. Moreover, they
+      // can give higher performance on smaller (sub)problems.
       if (left_side && lower) {
         TF_RETURN_IF_ERROR(TriangularSolveLeftLooking(sub.get(), a_param,
                                                       b_param, transpose_a,
                                                       conjugate_a)
                                .status());
+      } else if (!left_side && lower) {
+        TF_RETURN_IF_ERROR(TriangularSolveRightLooking(sub.get(), a_param,
+                                                       b_param, transpose_a,
+                                                       conjugate_a)
+                               .status());
       } else {
         TF_RETURN_IF_ERROR(TriangularSolve(sub.get(), a_param, b_param,
                                            left_side, lower, transpose_a,
@@ -169,7 +167,9 @@ xla::StatusOr<xla::XlaOp> TriangularSolve(xla::XlaBuilder* builder,
                             get_base_triangular_solve(k));
         update = builder->Call(*solve, {a_slice, b_slice});
       } else {
-        update = builder->Div(b_slice, maybe_conj(builder, a_slice));
+        TF_ASSIGN_OR_RETURN(auto a_slice_conj,
+                            MaybeConjugate(builder, a_slice, conjugate_a));
+        update = builder->Div(b_slice, a_slice_conj);
       }
       TF_ASSIGN_OR_RETURN(
           output, UpdateSliceInMinorDims(builder, output, update, {0, i}));
@@ -219,7 +219,9 @@ xla::StatusOr<xla::XlaOp> TriangularSolve(xla::XlaBuilder* builder,
                             get_base_triangular_solve(k));
         update = builder->Call(*solve, {a_slice, b_slice});
       } else {
-        update = builder->Div(b_slice, maybe_conj(builder, a_slice));
+        TF_ASSIGN_OR_RETURN(auto a_slice_conj,
+                            MaybeConjugate(builder, a_slice, conjugate_a));
+        update = builder->Div(b_slice, a_slice_conj);
       }
       TF_ASSIGN_OR_RETURN(
           output, UpdateSliceInMinorDims(builder, output, update, {i, 0}));
@@ -268,7 +270,9 @@ xla::StatusOr<xla::XlaOp> TriangularSolve(xla::XlaBuilder* builder,
                             get_base_triangular_solve(k));
         update = builder->Call(*solve, {a_slice, b_slice});
       } else {
-        update = builder->Div(b_slice, maybe_conj(builder, a_slice));
+        TF_ASSIGN_OR_RETURN(auto a_slice_conj,
+                            MaybeConjugate(builder, a_slice, conjugate_a));
+        update = builder->Div(b_slice, a_slice_conj);
       }
       TF_ASSIGN_OR_RETURN(
           output, UpdateSliceInMinorDims(builder, output, update, {0, i}));
@@ -318,7 +322,9 @@ xla::StatusOr<xla::XlaOp> TriangularSolve(xla::XlaBuilder* builder,
                             get_base_triangular_solve(k));
         update = builder->Call(*solve, {a_slice, b_slice});
       } else {
-        update = builder->Div(b_slice, maybe_conj(builder, a_slice));
+        TF_ASSIGN_OR_RETURN(auto a_slice_conj,
+                            MaybeConjugate(builder, a_slice, conjugate_a));
+        update = builder->Div(b_slice, a_slice_conj);
       }
       TF_ASSIGN_OR_RETURN(
           output, UpdateSliceInMinorDims(builder, output, update, {i, 0}));
@@ -371,11 +377,6 @@ xla::StatusOr<xla::XlaOp> TriangularSolveLeftLooking(xla::XlaBuilder* builder,
     batch_dimensions.push_back(a_size);
   }
 
-  auto maybe_conj = [&](xla::XlaBuilder* builder, xla::XlaOp x) {
-    auto perform_conj = a_shape.element_type() == xla::C64 && conjugate_a;
-    return perform_conj ? builder->Conj(x) : x;
-  };
-
   // The main computation is performed in a While loop.
 
   // Allocate the output and set its first or last row,
@@ -391,7 +392,9 @@ xla::StatusOr<xla::XlaOp> TriangularSolveLeftLooking(xla::XlaBuilder* builder,
                         SliceInMinorDims(builder, a, {i, i}, {i + 1, i + 1}));
     TF_ASSIGN_OR_RETURN(auto b_slice,
                         SliceInMinorDims(builder, b, {i, 0}, {i + 1, n}));
-    auto update = builder->Div(b_slice, maybe_conj(builder, a_slice));
+    TF_ASSIGN_OR_RETURN(auto a_slice_conj,
+                        MaybeConjugate(builder, a_slice, conjugate_a));
+    auto update = builder->Div(b_slice, a_slice_conj);
     TF_ASSIGN_OR_RETURN(
         output, UpdateSliceInMinorDims(builder, output, update, {i, 0}));
   }
@@ -493,7 +496,9 @@ xla::StatusOr<xla::XlaOp> TriangularSolveLeftLooking(xla::XlaBuilder* builder,
     // body_out[..., i:i+1, :] = result_row / a[..., i:i+1, i:i+1]
     TF_ASSIGN_OR_RETURN(auto a_elt, DynamicSliceInMinorDims(bodyb.get(), body_a,
                                                             {i, i}, {1, 1}));
-    auto div_result = bodyb->Div(result_row, maybe_conj(bodyb.get(), a_elt));
+    TF_ASSIGN_OR_RETURN(auto a_elt_conj,
+                        MaybeConjugate(bodyb.get(), a_elt, conjugate_a));
+    auto div_result = bodyb->Div(result_row, a_elt_conj);
     TF_ASSIGN_OR_RETURN(body_out,
                         DynamicUpdateSliceInMinorDims(bodyb.get(), body_out,
                                                       div_result, {i, zero}));
@@ -513,4 +518,130 @@ xla::StatusOr<xla::XlaOp> TriangularSolveLeftLooking(xla::XlaBuilder* builder,
   return builder->GetTupleElement(triangular_solve_left_looking_while, 1);
 }
 
+xla::StatusOr<xla::XlaOp> TriangularSolveRightLooking(xla::XlaBuilder* builder,
+                                                      const xla::XlaOp& a,
+                                                      const xla::XlaOp& b,
+                                                      bool transpose_a,
+                                                      bool conjugate_a) {
+  TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a));
+  TF_ASSIGN_OR_RETURN(xla::Shape b_shape, builder->GetShape(b));
+  const int64 m = xla::ShapeUtil::GetDimension(b_shape, -2);
+  const int64 n = xla::ShapeUtil::GetDimension(b_shape, -1);
+  const int64 ndims = xla::ShapeUtil::Rank(a_shape);
+
+  std::vector<int64> batch_dimensions;
+  for (int i = 0; i < ndims - 2; ++i) {
+    int64 a_size = a_shape.dimensions(i);
+    batch_dimensions.push_back(a_size);
+  }
+
+  // The main computation is performed in a While loop.
+  xla::XlaOp output = Zeros(builder, b_shape);
+
+  // Construct the initial loop carry tuple,
+  // if transpose_a:
+  //   init = (0, output, a, b)
+  // else:
+  //   init = (n-1, output, a, b)
+  std::vector<xla::Shape> tuple_shapes = {
+      // The loop iteration counter is a scalar, incremented each iteration.
+      xla::ShapeUtil::MakeShape(xla::S32, {}),
+      // The output has the shape of b, with one row updated each iteration.
+      b_shape,
+      // The coefficient matrix a is a loop invariant.
+      a_shape,
+      // The right-hand-side matrix b is a loop invariant.
+      b_shape};
+  xla::Shape tuple_shape = xla::ShapeUtil::MakeTupleShape(tuple_shapes);
+  auto init_i = builder->ConstantR0<int32>(transpose_a ? 0 : n - 1);
+  auto init = builder->Tuple({init_i, output, a, b});
+
+  // Construct the loop condition function,
+  // def cond_fun(loop_carry):
+  //   i, output, a, b = loop_carry
+  //   return i < n if transpose_a else i >= 0
+  std::unique_ptr<xla::XlaBuilder> condb =
+      builder->CreateSubBuilder("TriangularSolveRightLookingWhileCond");
+  {
+    auto i = condb->GetTupleElement(
+        condb->Parameter(0, tuple_shape,
+                         "TriangularSolveRightLookingWhileTuple"),
+        0);
+    if (transpose_a) {
+      condb->Lt(i, condb->ConstantR0<int32>(n));
+    } else {
+      condb->Ge(i, condb->ConstantR0<int32>(0));
+    }
+  }
+  TF_ASSIGN_OR_RETURN(auto cond, condb->Build());
+
+  // Construct the loop body function,
+  // def body_fun(loop_carry):
+  //   i, output, a, b = loop_carry
+  //   if transpose_a:
+  //     a_row = np.swapaxes(a[..., :, i:i+1], -1 -2)
+  //   else:
+  //     a_row = a[..., :, i:i+1]
+  //   result_row = b[..., :, i:i+1] - np.matmul(output, a_row)
+  //   output[..., :, i:i+1] = result_row / a[..., i:i+1, i:i+1]
+  //   if transpose_a:
+  //     return (i - 1, output, a, b)
+  //   else:
+  //     return (i + 1, output, a, b)
+  // We have to do some extra FLOPs propagating zeros in the matrix multiply
+  // because we can't have the size of its arguments depend on the loop counter.
+  std::unique_ptr<xla::XlaBuilder> bodyb =
+      builder->CreateSubBuilder("TriangularSolveRightLookingWhileBody");
+  {
+    auto input_tuple = bodyb->Parameter(
+        0, tuple_shape, "TriangularSolveRightLookingWhileTuple");
+
+    // i, output, a, b = loop_carry
+    auto i = bodyb->GetTupleElement(input_tuple, 0);
+    auto body_out = bodyb->GetTupleElement(input_tuple, 1);
+    auto body_a = bodyb->GetTupleElement(input_tuple, 2);
+    auto body_b = bodyb->GetTupleElement(input_tuple, 3);
+    auto zero = bodyb->ConstantR0<int32>(0);
+
+    // We'd like to implement b[..., :, i:i+1] - np.matmul(output, a[..., :,
+    // i:i+1]) But since we can't have intermediate array sizes depend on the
+    // loop counter, we instead exploit the fact that we initialized the output
+    // to all zeros and use that as zero-padding (doing unnecessary FLOPs).
+    TF_ASSIGN_OR_RETURN(auto b_update, BatchDot(bodyb.get(), body_out, body_a,
+                                                /*transpose_x=*/false,
+                                                /*transpose_y=*/transpose_a,
+                                                /*conjugate_x=*/false,
+                                                /*conjugate_y=*/conjugate_a));
+    // result = b - np.matmul(output, a)
+    auto result = bodyb->Sub(body_b, b_update);
+    // result_row = result[..., :, i:i+1]
+    TF_ASSIGN_OR_RETURN(
+        auto result_row,
+        DynamicSliceInMinorDims(bodyb.get(), result, {zero, i}, {m, 1}));
+
+    // body_out[..., :, i:i+1] = result_row / a[..., i:i+1, i:i+1]
+    TF_ASSIGN_OR_RETURN(auto a_ii, DynamicSliceInMinorDims(bodyb.get(), body_a,
+                                                           {i, i}, {1, 1}));
+    TF_ASSIGN_OR_RETURN(auto a_ii_conj,
+                        MaybeConjugate(bodyb.get(), a_ii, conjugate_a));
+    auto div_result = bodyb->Div(result_row, a_ii_conj);
+    TF_ASSIGN_OR_RETURN(body_out,
+                        DynamicUpdateSliceInMinorDims(bodyb.get(), body_out,
+                                                      div_result, {zero, i}));
+
+    // if transpose_a:
+    //   return (i + 1, body_out, a, b)
+    // else:
+    //   return (i - 1, body_out, a, b)
+    auto next_i = bodyb->Add(i, bodyb->ConstantR0<int32>(transpose_a ? 1 : -1));
+    bodyb->Tuple({next_i, body_out, body_a, body_b});
+  }
+  TF_ASSIGN_OR_RETURN(auto body, bodyb->Build());
+
+  // Construct the While loop and return the result,
+  // return while_loop(cond_fun, body_fun, init)[1]
+  auto triangular_solve_left_looking_while = builder->While(cond, body, init);
+  return builder->GetTupleElement(triangular_solve_left_looking_while, 1);
+}
+
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/tf2xla/lib/triangular_solve.h b/tensorflow/compiler/tf2xla/lib/triangular_solve.h
index fd8f2489d1..540c26b247 100644
--- a/tensorflow/compiler/tf2xla/lib/triangular_solve.h
+++ b/tensorflow/compiler/tf2xla/lib/triangular_solve.h
@@ -69,6 +69,12 @@ xla::StatusOr<xla::XlaOp> TriangularSolveLeftLooking(xla::XlaBuilder* builder,
                                                      bool transpose_a,
                                                      bool conjugate_a);
 
+xla::StatusOr<xla::XlaOp> TriangularSolveRightLooking(xla::XlaBuilder* builder,
+                                                      const xla::XlaOp& a,
+                                                      const xla::XlaOp& b,
+                                                      bool transpose_a,
+                                                      bool conjugate_a);
+
 }  // namespace tensorflow
 
 #endif  // TENSORFLOW_COMPILER_TF2XLA_LIB_TRIANGULAR_SOLVE_H_
diff --git a/tensorflow/compiler/tf2xla/lib/util.cc b/tensorflow/compiler/tf2xla/lib/util.cc
index cc7b13571c..d9ff7e6259 100644
--- a/tensorflow/compiler/tf2xla/lib/util.cc
+++ b/tensorflow/compiler/tf2xla/lib/util.cc
@@ -230,4 +230,11 @@ xla::StatusOr<xla::XlaOp> TransposeInMinorDims(xla::XlaBuilder* builder,
   return builder->Transpose(x, permutation);
 }
 
+xla::StatusOr<xla::XlaOp> MaybeConjugate(xla::XlaBuilder* builder,
+                                         const xla::XlaOp& x, bool conjugate) {
+  TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(x));
+  auto perform_conj = shape.element_type() == xla::C64 && conjugate;
+  return perform_conj ? builder->Conj(x) : x;
+}
+
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/tf2xla/lib/util.h b/tensorflow/compiler/tf2xla/lib/util.h
index 3df44ef035..3c120a2548 100644
--- a/tensorflow/compiler/tf2xla/lib/util.h
+++ b/tensorflow/compiler/tf2xla/lib/util.h
@@ -85,6 +85,11 @@ xla::StatusOr<xla::XlaOp> DynamicUpdateSliceInMinorDims(
 xla::StatusOr<xla::XlaOp> TransposeInMinorDims(xla::XlaBuilder* builder,
                                                const xla::XlaOp& x);
 
+// Applies a complex conjugation operation if `a` is complex and `conjugate_a`
+// is true, otherwise returns its argument.
+xla::StatusOr<xla::XlaOp> MaybeConjugate(xla::XlaBuilder* builder,
+                                         const xla::XlaOp& x, bool conjugate);
+
 }  // namespace tensorflow
 
 #endif  // TENSORFLOW_COMPILER_TF2XLA_LIB_UTIL_H_
-- 
GitLab


From 0237e86297087ba3e700ac9218f846e6e662c60f Mon Sep 17 00:00:00 2001
From: Jianwei Xie <xiejw@google.com>
Date: Wed, 2 May 2018 15:36:54 -0700
Subject: [PATCH 178/395] Adds the EvalListener support for run_local.

PiperOrigin-RevId: 195163507
---
 tensorflow/python/estimator/training.py      | 10 +++
 tensorflow/python/estimator/training_test.py | 67 ++++++++++++++++++++
 2 files changed, 77 insertions(+)

diff --git a/tensorflow/python/estimator/training.py b/tensorflow/python/estimator/training.py
index 534c357067..41ffa371aa 100644
--- a/tensorflow/python/estimator/training.py
+++ b/tensorflow/python/estimator/training.py
@@ -656,6 +656,11 @@ class _TrainingExecutor(object):
           max_steps=self._train_spec.max_steps,
           hooks=train_hooks)
 
+      if not self._continuous_eval_listener.before_eval():
+        logging.info('Exiting training and evaluation lopp, as requested by '
+                     '_ContinuousEvalListener.before_eval.')
+        break
+
       # Final export signal: For any eval result with global_step >= train
       # max_steps, the evaluator will send the final export signal. The
       # _should_stop_local_train will then end the while True as the stopping
@@ -669,6 +674,11 @@ class _TrainingExecutor(object):
         raise RuntimeError('There was no new checkpoint after the training. '
                            'Eval status: {}'.format(eval_result.status))
 
+      if not self._continuous_eval_listener.after_eval(eval_result):
+        logging.info('Exiting evaluation, as requested by '
+                     '_ContinuousEvalListener.after_eval.')
+        break
+
       if _should_stop_local_train(
           eval_result.metrics[ops.GraphKeys.GLOBAL_STEP]):
         break
diff --git a/tensorflow/python/estimator/training_test.py b/tensorflow/python/estimator/training_test.py
index c04905ae65..3b6f5e18cb 100644
--- a/tensorflow/python/estimator/training_test.py
+++ b/tensorflow/python/estimator/training_test.py
@@ -1628,6 +1628,73 @@ class TrainingExecutorRunLocalTest(test.TestCase):
     self.assertEqual(3, mock_est.times_export_was_called)
     self.assertEqual(1, mock_est.times_final_export_was_called)
 
+  def test_runs_with_eval_listener_before_eval(self):
+    mock_est = test.mock.Mock(spec=estimator_lib.Estimator, model_dir='path/')
+    mock_est.latest_checkpoint = self.unique_checkpoint_every_time_fn
+
+    train_spec = training.TrainSpec(input_fn=lambda: 1, max_steps=300)
+    eval_spec = training.EvalSpec(input_fn=lambda: 1, throttle_secs=100)
+    # should be called 2 times without the evallistener
+    mock_est.evaluate.side_effect = [{
+        _GLOBAL_STEP_KEY: train_spec.max_steps - 50
+    }, {
+        _GLOBAL_STEP_KEY: train_spec.max_steps
+    }]
+
+    class _Listener(training._ContinuousEvalListener):
+
+      def __init__(self):
+        self.call_count = 0
+
+      def before_eval(self):
+        self.call_count += 1
+        return False  # Will stop the run_local before first eval.
+
+    listener = _Listener()
+
+    executor = training._TrainingExecutor(
+        mock_est, train_spec, eval_spec, continuous_eval_listener=listener)
+    executor.run_local()
+
+    self.assertEqual(1, mock_est.train.call_count)
+    self.assertEqual(0, mock_est.evaluate.call_count)
+    self.assertEqual(1, listener.call_count)
+
+  def test_runs_with_eval_listener_after_eval(self):
+    mock_est = test.mock.Mock(spec=estimator_lib.Estimator, model_dir='path/')
+    mock_est.latest_checkpoint = self.unique_checkpoint_every_time_fn
+
+    train_spec = training.TrainSpec(input_fn=lambda: 1, max_steps=300)
+    eval_spec = training.EvalSpec(input_fn=lambda: 1, throttle_secs=100)
+    # should be called 2 times without the evallistener
+    mock_est.evaluate.side_effect = [{
+        _GLOBAL_STEP_KEY: train_spec.max_steps - 50
+    }, {
+        _GLOBAL_STEP_KEY: train_spec.max_steps
+    }]
+
+    class _Listener(training._ContinuousEvalListener):
+
+      def __init__(self, test_case):
+        self.call_count = 0
+        self._test_case = test_case
+
+      def after_eval(self, eval_result):
+        self.call_count += 1
+        self._test_case.assertEqual(
+            train_spec.max_steps - 50, eval_result.metrics[_GLOBAL_STEP_KEY])
+        return False  # Will stop the run_local after first eval.
+
+    listener = _Listener(test_case=self)
+
+    executor = training._TrainingExecutor(
+        mock_est, train_spec, eval_spec, continuous_eval_listener=listener)
+    executor.run_local()
+
+    self.assertEqual(1, mock_est.train.call_count)
+    self.assertEqual(1, mock_est.evaluate.call_count)
+    self.assertEqual(1, listener.call_count)
+
   def test_handles_no_new_checkpoint_found(self):
     mock_est = test.mock.Mock(spec=estimator_lib.Estimator, model_dir='path/')
     mock_est.latest_checkpoint.return_value = (
-- 
GitLab


From 49f2afe21e3cada8951205d00e877c873a33754c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 15:51:16 -0700
Subject: [PATCH 179/395] Allow evaluation and prediction through warm-starting
 (no current checkpoint / model_dir).

PiperOrigin-RevId: 195165732
---
 tensorflow/python/estimator/BUILD             |  1 +
 tensorflow/python/estimator/estimator.py      | 20 ++++-
 tensorflow/python/estimator/estimator_test.py | 89 +++++++++++++++++++
 3 files changed, 106 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/estimator/BUILD b/tensorflow/python/estimator/BUILD
index c6bb9b9be7..56dec1eaa1 100644
--- a/tensorflow/python/estimator/BUILD
+++ b/tensorflow/python/estimator/BUILD
@@ -478,6 +478,7 @@ py_library(
         "//tensorflow/python:util",
         "//tensorflow/python/data",
         "//tensorflow/python/saved_model:builder",
+        "//tensorflow/python/saved_model:constants",
         "//tensorflow/python/saved_model:tag_constants",
         "//third_party/py/numpy",
         "@six_archive//:six",
diff --git a/tensorflow/python/estimator/estimator.py b/tensorflow/python/estimator/estimator.py
index 946f093ba7..530a4a24ef 100644
--- a/tensorflow/python/estimator/estimator.py
+++ b/tensorflow/python/estimator/estimator.py
@@ -50,6 +50,7 @@ from tensorflow.python.ops import variables
 from tensorflow.python.platform import gfile
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.saved_model import builder as saved_model_builder
+from tensorflow.python.saved_model import constants
 from tensorflow.python.saved_model import tag_constants
 from tensorflow.python.summary import summary
 from tensorflow.python.summary.writer import writer_cache
@@ -493,7 +494,6 @@ class Estimator(object):
       if not checkpoint_path:
         logging.info('Could not find trained model in model_dir: {}, running '
                      'initialization to predict.'.format(self._model_dir))
-
       with ops.Graph().as_default() as g:
         random_seed.set_random_seed(self._config.tf_random_seed)
         self._create_and_assert_global_step(g)
@@ -501,6 +501,10 @@ class Estimator(object):
             input_fn, model_fn_lib.ModeKeys.PREDICT)
         estimator_spec = self._call_model_fn(
             features, None, model_fn_lib.ModeKeys.PREDICT, self.config)
+
+        # Call to warm_start has to be after model_fn is called.
+        self._maybe_warm_start(checkpoint_path)
+
         predictions = self._extract_keys(
             estimator_spec.predictions, predict_keys)
         all_hooks = list(input_hooks)
@@ -982,9 +986,7 @@ class Estimator(object):
     if self._warm_start_settings:
       logging.info('Warm-starting with WarmStartSettings: %s' %
                    (self._warm_start_settings,))
-      # pylint: disable=protected-access
       warm_starting_util.warm_start(*self._warm_start_settings)
-      # pylint: enable=protected-access
     # Check if the user created a loss summary, and add one if they didn't.
     # We assume here that the summary is called 'loss'. If it is not, we will
     # make another one with the name 'loss' to ensure it shows up in the right
@@ -1089,6 +1091,9 @@ class Estimator(object):
       estimator_spec = self._call_model_fn(
           features, labels, model_fn_lib.ModeKeys.EVAL, self.config)
 
+      # Call to warm_start has to be after model_fn is called.
+      self._maybe_warm_start(checkpoint_path)
+
       if model_fn_lib.LOSS_METRIC_KEY in estimator_spec.eval_metric_ops:
         raise ValueError(
             'Metric with name "%s" is not allowed, because Estimator ' % (
@@ -1126,6 +1131,12 @@ class Estimator(object):
 
     return eval_results
 
+  def _maybe_warm_start(self, checkpoint_path):
+    if not checkpoint_path and self._warm_start_settings:
+      logging.info('Warm-starting with WarmStartSettings: %s' %
+                   (self._warm_start_settings,))
+      warm_starting_util.warm_start(*self._warm_start_settings)
+
 
 def create_per_tower_ready_op(scaffold):
   """Create a Scaffold.ready_op inside a tower."""
@@ -1525,7 +1536,8 @@ def _get_default_warm_start_settings(warm_start_from):
       logging.info('Warm-starting from a SavedModel')
       return WarmStartSettings(ckpt_to_initialize_from=os.path.join(
           compat.as_bytes(warm_start_from),
-          compat.as_bytes('variables/variables')))
+          compat.as_bytes('{}/{}'.format(constants.VARIABLES_DIRECTORY,
+                                         constants.VARIABLES_FILENAME))))
     return WarmStartSettings(ckpt_to_initialize_from=warm_start_from)
   elif isinstance(warm_start_from, WarmStartSettings):
     return warm_start_from
diff --git a/tensorflow/python/estimator/estimator_test.py b/tensorflow/python/estimator/estimator_test.py
index 4d958f8b43..76b45b7f57 100644
--- a/tensorflow/python/estimator/estimator_test.py
+++ b/tensorflow/python/estimator/estimator_test.py
@@ -1116,6 +1116,52 @@ class EstimatorEvaluateTest(test.TestCase):
     # initialized (since there is no checkpoint).
     self.assertEqual(3., metrics['metric'])
 
+  def test_no_checkpoint_uses_init_with_warm_starting(self):
+    def _make_model_fn(x):
+      def _variable_creating_and_export_model_fn(features, labels, mode):
+        _, _ = features, labels
+        x_var = variable_scope.get_variable('x', initializer=x)
+        global_step = training.get_global_step()
+        return model_fn_lib.EstimatorSpec(
+            mode,
+            predictions={'y': constant_op.constant(1.0)},
+            loss=constant_op.constant(1.),
+            eval_metric_ops={'metric': metrics_lib.mean(x_var + 1)},
+            train_op=state_ops.assign_add(global_step, 1),
+            export_outputs={'test': export_output.ClassificationOutput(
+                constant_op.constant([4.2]), constant_op.constant(['label']))})
+      return _variable_creating_and_export_model_fn
+
+    first_est = estimator.Estimator(model_fn=_make_model_fn(42.))
+    first_est.train(dummy_input_fn, steps=10)
+    feature_spec = {'x': parsing_ops.VarLenFeature(dtype=dtypes.int64),
+                    'y': parsing_ops.VarLenFeature(dtype=dtypes.int64)}
+    serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
+        feature_spec)
+    tmpdir = tempfile.mkdtemp()
+    export_dir_base = os.path.join(
+        compat.as_bytes(tmpdir), compat.as_bytes('export'))
+    exported_path = first_est.export_savedmodel(export_dir_base,
+                                                serving_input_receiver_fn)
+
+    # Test that we can pass either warm_start_from as an external checkpoint
+    # or an exported SavedModel.
+    est = estimator.Estimator(model_fn=_make_model_fn(52.),
+                              warm_start_from=exported_path)
+    metrics = est.evaluate(dummy_input_fn, steps=1)
+    # Metric value here is set to 1 + the value of the Variable that is
+    # warm-started from the SavedModel of the first model (42.), as opposed to
+    # the initialization in the new model_fn (52.).
+    self.assertEqual(43., metrics['metric'])
+
+    est = estimator.Estimator(model_fn=_make_model_fn(62.),
+                              warm_start_from=first_est.model_dir)
+    metrics = est.evaluate(dummy_input_fn, steps=1)
+    # Metric value here is set to 1 + the value of the Variable that is
+    # warm-started from a checkpoint of the first model (42.), as opposed to
+    # the initialization in the new model_fn (52.).
+    self.assertEqual(43., metrics['metric'])
+
   def test_scores(self):
     est = estimator.Estimator(
         model_fn=_model_fn_with_eval_metric_ops,
@@ -1384,6 +1430,49 @@ class EstimatorPredictTest(test.TestCase):
     # initialized (since there is no checkpoint).
     self.assertEqual(4., next(est.predict(dummy_input_fn)))
 
+  def test_no_checkpoint_uses_init_with_warm_starting(self):
+    def _make_model_fn(x):
+      def _variable_creating_and_export_model_fn(features, labels, mode):
+        _, _ = features, labels
+        x_var = variables.Variable([[x]], name='x')
+        return model_fn_lib.EstimatorSpec(
+            mode,
+            predictions=math_ops.add(x_var, 1.),
+            loss=constant_op.constant(1.),
+            train_op=state_ops.assign_add(training.get_global_step(), 1),
+            export_outputs={'test': export_output.ClassificationOutput(
+                constant_op.constant([4.2]),
+                constant_op.constant(['label']))})
+      return _variable_creating_and_export_model_fn
+
+    first_est = estimator.Estimator(model_fn=_make_model_fn(3.))
+    first_est.train(dummy_input_fn, steps=10)
+    feature_spec = {'x': parsing_ops.VarLenFeature(dtype=dtypes.int64),
+                    'y': parsing_ops.VarLenFeature(dtype=dtypes.int64)}
+    serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
+        feature_spec)
+    tmpdir = tempfile.mkdtemp()
+    export_dir_base = os.path.join(
+        compat.as_bytes(tmpdir), compat.as_bytes('export'))
+    exported_path = first_est.export_savedmodel(export_dir_base,
+                                                serving_input_receiver_fn)
+
+    # Test that we can pass either warm_start_from as an external checkpoint
+    # or an exported SavedModel.
+    est = estimator.Estimator(model_fn=_make_model_fn(30.),
+                              warm_start_from=exported_path)
+    # Prediction here is set to 1 + the value of the Variable that is
+    # warm-started from the SavedModel of the first model (3.), as opposed to
+    # the initialization in the new model_fn (30.).
+    self.assertEqual(4., next(est.predict(dummy_input_fn)))
+
+    est = estimator.Estimator(model_fn=_make_model_fn(40.),
+                              warm_start_from=first_est.model_dir)
+    # Prediction here is set to 1 + the value of the Variable that is
+    # warm-started from a checkpoint of the first model (3.), as opposed to
+    # the initialization in the new model_fn (40.).
+    self.assertEqual(4., next(est.predict(dummy_input_fn)))
+
   def test_no_trained_model_invalid_checkpoint_path(self):
     est = estimator.Estimator(model_fn=model_fn_global_step_incrementer)
     with self.assertRaises(ValueError):
-- 
GitLab


From 30927ec6b625121bae1b89b07f9faeaebaed321f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 16:04:09 -0700
Subject: [PATCH 180/395] Mark all nodes processed by AddOpsRewrite/MinBCast
 stages with a tag.

PiperOrigin-RevId: 195167597
---
 ...direct_session_with_tracking_alloc_test.cc |  4 +-
 .../optimizers/arithmetic_optimizer.cc        | 77 +++++++++++--------
 .../grappler/optimizers/meta_optimizer.cc     |  2 +-
 .../python/grappler/layout_optimizer_test.py  |  8 +-
 4 files changed, 53 insertions(+), 38 deletions(-)

diff --git a/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc b/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc
index b4dd521bbc..695423b2cb 100644
--- a/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc
+++ b/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc
@@ -102,9 +102,9 @@ TEST(DirectSessionWithTrackingAllocTest, CostModelTest) {
         EXPECT_EQ(2, shape.dim(0).size());
         EXPECT_EQ(1, shape.dim(1).size());
         if (node->name() == y->name()) {
-          EXPECT_EQ(7, cm->AllocationId(node, 0));
+          EXPECT_EQ(9, cm->AllocationId(node, 0));
         } else {
-          EXPECT_EQ(8, cm->AllocationId(node, 0));
+          EXPECT_EQ(10, cm->AllocationId(node, 0));
         }
       }
       EXPECT_LE(0, cm->MaxExecutionTime(node));
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
index bf59b25449..d6510ba681 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
@@ -24,6 +24,7 @@ limitations under the License.
 
 #include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/tensor_shape.pb.h"
 #include "tensorflow/core/framework/types.h"
@@ -49,6 +50,12 @@ namespace tensorflow {
 namespace grappler {
 namespace {
 
+// Mark nodes created or optimized by a stage with a tag.
+constexpr char kAddOpsRewriteTag[] =
+    "_grappler:ArithmeticOptimizer:AddOpsRewriteStage";
+constexpr char kMinimizeBroadcastsTag[] =
+    "_grappler:ArithmeticOptimizer:MinimizeBroadcasts";
+
 // Extract values from a Const op to `values`. Returns true if succeeds.
 template <typename T>
 bool ValuesFromConstNode(const NodeDef& node, std::vector<T>* values) {
@@ -142,18 +149,6 @@ bool MaybeAddControlInput(const string& new_input, NodeDef* node,
   return !already_exists;
 }
 
-int CopyControlInputs(const NodeDef& from, NodeDef* to, GraphDef* graph,
-                      NodeMap* node_map) {
-  int num_copied = 0;
-  for (const string& input : from.input()) {
-    if (IsControlInput(input) &&
-        MaybeAddControlInput(input, to, graph, node_map)) {
-      ++num_copied;
-    }
-  }
-  return num_copied;
-}
-
 void SetDataTypeToAttr(DataType dtype, const string& attr_name, NodeDef* node) {
   (*node->mutable_attr())[attr_name].set_type(dtype);
 }
@@ -326,7 +321,7 @@ class ArithmeticNodesGroupOptimizerStage : public ArithmeticOptimizerStage {
   explicit ArithmeticNodesGroupOptimizerStage(
       const string& name, const GraphOptimizerContext& ctx,
       const ArithmeticOptimizerContext ctx_ext)
-      : ArithmeticOptimizerStage(name, ctx, ctx_ext), optimized_nodes_{} {}
+      : ArithmeticOptimizerStage(name, ctx, ctx_ext) {}
   ~ArithmeticNodesGroupOptimizerStage() override = default;
 
   // Input name with a statically inferred shape from GraphProperties
@@ -465,13 +460,16 @@ class ArithmeticNodesGroupOptimizerStage : public ArithmeticOptimizerStage {
     return signature;
   }
 
-  void AddToOptimizedNodes(const NodeDef* node) {
-    optimized_nodes_.insert(node->name());
+  void MarkWithTag(const StringPiece tag, NodeDef* node) {
+    AddNodeAttr(tag, true, node);
   }
 
-  void AddAllMembersToOptimizedNodes(const OptimizedNodesGroup& group) {
-    AddToOptimizedNodes(group.root_node);
-    for (const NodeDef* opt : group.optimized_nodes) AddToOptimizedNodes(opt);
+  void MarkAllMembersWithTag(const OptimizedNodesGroup& group,
+                             const StringPiece tag) const {
+    AddNodeAttr(tag, true, group.root_node);
+    for (NodeDef* optimized_node : group.optimized_nodes) {
+      AddNodeAttr(tag, true, optimized_node);
+    }
   }
 
   bool IsOnTheSameDevice(const OptimizedNodesGroup& group,
@@ -479,13 +477,19 @@ class ArithmeticNodesGroupOptimizerStage : public ArithmeticOptimizerStage {
     return group.root_node->device() == node.device();
   }
 
-  bool IsAlreadyOptimized(const NodeDef& node) const {
-    return optimized_nodes_.find(node.name()) != optimized_nodes_.end();
+  bool IsInPreserveSet(const NodeDef& node) const {
+    return ctx().nodes_to_preserve->find(node.name()) !=
+           ctx().nodes_to_preserve->end();
   }
 
- private:
-  // set of nodes already processed by this optimizer stage
-  std::unordered_set<string> optimized_nodes_;
+  bool IsMarkedWithTag(const NodeDef& node, const StringPiece tag) const {
+    return HasNodeAttr(node, tag);
+  }
+
+  bool IsMarkedWithAnyTag(const NodeDef& node, const StringPiece tag1,
+                          const StringPiece tag2) const {
+    return IsMarkedWithTag(node, tag1) || IsMarkedWithTag(node, tag2);
+  }
 };
 
 // Rewrite a tree of Add/AddN with a single AddN operation, consuming all the
@@ -561,7 +565,7 @@ class AddOpsRewriteStage : public ArithmeticNodesGroupOptimizerStage {
     if (!IsAdd(node) && !IsAddN(node)) {
       return false;
     }
-    if (IsInPreserveSet(node) || IsAlreadyOptimized(node)) {
+    if (IsInPreserveSet(node) || IsMarkedWithTag(node, kAddOpsRewriteTag)) {
       return false;
     }
     // TODO(ezhulenev): relax this condition for root node
@@ -579,7 +583,7 @@ class AddOpsRewriteStage : public ArithmeticNodesGroupOptimizerStage {
             << " num_inputs=" << group.inputs.size();
 
     // Do not optimize any of the nodes that are part of this group.
-    AddAllMembersToOptimizedNodes(group);
+    MarkAllMembersWithTag(group, kAddOpsRewriteTag);
 
     // All new nodes will be placed under the scope of a root node.
     auto root_scope_and_name = ParseNodeScopeAndName(group.root_node->name());
@@ -688,7 +692,7 @@ class AddOpsRewriteStage : public ArithmeticNodesGroupOptimizerStage {
       node->add_input(inputAndShape.input);
     }
 
-    AddToOptimizedNodes(node);
+    MarkWithTag(kAddOpsRewriteTag, node);
     return InputAndShape(node_name, shape);
   }
 
@@ -705,14 +709,13 @@ class AddOpsRewriteStage : public ArithmeticNodesGroupOptimizerStage {
     node->set_op("Add");
     node->set_device(root_node.device());
     (*node->mutable_attr())["T"].set_type(dtype);
+    node->add_input(left.input);
+    node->add_input(right.input);
 
     ctx().node_map->AddOutput(left.input, node_name);
     ctx().node_map->AddOutput(right.input, node_name);
 
-    node->add_input(left.input);
-    node->add_input(right.input);
-
-    AddToOptimizedNodes(node);
+    MarkWithTag(kAddOpsRewriteTag, node);
     return InputAndShape(
         node_name, TensorShapeProto());  // shape is not important at this point
   }
@@ -960,7 +963,9 @@ class MinimizeBroadcasts : public ArithmeticNodesGroupOptimizerStage {
 
   bool IsSupported(const NodeDef* node) const override {
     if (!IsBinaryAssociative(*node)) return false;
-    if (IsAlreadyOptimized(*node)) return false;
+
+    if (IsMarkedWithAnyTag(*node, kMinimizeBroadcastsTag, kAddOpsRewriteTag))
+      return false;
 
     // has a symbolically defined shape with broadcastable inputs
     OpInfo::TensorProperties properties;
@@ -984,7 +989,11 @@ class MinimizeBroadcasts : public ArithmeticNodesGroupOptimizerStage {
     if (!IsSameOp(group, node)) {
       return false;
     }
-    if (IsInPreserveSet(node) || IsAlreadyOptimized(node)) {
+    if (IsInPreserveSet(node)) {
+      return false;
+    }
+    // Nodes optimized by AddOpsRewrite already have optimal broadcasts.
+    if (IsMarkedWithAnyTag(node, kMinimizeBroadcastsTag, kAddOpsRewriteTag)) {
       return false;
     }
     if (IsDrivenByControlDependency(node) || DrivesControlDependency(node)) {
@@ -1019,7 +1028,7 @@ class MinimizeBroadcasts : public ArithmeticNodesGroupOptimizerStage {
             << " num_optimized_nodes=" << group.optimized_nodes.size();
 
     // Do not optimize any of the nodes that are part of this group.
-    AddAllMembersToOptimizedNodes(group);
+    MarkAllMembersWithTag(group, kMinimizeBroadcastsTag);
 
     if (CountUniqueShapes(group.inputs) <= 1) {
       VLOG(3) << "Skip min-bcast group with single unique shape";
@@ -1905,6 +1914,8 @@ void ArithmeticOptimizer::DedupComputations() {
           FeedsInPlaceOp(graph_view, *node)) {
         continue;
       }
+      VLOG(3) << "Remove duplicated node: node=" << node->name()
+              << " representative=" << rep->name();
       const std::set<NodeDef*>& fanouts = node_map_->GetOutputs(node->name());
       for (NodeDef* fanout : fanouts) {
         for (int i = 0; i < fanout->input_size(); ++i) {
diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer.cc b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
index 5230177dca..0c8e18d7ab 100644
--- a/tensorflow/core/grappler/optimizers/meta_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
@@ -65,7 +65,7 @@ int NumIterations(const RewriterConfig& cfg) {
 // Check if optimizer is allowed to run only once.
 bool IsRunOnceOptimizer(const string& name) {
   return name == "layout" || name == "memory_optimizer" ||
-         name == "arithmetic_optimizer" || name == "loop_optimizer";
+         name == "loop_optimizer";
 }
 
 }  // namespace
diff --git a/tensorflow/python/grappler/layout_optimizer_test.py b/tensorflow/python/grappler/layout_optimizer_test.py
index e3dd4b0bdf..2d6925d1a8 100644
--- a/tensorflow/python/grappler/layout_optimizer_test.py
+++ b/tensorflow/python/grappler/layout_optimizer_test.py
@@ -150,10 +150,14 @@ def _loop_with_vec_and_4d():
 def _get_config(layout_optimizer=True):
   if layout_optimizer:
     rewrite_options = rewriter_config_pb2.RewriterConfig(
-        layout_optimizer=rewriter_config_pb2.RewriterConfig.ON)
+        layout_optimizer=rewriter_config_pb2.RewriterConfig.ON,
+        # do not remove duplicated nodes
+        arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF)
   else:
     rewrite_options = rewriter_config_pb2.RewriterConfig(
-        layout_optimizer=rewriter_config_pb2.RewriterConfig.OFF)
+        layout_optimizer=rewriter_config_pb2.RewriterConfig.OFF,
+        # do not remove duplicated nodes
+        arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF)
   graph_options = config_pb2.GraphOptions(
       rewrite_options=rewrite_options, build_cost_model=1)
   config = config_pb2.ConfigProto(graph_options=graph_options)
-- 
GitLab


From 1f4efb78320e1406c0cc9ce4b8753f3d2511048e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 16:05:43 -0700
Subject: [PATCH 181/395] Add RNNEstimator which takes in arbitrary heads.

PiperOrigin-RevId: 195167853
---
 tensorflow/contrib/estimator/BUILD            |   4 +
 tensorflow/contrib/estimator/__init__.py      |   1 +
 .../contrib/estimator/python/estimator/rnn.py | 164 ++++++++++++++++--
 .../estimator/python/estimator/rnn_test.py    | 119 ++++++++-----
 4 files changed, 237 insertions(+), 51 deletions(-)

diff --git a/tensorflow/contrib/estimator/BUILD b/tensorflow/contrib/estimator/BUILD
index b473de86ee..41a817673d 100644
--- a/tensorflow/contrib/estimator/BUILD
+++ b/tensorflow/contrib/estimator/BUILD
@@ -452,18 +452,22 @@ py_test(
         "notsan",
     ],
     deps = [
+        ":head",
         ":rnn",
+        "//tensorflow/contrib/data",
         "//tensorflow/core:protos_all_py",
         "//tensorflow/python:check_ops",
         "//tensorflow/python:client_testlib",
         "//tensorflow/python:dtypes",
         "//tensorflow/python:framework_ops",
+        "//tensorflow/python:lib",
         "//tensorflow/python:math_ops",
         "//tensorflow/python:state_ops",
         "//tensorflow/python:summary",
         "//tensorflow/python:training",
         "//tensorflow/python:variables",
         "//tensorflow/python/estimator:numpy_io",
+        "//tensorflow/python/estimator:parsing_utils",
         "//tensorflow/python/feature_column",
         "//third_party/py/numpy",
         "@six_archive//:six",
diff --git a/tensorflow/contrib/estimator/__init__.py b/tensorflow/contrib/estimator/__init__.py
index f66d844660..d43b3ea6bf 100644
--- a/tensorflow/contrib/estimator/__init__.py
+++ b/tensorflow/contrib/estimator/__init__.py
@@ -55,6 +55,7 @@ _allowed_symbols = [
     'replicate_model_fn',
     'TowerOptimizer',
     'RNNClassifier',
+    'RNNEstimator',
 ]
 
 remove_undocumented(__name__, allowed_exception_list=_allowed_symbols)
diff --git a/tensorflow/contrib/estimator/python/estimator/rnn.py b/tensorflow/contrib/estimator/python/estimator/rnn.py
index b475c12f5a..7f385fd76e 100644
--- a/tensorflow/contrib/estimator/python/estimator/rnn.py
+++ b/tensorflow/contrib/estimator/python/estimator/rnn.py
@@ -328,6 +328,19 @@ def _rnn_model_fn(features,
         logits=logits)
 
 
+def _assert_rnn_cell_fn(rnn_cell_fn, num_units, cell_type):
+  """Assert arguments are valid and return rnn_cell_fn."""
+  if rnn_cell_fn and (num_units or cell_type != USE_DEFAULT):
+    raise ValueError(
+        'num_units and cell_type must not be specified when using rnn_cell_fn'
+    )
+  if not rnn_cell_fn:
+    if cell_type == USE_DEFAULT:
+      cell_type = 'basic_rnn'
+    rnn_cell_fn = _make_rnn_cell_fn(num_units, cell_type)
+  return rnn_cell_fn
+
+
 class RNNClassifier(estimator.Estimator):
   """A classifier for TensorFlow RNN models.
 
@@ -341,8 +354,8 @@ class RNNClassifier(estimator.Estimator):
   token_emb = embedding_column(categorical_column=token_sequence, ...)
 
   estimator = RNNClassifier(
-      num_units=[32, 16], cell_type='lstm',
-      sequence_feature_columns=[token_emb])
+      sequence_feature_columns=[token_emb],
+      num_units=[32, 16], cell_type='lstm')
 
   # Input builders
   def input_fn_train: # returns x, y
@@ -438,8 +451,8 @@ class RNNClassifier(estimator.Estimator):
         encoded as integer values in {0, 1,..., n_classes-1} for `n_classes`>2 .
         Also there will be errors if vocabulary is not provided and labels are
         string.
-      optimizer: An instance of `tf.Optimizer` used to train the model. Defaults
-        to Adagrad optimizer.
+      optimizer: An instance of `tf.Optimizer` or string specifying optimizer
+        type. Defaults to Adagrad optimizer.
       input_layer_partitioner: Optional. Partitioner for input layer. Defaults
         to `min_max_variable_partitioner` with `min_slice_size` 64 << 20.
       config: `RunConfig` object to configure the runtime settings.
@@ -448,14 +461,7 @@ class RNNClassifier(estimator.Estimator):
       ValueError: If `num_units`, `cell_type`, and `rnn_cell_fn` are not
         compatible.
     """
-    if rnn_cell_fn and (num_units or cell_type != USE_DEFAULT):
-      raise ValueError(
-          'num_units and cell_type must not be specified when using rnn_cell_fn'
-      )
-    if not rnn_cell_fn:
-      if cell_type == USE_DEFAULT:
-        cell_type = 'basic_rnn'
-      rnn_cell_fn = _make_rnn_cell_fn(num_units, cell_type)
+    rnn_cell_fn = _assert_rnn_cell_fn(rnn_cell_fn, num_units, cell_type)
 
     if n_classes == 2:
       head = head_lib._binary_logistic_head_with_sigmoid_cross_entropy_loss(  # pylint: disable=protected-access
@@ -479,3 +485,137 @@ class RNNClassifier(estimator.Estimator):
           config=config)
     super(RNNClassifier, self).__init__(
         model_fn=_model_fn, model_dir=model_dir, config=config)
+
+
+class RNNEstimator(estimator.Estimator):
+  """An Estimator for TensorFlow RNN models with user-specified head.
+
+  Example:
+
+  ```python
+  token_sequence = sequence_categorical_column_with_hash_bucket(...)
+  token_emb = embedding_column(categorical_column=token_sequence, ...)
+
+  estimator = RNNEstimator(
+      head=tf.contrib.estimator.regression_head(),
+      sequence_feature_columns=[token_emb],
+      num_units=[32, 16], cell_type='lstm')
+
+  # Or with custom RNN cell:
+  def rnn_cell_fn(mode):
+    cells = [ tf.contrib.rnn.LSTMCell(size) for size in [32, 16] ]
+    if mode == tf.estimator.ModeKeys.TRAIN:
+      cells = [ tf.contrib.rnn.DropoutWrapper(cell, input_keep_prob=0.5)
+                    for cell in cells ]
+    return tf.contrib.rnn.MultiRNNCell(cells)
+
+  estimator = RNNEstimator(
+      head=tf.contrib.estimator.regression_head(),
+      sequence_feature_columns=[token_emb],
+      rnn_cell_fn=rnn_cell_fn)
+
+  # Input builders
+  def input_fn_train: # returns x, y
+    pass
+  estimator.train(input_fn=input_fn_train, steps=100)
+
+  def input_fn_eval: # returns x, y
+    pass
+  metrics = estimator.evaluate(input_fn=input_fn_eval, steps=10)
+  def input_fn_predict: # returns x, None
+    pass
+  predictions = estimator.predict(input_fn=input_fn_predict)
+  ```
+
+  Input of `train` and `evaluate` should have following features,
+  otherwise there will be a `KeyError`:
+
+  * if the head's `weight_column` is not `None`, a feature with
+    `key=weight_column` whose value is a `Tensor`.
+  * for each `column` in `sequence_feature_columns`:
+    - a feature with `key=column.name` whose `value` is a `SparseTensor`.
+  * for each `column` in `context_feature_columns`:
+    - if `column` is a `_CategoricalColumn`, a feature with `key=column.name`
+      whose `value` is a `SparseTensor`.
+    - if `column` is a `_WeightedCategoricalColumn`, two features: the first
+      with `key` the id column name, the second with `key` the weight column
+      name. Both features' `value` must be a `SparseTensor`.
+    - if `column` is a `_DenseColumn`, a feature with `key=column.name`
+      whose `value` is a `Tensor`.
+
+  Loss and predicted output are determined by the specified head.
+
+  @compatibility(eager)
+  Estimators are not compatible with eager execution.
+  @end_compatibility
+  """
+
+  def __init__(self,
+               head,
+               sequence_feature_columns,
+               context_feature_columns=None,
+               num_units=None,
+               cell_type=USE_DEFAULT,
+               rnn_cell_fn=None,
+               model_dir=None,
+               optimizer='Adagrad',
+               input_layer_partitioner=None,
+               config=None):
+    """Initializes a `RNNClassifier` instance.
+
+    Args:
+      head: A `_Head` instance constructed with a method such as
+        `tf.contrib.estimator.multi_label_head`. This specifies the model's
+        output and loss function to be optimized.
+      sequence_feature_columns: An iterable containing the `FeatureColumn`s
+        that represent sequential input. All items in the set should either be
+        sequence columns (e.g. `sequence_numeric_column`) or constructed from
+        one (e.g. `embedding_column` with `sequence_categorical_column_*` as
+        input).
+      context_feature_columns: An iterable containing the `FeatureColumn`s
+        for contextual input. The data represented by these columns will be
+        replicated and given to the RNN at each timestep. These columns must be
+        instances of classes derived from `_DenseColumn` such as
+        `numeric_column`, not the sequential variants.
+      num_units: Iterable of integer number of hidden units per RNN layer. If
+        set, `cell_type` must also be specified and `rnn_cell_fn` must be
+        `None`.
+      cell_type: A subclass of `tf.nn.rnn_cell.RNNCell` or a string specifying
+        the cell type. Supported strings are: `'basic_rnn'`, `'lstm'`, and
+        `'gru'`. If set, `num_units` must also be specified and `rnn_cell_fn`
+        must be `None`.
+      rnn_cell_fn: A function with one argument, a `tf.estimator.ModeKeys`, and
+        returns an object of type `tf.nn.rnn_cell.RNNCell` that will be used to
+        construct the RNN. If set, `num_units` and `cell_type` cannot be set.
+        This is for advanced users who need additional customization beyond
+        `num_units` and `cell_type`. Note that `tf.nn.rnn_cell.MultiRNNCell` is
+        needed for stacked RNNs.
+      model_dir: Directory to save model parameters, graph and etc. This can
+        also be used to load checkpoints from the directory into a estimator to
+        continue training a previously saved model.
+      optimizer: An instance of `tf.Optimizer` or string specifying optimizer
+        type. Defaults to Adagrad optimizer.
+      input_layer_partitioner: Optional. Partitioner for input layer. Defaults
+        to `min_max_variable_partitioner` with `min_slice_size` 64 << 20.
+      config: `RunConfig` object to configure the runtime settings.
+
+    Raises:
+      ValueError: If `num_units`, `cell_type`, and `rnn_cell_fn` are not
+        compatible.
+    """
+    rnn_cell_fn = _assert_rnn_cell_fn(rnn_cell_fn, num_units, cell_type)
+
+    def _model_fn(features, labels, mode, config):
+      return _rnn_model_fn(
+          features=features,
+          labels=labels,
+          mode=mode,
+          head=head,
+          rnn_cell_fn=rnn_cell_fn,
+          sequence_feature_columns=tuple(sequence_feature_columns or []),
+          context_feature_columns=tuple(context_feature_columns or []),
+          optimizer=optimizer,
+          input_layer_partitioner=input_layer_partitioner,
+          config=config)
+    super(RNNEstimator, self).__init__(
+        model_fn=_model_fn, model_dir=model_dir, config=config)
diff --git a/tensorflow/contrib/estimator/python/estimator/rnn_test.py b/tensorflow/contrib/estimator/python/estimator/rnn_test.py
index 393f94f5c7..959b40371a 100644
--- a/tensorflow/contrib/estimator/python/estimator/rnn_test.py
+++ b/tensorflow/contrib/estimator/python/estimator/rnn_test.py
@@ -25,12 +25,15 @@ import tempfile
 import numpy as np
 import six
 
+from tensorflow.contrib.data.python.ops import readers
+from tensorflow.contrib.estimator.python.estimator import head as head_lib
 from tensorflow.contrib.estimator.python.estimator import rnn
 from tensorflow.contrib.feature_column.python.feature_column import sequence_feature_column as seq_fc
 from tensorflow.core.example import example_pb2
 from tensorflow.core.example import feature_pb2
 from tensorflow.python.estimator import model_fn
 from tensorflow.python.estimator.canned import metric_keys
+from tensorflow.python.estimator.canned import parsing_utils
 from tensorflow.python.estimator.canned import prediction_keys
 from tensorflow.python.estimator.export import export
 from tensorflow.python.estimator.inputs import numpy_io
@@ -38,9 +41,9 @@ from tensorflow.python.feature_column import feature_column as fc
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.lib.io import python_io
 from tensorflow.python.ops import check_ops
 from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import parsing_ops
 from tensorflow.python.ops import partitioned_variables
 from tensorflow.python.ops import rnn_cell
 from tensorflow.python.ops import state_ops
@@ -50,7 +53,6 @@ from tensorflow.python.platform import gfile
 from tensorflow.python.platform import test
 from tensorflow.python.summary.writer import writer_cache
 from tensorflow.python.training import checkpoint_utils
-from tensorflow.python.training import input as input_lib
 from tensorflow.python.training import monitored_session
 from tensorflow.python.training import optimizer
 from tensorflow.python.training import training_util
@@ -984,7 +986,10 @@ class RNNClassifierPredictionTest(test.TestCase):
                      predictions[prediction_keys.PredictionKeys.CLASSES])
 
 
-class RNNClassifierIntegrationTest(test.TestCase):
+class BaseRNNClassificationIntegrationTest(object):
+
+  def __init__(self, _create_estimator_fn):
+    self._create_estimator_fn = _create_estimator_fn
 
   def setUp(self):
     self._model_dir = tempfile.mkdtemp()
@@ -994,20 +999,11 @@ class RNNClassifierIntegrationTest(test.TestCase):
       writer_cache.FileWriterCache.clear()
       shutil.rmtree(self._model_dir)
 
-  def _test_complete_flow(
-      self, train_input_fn, eval_input_fn, predict_input_fn, n_classes,
-      batch_size):
-    col = seq_fc.sequence_categorical_column_with_hash_bucket(
-        'tokens', hash_bucket_size=10)
-    embed = fc.embedding_column(col, dimension=2)
-    feature_columns = [embed]
-
+  def _test_complete_flow(self, feature_columns, train_input_fn, eval_input_fn,
+                          predict_input_fn, n_classes, batch_size):
     cell_units = [4, 2]
-    est = rnn.RNNClassifier(
-        num_units=cell_units,
-        sequence_feature_columns=feature_columns,
-        n_classes=n_classes,
-        model_dir=self._model_dir)
+    est = self._create_estimator_fn(feature_columns, n_classes, cell_units,
+                                    self._model_dir)
 
     # TRAIN
     num_steps = 10
@@ -1026,10 +1022,10 @@ class RNNClassifierIntegrationTest(test.TestCase):
     self.assertAllEqual((batch_size, n_classes), predicted_proba.shape)
 
     # EXPORT
-    feature_spec = {
-        'tokens': parsing_ops.VarLenFeature(dtypes.string),
-        'label': parsing_ops.FixedLenFeature([1], dtypes.int64),
-    }
+    feature_spec = parsing_utils.classifier_parse_example_spec(
+        feature_columns,
+        label_key='label',
+        label_dtype=dtypes.int64)
     serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
         feature_spec)
     export_dir = est.export_savedmodel(tempfile.mkdtemp(),
@@ -1069,7 +1065,13 @@ class RNNClassifierIntegrationTest(test.TestCase):
         batch_size=batch_size,
         shuffle=False)
 
+    col = seq_fc.sequence_categorical_column_with_hash_bucket(
+        'tokens', hash_bucket_size=10)
+    embed = fc.embedding_column(col, dimension=2)
+    feature_columns = [embed]
+
     self._test_complete_flow(
+        feature_columns=feature_columns,
         train_input_fn=train_input_fn,
         eval_input_fn=eval_input_fn,
         predict_input_fn=predict_input_fn,
@@ -1082,7 +1084,8 @@ class RNNClassifierIntegrationTest(test.TestCase):
     batch_size = 10
     words = [b'dog', b'cat', b'bird', b'the', b'a', b'sat', b'flew', b'slept']
 
-    serialized_examples = []
+    _, examples_file = tempfile.mkstemp()
+    writer = python_io.TFRecordWriter(examples_file)
     for _ in range(batch_size):
       sequence_length = random.randint(1, len(words))
       sentence = random.sample(words, sequence_length)
@@ -1096,30 +1099,36 @@ class RNNClassifierIntegrationTest(test.TestCase):
                   feature_pb2.Feature(int64_list=feature_pb2.Int64List(
                       value=[label])),
           }))
-      serialized_examples.append(example.SerializeToString())
+      writer.write(example.SerializeToString())
+    writer.close()
+
+    col = seq_fc.sequence_categorical_column_with_hash_bucket(
+        'tokens', hash_bucket_size=10)
+    embed = fc.embedding_column(col, dimension=2)
+    feature_columns = [embed]
+    feature_spec = parsing_utils.classifier_parse_example_spec(
+        feature_columns,
+        label_key='label',
+        label_dtype=dtypes.int64)
 
-    feature_spec = {
-        'tokens': parsing_ops.VarLenFeature(dtypes.string),
-        'label': parsing_ops.FixedLenFeature([1], dtypes.int64),
-    }
     def _train_input_fn():
-      features = parsing_ops.parse_example(serialized_examples, feature_spec)
-      labels = features.pop('label')
-      return features, labels
+      dataset = readers.make_batched_features_dataset(
+          examples_file, batch_size, feature_spec)
+      return dataset.map(lambda features: (features, features.pop('label')))
     def _eval_input_fn():
-      features = parsing_ops.parse_example(
-          input_lib.limit_epochs(serialized_examples, num_epochs=1),
-          feature_spec)
-      labels = features.pop('label')
-      return features, labels
+      dataset = readers.make_batched_features_dataset(
+          examples_file, batch_size, feature_spec, num_epochs=1)
+      return dataset.map(lambda features: (features, features.pop('label')))
     def _predict_input_fn():
-      features = parsing_ops.parse_example(
-          input_lib.limit_epochs(serialized_examples, num_epochs=1),
-          feature_spec)
-      features.pop('label')
-      return features, None
+      dataset = readers.make_batched_features_dataset(
+          examples_file, batch_size, feature_spec, num_epochs=1)
+      def features_fn(features):
+        features.pop('label')
+        return features
+      return dataset.map(features_fn)
 
     self._test_complete_flow(
+        feature_columns=feature_columns,
         train_input_fn=_train_input_fn,
         eval_input_fn=_eval_input_fn,
         predict_input_fn=_predict_input_fn,
@@ -1127,5 +1136,37 @@ class RNNClassifierIntegrationTest(test.TestCase):
         batch_size=batch_size)
 
 
+def _rnn_classifier_fn(feature_columns, n_classes, cell_units, model_dir):
+  return rnn.RNNClassifier(
+      num_units=cell_units,
+      sequence_feature_columns=feature_columns,
+      n_classes=n_classes,
+      model_dir=model_dir)
+
+
+class RNNClassifierIntegrationTest(BaseRNNClassificationIntegrationTest,
+                                   test.TestCase):
+
+  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
+    test.TestCase.__init__(self, methodName)
+    BaseRNNClassificationIntegrationTest.__init__(self, _rnn_classifier_fn)
+
+
+def _rnn_estimator_fn(feature_columns, n_classes, cell_units, model_dir):
+  return rnn.RNNEstimator(
+      head=head_lib.multi_class_head(n_classes=n_classes),
+      num_units=cell_units,
+      sequence_feature_columns=feature_columns,
+      model_dir=model_dir)
+
+
+class RNNEstimatorIntegrationTest(BaseRNNClassificationIntegrationTest,
+                                  test.TestCase):
+
+  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
+    test.TestCase.__init__(self, methodName)
+    BaseRNNClassificationIntegrationTest.__init__(self, _rnn_estimator_fn)
+
+
 if __name__ == '__main__':
   test.main()
-- 
GitLab


From c7a5787fef8daf3e44313cbd48591464f9643f56 Mon Sep 17 00:00:00 2001
From: Ayush Dubey <ayushd@google.com>
Date: Wed, 2 May 2018 16:13:06 -0700
Subject: [PATCH 182/395] Enable reshape of _ScopedAllocatorConcat output.

The _ScopedAllocatorConcat kernel outputs the backing tensor after performing
runtime bounds checks.  However, the shape of the backing tensor may not match
the desired output shape of the concat operation.

This change adds a "reshape" boolean attribute to _ScopedAllocatorConcat kernel.
When this attribute is set to true, the kernel outputs a reshaped backing tensor
according to the "shape" attribute.

PiperOrigin-RevId: 195169105
---
 .../core/kernels/scoped_allocator_ops.cc      | 39 +++++++----
 .../core/kernels/scoped_allocator_ops_test.cc | 64 ++++++++++++++++---
 tensorflow/core/ops/scoped_allocator_ops.cc   | 11 ++--
 3 files changed, 89 insertions(+), 25 deletions(-)

diff --git a/tensorflow/core/kernels/scoped_allocator_ops.cc b/tensorflow/core/kernels/scoped_allocator_ops.cc
index d7b25ffad0..1800ee8c1f 100644
--- a/tensorflow/core/kernels/scoped_allocator_ops.cc
+++ b/tensorflow/core/kernels/scoped_allocator_ops.cc
@@ -94,7 +94,8 @@ class ScopedAllocatorConcatOp : public OpKernel {
       : OpKernel(context) {
     OP_REQUIRES_OK(context, context->GetAttr("shape", &shape_));
     OP_REQUIRES_OK(context, context->GetAttr("T", &dtype_));
-    // This stuff is just for debugging
+    OP_REQUIRES_OK(context, context->GetAttr("reshape", &reshape_));
+    // These attributes are just for debugging.
     OP_REQUIRES_OK(context, context->GetAttr("sa_name", &name_));
     OP_REQUIRES_OK(context, context->GetAttr("id", &id_));
     device_ = context->device();
@@ -114,11 +115,14 @@ class ScopedAllocatorConcatOp : public OpKernel {
                                         backing_tensor.NumElements(),
                                         " is not equal to expected ",
                                         shape_.num_elements()));
-    VLOG(1) << "_ScopedAllocatorConcatOp outputting backing tensor at "
-            << DMAHelper::base(&backing_tensor);
-    Tensor backing_copy(backing_tensor);
-    context->set_output(0, backing_copy);
-    const TensorBuffer* backing_buf = DMAHelper::buffer(&backing_copy);
+    Tensor output(dtype_);
+    if (reshape_) {
+      CHECK(output.CopyFrom(backing_tensor, shape_));
+    } else {
+      CHECK(output.CopyFrom(backing_tensor, backing_tensor.shape()));
+    }
+    context->set_output(0, output);
+    const TensorBuffer* backing_buf = DMAHelper::buffer(&output);
     const void* backing_tensor_lb = backing_buf->data();
     const void* backing_tensor_ub = static_cast<const void*>(
         static_cast<const char*>(backing_tensor_lb) + backing_buf->size());
@@ -126,17 +130,27 @@ class ScopedAllocatorConcatOp : public OpKernel {
     for (int i = 1; i < context->num_inputs(); ++i) {
       const TensorBuffer* input_buf = DMAHelper::buffer(&context->input(i));
       const void* input_lb = input_buf->data();
-      OP_REQUIRES(
-          context, input_lb >= backing_tensor_lb,
-          errors::InvalidArgument("Lower bound check fail for input ", i,
-                                  " to node ", context->op_kernel().name()));
       const void* input_ub = static_cast<const void*>(
           static_cast<const char*>(input_lb) + input_buf->size());
+      OP_REQUIRES(
+          context, input_lb >= backing_tensor_lb,
+          errors::InvalidArgument(
+              "Lower bound check fail for input ", i, " from node ",
+              context->op_kernel().requested_input(i), " to node ",
+              context->op_kernel().name(), " input bounds = [", input_lb, ", ",
+              input_ub, "]", " backing_tensor bounds = [", backing_tensor_lb,
+              ", ", backing_tensor_ub, "]"));
       OP_REQUIRES(
           context, input_ub <= backing_tensor_ub,
-          errors::InvalidArgument("Upper bound check fail for input ", i,
-                                  " to node ", context->op_kernel().name()));
+          errors::InvalidArgument(
+              "Upper bound check fail for input ", i, " from node ",
+              context->op_kernel().requested_input(i), " to node ",
+              context->op_kernel().name(), " input bounds = [", input_lb, ", ",
+              input_ub, "]", " backing_tensor bounds = [", backing_tensor_lb,
+              ", ", backing_tensor_ub, "]"));
     }
+    VLOG(1) << "_ScopedAllocatorConcatOp outputting backing tensor at "
+            << backing_buf;
   }
 
  private:
@@ -144,6 +158,7 @@ class ScopedAllocatorConcatOp : public OpKernel {
   DataType dtype_;
   string name_;
   int32 id_;
+  bool reshape_;
   DeviceBase* device_;
 };
 
diff --git a/tensorflow/core/kernels/scoped_allocator_ops_test.cc b/tensorflow/core/kernels/scoped_allocator_ops_test.cc
index 3d36c8b7d4..019c6619ee 100644
--- a/tensorflow/core/kernels/scoped_allocator_ops_test.cc
+++ b/tensorflow/core/kernels/scoped_allocator_ops_test.cc
@@ -120,18 +120,43 @@ void PrepOp(DataType dtype, int32 id,
 
 class ScopedAllocatorConcatOpTest : public OpsTestBase {
  protected:
-  void MakeOp(const TensorShape& shape, DataType dtype, const string& name,
-              int32 id, int32 num_tensors) {
+  void BuildNodeDef(const TensorShape& shape, DataType dtype,
+                    const string& name, int32 id, int32 num_tensors) {
+    TF_EXPECT_OK(
+        NodeDefBuilder("scoped_allocator_concat_op", "_ScopedAllocatorConcat")
+            .Attr("shape", shape)
+            .Attr("T", dtype)
+            .Attr("N", num_tensors)
+            .Attr("sa_name", name)
+            .Attr("id", id)
+            .Input(FakeInput(dtype))               // backing tensor
+            .Input(FakeInput(num_tensors, dtype))  // list of tensors
+            .Finalize(node_def()));
+    shape_ = shape;
+    reshape_ = false;
+  }
+
+  void BuildNodeDefWithReshape(const TensorShape& shape, DataType dtype,
+                               bool reshape, const string& name, int32 id,
+                               int32 num_tensors) {
     TF_EXPECT_OK(
         NodeDefBuilder("scoped_allocator_concat_op", "_ScopedAllocatorConcat")
             .Attr("shape", shape)
             .Attr("T", dtype)
+            .Attr("reshape", reshape)
             .Attr("N", num_tensors)
             .Attr("sa_name", name)
             .Attr("id", id)
             .Input(FakeInput(dtype))               // backing tensor
             .Input(FakeInput(num_tensors, dtype))  // list of tensors
             .Finalize(node_def()));
+    shape_ = shape;
+    reshape_ = reshape;
+  }
+
+  void MakeOp(const TensorShape& shape, DataType dtype, bool reshape,
+              const string& name, int32 id, int32 num_tensors) {
+    BuildNodeDefWithReshape(shape, dtype, reshape, name, id, num_tensors);
     TF_EXPECT_OK(InitOp());
   }
 
@@ -141,7 +166,7 @@ class ScopedAllocatorConcatOpTest : public OpsTestBase {
     std::vector<Tensor> tensors;
     std::vector<ScopedAllocator::Field> fields;
     PrepOp(dtype, id, fields_shapes, &fields, &backing_tensor, allocator(),
-           device_->GetScopedAllocatorMgr(), "split", &tensors, &inputs_,
+           device_->GetScopedAllocatorMgr(), "concat", &tensors, &inputs_,
            input_types_);
 
     TF_ASSERT_OK(RunOpKernel());
@@ -155,34 +180,55 @@ class ScopedAllocatorConcatOpTest : public OpsTestBase {
     CHECK_EQ(DMAHelper::base(&input), DMAHelper::base(&output));
     CHECK_EQ(input.dtype(), output.dtype());
     CHECK_EQ(input.NumElements(), output.NumElements());
+    if (reshape_) {
+      CHECK_EQ(shape_, output.shape());
+    } else {
+      TensorShape expected_shape({input.NumElements()});
+      CHECK_EQ(expected_shape, output.shape());
+    }
 
     // Free the backing tensor which was allocated in PrepOp.
     delete backing_tensor;
   }
+
+ private:
+  TensorShape shape_;
+  bool reshape_;
 };
 
 TEST_F(ScopedAllocatorConcatOpTest, Success1) {
-  MakeOp({32}, DT_FLOAT, "test", 120, 2);
+  MakeOp({32}, DT_FLOAT, false, "test", 120, 2);
   ExecOp(DT_FLOAT, 120, {{16}, {16}});
 }
 
 TEST_F(ScopedAllocatorConcatOpTest, Success2) {
-  MakeOp({2, 2, 2}, DT_DOUBLE, "test", 120, 2);
+  MakeOp({2, 2, 2}, DT_DOUBLE, false, "test", 120, 2);
   ExecOp(DT_DOUBLE, 120, {{2, 2}, {2, 2}});
 }
 
 TEST_F(ScopedAllocatorConcatOpTest, Success3) {
-  MakeOp({3, 3, 3}, DT_HALF, "test", 120, 3);
+  MakeOp({3, 3, 3}, DT_HALF, false, "test", 120, 3);
   ExecOp(DT_HALF, 120, {{3, 3}, {3, 3}, {3, 3}});
 }
 
+TEST_F(ScopedAllocatorConcatOpTest, Reshape) {
+  MakeOp({2, 2, 2}, DT_DOUBLE, true, "test", 120, 2);
+  ExecOp(DT_DOUBLE, 120, {{2, 2}, {2, 2}});
+}
+
+TEST_F(ScopedAllocatorConcatOpTest, NoReshapeAttr) {
+  BuildNodeDef({3, 4, 4}, DT_HALF, "test", 120, 3);
+  TF_EXPECT_OK(InitOp());
+  ExecOp(DT_HALF, 120, {{4, 4}, {4, 4}, {4, 4}});
+}
+
 TEST_F(ScopedAllocatorConcatOpTest, FailDtypeCheck) {
-  MakeOp({8}, DT_FLOAT, "test", 120, 2);
+  MakeOp({8}, DT_FLOAT, false, "test", 120, 2);
   EXPECT_DEATH(ExecOp(DT_DOUBLE, 120, {{4}, {4}}), "");
 }
 
 TEST_F(ScopedAllocatorConcatOpTest, FailNumElementsCheck) {
-  MakeOp({32}, DT_FLOAT, "test", 120, 2);
+  MakeOp({32}, DT_FLOAT, false, "test", 120, 2);
   AddInputFromArray<float>({8}, {0, 1, 2, 3, 4, 5, 6, 7});
   AddInputFromArray<float>({4}, {0, 1, 2, 3});
   AddInputFromArray<float>({4}, {4, 5, 6, 7});
@@ -193,7 +239,7 @@ TEST_F(ScopedAllocatorConcatOpTest, FailNumElementsCheck) {
 // This test should fail because the backing tensor and the input tensors are
 // unrelated, i.e. the inputs are not slices of the backing tensor.
 TEST_F(ScopedAllocatorConcatOpTest, FailBounds) {
-  MakeOp({8}, DT_DOUBLE, "test", 120, 2);
+  MakeOp({8}, DT_DOUBLE, false, "test", 120, 2);
   AddInputFromArray<double>({8}, {0, 1, 2, 3, 4, 5, 6, 7});
   AddInputFromArray<double>({4}, {0, 1, 2, 3});
   AddInputFromArray<double>({4}, {4, 5, 6, 7});
diff --git a/tensorflow/core/ops/scoped_allocator_ops.cc b/tensorflow/core/ops/scoped_allocator_ops.cc
index f053a53f4c..1e0dcdac96 100644
--- a/tensorflow/core/ops/scoped_allocator_ops.cc
+++ b/tensorflow/core/ops/scoped_allocator_ops.cc
@@ -43,6 +43,7 @@ REGISTER_OP("_ScopedAllocatorConcat")
     .Input("inputs: N * T")
     .Attr("shape: shape")
     .Attr("T: type")
+    .Attr("reshape: bool = false")
     .Attr("sa_name: string")
     .Attr("id: int")
     .Attr("N: int >= 2")
@@ -69,10 +70,12 @@ REGISTER_OP("_ScopedAllocatorSplit")
     .SetIsStateful()
     .SetShapeFn(shape_inference::ExplicitShape)
     .Doc(R"doc(
-Acts like a Concat Op that merges multple tensors into one, however it must
-only be used in conjunction with a ScopedAllocator which is backing the memory
-of all of its input tensors so that actually it just outputs a read-only
-reference to that ScopedAllocator's backing tensor.
+Acts roughly like a SplitV Op that splits one tensor into multiple tensors
+but must only be used in conjunction with corresponding ScopedAllocator
+and ScopedAllocatorConcat instances.  In practice it is provided as inputs
+the backing tensor as first input, which contains the concatenated values,
+and a list of alias tensors as its other input and it simply outputs that
+second list.
 
 This is an experimental op for internal use only.  It is possible to use this
 op in unsafe ways.
-- 
GitLab


From 8a022d3d0e1bc521ccfee74174e75821bdc1bfa9 Mon Sep 17 00:00:00 2001
From: Francois Chollet <fchollet@google.com>
Date: Wed, 2 May 2018 16:58:35 -0700
Subject: [PATCH 183/395] Allow `Layer.add_loss` to receive non-tensor; fixes
 error triggered when using a weight regularizer of factor 0.

PiperOrigin-RevId: 195175909
---
 tensorflow/python/keras/_impl/keras/engine/base_layer.py | 3 +++
 tensorflow/python/keras/_impl/keras/regularizers_test.py | 5 +++++
 2 files changed, 8 insertions(+)

diff --git a/tensorflow/python/keras/_impl/keras/engine/base_layer.py b/tensorflow/python/keras/_impl/keras/engine/base_layer.py
index a3e78c95dc..3af4eaabe9 100644
--- a/tensorflow/python/keras/_impl/keras/engine/base_layer.py
+++ b/tensorflow/python/keras/_impl/keras/engine/base_layer.py
@@ -29,6 +29,7 @@ from tensorflow.python.estimator import util as estimator_util
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_util
 from tensorflow.python.keras._impl.keras import backend
 from tensorflow.python.keras._impl.keras import constraints
 from tensorflow.python.keras._impl.keras import initializers
@@ -390,6 +391,8 @@ class Layer(checkpointable.CheckpointableBase):
       raise RuntimeError('Layer.add_loss not supported in Eager mode.')
 
     losses = generic_utils.to_list(losses)
+    losses = [ops.convert_to_tensor(loss, dtype=backend.floatx())
+              if not tensor_util.is_tensor(loss) else loss for loss in losses]
     self._losses += losses
     if inputs is None:
       for loss in losses:
diff --git a/tensorflow/python/keras/_impl/keras/regularizers_test.py b/tensorflow/python/keras/_impl/keras/regularizers_test.py
index 9a1612b777..c4f04833ba 100644
--- a/tensorflow/python/keras/_impl/keras/regularizers_test.py
+++ b/tensorflow/python/keras/_impl/keras/regularizers_test.py
@@ -71,6 +71,11 @@ class KerasRegularizersTest(test.TestCase):
         model.fit(x_train, y_train, batch_size=10,
                   epochs=1, verbose=0)
 
+  def test_zero_regularization(self):
+    inputs = keras.backend.ones(shape=(10, 10))
+    layer = keras.layers.Dense(3, kernel_regularizer=keras.regularizers.l2(0))
+    layer(inputs)
+
 
 if __name__ == '__main__':
   test.main()
-- 
GitLab


From dde83d4bee2c524cb5bd0adc4f702c9fc5ac6f3f Mon Sep 17 00:00:00 2001
From: Benoit Steiner <bsteiner@google.com>
Date: Wed, 2 May 2018 17:00:16 -0700
Subject: [PATCH 184/395] Handle negative values when slicing symbolic shapes

PiperOrigin-RevId: 195176133
---
 tensorflow/core/grappler/costs/graph_properties.cc | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/grappler/costs/graph_properties.cc b/tensorflow/core/grappler/costs/graph_properties.cc
index 23d25cba8d..eaf7634daa 100644
--- a/tensorflow/core/grappler/costs/graph_properties.cc
+++ b/tensorflow/core/grappler/costs/graph_properties.cc
@@ -804,11 +804,16 @@ class SymbolicShapeRefiner {
           int64 start = slice_offset->dtype() == DT_INT32
                             ? slice_offset->flat<int32>()(0)
                             : slice_offset->flat<int64>()(0);
-          int64 end = start + (slice_size->dtype() == DT_INT32
-                                   ? slice_size->flat<int32>()(0)
-                                   : slice_size->flat<int64>()(0));
+          int64 size =
+              (slice_size->dtype() == DT_INT32 ? slice_size->flat<int32>()(0)
+                                               : slice_size->flat<int64>()(0));
           ShapeHandle result;
-          TF_RETURN_IF_ERROR(ic->Subshape(input, start, end, &result));
+          if (size == -1) {
+            TF_RETURN_IF_ERROR(ic->Subshape(input, start, &result));
+          } else {
+            int64 end = start + size;
+            TF_RETURN_IF_ERROR(ic->Subshape(input, start, end, &result));
+          }
           c->output_tensors_as_shapes.resize(1);
           c->output_tensors_as_shapes[0] = result;
         }
-- 
GitLab


From a1ef905926d12b0362c0dcf6d669e1c3d2ffcf70 Mon Sep 17 00:00:00 2001
From: Jeremy Lau <lauj@google.com>
Date: Wed, 2 May 2018 17:25:10 -0700
Subject: [PATCH 185/395] BufferValue is a new base class for LogicalBuffer and
 HloValue. This makes it easier to migrate from
 TuplePointsToAnalysis/LogicalBuffer to HloDataflowAnalysis/HloValue. No
 functional changes.

PiperOrigin-RevId: 195179676
---
 tensorflow/compiler/xla/service/BUILD         |  22 +++
 .../xla/service/buffer_assignment_test.cc     |   3 +-
 .../compiler/xla/service/buffer_value.cc      |  66 +++++++
 .../compiler/xla/service/buffer_value.h       | 177 ++++++++++++++++++
 tensorflow/compiler/xla/service/compiler.h    |   5 +-
 tensorflow/compiler/xla/service/gpu/BUILD     |   1 +
 .../compiler/xla/service/gpu/hlo_schedule.cc  |   3 +-
 .../xla/service/heap_simulator_test.cc        |   5 +-
 .../xla/service/hlo_rematerialization.cc      |   3 +-
 .../xla/service/hlo_scheduling_test.cc        |   6 +-
 tensorflow/compiler/xla/service/hlo_value.cc  |   8 +-
 tensorflow/compiler/xla/service/hlo_value.h   |  51 ++---
 .../compiler/xla/service/logical_buffer.cc    |  42 +----
 .../compiler/xla/service/logical_buffer.h     | 123 +-----------
 14 files changed, 318 insertions(+), 197 deletions(-)
 create mode 100644 tensorflow/compiler/xla/service/buffer_value.cc
 create mode 100644 tensorflow/compiler/xla/service/buffer_value.h

diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD
index 17964cdd59..0b8b22b44c 100644
--- a/tensorflow/compiler/xla/service/BUILD
+++ b/tensorflow/compiler/xla/service/BUILD
@@ -780,6 +780,7 @@ cc_library(
     srcs = ["compiler.cc"],
     hdrs = ["compiler.h"],
     deps = [
+        ":buffer_value",
         ":executable",
         ":hlo",
         ":hlo_module_config",
@@ -1014,6 +1015,7 @@ tf_cc_test(
     srcs = ["buffer_assignment_test.cc"],
     deps = [
         ":buffer_assignment",
+        ":buffer_value",
         ":call_graph",
         ":computation_tracker",
         ":copy_insertion",
@@ -1095,6 +1097,7 @@ tf_cc_test(
     name = "heap_simulator_test",
     srcs = ["heap_simulator_test.cc"],
     deps = [
+        ":buffer_value",
         ":heap_simulator",
         ":hlo",
         ":hlo_ordering",
@@ -1163,6 +1166,7 @@ tf_cc_test(
     name = "hlo_scheduling_test",
     srcs = ["hlo_scheduling_test.cc"],
     deps = [
+        ":buffer_value",
         ":hlo",
         ":hlo_ordering",
         ":hlo_scheduling",
@@ -1749,11 +1753,27 @@ tf_cc_test(
     ],
 )
 
+cc_library(
+    name = "buffer_value",
+    srcs = ["buffer_value.cc"],
+    hdrs = ["buffer_value.h"],
+    deps = [
+        ":hlo",
+        ":hlo_proto",
+        "//tensorflow/compiler/xla:shape_util",
+        "//tensorflow/compiler/xla:types",
+        "//tensorflow/compiler/xla:xla_data_proto",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+    ],
+)
+
 cc_library(
     name = "logical_buffer",
     srcs = ["logical_buffer.cc"],
     hdrs = ["logical_buffer.h"],
     deps = [
+        ":buffer_value",
         ":hlo",
         ":hlo_proto",
         "//tensorflow/compiler/xla:shape_util",
@@ -1769,6 +1789,7 @@ cc_library(
     srcs = ["hlo_value.cc"],
     hdrs = ["hlo_value.h"],
     deps = [
+        ":buffer_value",
         ":hlo",
         "//tensorflow/compiler/xla:shape_tree",
         "//tensorflow/compiler/xla:shape_util",
@@ -2066,6 +2087,7 @@ cc_library(
     hdrs = ["hlo_rematerialization.h"],
     deps = [
         ":buffer_liveness",
+        ":buffer_value",
         ":call_graph",
         ":flatten_call_graph",
         ":hlo",
diff --git a/tensorflow/compiler/xla/service/buffer_assignment_test.cc b/tensorflow/compiler/xla/service/buffer_assignment_test.cc
index f6d6b5c36a..a4fb0eefac 100644
--- a/tensorflow/compiler/xla/service/buffer_assignment_test.cc
+++ b/tensorflow/compiler/xla/service/buffer_assignment_test.cc
@@ -23,6 +23,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/compiler/xla/ptr_util.h"
+#include "tensorflow/compiler/xla/service/buffer_value.h"
 #include "tensorflow/compiler/xla/service/call_graph.h"
 #include "tensorflow/compiler/xla/service/computation_tracker.h"
 #include "tensorflow/compiler/xla/service/copy_insertion.h"
@@ -1684,7 +1685,7 @@ class WhileBufferAssignmentTest : public HloTestBase {
         .ConsumeValueOrDie();
   }
 
-  static int64 ByteSizeOf(const LogicalBuffer& buffer) {
+  static int64 ByteSizeOf(const BufferValue& buffer) {
     return ShapeUtil::ByteSizeOf(buffer.shape(), sizeof(void*));
   }
 
diff --git a/tensorflow/compiler/xla/service/buffer_value.cc b/tensorflow/compiler/xla/service/buffer_value.cc
new file mode 100644
index 0000000000..df1a5ca435
--- /dev/null
+++ b/tensorflow/compiler/xla/service/buffer_value.cc
@@ -0,0 +1,66 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/buffer_value.h"
+
+#include <iosfwd>
+
+#include "tensorflow/compiler/xla/service/hlo_computation.h"
+#include "tensorflow/compiler/xla/service/hlo_instruction.h"
+#include "tensorflow/compiler/xla/types.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace xla {
+
+BufferValue::BufferValue(HloInstruction* instruction, const ShapeIndex& index,
+                         Id id)
+    : id_(id) {
+  const Shape& shape = ShapeUtil::GetSubshape(instruction->shape(), index);
+  is_array_ = ShapeUtil::IsArray(shape);
+  is_tuple_ = ShapeUtil::IsTuple(shape);
+}
+
+BufferValue::~BufferValue() {}
+
+std::ostream& operator<<(std::ostream& out, const BufferValue& buffer) {
+  out << buffer.ToString();
+  return out;
+}
+
+/*static*/ LogicalBufferProto::Location BufferValue::ToLocationProto(
+    const HloInstruction& instruction, const ShapeIndex& index) {
+  LogicalBufferProto::Location proto;
+  proto.set_computation_name(instruction.parent()->name());
+  proto.set_instruction_name(instruction.name());
+  for (const int64 index_entry : index) {
+    proto.add_shape_index(index_entry);
+  }
+  return proto;
+}
+
+LogicalBufferProto BufferValue::ToProto(const SizeFunction& size_fn) const {
+  LogicalBufferProto proto;
+  proto.set_id(id());
+  proto.set_size(size_fn(*this));
+  LogicalBufferProto::Location proto_location =
+      ToLocationProto(*instruction(), index());
+  proto.mutable_defined_at()->Swap(&proto_location);
+  proto.set_color(color().value());
+  return proto;
+}
+
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/buffer_value.h b/tensorflow/compiler/xla/service/buffer_value.h
new file mode 100644
index 0000000000..f4be16e084
--- /dev/null
+++ b/tensorflow/compiler/xla/service/buffer_value.h
@@ -0,0 +1,177 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_BUFFER_VALUE_H_
+#define TENSORFLOW_COMPILER_XLA_SERVICE_BUFFER_VALUE_H_
+
+#include <functional>
+#include <string>
+
+#include "tensorflow/compiler/xla/service/hlo.pb.h"
+#include "tensorflow/compiler/xla/service/hlo_instruction.h"
+#include "tensorflow/compiler/xla/shape_util.h"
+#include "tensorflow/compiler/xla/types.h"
+#include "tensorflow/compiler/xla/xla_data.pb.h"
+#include "tensorflow/core/lib/gtl/array_slice.h"
+#include "tensorflow/core/lib/gtl/int_type.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace xla {
+
+// Abstract class describing a value used by one of the dataflow analyses -
+// TuplePointsToAnalysis or HloDataflowAnalysis.
+// TODO(b/78906445) Delete this class when TuplePointsToAnalysis is unused.
+//
+// XLA arrays are trivially a single BufferValue. Tuples are made up of more
+// than one BufferValue: an BufferValue for the pointer vector, and an
+// BufferValue for each child element.
+//
+// Every BufferValue is defined by a particular instruction and most
+// instructions define only a single BufferValue. Instructions which define a
+// single BufferValue include array-shaped instructions such as Add but also
+// includes Tuple-shaped instructions such as Tuple. The Tuple instruction
+// defines a single BufferValue which is a vector of pointers to the values
+// containing the Tuple instruction's operands. Though the result of the Tuple
+// instruction includes multiple values only the top-level BufferValue (the
+// vector of pointers) is defined by the Tuple instruction. The values
+// containing the tuple elements are defined by earlier instructions, usually
+// the operands of the Tuple instruction.
+//
+// Instructions which construct both the tuple *and* the tuple elements define
+// more than one BufferValue. This includes (at least) tuple-shaped Constant,
+// Parameter, Infeed and While instructions. These tuple-shaped instructions do
+// not assemble a tuple from existing BufferValues like the Tuple instruction
+// does, but rather define all the BufferValues in the tuple.
+//
+// Some instructions, such as Bitcast, define no buffers. These instructions
+// simply forward buffers from their operands.
+//
+// The BufferValue object describes which HLO instruction defines a buffer and
+// where within that instruction's output shape the buffer is defined. The
+// location within the output shape is indicated by BufferValue::index() which
+// is defined identically to the index used in ShapeUtil::GetSubshape().
+// Examples:
+//
+// %add = Add(%foo, %bar)
+// %tuple_constant = Constant({1, {42, 43}})
+//
+// %add defines a single array-shaped buffer BufferValue(%add, {}) which holds
+// the array result of the add operation. The nested-tuple-shaped
+// %tuple_constant defines 5 buffers described by the following BufferValue
+// objects:
+//
+//   BufferValue(%tuple_constant, {})      // "Top-level" buffer: vector of
+//                                         //  pointers to BufferValues at
+//                                         //  indices {0} and {1}
+//   BufferValue(%tuple_constant, {0})     // Holds value "1"
+//   BufferValue(%tuple_constant, {1})     // Holds nested tuple: vector of
+//                                         //  pointers to BufferValues at
+//                                         //  indices {1, 0} and {1, 1}
+//   BufferValue(%tuple_constant, {1, 0})  // Holds value "42"
+//   BufferValue(%tuple_constant, {1, 1})  // Holds value "43"
+
+class BufferValue {
+ public:
+  TF_LIB_GTL_DEFINE_INT_TYPE(Color, int64);
+
+  // Id is a unique identifier for the BufferValue to facilitate efficient
+  // collections of BufferValues with stable iteration order.
+  using Id = int64;
+
+  // Functions which return the size and alignment of a logical buffer in bytes.
+  using SizeFunction = std::function<int64(const BufferValue&)>;
+  using AlignmentFunction = std::function<int64(BufferValue::Color)>;
+
+  virtual ~BufferValue();
+
+  Id id() const { return id_; }
+
+  // Return the instruction that defines the buffer.
+  virtual HloInstruction* instruction() const = 0;
+
+  // Return the index within the output of the instruction where the buffer is
+  // defined. Index used defined as in ShapeUtil::GetSubshape()
+  virtual const ShapeIndex& index() const = 0;
+
+  // Return the color of the BufferValue. Differently colored buffers can not be
+  // parts of the same allocation.
+  Color color() const {
+    CHECK_NE(color_, kInvalidColor)
+        << "Should not query the color of a buffer that was never colored";
+    return color_;
+  }
+
+  void set_color(Color color) {
+    CHECK_NE(color, kInvalidColor)
+        << "Should not set the color of a buffer to the invalid color";
+    color_ = color;
+  }
+
+  bool has_color() const { return color_ != kInvalidColor; }
+
+  // Return the shape of the buffer. This reference points into the shape field
+  // of the instruction defining the buffer.  Therefore, the returned shape will
+  // contain the layout of instruction, if any.
+  virtual const Shape& shape() const = 0;
+
+  // Returns true if this buffer is the top-level output buffer of the defining
+  // HLO instruction. This is equivalent to index == {}.
+  bool IsTopLevel() const { return index().empty(); }
+
+  // Whether this buffer contains a tuple.
+  bool IsTuple() const { return is_tuple_; }
+
+  // Whether this buffer contains an array.
+  bool IsArray() const { return is_array_; }
+
+  // operator< is required for std::set.
+  bool operator<(const BufferValue& other) const { return id_ < other.id_; }
+
+  virtual string ToString() const = 0;
+
+  // TODO(lauj) rename LogicalBufferProto to BufferValueProto.
+  LogicalBufferProto ToProto(const SizeFunction& size_fn) const;
+
+  // Returns the LogicalBufferProto::Location that serializes the given
+  // instruction and index.
+  static LogicalBufferProto::Location ToLocationProto(
+      const HloInstruction& instruction, const ShapeIndex& index);
+
+  const Color kInvalidColor = Color(-1);
+
+ protected:
+  BufferValue(HloInstruction* instruction, const ShapeIndex& index, Id id);
+
+ private:
+  // The definining instruction and index are not stored here; they can be found
+  // in the LogicalBuffer and HloValue subclasses. This class exists only to
+  // support migrations from TuplePointsToAnalysis to HloDataflowAnalysis, by
+  // allowing abstract use of LogicalBuffer or HloValue. After those migrations
+  // are complete, this class should be deleted (b/78906445). Because we plan to
+  // delete LogicalBuffer and this class, we don't refactor all the shared
+  // features from LogicalBuffer and HloValue into this class.
+  Id id_ : 62;
+  bool is_array_ : 1;
+  bool is_tuple_ : 1;
+  Color color_ = kInvalidColor;
+};
+
+std::ostream& operator<<(std::ostream& out, const BufferValue& buffer);
+
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_XLA_SERVICE_BUFFER_VALUE_H_
diff --git a/tensorflow/compiler/xla/service/compiler.h b/tensorflow/compiler/xla/service/compiler.h
index 5c14591d93..a4b59d1ba9 100644
--- a/tensorflow/compiler/xla/service/compiler.h
+++ b/tensorflow/compiler/xla/service/compiler.h
@@ -25,6 +25,7 @@ limitations under the License.
 #include <memory>
 #include <string>
 
+#include "tensorflow/compiler/xla/service/buffer_value.h"
 #include "tensorflow/compiler/xla/service/executable.h"
 #include "tensorflow/compiler/xla/service/hlo_module.h"
 #include "tensorflow/compiler/xla/service/hlo_module_config.h"
@@ -181,9 +182,9 @@ class Compiler {
 
   // Returns a function that computes the size in bytes of a given
   // logical buffer.
-  std::function<int64(const LogicalBuffer&)> BufferSizeBytesFunction() {
+  std::function<int64(const BufferValue&)> BufferSizeBytesFunction() {
     HloCostAnalysis::ShapeSizeFunction shape_size = ShapeSizeBytesFunction();
-    return [shape_size](const LogicalBuffer& buffer) {
+    return [shape_size](const BufferValue& buffer) {
       return shape_size(buffer.shape());
     };
   }
diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD
index f1707442fe..7cb7f55073 100644
--- a/tensorflow/compiler/xla/service/gpu/BUILD
+++ b/tensorflow/compiler/xla/service/gpu/BUILD
@@ -620,6 +620,7 @@ cc_library(
         "//tensorflow/compiler/xla:statusor",
         "//tensorflow/compiler/xla:types",
         "//tensorflow/compiler/xla:util",
+        "//tensorflow/compiler/xla/service:buffer_value",
         "//tensorflow/compiler/xla/service:hlo",
         "//tensorflow/compiler/xla/service:hlo_ordering",
         "//tensorflow/compiler/xla/service:hlo_reachability",
diff --git a/tensorflow/compiler/xla/service/gpu/hlo_schedule.cc b/tensorflow/compiler/xla/service/gpu/hlo_schedule.cc
index 42c1539e86..f766f96882 100644
--- a/tensorflow/compiler/xla/service/gpu/hlo_schedule.cc
+++ b/tensorflow/compiler/xla/service/gpu/hlo_schedule.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/service/gpu/hlo_schedule.h"
 
 #include "tensorflow/compiler/xla/ptr_util.h"
+#include "tensorflow/compiler/xla/service/buffer_value.h"
 #include "tensorflow/compiler/xla/service/hlo_reachability.h"
 #include "tensorflow/compiler/xla/service/hlo_scheduling.h"
 #include "tensorflow/compiler/xla/types.h"
@@ -199,7 +200,7 @@ StatusOr<std::unique_ptr<HloSchedule>> HloSchedule::Build(
     TF_ASSIGN_OR_RETURN(
         schedule->thunk_launch_order_,
         CreateMemoryMinimizingSequence(
-            *entry_computation, [pointer_size](const LogicalBuffer& buffer) {
+            *entry_computation, [pointer_size](const BufferValue& buffer) {
               return ShapeUtil::ByteSizeOf(buffer.shape(), pointer_size);
             }));
   } else {
diff --git a/tensorflow/compiler/xla/service/heap_simulator_test.cc b/tensorflow/compiler/xla/service/heap_simulator_test.cc
index e983fd11d4..fd56a603bb 100644
--- a/tensorflow/compiler/xla/service/heap_simulator_test.cc
+++ b/tensorflow/compiler/xla/service/heap_simulator_test.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <vector>
 
 #include "tensorflow/compiler/xla/literal_util.h"
+#include "tensorflow/compiler/xla/service/buffer_value.h"
 #include "tensorflow/compiler/xla/service/hlo_computation.h"
 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
 #include "tensorflow/compiler/xla/service/hlo_module.h"
@@ -85,7 +86,7 @@ class HeapSimulatorTracker {
     // size of the buffers doesn't matter, so we always return 0.  We rely on
     // the secondary sorting criteria of DecreasingSizeRunsHeap to sort calls by
     // buffer id, for determinism in the tests.
-    auto zero_size = [](const LogicalBuffer& buffer) { return 0; };
+    auto zero_size = [](const BufferValue& buffer) { return 0; };
     auto algorithm = MakeUnique<DecreasingSizeRunsHeap>(
         MakeUnique<HeapCallRecorder>(&actual_calls_));
     result_ = HeapSimulator::Run(
@@ -119,7 +120,7 @@ class HeapSimulatorTracker {
     // the sequence. This lets us ensure the Alloc calls are in the sequence
     // order. The Free calls are sorted by LogicalBuffer.id, which is at least
     // deterministic.
-    auto size_fn = [&reverse_position](const LogicalBuffer& buffer) {
+    auto size_fn = [&reverse_position](const BufferValue& buffer) {
       return reverse_position[buffer.instruction()];
     };
     auto algorithm = MakeUnique<DecreasingSizeRunsHeap>(
diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization.cc b/tensorflow/compiler/xla/service/hlo_rematerialization.cc
index b063244893..b171d41a31 100644
--- a/tensorflow/compiler/xla/service/hlo_rematerialization.cc
+++ b/tensorflow/compiler/xla/service/hlo_rematerialization.cc
@@ -22,6 +22,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/map_util.h"
 #include "tensorflow/compiler/xla/primitive_util.h"
+#include "tensorflow/compiler/xla/service/buffer_value.h"
 #include "tensorflow/compiler/xla/service/flatten_call_graph.h"
 #include "tensorflow/compiler/xla/service/hlo_computation.h"
 #include "tensorflow/compiler/xla/service/hlo_dce.h"
@@ -1216,7 +1217,7 @@ StatusOr<bool> HloRematerialization::Run(
   // Create initial sequence of HLO instructions.
   TF_ASSIGN_OR_RETURN(*sequence, CreateMemoryMinimizingSequence(
                                      *module,
-                                     [this](const LogicalBuffer& buffer) {
+                                     [this](const BufferValue& buffer) {
                                        return size_function_(buffer.shape());
                                      },
                                      scheduler_algorithm_));
diff --git a/tensorflow/compiler/xla/service/hlo_scheduling_test.cc b/tensorflow/compiler/xla/service/hlo_scheduling_test.cc
index 74544c4a67..92df7c1427 100644
--- a/tensorflow/compiler/xla/service/hlo_scheduling_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_scheduling_test.cc
@@ -77,7 +77,7 @@ TEST_F(MinimumMemoryForSequenceTest, MultiComputation) {
   HloComputation* entry_computation =
       module->AddEntryComputation(builder.Build());
 
-  auto size_fn = [](const LogicalBuffer& buffer) {
+  auto size_fn = [](const BufferValue& buffer) {
     return ShapeUtil::ByteSizeOf(buffer.shape(), /*pointer_size=*/8);
   };
 
@@ -124,7 +124,7 @@ TEST_F(HloSchedulingTest, LastUseScheduledFirst) {
 
   TF_ASSERT_OK_AND_ASSIGN(
       SequentialHloOrdering::HloModuleSequence sequence,
-      CreateMemoryMinimizingSequence(*module, [](const LogicalBuffer& buffer) {
+      CreateMemoryMinimizingSequence(*module, [](const BufferValue& buffer) {
         return ShapeUtil::ByteSizeOf(buffer.shape());
       }));
   // Verify that all instructions are in the sequence.
@@ -160,7 +160,7 @@ ENTRY root {
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                           tools::Parse(module_str));
 
-  auto size_fn = [](const LogicalBuffer& buffer) {
+  auto size_fn = [](const BufferValue& buffer) {
     return ShapeUtil::ByteSizeOf(buffer.shape(), /*pointer_size=*/8);
   };
   TF_ASSERT_OK_AND_ASSIGN(
diff --git a/tensorflow/compiler/xla/service/hlo_value.cc b/tensorflow/compiler/xla/service/hlo_value.cc
index 05b7dce3d1..7b27dbfec3 100644
--- a/tensorflow/compiler/xla/service/hlo_value.cc
+++ b/tensorflow/compiler/xla/service/hlo_value.cc
@@ -29,9 +29,11 @@ limitations under the License.
 #include "tensorflow/compiler/xla/types.h"
 #include "tensorflow/compiler/xla/util.h"
 #include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/gtl/flatset.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/types.h"
 
 namespace xla {
 
@@ -69,7 +71,7 @@ std::ostream& operator<<(std::ostream& out, const HloUse& use) {
 
 HloValue::HloValue(HloValue::Id id, HloInstruction* instruction,
                    const ShapeIndex& index, bool is_phi)
-    : id_(id), is_phi_(is_phi) {
+    : BufferValue(instruction, index, id), is_phi_(is_phi) {
   // The defining position is always the first element in the positions_ vector.
   positions_.push_back(HloPosition{instruction, index});
 }
@@ -90,8 +92,8 @@ string HloValue::ToShortString() const {
   string index_str = ShapeUtil::IsTuple(defining_instruction()->shape())
                          ? defining_index().ToString()
                          : "";
-  return StrCat(id_, " ", is_phi_ ? "PHI " : "", defining_instruction()->name(),
-                index_str);
+  return StrCat(id(), " ", is_phi_ ? "PHI " : "",
+                defining_instruction()->name(), index_str);
 }
 
 string HloValue::ToString(int indent) const {
diff --git a/tensorflow/compiler/xla/service/hlo_value.h b/tensorflow/compiler/xla/service/hlo_value.h
index 2a711e8b42..a1151f65e0 100644
--- a/tensorflow/compiler/xla/service/hlo_value.h
+++ b/tensorflow/compiler/xla/service/hlo_value.h
@@ -16,16 +16,20 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VALUE_H_
 #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VALUE_H_
 
-#include <ostream>
+#include <stddef.h>
 #include <string>
 #include <vector>
 
+#include "tensorflow/compiler/xla/service/buffer_value.h"
 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
 #include "tensorflow/compiler/xla/shape_tree.h"
+#include "tensorflow/compiler/xla/shape_util.h"
 #include "tensorflow/compiler/xla/types.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
 #include "tensorflow/core/lib/gtl/array_slice.h"
+#include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
 
 namespace xla {
 
@@ -80,30 +84,9 @@ struct HloUse {
 
 std::ostream& operator<<(std::ostream& out, const HloUse& use);
 
-// Class describing a value used by the dataflow analysis. XLA arrays are
-// trivially a single HloValue. Tuples are made up of more than one HloValue: an
-// HloValue for the pointer vector, and an HloValue for each child element.
-//
-// Every HloValue is defined by a particular instruction and most instructions
-// define only a single HloValue. Instructions which define a single HloValue
-// include array-shaped instructions such as Add but also includes Tuple-shaped
-// instructions such as Tuple. The Tuple instruction defines a single HloValue
-// which is a vector of pointers to the values containing the Tuple
-// instruction's operands. Though the result of the Tuple instruction includes
-// multiple values only the top-level HloValue (the vector of pointers) is
-// defined by the Tuple instruction. The values containing the tuple elements
-// are defined by earlier instructions, usually the operands of the Tuple
-// instruction.
-//
-// Instructions which construct both the tuple *and* the tuple elements define
-// more than one HloValue. This includes (at least) tuple-shaped Constant,
-// Parameter, Infeed and While instructions. These tuple-shaped instructions do
-// not assemble a tuple from existing HloValues like the Tuple instruction does,
-// but rather define all the HloValues in the tuple.
-class HloValue {
+// HloDataflowAnalysis uses this subclass of BufferValue.
+class HloValue : public BufferValue {
  public:
-  using Id = int64;
-
   // Predicate comparing HloValues by increasing id, useful for std::sort.
   static bool IdLessThan(const HloValue* a, const HloValue* b) {
     return a->id() < b->id();
@@ -120,6 +103,7 @@ class HloValue {
   // dataflow analysis (HloDataflowAnalysis::ssa_form_ is true).
   HloValue(Id id, HloInstruction* instruction, const ShapeIndex& index,
            bool is_phi = false);
+  ~HloValue() override {}
 
   // Sets the positions in the module at which the HloValue appears. Updates
   // uses. Should be called once and only once. The defining position should not
@@ -127,10 +111,6 @@ class HloValue {
   void SetPositionsAndComputeUses(
       tensorflow::gtl::ArraySlice<HloPosition> positions);
 
-  // Return a unique identifier for this HloValue. This value is used for stable
-  // sorting and iteration
-  Id id() const { return id_; }
-
   // Returns whether this value is a phi value.
   bool is_phi() const { return is_phi_; }
 
@@ -142,12 +122,18 @@ class HloValue {
     return defining_position().instruction;
   }
 
+  HloInstruction* instruction() const override {
+    return defining_instruction();
+  }
+
   // Return the shape index at which this HloValue is defined in the output of
   // its defining instruction.
   const ShapeIndex& defining_index() const { return defining_position().index; }
 
+  const ShapeIndex& index() const override { return defining_index(); }
+
   // Return the shape of this HloValue.
-  const Shape& shape() const { return defining_position().shape(); }
+  const Shape& shape() const override { return defining_position().shape(); }
 
   // Return all positions of the HloValue in the module.
   const std::vector<HloPosition>& positions() const { return positions_; }
@@ -164,12 +150,11 @@ class HloValue {
   // Return a single-line string representation of the value.
   string ToShortString() const;
 
-  string ToString(int indent = 0) const;
+  string ToString(int indent) const;
 
- private:
-  // Unique identifier for this HloValue. Used for stable sorting and iteration.
-  const Id id_;
+  string ToString() const override { return ToString(0); }
 
+ private:
   // Whether this instruction is a phi value.
   const bool is_phi_;
 
diff --git a/tensorflow/compiler/xla/service/logical_buffer.cc b/tensorflow/compiler/xla/service/logical_buffer.cc
index 68553bed12..1b3de8ad17 100644
--- a/tensorflow/compiler/xla/service/logical_buffer.cc
+++ b/tensorflow/compiler/xla/service/logical_buffer.cc
@@ -15,9 +15,6 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/service/logical_buffer.h"
 
-#include <ostream>
-#include <vector>
-
 #include "tensorflow/compiler/xla/service/hlo_computation.h"
 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
 #include "tensorflow/compiler/xla/types.h"
@@ -28,43 +25,16 @@ namespace xla {
 
 LogicalBuffer::LogicalBuffer(HloInstruction* instruction,
                              const ShapeIndex& index, Id id)
-    : instruction_(instruction), id_(id), color_(kInvalidColor), index_(index) {
-  const auto& s = shape();
-  is_array_ = ShapeUtil::IsArray(s);
-  is_tuple_ = ShapeUtil::IsTuple(s);
-}
+    : BufferValue(instruction, index, id),
+      instruction_(instruction),
+      index_(index) {}
+
+LogicalBuffer::~LogicalBuffer() {}
 
 string LogicalBuffer::ToString() const {
   return tensorflow::strings::StrCat(instruction_->name(), "[",
                                      tensorflow::str_util::Join(index_, ","),
-                                     "](#", id_, " @", color_.value(), ")");
-}
-
-std::ostream& operator<<(std::ostream& out, const LogicalBuffer& buffer) {
-  out << buffer.ToString();
-  return out;
-}
-
-/*static*/ LogicalBufferProto::Location LogicalBuffer::ToLocationProto(
-    const HloInstruction& instruction, const ShapeIndex& index) {
-  LogicalBufferProto::Location proto;
-  proto.set_computation_name(instruction.parent()->name());
-  proto.set_instruction_name(instruction.name());
-  for (const int64 index_entry : index) {
-    proto.add_shape_index(index_entry);
-  }
-  return proto;
-}
-
-LogicalBufferProto LogicalBuffer::ToProto(const SizeFunction& size_fn) const {
-  LogicalBufferProto proto;
-  proto.set_id(id_);
-  proto.set_size(size_fn(*this));
-  LogicalBufferProto::Location proto_location =
-      ToLocationProto(*instruction_, index_);
-  proto.mutable_defined_at()->Swap(&proto_location);
-  proto.set_color(color_.value());
-  return proto;
+                                     "](#", id(), " @", color().value(), ")");
 }
 
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/service/logical_buffer.h b/tensorflow/compiler/xla/service/logical_buffer.h
index 67b205e289..f9ba5a5547 100644
--- a/tensorflow/compiler/xla/service/logical_buffer.h
+++ b/tensorflow/compiler/xla/service/logical_buffer.h
@@ -16,11 +16,9 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_H_
 #define TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_H_
 
-#include <functional>
-#include <iosfwd>
 #include <string>
-#include <vector>
 
+#include "tensorflow/compiler/xla/service/buffer_value.h"
 #include "tensorflow/compiler/xla/service/hlo.pb.h"
 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
 #include "tensorflow/compiler/xla/shape_util.h"
@@ -33,133 +31,30 @@ limitations under the License.
 
 namespace xla {
 
-// Class describing a contiguous sequence of elements (ie, C array) which form
-// the components of Shaped values in XLA. XLA arrays are trivially a
-// single LogicalBuffer. Tuple values are made up of more than one
-// LogicalBuffer: a LogicalBuffer for the pointers to elements, and a
-// LogicalBuffer for each child element.
-//
-// Every buffer is defined by a particular instruction and most instructions
-// define only a single buffer. Instructions which define a single buffer
-// include array-shaped instructions such as Add but also includes Tuple-shaped
-// instructions such as Tuple. The Tuple instruction defines a single buffer
-// which is a vector of pointers to the buffers containing the Tuple
-// instruction's operands. Though the result of the Tuple instruction includes
-// multiple buffers only the top-level buffer (the vector of pointers) is
-// defined by the Tuple instruction. The buffers containing the tuple elements
-// are defined by earlier instructions, usually the operands of the Tuple
-// instruction.
-//
-// Instructions which construct both the tuple *and* the tuple elements define
-// more than one buffer. This includes (at least) tuple-shaped Constant,
-// Parameter, Infeed and While instructions. The tuple-shaped instructions do
-// not assemble a tuple from existing buffers like the Tuple instruction does,
-// but rather define the entire tuple.
-//
-// Some instructions, such as Bitcast, define no buffers. These instructions
-// simply forward buffers from their operands.
-//
-// The LogicalBuffer object describes which HLO instruction defines a buffer and
-// where within that instruction's output shape the buffer is defined. The
-// location within the output shape is indicated by LogicalBuffer::index() which
-// is defined identically to the index used in
-// ShapeUtil::GetSubshape(). Examples:
-//
-// %add = Add(%foo, %bar)
-// %tuple_constant = Constant({1, {42, 43}})
-//
-// %add defines a single array-shaped buffer LogicalBuffer(%add, {}) which holds
-// the array result of the add operation. The nested-tuple-shaped
-// %tuple_constant defines 5 buffers described by the following LogicalBuffer
-// objects:
-//
-//   LogicalBuffer(%tuple_constant, {})      // "Top-level" buffer: vector of
-//                                           //  pointers to LogicalBuffers at
-//                                           //  indices {0} and {1}
-//   LogicalBuffer(%tuple_constant, {0})     // Holds value "1"
-//   LogicalBuffer(%tuple_constant, {1})     // Holds nested tuple: vector of
-//                                           //  pointers to LogicalBuffers at
-//                                           //  indices {1, 0} and {1, 1}
-//   LogicalBuffer(%tuple_constant, {1, 0})  // Holds value "42"
-//   LogicalBuffer(%tuple_constant, {1, 1})  // Holds value "43"
-class LogicalBuffer {
+// TuplePointsToAnalysis uses this subclass of BufferValue.
+class LogicalBuffer : public BufferValue {
  public:
-  TF_LIB_GTL_DEFINE_INT_TYPE(Color, int64);
-
-  // Id is a unique identifier for the LogicalBuffer to facilitate efficient
-  // collections of LogicalBuffers with stable iteration order.
-  // LogicalBuffers are typically created and accessed through
-  // TuplePointsToAnalysis, and points-to analysis assigns each LogicalBuffer a
-  // unique value.
-  using Id = int64;
-
-  // Functions which return the size and alignment of a logical buffer in bytes.
-  using SizeFunction = std::function<int64(const LogicalBuffer&)>;
-  using AlignmentFunction = std::function<int64(LogicalBuffer::Color)>;
-
   LogicalBuffer(HloInstruction* instruction, const ShapeIndex& index, Id id);
-
-  Id id() const { return id_; }
+  ~LogicalBuffer() override;
 
   // Return the instruction that defines the buffer.
-  HloInstruction* instruction() const { return instruction_; }
+  HloInstruction* instruction() const override { return instruction_; }
 
   // Return the index within the output of the instruction where the buffer is
   // defined. Index used defined as in ShapeUtil::GetSubshape()
-  const ShapeIndex& index() const { return index_; }
-
-  // Return the color of the logical buffer. Differently colored buffers can
-  // not be parts of the same allocation.
-  Color color() const {
-    CHECK_NE(color_, kInvalidColor)
-        << "Should not query the color of a buffer that was never colored";
-    return color_;
-  }
-
-  void set_color(Color color) {
-    CHECK_NE(color, kInvalidColor)
-        << "Should not set the color of a buffer to the invalid color";
-    color_ = color;
-  }
-
-  bool has_color() const { return color_ != kInvalidColor; }
+  const ShapeIndex& index() const override { return index_; }
 
   // Return the shape of the buffer. This reference points into the shape field
   // of the instruction defining the buffer.  Therefore, the returned shape will
   // contain the layout of instruction, if any.
-  const Shape& shape() const {
+  const Shape& shape() const override {
     return ShapeUtil::GetSubshape(instruction_->shape(), index_);
   }
 
-  // Returns true if this buffer is the top-level output buffer of the defining
-  // HLO instruction. This is equivalent to index == {}.
-  bool IsTopLevel() const { return index_.empty(); }
-
-  // Whether this buffer contains a tuple.
-  bool IsTuple() const { return is_tuple_; }
-
-  // Whether this buffer contains an array.
-  bool IsArray() const { return is_array_; }
-
-  // operator< is required for std::set.
-  bool operator<(const LogicalBuffer& other) const { return id_ < other.id_; }
-
-  string ToString() const;
-  LogicalBufferProto ToProto(const SizeFunction& size_fn) const;
-
-  // Returns the LogicalBufferProto::Location that serializes the given
-  // instruction and index.
-  static LogicalBufferProto::Location ToLocationProto(
-      const HloInstruction& instruction, const ShapeIndex& index);
-
-  const Color kInvalidColor = Color(-1);
+  string ToString() const override;
 
  private:
   HloInstruction* instruction_;
-  Id id_ : 62;
-  bool is_array_ : 1;
-  bool is_tuple_ : 1;
-  Color color_;
   ShapeIndex index_;
 
   // Similar to HLO constructs (HloInstruction, etc), pointers are used for
@@ -167,8 +62,6 @@ class LogicalBuffer {
   TF_DISALLOW_COPY_AND_ASSIGN(LogicalBuffer);
 };
 
-std::ostream& operator<<(std::ostream& out, const LogicalBuffer& buffer);
-
 }  // namespace xla
 
 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_H_
-- 
GitLab


From 7833890a0da5226e4c409b1020155f1718c0edb2 Mon Sep 17 00:00:00 2001
From: Anna R <annarev@google.com>
Date: Wed, 2 May 2018 17:41:26 -0700
Subject: [PATCH 186/395] Add a collect_trace option to run_op_benchmark for
 cases when callers just want to pass RunOptions.FULL_TRACE but don't want to
 store trace in extras.

PiperOrigin-RevId: 195181533
---
 tensorflow/python/kernel_tests/benchmark_test.py | 12 ++++++++----
 tensorflow/python/platform/benchmark.py          | 14 ++++++++++----
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/kernel_tests/benchmark_test.py b/tensorflow/python/kernel_tests/benchmark_test.py
index 623343602d..78b6e38d94 100644
--- a/tensorflow/python/kernel_tests/benchmark_test.py
+++ b/tensorflow/python/kernel_tests/benchmark_test.py
@@ -67,7 +67,7 @@ class TestReportingBenchmark(test.Benchmark):
     with session.Session() as sess:
       a = constant_op.constant(0.0)
       a_plus_a = a + a
-      self.run_op_benchmark(
+      return self.run_op_benchmark(
           sess, a_plus_a, min_iters=1000, store_trace=True, name="op_benchmark")
 
 
@@ -148,7 +148,7 @@ class BenchmarkTest(test.TestCase):
       reporting = TestReportingBenchmark()
       reporting.benchmarkReport1()  # This should write
       reporting.benchmarkReport2()  # This should write
-      reporting.benchmark_times_an_op()  # This should write
+      benchmark_values3 = reporting.benchmark_times_an_op()  # This should write
 
       # Check the files were written
       self.assertTrue(gfile.Exists(expected_output_file))
@@ -186,8 +186,12 @@ class BenchmarkTest(test.TestCase):
       self.assertEquals(expected_3.name, read_benchmark_3.name)
       self.assertEquals(expected_3.iters, read_benchmark_3.iters)
       self.assertGreater(read_benchmark_3.wall_time, 0)
-      full_trace = read_benchmark_3.extras["full_trace_chrome_format"]
-      json_trace = json.loads(full_trace.string_value)
+
+      # Trace is not stored in benchmark entry. Instead we get it from
+      # return value of `run_op_benchmark` call.
+      full_trace = benchmark_values3["extras"]["full_trace_chrome_format"]
+      json_trace = json.loads(full_trace)
+
       self.assertTrue(isinstance(json_trace, dict))
       self.assertTrue("traceEvents" in json_trace.keys())
       allocator_keys = [k for k in read_benchmark_3.extras.keys()
diff --git a/tensorflow/python/platform/benchmark.py b/tensorflow/python/platform/benchmark.py
index 12dae94a64..eba2baaf6f 100644
--- a/tensorflow/python/platform/benchmark.py
+++ b/tensorflow/python/platform/benchmark.py
@@ -213,9 +213,10 @@ class TensorFlowBenchmark(Benchmark):
       burn_iters: Number of burn-in iterations to run.
       min_iters: Minimum number of iterations to use for timing.
       store_trace: Boolean, whether to run an extra untimed iteration and
-        store the trace of iteration in the benchmark report.
+        store the trace of iteration in returned extras.
         The trace will be stored as a string in Google Chrome trace format
-        in the extras field "full_trace_chrome_format".
+        in the extras field "full_trace_chrome_format". Note that trace
+        will not be stored in test_log_pb2.TestResults proto.
       store_memory_usage: Boolean, whether to run an extra untimed iteration,
         calculate memory usage, and store that in extras fields.
       name: (optional) Override the BenchmarkEntry name with `name`.
@@ -227,7 +228,9 @@ class TensorFlowBenchmark(Benchmark):
 
     Returns:
       A `dict` containing the key-value pairs that were passed to
-      `report_benchmark`.
+      `report_benchmark`. If `store_trace` option is used, then
+      `full_chrome_trace_format` will be included in return dictionary even
+      though it is not passed to `report_benchmark` with `extras`.
     """
     for _ in range(burn_iters):
       sess.run(op_or_tensor, feed_dict=feed_dict)
@@ -242,6 +245,7 @@ class TensorFlowBenchmark(Benchmark):
       deltas[i] = delta
 
     extras = extras if extras is not None else {}
+    unreported_extras = {}
     if store_trace or store_memory_usage:
       run_options = config_pb2.RunOptions(
           trace_level=config_pb2.RunOptions.FULL_TRACE)
@@ -251,7 +255,8 @@ class TensorFlowBenchmark(Benchmark):
       tl = timeline.Timeline(run_metadata.step_stats)
 
       if store_trace:
-        extras["full_trace_chrome_format"] = tl.generate_chrome_trace_format()
+        unreported_extras["full_trace_chrome_format"] = (
+            tl.generate_chrome_trace_format())
 
       if store_memory_usage:
         step_stats_analysis = tl.analyze_step_stats(show_memory=True)
@@ -277,6 +282,7 @@ class TensorFlowBenchmark(Benchmark):
         "throughput": mbs / median_delta
     }
     self.report_benchmark(**benchmark_values)
+    benchmark_values["extras"].update(unreported_extras)
     return benchmark_values
 
 
-- 
GitLab


From a44996a84b24c43cca40c685a009fd59275755ab Mon Sep 17 00:00:00 2001
From: Yong Tang <yong.tang.github@outlook.com>
Date: Thu, 3 May 2018 02:45:29 +0200
Subject: [PATCH 187/395] Add go_package to proto definition files (#17262)

* Add go_package to proto definition files

This fix tries to address the issue raised in 16282 by
add go_package to proto files, so that generated go files
have correct path.

This fix fixes 16282.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

* Add go_package to proto definition in tensorflow/core/framework

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

* Add go_package to proto definition in tensorflow/core/example

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

* Add go_package to proto definition in tensorflow/core/example

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
---
 tensorflow/core/example/example.proto                           | 2 +-
 tensorflow/core/example/example_parser_configuration.proto      | 1 +
 tensorflow/core/example/feature.proto                           | 2 +-
 tensorflow/core/framework/allocation_description.proto          | 1 +
 tensorflow/core/framework/api_def.proto                         | 1 +
 tensorflow/core/framework/attr_value.proto                      | 2 +-
 tensorflow/core/framework/cost_graph.proto                      | 2 +-
 tensorflow/core/framework/device_attributes.proto               | 1 +
 tensorflow/core/framework/function.proto                        | 2 +-
 tensorflow/core/framework/graph.proto                           | 2 +-
 tensorflow/core/framework/graph_transfer_info.proto             | 2 +-
 tensorflow/core/framework/iterator.proto                        | 1 +
 tensorflow/core/framework/kernel_def.proto                      | 2 +-
 tensorflow/core/framework/log_memory.proto                      | 2 +-
 tensorflow/core/framework/node_def.proto                        | 2 +-
 tensorflow/core/framework/op_def.proto                          | 2 +-
 tensorflow/core/framework/reader_base.proto                     | 1 +
 tensorflow/core/framework/remote_fused_graph_execute_info.proto | 2 +-
 tensorflow/core/framework/resource_handle.proto                 | 1 +
 tensorflow/core/framework/step_stats.proto                      | 2 +-
 tensorflow/core/framework/summary.proto                         | 2 +-
 tensorflow/core/framework/tensor.proto                          | 2 +-
 tensorflow/core/framework/tensor_description.proto              | 2 +-
 tensorflow/core/framework/tensor_shape.proto                    | 1 +
 tensorflow/core/framework/tensor_slice.proto                    | 1 +
 tensorflow/core/framework/types.proto                           | 1 +
 tensorflow/core/framework/variable.proto                        | 1 +
 tensorflow/core/framework/versions.proto                        | 1 +
 tensorflow/core/lib/core/error_codes.proto                      | 1 +
 tensorflow/core/protobuf/cluster.proto                          | 1 +
 tensorflow/core/protobuf/config.proto                           | 2 +-
 tensorflow/core/protobuf/control_flow.proto                     | 1 +
 tensorflow/core/protobuf/critical_section.proto                 | 1 +
 tensorflow/core/protobuf/debug.proto                            | 1 +
 tensorflow/core/protobuf/device_properties.proto                | 1 +
 tensorflow/core/protobuf/master.proto                           | 2 +-
 tensorflow/core/protobuf/master_service.proto                   | 2 +-
 tensorflow/core/protobuf/meta_graph.proto                       | 2 +-
 tensorflow/core/protobuf/named_tensor.proto                     | 2 +-
 tensorflow/core/protobuf/queue_runner.proto                     | 2 +-
 tensorflow/core/protobuf/rewriter_config.proto                  | 1 +
 tensorflow/core/protobuf/saved_model.proto                      | 2 +-
 tensorflow/core/protobuf/saver.proto                            | 1 +
 tensorflow/core/protobuf/tensor_bundle.proto                    | 2 +-
 tensorflow/core/protobuf/tensorflow_server.proto                | 2 +-
 tensorflow/core/protobuf/worker.proto                           | 2 +-
 tensorflow/core/protobuf/worker_service.proto                   | 2 +-
 47 files changed, 47 insertions(+), 27 deletions(-)

diff --git a/tensorflow/core/example/example.proto b/tensorflow/core/example/example.proto
index b2b723278b..e7142a4ef9 100644
--- a/tensorflow/core/example/example.proto
+++ b/tensorflow/core/example/example.proto
@@ -7,7 +7,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "ExampleProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.example";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/example";
 package tensorflow;
 
 // An Example is a mostly-normalized data format for storing data for
diff --git a/tensorflow/core/example/example_parser_configuration.proto b/tensorflow/core/example/example_parser_configuration.proto
index 15846c0e30..b2c115d80e 100644
--- a/tensorflow/core/example/example_parser_configuration.proto
+++ b/tensorflow/core/example/example_parser_configuration.proto
@@ -6,6 +6,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "ExampleParserConfigurationProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.example";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/example";
 package tensorflow;
 
 import "tensorflow/core/framework/tensor_shape.proto";
diff --git a/tensorflow/core/example/feature.proto b/tensorflow/core/example/feature.proto
index da3dc59a12..6d81974aac 100644
--- a/tensorflow/core/example/feature.proto
+++ b/tensorflow/core/example/feature.proto
@@ -58,7 +58,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "FeatureProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.example";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/example";
 package tensorflow;
 
 // Containers to hold repeated fundamental values.
diff --git a/tensorflow/core/framework/allocation_description.proto b/tensorflow/core/framework/allocation_description.proto
index bb1037c2df..64133b05e1 100644
--- a/tensorflow/core/framework/allocation_description.proto
+++ b/tensorflow/core/framework/allocation_description.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "AllocationDescriptionProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 
 message AllocationDescription {
   // Total number of bytes requested
diff --git a/tensorflow/core/framework/api_def.proto b/tensorflow/core/framework/api_def.proto
index 98c38efc0e..cce02d84b2 100644
--- a/tensorflow/core/framework/api_def.proto
+++ b/tensorflow/core/framework/api_def.proto
@@ -8,6 +8,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "ApiDefProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/attr_value.proto";
 
 // Used to specify and override the default API & behavior in the
diff --git a/tensorflow/core/framework/attr_value.proto b/tensorflow/core/framework/attr_value.proto
index 62f0a9050f..054e3ec97c 100644
--- a/tensorflow/core/framework/attr_value.proto
+++ b/tensorflow/core/framework/attr_value.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "AttrValueProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/tensor.proto";
 import "tensorflow/core/framework/tensor_shape.proto";
 import "tensorflow/core/framework/types.proto";
diff --git a/tensorflow/core/framework/cost_graph.proto b/tensorflow/core/framework/cost_graph.proto
index 7885b0171a..19d765cd32 100644
--- a/tensorflow/core/framework/cost_graph.proto
+++ b/tensorflow/core/framework/cost_graph.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "CostGraphProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/tensor_shape.proto";
 import "tensorflow/core/framework/types.proto";
 
diff --git a/tensorflow/core/framework/device_attributes.proto b/tensorflow/core/framework/device_attributes.proto
index 0b3c0d5bdf..44236ca979 100644
--- a/tensorflow/core/framework/device_attributes.proto
+++ b/tensorflow/core/framework/device_attributes.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "DeviceAttributesProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 
 message InterconnectLink {
   int32 device_id = 1;
diff --git a/tensorflow/core/framework/function.proto b/tensorflow/core/framework/function.proto
index 72e3c43831..e69d3938d9 100644
--- a/tensorflow/core/framework/function.proto
+++ b/tensorflow/core/framework/function.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "FunctionProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/attr_value.proto";
 import "tensorflow/core/framework/node_def.proto";
 import "tensorflow/core/framework/op_def.proto";
diff --git a/tensorflow/core/framework/graph.proto b/tensorflow/core/framework/graph.proto
index 7d6e16d5c1..76d358971d 100644
--- a/tensorflow/core/framework/graph.proto
+++ b/tensorflow/core/framework/graph.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "GraphProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/node_def.proto";
 import "tensorflow/core/framework/function.proto";
 import "tensorflow/core/framework/versions.proto";
diff --git a/tensorflow/core/framework/graph_transfer_info.proto b/tensorflow/core/framework/graph_transfer_info.proto
index 41dd54d78c..232297d460 100644
--- a/tensorflow/core/framework/graph_transfer_info.proto
+++ b/tensorflow/core/framework/graph_transfer_info.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "GraphTransferInfoProto";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/types.proto";
 
 message GraphTransferNodeInput {
diff --git a/tensorflow/core/framework/iterator.proto b/tensorflow/core/framework/iterator.proto
index 7e5f5ea2e0..f015342e13 100644
--- a/tensorflow/core/framework/iterator.proto
+++ b/tensorflow/core/framework/iterator.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "IteratorProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.util";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 
 // Protocol buffer representing the metadata for an iterator's state stored
 // as a Variant tensor.
diff --git a/tensorflow/core/framework/kernel_def.proto b/tensorflow/core/framework/kernel_def.proto
index 65e9ef04a0..a17b9c8492 100644
--- a/tensorflow/core/framework/kernel_def.proto
+++ b/tensorflow/core/framework/kernel_def.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "KernelDefProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/attr_value.proto";
 
 message KernelDef {
diff --git a/tensorflow/core/framework/log_memory.proto b/tensorflow/core/framework/log_memory.proto
index d1e126330d..7f37eadc3b 100644
--- a/tensorflow/core/framework/log_memory.proto
+++ b/tensorflow/core/framework/log_memory.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "LogMemoryProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/tensor_description.proto";
 
 message MemoryLogStep {
diff --git a/tensorflow/core/framework/node_def.proto b/tensorflow/core/framework/node_def.proto
index 8fcee32e29..0a095f903f 100644
--- a/tensorflow/core/framework/node_def.proto
+++ b/tensorflow/core/framework/node_def.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "NodeProto";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/attr_value.proto";
 
 message NodeDef {
diff --git a/tensorflow/core/framework/op_def.proto b/tensorflow/core/framework/op_def.proto
index ca0e5e7133..aea2d2bb09 100644
--- a/tensorflow/core/framework/op_def.proto
+++ b/tensorflow/core/framework/op_def.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "OpDefProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/attr_value.proto";
 import "tensorflow/core/framework/types.proto";
 
diff --git a/tensorflow/core/framework/reader_base.proto b/tensorflow/core/framework/reader_base.proto
index 1b8b965ee1..9e187cfa79 100644
--- a/tensorflow/core/framework/reader_base.proto
+++ b/tensorflow/core/framework/reader_base.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "ReaderBaseProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 
 // For serializing and restoring the state of ReaderBase, see
 // reader_base.h for details.
diff --git a/tensorflow/core/framework/remote_fused_graph_execute_info.proto b/tensorflow/core/framework/remote_fused_graph_execute_info.proto
index 946da40d0e..10072724d2 100644
--- a/tensorflow/core/framework/remote_fused_graph_execute_info.proto
+++ b/tensorflow/core/framework/remote_fused_graph_execute_info.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "RemoteFusedGraphExecuteInfoProto";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/graph.proto";
 import "tensorflow/core/framework/tensor_shape.proto";
 import "tensorflow/core/framework/types.proto";
diff --git a/tensorflow/core/framework/resource_handle.proto b/tensorflow/core/framework/resource_handle.proto
index b1921337f5..a54d3d906c 100644
--- a/tensorflow/core/framework/resource_handle.proto
+++ b/tensorflow/core/framework/resource_handle.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "ResourceHandle";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 
 // Protocol buffer representing a handle to a tensorflow resource. Handles are
 // not valid across executions, but can be serialized back and forth from within
diff --git a/tensorflow/core/framework/step_stats.proto b/tensorflow/core/framework/step_stats.proto
index 65c8089d51..d98999cb54 100644
--- a/tensorflow/core/framework/step_stats.proto
+++ b/tensorflow/core/framework/step_stats.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "StepStatsProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/allocation_description.proto";
 import "tensorflow/core/framework/tensor_description.proto";
 
diff --git a/tensorflow/core/framework/summary.proto b/tensorflow/core/framework/summary.proto
index 55879f8783..532e4fcd87 100644
--- a/tensorflow/core/framework/summary.proto
+++ b/tensorflow/core/framework/summary.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "SummaryProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/tensor.proto";
 
 // Metadata associated with a series of Summary data
diff --git a/tensorflow/core/framework/tensor.proto b/tensorflow/core/framework/tensor.proto
index abbf16e810..55921af1d0 100644
--- a/tensorflow/core/framework/tensor.proto
+++ b/tensorflow/core/framework/tensor.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "TensorProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/resource_handle.proto";
 import "tensorflow/core/framework/tensor_shape.proto";
 import "tensorflow/core/framework/types.proto";
diff --git a/tensorflow/core/framework/tensor_description.proto b/tensorflow/core/framework/tensor_description.proto
index 6ac3c1b881..4c23c7e620 100644
--- a/tensorflow/core/framework/tensor_description.proto
+++ b/tensorflow/core/framework/tensor_description.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "TensorDescriptionProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 import "tensorflow/core/framework/types.proto";
 import "tensorflow/core/framework/tensor_shape.proto";
 import "tensorflow/core/framework/allocation_description.proto";
diff --git a/tensorflow/core/framework/tensor_shape.proto b/tensorflow/core/framework/tensor_shape.proto
index 1ec3c5323c..286156a012 100644
--- a/tensorflow/core/framework/tensor_shape.proto
+++ b/tensorflow/core/framework/tensor_shape.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "TensorShapeProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 
 package tensorflow;
 
diff --git a/tensorflow/core/framework/tensor_slice.proto b/tensorflow/core/framework/tensor_slice.proto
index 24b01661dc..a5c366ed60 100644
--- a/tensorflow/core/framework/tensor_slice.proto
+++ b/tensorflow/core/framework/tensor_slice.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "TensorSliceProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 
 package tensorflow;
 
diff --git a/tensorflow/core/framework/types.proto b/tensorflow/core/framework/types.proto
index e003fd0010..03835d1b92 100644
--- a/tensorflow/core/framework/types.proto
+++ b/tensorflow/core/framework/types.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "TypesProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 
 // LINT.IfChange
 enum DataType {
diff --git a/tensorflow/core/framework/variable.proto b/tensorflow/core/framework/variable.proto
index e0df01cc9b..93ae423bab 100644
--- a/tensorflow/core/framework/variable.proto
+++ b/tensorflow/core/framework/variable.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "VariableProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 
 // Protocol buffer representing a Variable.
 message VariableDef {
diff --git a/tensorflow/core/framework/versions.proto b/tensorflow/core/framework/versions.proto
index 7d5e58ae7d..dd2ec55238 100644
--- a/tensorflow/core/framework/versions.proto
+++ b/tensorflow/core/framework/versions.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "VersionsProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framework";
 
 // Version information for a piece of serialized data
 //
diff --git a/tensorflow/core/lib/core/error_codes.proto b/tensorflow/core/lib/core/error_codes.proto
index b82d389146..5ced65a973 100644
--- a/tensorflow/core/lib/core/error_codes.proto
+++ b/tensorflow/core/lib/core/error_codes.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "ErrorCodesProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/lib/core";
 
 // The canonical error codes for TensorFlow APIs.
 //
diff --git a/tensorflow/core/protobuf/cluster.proto b/tensorflow/core/protobuf/cluster.proto
index 33c87eefe0..c696d345e0 100644
--- a/tensorflow/core/protobuf/cluster.proto
+++ b/tensorflow/core/protobuf/cluster.proto
@@ -20,6 +20,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "ClusterProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.distruntime";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 
 // This file contains protos to be used when defining a TensorFlow
 // cluster.
diff --git a/tensorflow/core/protobuf/config.proto b/tensorflow/core/protobuf/config.proto
index c1a0075b64..078e76e7dc 100644
--- a/tensorflow/core/protobuf/config.proto
+++ b/tensorflow/core/protobuf/config.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "ConfigProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 import "tensorflow/core/framework/cost_graph.proto";
 import "tensorflow/core/framework/graph.proto";
 import "tensorflow/core/framework/step_stats.proto";
diff --git a/tensorflow/core/protobuf/control_flow.proto b/tensorflow/core/protobuf/control_flow.proto
index 3c05b4f0e2..5f44878c44 100644
--- a/tensorflow/core/protobuf/control_flow.proto
+++ b/tensorflow/core/protobuf/control_flow.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "ControlFlowProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 
 // Control flow context related protocol buffers.
 
diff --git a/tensorflow/core/protobuf/critical_section.proto b/tensorflow/core/protobuf/critical_section.proto
index 0b3f531e6d..7954e7ba87 100644
--- a/tensorflow/core/protobuf/critical_section.proto
+++ b/tensorflow/core/protobuf/critical_section.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "CriticalSectionProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 
 // Protocol buffer representing a CriticalSection.
 message CriticalSectionDef {
diff --git a/tensorflow/core/protobuf/debug.proto b/tensorflow/core/protobuf/debug.proto
index 56983f3b7d..499900f965 100644
--- a/tensorflow/core/protobuf/debug.proto
+++ b/tensorflow/core/protobuf/debug.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "DebugProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 
 // EXPERIMENTAL. Option for watching a node.
 message DebugTensorWatch {
diff --git a/tensorflow/core/protobuf/device_properties.proto b/tensorflow/core/protobuf/device_properties.proto
index 3bd3015900..11e1258e75 100644
--- a/tensorflow/core/protobuf/device_properties.proto
+++ b/tensorflow/core/protobuf/device_properties.proto
@@ -18,6 +18,7 @@ syntax = "proto3";
 package tensorflow;
 option cc_enable_arenas = true;
 option java_outer_classname = "DevicePropertiesProtos";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 
 message DeviceProperties {
   // Device type (CPU, GPU, ...)
diff --git a/tensorflow/core/protobuf/master.proto b/tensorflow/core/protobuf/master.proto
index 96c91536f7..03022875e6 100644
--- a/tensorflow/core/protobuf/master.proto
+++ b/tensorflow/core/protobuf/master.proto
@@ -20,7 +20,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "DistributedRuntimeProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.distruntime";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 import "tensorflow/core/framework/device_attributes.proto";
 import "tensorflow/core/framework/graph.proto";
 import "tensorflow/core/framework/tensor.proto";
diff --git a/tensorflow/core/protobuf/master_service.proto b/tensorflow/core/protobuf/master_service.proto
index 1170611f37..ce0e4f6435 100644
--- a/tensorflow/core/protobuf/master_service.proto
+++ b/tensorflow/core/protobuf/master_service.proto
@@ -19,7 +19,7 @@ package tensorflow.grpc;
 option java_outer_classname = "MasterServiceProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.distruntime";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 import "tensorflow/core/protobuf/master.proto";
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/tensorflow/core/protobuf/meta_graph.proto b/tensorflow/core/protobuf/meta_graph.proto
index fd86c0da12..75a2a88ed7 100644
--- a/tensorflow/core/protobuf/meta_graph.proto
+++ b/tensorflow/core/protobuf/meta_graph.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "MetaGraphProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 import "google/protobuf/any.proto";
 
 import "tensorflow/core/framework/graph.proto";
diff --git a/tensorflow/core/protobuf/named_tensor.proto b/tensorflow/core/protobuf/named_tensor.proto
index dd4976e354..6e2f7feee2 100644
--- a/tensorflow/core/protobuf/named_tensor.proto
+++ b/tensorflow/core/protobuf/named_tensor.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "NamedTensorProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 import "tensorflow/core/framework/tensor.proto";
 
 // A pair of tensor name and tensor values.
diff --git a/tensorflow/core/protobuf/queue_runner.proto b/tensorflow/core/protobuf/queue_runner.proto
index 05a48d0acf..f4df649f7d 100644
--- a/tensorflow/core/protobuf/queue_runner.proto
+++ b/tensorflow/core/protobuf/queue_runner.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "QueueRunnerProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 import "tensorflow/core/lib/core/error_codes.proto";
 
 // Protocol buffer representing a QueueRunner.
diff --git a/tensorflow/core/protobuf/rewriter_config.proto b/tensorflow/core/protobuf/rewriter_config.proto
index 029b27cd04..a15ccdfd87 100644
--- a/tensorflow/core/protobuf/rewriter_config.proto
+++ b/tensorflow/core/protobuf/rewriter_config.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "RewriterConfigProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 
 import "tensorflow/core/framework/attr_value.proto";
 
diff --git a/tensorflow/core/protobuf/saved_model.proto b/tensorflow/core/protobuf/saved_model.proto
index c2595ddf88..03789d3df7 100644
--- a/tensorflow/core/protobuf/saved_model.proto
+++ b/tensorflow/core/protobuf/saved_model.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "SavedModelProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 import "tensorflow/core/protobuf/meta_graph.proto";
 
 // SavedModel is the high level serialization format for TensorFlow Models.
diff --git a/tensorflow/core/protobuf/saver.proto b/tensorflow/core/protobuf/saver.proto
index a757d3f756..4245386145 100644
--- a/tensorflow/core/protobuf/saver.proto
+++ b/tensorflow/core/protobuf/saver.proto
@@ -5,6 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "SaverProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.util";
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 
 // Protocol buffer representing the configuration of a Saver.
 message SaverDef {
diff --git a/tensorflow/core/protobuf/tensor_bundle.proto b/tensorflow/core/protobuf/tensor_bundle.proto
index 80e87f14f9..681c01bbbd 100644
--- a/tensorflow/core/protobuf/tensor_bundle.proto
+++ b/tensorflow/core/protobuf/tensor_bundle.proto
@@ -5,7 +5,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "TensorBundleProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.util";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 import "tensorflow/core/framework/tensor_shape.proto";
 import "tensorflow/core/framework/tensor_slice.proto";
 import "tensorflow/core/framework/types.proto";
diff --git a/tensorflow/core/protobuf/tensorflow_server.proto b/tensorflow/core/protobuf/tensorflow_server.proto
index 6199e707e5..be25804a1b 100644
--- a/tensorflow/core/protobuf/tensorflow_server.proto
+++ b/tensorflow/core/protobuf/tensorflow_server.proto
@@ -23,7 +23,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "ServerProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.distruntime";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 // Defines the configuration of a single TensorFlow server.
 message ServerDef {
   // The cluster of which this server is a member.
diff --git a/tensorflow/core/protobuf/worker.proto b/tensorflow/core/protobuf/worker.proto
index 602f6a1ef1..d714d85ce6 100644
--- a/tensorflow/core/protobuf/worker.proto
+++ b/tensorflow/core/protobuf/worker.proto
@@ -20,7 +20,7 @@ option cc_enable_arenas = true;
 option java_outer_classname = "WorkerProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.distruntime";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 import "google/protobuf/any.proto";
 import "tensorflow/core/framework/cost_graph.proto";
 import "tensorflow/core/framework/step_stats.proto";
diff --git a/tensorflow/core/protobuf/worker_service.proto b/tensorflow/core/protobuf/worker_service.proto
index 01c76c01a9..025fa7ca59 100644
--- a/tensorflow/core/protobuf/worker_service.proto
+++ b/tensorflow/core/protobuf/worker_service.proto
@@ -19,7 +19,7 @@ package tensorflow.grpc;
 option java_outer_classname = "WorkerServiceProtos";
 option java_multiple_files = true;
 option java_package = "org.tensorflow.distruntime";
-
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
 import "tensorflow/core/protobuf/worker.proto";
 
 ////////////////////////////////////////////////////////////////////////////////
-- 
GitLab


From df5ae5ac2a58131737a11e417ac34a663efb3574 Mon Sep 17 00:00:00 2001
From: Sunitha Kambhampati <skambha@us.ibm.com>
Date: Wed, 2 May 2018 17:52:38 -0700
Subject: [PATCH 188/395] Add some todo's

---
 tensorflow/contrib/tensorboard/db/summary_db_writer.cc      | 1 +
 tensorflow/contrib/tensorboard/db/summary_db_writer_test.cc | 1 +
 2 files changed, 2 insertions(+)

diff --git a/tensorflow/contrib/tensorboard/db/summary_db_writer.cc b/tensorflow/contrib/tensorboard/db/summary_db_writer.cc
index 046a2d3884..630c0607ae 100644
--- a/tensorflow/contrib/tensorboard/db/summary_db_writer.cc
+++ b/tensorflow/contrib/tensorboard/db/summary_db_writer.cc
@@ -1183,6 +1183,7 @@ class SummaryDbWriter : public SummaryWriterInterface {
     Tensor t{DT_DOUBLE, {k, 3}};
     auto data = t.flat<double>();
     for (int i = 0, j = 0; i < k; ++i) {
+      // TODO(nickfelt): reconcile with TensorBoard's data_compat.py
       // From summary.proto
       // Parallel arrays encoding the bucket boundaries and the bucket values.
       // bucket(i) is the count for the bucket i.  The range for
diff --git a/tensorflow/contrib/tensorboard/db/summary_db_writer_test.cc b/tensorflow/contrib/tensorboard/db/summary_db_writer_test.cc
index cb51325d15..2044692b6e 100644
--- a/tensorflow/contrib/tensorboard/db/summary_db_writer_test.cc
+++ b/tensorflow/contrib/tensorboard/db/summary_db_writer_test.cc
@@ -131,6 +131,7 @@ TEST_F(SummaryDbWriterTest, WriteHistogram_VerifyTensorValues) {
   writer_->Unref();
   writer_ = nullptr;
 
+  // TODO(nickfelt): implement QueryTensor() to encapsulate this
   // Verify the data
   string result = QueryString("SELECT data FROM Tensors");
   const double* val = reinterpret_cast<const double*>(result.data());
-- 
GitLab


From 8f0a90b711480c12716d1a3b1094cc8b34939f2d Mon Sep 17 00:00:00 2001
From: RJ Ryan <rjryan@google.com>
Date: Wed, 2 May 2018 17:57:27 -0700
Subject: [PATCH 189/395] Add complex128 support to FFT, FFT2D, FFT3D, IFFT,
 IFFT2D, and IFFT3D.

NumPy automatically upcasts to complex128 when computing FFTs, leading to issues like:
#10749

This change allows users to choose between 32-bit and 64-bit precision FFTs on CPU and GPU.

PiperOrigin-RevId: 195183206
---
 tensorflow/compiler/tf2xla/kernels/fft_ops.cc |  17 +-
 tensorflow/core/kernels/fft_ops.cc            |  78 +++++++---
 tensorflow/core/ops/spectral_ops.cc           |  30 ++--
 .../python/kernel_tests/fft_ops_test.py       | 145 +++++++++++-------
 .../linalg/linear_operator_circulant_test.py  |   6 +-
 tensorflow/python/ops/spectral_grad.py        |  30 ++--
 6 files changed, 196 insertions(+), 110 deletions(-)

diff --git a/tensorflow/compiler/tf2xla/kernels/fft_ops.cc b/tensorflow/compiler/tf2xla/kernels/fft_ops.cc
index fcb927dab0..933924cad1 100644
--- a/tensorflow/compiler/tf2xla/kernels/fft_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/fft_ops.cc
@@ -81,9 +81,11 @@ class FFTOp : public GenericFftOp {
   explicit FFTOp(OpKernelConstruction* ctx)
       : GenericFftOp(ctx, /*fft_type=*/FftType::FFT, /*fft_rank=*/FFTRank) {}
 };
-REGISTER_XLA_OP(Name("FFT"), FFTOp<1>);
-REGISTER_XLA_OP(Name("FFT2D"), FFTOp<2>);
-REGISTER_XLA_OP(Name("FFT3D"), FFTOp<3>);
+REGISTER_XLA_OP(Name("FFT").TypeConstraint("Tcomplex", DT_COMPLEX64), FFTOp<1>);
+REGISTER_XLA_OP(Name("FFT2D").TypeConstraint("Tcomplex", DT_COMPLEX64),
+                FFTOp<2>);
+REGISTER_XLA_OP(Name("FFT3D").TypeConstraint("Tcomplex", DT_COMPLEX64),
+                FFTOp<3>);
 
 template <int FFTRank>
 class IFFTOp : public GenericFftOp {
@@ -91,9 +93,12 @@ class IFFTOp : public GenericFftOp {
   explicit IFFTOp(OpKernelConstruction* ctx)
       : GenericFftOp(ctx, /*fft_type=*/FftType::IFFT, /*fft_rank=*/FFTRank) {}
 };
-REGISTER_XLA_OP(Name("IFFT"), IFFTOp<1>);
-REGISTER_XLA_OP(Name("IFFT2D"), IFFTOp<2>);
-REGISTER_XLA_OP(Name("IFFT3D"), IFFTOp<3>);
+REGISTER_XLA_OP(Name("IFFT").TypeConstraint("Tcomplex", DT_COMPLEX64),
+                IFFTOp<1>);
+REGISTER_XLA_OP(Name("IFFT2D").TypeConstraint("Tcomplex", DT_COMPLEX64),
+                IFFTOp<2>);
+REGISTER_XLA_OP(Name("IFFT3D").TypeConstraint("Tcomplex", DT_COMPLEX64),
+                IFFTOp<3>);
 
 template <int FFTRank>
 class RFFTOp : public GenericFftOp {
diff --git a/tensorflow/core/kernels/fft_ops.cc b/tensorflow/core/kernels/fft_ops.cc
index 661bf5fc5f..d7105a71bb 100644
--- a/tensorflow/core/kernels/fft_ops.cc
+++ b/tensorflow/core/kernels/fft_ops.cc
@@ -129,13 +129,23 @@ class FFTCPU : public FFTBase {
     auto device = ctx->eigen_device<CPUDevice>();
 
     if (!IsReal()) {
-      auto input = Tensor(in).flat_inner_dims<complex64, FFTRank + 1>();
-      // Compute the FFT using eigen.
-      auto output = out->flat_inner_dims<complex64, FFTRank + 1>();
+      // Compute the FFT using Eigen.
       constexpr auto direction =
           Forward ? Eigen::FFT_FORWARD : Eigen::FFT_REVERSE;
-      output.device(device) =
-          input.template fft<Eigen::BothParts, direction>(axes);
+      if (in.dtype() == DT_COMPLEX64) {
+        DCHECK_EQ(out->dtype(), DT_COMPLEX64);
+        auto input = Tensor(in).flat_inner_dims<complex64, FFTRank + 1>();
+        auto output = out->flat_inner_dims<complex64, FFTRank + 1>();
+        output.device(device) =
+            input.template fft<Eigen::BothParts, direction>(axes);
+      } else {
+        DCHECK_EQ(DT_COMPLEX128, in.dtype());
+        DCHECK_EQ(DT_COMPLEX128, out->dtype());
+        auto input = Tensor(in).flat_inner_dims<complex128, FFTRank + 1>();
+        auto output = out->flat_inner_dims<complex128, FFTRank + 1>();
+        output.device(device) =
+            input.template fft<Eigen::BothParts, direction>(axes);
+      }
     } else {
       if (IsForward()) {
         auto input = Tensor(in).flat_inner_dims<float, FFTRank + 1>();
@@ -392,10 +402,16 @@ class FFTGPUBase : public FFTBase {
     }
 
     constexpr bool kInPlaceFft = false;
+    const bool is_complex128 = in.dtype() == DT_COMPLEX128;
+    // complex128 real FFT is not supported yet.
+    DCHECK(!IsReal() || !is_complex128);
+
     const auto kFftType =
         IsReal() ? (IsForward() ? se::fft::Type::kR2C : se::fft::Type::kC2R)
-                 : (IsForward() ? se::fft::Type::kC2CForward
-                                : se::fft::Type::kC2CInverse);
+                 : (IsForward() ? (is_complex128 ? se::fft::Type::kZ2ZForward
+                                                 : se::fft::Type::kC2CForward)
+                                : (is_complex128 ? se::fft::Type::kZ2ZInverse
+                                                 : se::fft::Type::kC2CInverse));
 
     CufftScratchAllocator scratch_allocator(CufftScratchSize, ctx);
     auto plan =
@@ -428,20 +444,42 @@ class FFTGPUBase : public FFTBase {
                              input_shape.DebugString()));
       }
     } else {
-      auto src = AsDeviceMemory<complex64>(in.flat<complex64>().data());
-      auto dst = AsDeviceMemory<complex64>(out->flat<complex64>().data());
-      OP_REQUIRES(
-          ctx, stream->ThenFft(plan.get(), src, &dst).ok(),
-          errors::Internal("fft failed : type=", static_cast<int>(kFftType),
-                           " in.shape=", input_shape.DebugString()));
-      if (!IsForward()) {
-        auto alpha = complex64(1.f / output_distance);
+      if (!is_complex128) {
+        DCHECK_EQ(in.dtype(), DT_COMPLEX64);
+        DCHECK_EQ(out->dtype(), DT_COMPLEX64);
+        auto src = AsDeviceMemory<complex64>(in.flat<complex64>().data());
+        auto dst = AsDeviceMemory<complex64>(out->flat<complex64>().data());
         OP_REQUIRES(
-            ctx,
-            stream->ThenBlasScal(output_shape.num_elements(), alpha, &dst, 1)
-                .ok(),
-            errors::Internal("BlasScal failed : in.shape=",
-                             input_shape.DebugString()));
+            ctx, stream->ThenFft(plan.get(), src, &dst).ok(),
+            errors::Internal("fft failed : type=", static_cast<int>(kFftType),
+                             " in.shape=", input_shape.DebugString()));
+        if (!IsForward()) {
+          float alpha = 1.f / output_distance;
+          OP_REQUIRES(
+              ctx,
+              stream->ThenBlasScal(output_shape.num_elements(), alpha, &dst, 1)
+                  .ok(),
+              errors::Internal("BlasScal failed : in.shape=",
+                               input_shape.DebugString()));
+        }
+      } else {
+        DCHECK_EQ(in.dtype(), DT_COMPLEX128);
+        DCHECK_EQ(out->dtype(), DT_COMPLEX128);
+        auto src = AsDeviceMemory<complex128>(in.flat<complex128>().data());
+        auto dst = AsDeviceMemory<complex128>(out->flat<complex128>().data());
+        OP_REQUIRES(
+            ctx, stream->ThenFft(plan.get(), src, &dst).ok(),
+            errors::Internal("fft failed : type=", static_cast<int>(kFftType),
+                             " in.shape=", input_shape.DebugString()));
+        if (!IsForward()) {
+          double alpha = 1.0 / output_distance;
+          OP_REQUIRES(
+              ctx,
+              stream->ThenBlasScal(output_shape.num_elements(), alpha, &dst, 1)
+                  .ok(),
+              errors::Internal("BlasScal failed : in.shape=",
+                               input_shape.DebugString()));
+        }
       }
     }
   }
diff --git a/tensorflow/core/ops/spectral_ops.cc b/tensorflow/core/ops/spectral_ops.cc
index 2790aee37e..b1ae7040f0 100644
--- a/tensorflow/core/ops/spectral_ops.cc
+++ b/tensorflow/core/ops/spectral_ops.cc
@@ -25,43 +25,49 @@ using shape_inference::InferenceContext;
 using shape_inference::ShapeHandle;
 
 REGISTER_OP("FFT")
-    .Input("input: complex64")
-    .Output("output: complex64")
+    .Input("input: Tcomplex")
+    .Output("output: Tcomplex")
+    .Attr("Tcomplex: {complex64, complex128} = DT_COMPLEX64")
     .SetShapeFn([](InferenceContext* c) {
       return shape_inference::UnchangedShapeWithRankAtLeast(c, 1);
     });
 
 REGISTER_OP("IFFT")
-    .Input("input: complex64")
-    .Output("output: complex64")
+    .Input("input: Tcomplex")
+    .Output("output: Tcomplex")
+    .Attr("Tcomplex: {complex64, complex128} = DT_COMPLEX64")
     .SetShapeFn([](InferenceContext* c) {
       return shape_inference::UnchangedShapeWithRankAtLeast(c, 1);
     });
 
 REGISTER_OP("FFT2D")
-    .Input("input: complex64")
-    .Output("output: complex64")
+    .Input("input: Tcomplex")
+    .Output("output: Tcomplex")
+    .Attr("Tcomplex: {complex64, complex128} = DT_COMPLEX64")
     .SetShapeFn([](InferenceContext* c) {
       return shape_inference::UnchangedShapeWithRankAtLeast(c, 2);
     });
 
 REGISTER_OP("IFFT2D")
-    .Input("input: complex64")
-    .Output("output: complex64")
+    .Input("input: Tcomplex")
+    .Output("output: Tcomplex")
+    .Attr("Tcomplex: {complex64, complex128} = DT_COMPLEX64")
     .SetShapeFn([](InferenceContext* c) {
       return shape_inference::UnchangedShapeWithRankAtLeast(c, 2);
     });
 
 REGISTER_OP("FFT3D")
-    .Input("input: complex64")
-    .Output("output: complex64")
+    .Input("input: Tcomplex")
+    .Output("output: Tcomplex")
+    .Attr("Tcomplex: {complex64, complex128} = DT_COMPLEX64")
     .SetShapeFn([](InferenceContext* c) {
       return shape_inference::UnchangedShapeWithRankAtLeast(c, 3);
     });
 
 REGISTER_OP("IFFT3D")
-    .Input("input: complex64")
-    .Output("output: complex64")
+    .Input("input: Tcomplex")
+    .Output("output: Tcomplex")
+    .Attr("Tcomplex: {complex64, complex128} = DT_COMPLEX64")
     .SetShapeFn([](InferenceContext* c) {
       return shape_inference::UnchangedShapeWithRankAtLeast(c, 3);
     });
diff --git a/tensorflow/python/kernel_tests/fft_ops_test.py b/tensorflow/python/kernel_tests/fft_ops_test.py
index b9e2aa1f3a..629acedda5 100644
--- a/tensorflow/python/kernel_tests/fft_ops_test.py
+++ b/tensorflow/python/kernel_tests/fft_ops_test.py
@@ -38,11 +38,13 @@ VALID_FFT_RANKS = (1, 2, 3)
 
 class BaseFFTOpsTest(test.TestCase):
 
-  def _compare(self, x, rank, fft_length=None, use_placeholder=False):
-    self._compareForward(x, rank, fft_length, use_placeholder)
-    self._compareBackward(x, rank, fft_length, use_placeholder)
+  def _compare(self, x, rank, fft_length=None, use_placeholder=False,
+               rtol=1e-4, atol=1e-4):
+    self._compareForward(x, rank, fft_length, use_placeholder, rtol, atol)
+    self._compareBackward(x, rank, fft_length, use_placeholder, rtol, atol)
 
-  def _compareForward(self, x, rank, fft_length=None, use_placeholder=False):
+  def _compareForward(self, x, rank, fft_length=None, use_placeholder=False,
+                      rtol=1e-4, atol=1e-4):
     x_np = self._npFFT(x, rank, fft_length)
     if use_placeholder:
       x_ph = array_ops.placeholder(dtype=dtypes.as_dtype(x.dtype))
@@ -50,9 +52,10 @@ class BaseFFTOpsTest(test.TestCase):
     else:
       x_tf = self._tfFFT(x, rank, fft_length)
 
-    self.assertAllClose(x_np, x_tf, rtol=1e-4, atol=1e-4)
+    self.assertAllClose(x_np, x_tf, rtol=rtol, atol=atol)
 
-  def _compareBackward(self, x, rank, fft_length=None, use_placeholder=False):
+  def _compareBackward(self, x, rank, fft_length=None, use_placeholder=False,
+                       rtol=1e-4, atol=1e-4):
     x_np = self._npIFFT(x, rank, fft_length)
     if use_placeholder:
       x_ph = array_ops.placeholder(dtype=dtypes.as_dtype(x.dtype))
@@ -60,7 +63,7 @@ class BaseFFTOpsTest(test.TestCase):
     else:
       x_tf = self._tfIFFT(x, rank, fft_length)
 
-    self.assertAllClose(x_np, x_tf, rtol=1e-4, atol=1e-4)
+    self.assertAllClose(x_np, x_tf, rtol=rtol, atol=atol)
 
   def _checkMemoryFail(self, x, rank):
     config = config_pb2.ConfigProto()
@@ -68,7 +71,8 @@ class BaseFFTOpsTest(test.TestCase):
     with self.test_session(config=config, force_gpu=True):
       self._tfFFT(x, rank, fft_length=None)
 
-  def _checkGradComplex(self, func, x, y, result_is_complex=True):
+  def _checkGradComplex(self, func, x, y, result_is_complex=True,
+                        rtol=1e-2, atol=1e-2):
     with self.test_session(use_gpu=True):
       inx = ops.convert_to_tensor(x)
       iny = ops.convert_to_tensor(y)
@@ -85,10 +89,10 @@ class BaseFFTOpsTest(test.TestCase):
            x_init_value=[x, y],
            delta=1e-2)
 
-    self.assertAllClose(x_jacob_t, x_jacob_n, rtol=1e-2, atol=1e-2)
-    self.assertAllClose(y_jacob_t, y_jacob_n, rtol=1e-2, atol=1e-2)
+    self.assertAllClose(x_jacob_t, x_jacob_n, rtol=rtol, atol=atol)
+    self.assertAllClose(y_jacob_t, y_jacob_n, rtol=rtol, atol=atol)
 
-  def _checkGradReal(self, func, x):
+  def _checkGradReal(self, func, x, rtol=1e-2, atol=1e-2):
     with self.test_session(use_gpu=True):
       inx = ops.convert_to_tensor(x)
       # func is a forward RFFT function (batched or unbatched).
@@ -98,7 +102,7 @@ class BaseFFTOpsTest(test.TestCase):
       x_jacob_t, x_jacob_n = test.compute_gradient(
           inx, list(x.shape), loss, [1], x_init_value=x, delta=1e-2)
 
-    self.assertAllClose(x_jacob_t, x_jacob_n, rtol=1e-2, atol=1e-2)
+    self.assertAllClose(x_jacob_t, x_jacob_n, rtol=rtol, atol=atol)
 
 
 class FFTOpsTest(BaseFFTOpsTest):
@@ -155,27 +159,30 @@ class FFTOpsTest(BaseFFTOpsTest):
 
   def testEmpty(self):
     with spectral_ops_test_util.fft_kernel_label_map():
-      for rank in VALID_FFT_RANKS:
-        for dims in xrange(rank, rank + 3):
-          x = np.zeros((0,) * dims).astype(np.complex64)
-          self.assertEqual(x.shape, self._tfFFT(x, rank).shape)
-          self.assertEqual(x.shape, self._tfIFFT(x, rank).shape)
+      for np_type in (np.complex64, np.complex128):
+        for rank in VALID_FFT_RANKS:
+          for dims in xrange(rank, rank + 3):
+            x = np.zeros((0,) * dims).astype(np_type)
+            self.assertEqual(x.shape, self._tfFFT(x, rank).shape)
+            self.assertEqual(x.shape, self._tfIFFT(x, rank).shape)
 
   def testBasic(self):
     with spectral_ops_test_util.fft_kernel_label_map():
-      for rank in VALID_FFT_RANKS:
-        for dims in xrange(rank, rank + 3):
-          self._compare(
-              np.mod(np.arange(np.power(4, dims)), 10).reshape(
-                  (4,) * dims).astype(np.complex64), rank)
+      for np_type, tol in ((np.complex64, 1e-4), (np.complex128, 1e-8)):
+        for rank in VALID_FFT_RANKS:
+          for dims in xrange(rank, rank + 3):
+            self._compare(
+                np.mod(np.arange(np.power(4, dims)), 10).reshape(
+                    (4,) * dims).astype(np_type), rank, rtol=tol, atol=tol)
 
   def testLargeBatch(self):
     if test.is_gpu_available(cuda_only=True):
       rank = 1
       for dims in xrange(rank, rank + 3):
-        self._compare(
-            np.mod(np.arange(np.power(128, dims)), 10).reshape(
-                (128,) * dims).astype(np.complex64), rank)
+        for np_type, tol in ((np.complex64, 1e-4), (np.complex128, 1e-5)):
+          self._compare(
+              np.mod(np.arange(np.power(128, dims)), 10).reshape(
+                  (128,) * dims).astype(np_type), rank, rtol=tol, atol=tol)
 
   # TODO(yangzihao): Disable before we can figure out a way to
   # properly test memory fail for large batch fft.
@@ -189,27 +196,49 @@ class FFTOpsTest(BaseFFTOpsTest):
 
   def testBasicPlaceholder(self):
     with spectral_ops_test_util.fft_kernel_label_map():
-      for rank in VALID_FFT_RANKS:
-        for dims in xrange(rank, rank + 3):
-          self._compare(
-              np.mod(np.arange(np.power(4, dims)), 10).reshape(
-                  (4,) * dims).astype(np.complex64),
-              rank,
-              use_placeholder=True)
+      for np_type, tol in ((np.complex64, 1e-4), (np.complex128, 1e-8)):
+        for rank in VALID_FFT_RANKS:
+          for dims in xrange(rank, rank + 3):
+            self._compare(
+                np.mod(np.arange(np.power(4, dims)), 10).reshape(
+                    (4,) * dims).astype(np_type),
+                rank, use_placeholder=True, rtol=tol, atol=tol)
 
   def testRandom(self):
     with spectral_ops_test_util.fft_kernel_label_map():
-      np.random.seed(12345)
+      for np_type, tol in ((np.complex64, 1e-4), (np.complex128, 5e-6)):
+        def gen(shape):
+          n = np.prod(shape)
+          re = np.random.uniform(size=n)
+          im = np.random.uniform(size=n)
+          return (re + im * 1j).reshape(shape)
 
-      def gen(shape):
-        n = np.prod(shape)
-        re = np.random.uniform(size=n)
-        im = np.random.uniform(size=n)
-        return (re + im * 1j).reshape(shape)
+        for rank in VALID_FFT_RANKS:
+          for dims in xrange(rank, rank + 3):
+            self._compare(gen((4,) * dims).astype(np_type), rank,
+                          rtol=tol, atol=tol)
 
-      for rank in VALID_FFT_RANKS:
-        for dims in xrange(rank, rank + 3):
-          self._compare(gen((4,) * dims), rank)
+  def testRandom1D(self):
+    with spectral_ops_test_util.fft_kernel_label_map():
+      for np_type in (np.complex64, np.complex128):
+        has_gpu = test.is_gpu_available(cuda_only=True)
+        tol = {(np.complex64, True): 1e-4,
+               (np.complex64, False): 1e-2,
+               (np.complex128, True): 1e-4,
+               (np.complex128, False): 1e-2}[(np_type, has_gpu)]
+        def gen(shape):
+          n = np.prod(shape)
+          re = np.random.uniform(size=n)
+          im = np.random.uniform(size=n)
+          return (re + im * 1j).reshape(shape)
+
+        # Check a variety of power-of-2 FFT sizes.
+        for dim in (128, 256, 512, 1024):
+          self._compare(gen((dim,)).astype(np_type), 1, rtol=tol, atol=tol)
+
+        # Check a variety of non-power-of-2 FFT sizes.
+        for dim in (127, 255, 511, 1023):
+          self._compare(gen((dim,)).astype(np_type), 1, rtol=tol, atol=tol)
 
   def testError(self):
     for rank in VALID_FFT_RANKS:
@@ -224,22 +253,27 @@ class FFTOpsTest(BaseFFTOpsTest):
 
   def testGrad_Simple(self):
     with spectral_ops_test_util.fft_kernel_label_map():
-      for rank in VALID_FFT_RANKS:
-        for dims in xrange(rank, rank + 2):
-          re = np.ones(shape=(4,) * dims, dtype=np.float32) / 10.0
-          im = np.zeros(shape=(4,) * dims, dtype=np.float32)
-          self._checkGradComplex(self._tfFFTForRank(rank), re, im)
-          self._checkGradComplex(self._tfIFFTForRank(rank), re, im)
+      for np_type, tol in ((np.float32, 1e-4), (np.float64, 1e-10)):
+        for rank in VALID_FFT_RANKS:
+          for dims in xrange(rank, rank + 2):
+            re = np.ones(shape=(4,) * dims, dtype=np_type) / 10.0
+            im = np.zeros(shape=(4,) * dims, dtype=np_type)
+            self._checkGradComplex(self._tfFFTForRank(rank), re, im,
+                                   rtol=tol, atol=tol)
+            self._checkGradComplex(self._tfIFFTForRank(rank), re, im,
+                                   rtol=tol, atol=tol)
 
   def testGrad_Random(self):
     with spectral_ops_test_util.fft_kernel_label_map():
-      np.random.seed(54321)
-      for rank in VALID_FFT_RANKS:
-        for dims in xrange(rank, rank + 2):
-          re = np.random.rand(*((3,) * dims)).astype(np.float32) * 2 - 1
-          im = np.random.rand(*((3,) * dims)).astype(np.float32) * 2 - 1
-          self._checkGradComplex(self._tfFFTForRank(rank), re, im)
-          self._checkGradComplex(self._tfIFFTForRank(rank), re, im)
+      for np_type, tol in ((np.float32, 1e-2), (np.float64, 1e-10)):
+        for rank in VALID_FFT_RANKS:
+          for dims in xrange(rank, rank + 2):
+            re = np.random.rand(*((3,) * dims)).astype(np_type) * 2 - 1
+            im = np.random.rand(*((3,) * dims)).astype(np_type) * 2 - 1
+            self._checkGradComplex(self._tfFFTForRank(rank), re, im,
+                                   rtol=tol, atol=tol)
+            self._checkGradComplex(self._tfIFFTForRank(rank), re, im,
+                                   rtol=tol, atol=tol)
 
 
 class RFFTOpsTest(BaseFFTOpsTest):
@@ -395,8 +429,6 @@ class RFFTOpsTest(BaseFFTOpsTest):
 
   def testRandom(self):
     with spectral_ops_test_util.fft_kernel_label_map():
-      np.random.seed(12345)
-
       def gen_real(shape):
         n = np.prod(shape)
         re = np.random.uniform(size=n)
@@ -491,7 +523,6 @@ class RFFTOpsTest(BaseFFTOpsTest):
 
   def testGrad_Random(self):
     with spectral_ops_test_util.fft_kernel_label_map():
-      np.random.seed(54321)
       for rank in VALID_FFT_RANKS:
         # rfft3d/irfft3d do not have gradients yet.
         if rank == 3:
diff --git a/tensorflow/python/kernel_tests/linalg/linear_operator_circulant_test.py b/tensorflow/python/kernel_tests/linalg/linear_operator_circulant_test.py
index e7f2f1c12b..5713d16969 100644
--- a/tensorflow/python/kernel_tests/linalg/linear_operator_circulant_test.py
+++ b/tensorflow/python/kernel_tests/linalg/linear_operator_circulant_test.py
@@ -73,7 +73,7 @@ class LinearOperatorCirculantBaseTest(object):
       x = np.zeros([domain_dimension])
       # x is a basis vector.
       x[m] = 1.0
-      fft_x = math_ops.fft(x)
+      fft_x = math_ops.fft(x.astype(np.complex64))
       h_convolve_x = math_ops.ifft(spectrum * fft_x)
       matrix_rows.append(h_convolve_x)
     matrix = array_ops.stack(matrix_rows, axis=-1)
@@ -91,7 +91,7 @@ class LinearOperatorCirculantTestSelfAdjointOperator(
 
   @property
   def _dtypes_to_test(self):
-    # This operator will always be complex because, although the specturm is
+    # This operator will always be complex because, although the spectrum is
     # real, the matrix will not be real.
     return [dtypes.complex64]
 
@@ -408,7 +408,7 @@ class LinearOperatorCirculant2DBaseTest(object):
         x = np.zeros(block_shape)
         # x is a basis vector.
         x[n0, n1] = 1.0
-        fft_x = math_ops.fft2d(x)
+        fft_x = math_ops.fft2d(x.astype(np.complex64))
         h_convolve_x = math_ops.ifft2d(spectrum * fft_x)
         # We want the flat version of the action of the operator on a basis
         # vector, not the block version.
diff --git a/tensorflow/python/ops/spectral_grad.py b/tensorflow/python/ops/spectral_grad.py
index deb0a57178..0af24114ac 100644
--- a/tensorflow/python/ops/spectral_grad.py
+++ b/tensorflow/python/ops/spectral_grad.py
@@ -32,38 +32,44 @@ def _FFTSizeForGrad(grad, rank):
 
 @ops.RegisterGradient("FFT")
 def _FFTGrad(_, grad):
-  size = math_ops.cast(_FFTSizeForGrad(grad, 1), dtypes.float32)
-  return spectral_ops.ifft(grad) * math_ops.complex(size, 0.)
+  size = math_ops.cast(_FFTSizeForGrad(grad, 1), grad.dtype)
+  return spectral_ops.ifft(grad) * size
 
 
 @ops.RegisterGradient("IFFT")
 def _IFFTGrad(_, grad):
-  rsize = 1. / math_ops.cast(_FFTSizeForGrad(grad, 1), dtypes.float32)
-  return spectral_ops.fft(grad) * math_ops.complex(rsize, 0.)
+  rsize = math_ops.cast(
+      1. / math_ops.cast(_FFTSizeForGrad(grad, 1), grad.dtype.real_dtype),
+      grad.dtype)
+  return spectral_ops.fft(grad) * rsize
 
 
 @ops.RegisterGradient("FFT2D")
 def _FFT2DGrad(_, grad):
-  size = math_ops.cast(_FFTSizeForGrad(grad, 2), dtypes.float32)
-  return spectral_ops.ifft2d(grad) * math_ops.complex(size, 0.)
+  size = math_ops.cast(_FFTSizeForGrad(grad, 2), grad.dtype)
+  return spectral_ops.ifft2d(grad) * size
 
 
 @ops.RegisterGradient("IFFT2D")
 def _IFFT2DGrad(_, grad):
-  rsize = 1. / math_ops.cast(_FFTSizeForGrad(grad, 2), dtypes.float32)
-  return spectral_ops.fft2d(grad) * math_ops.complex(rsize, 0.)
+  rsize = math_ops.cast(
+      1. / math_ops.cast(_FFTSizeForGrad(grad, 2), grad.dtype.real_dtype),
+      grad.dtype)
+  return spectral_ops.fft2d(grad) * rsize
 
 
 @ops.RegisterGradient("FFT3D")
 def _FFT3DGrad(_, grad):
-  size = math_ops.cast(_FFTSizeForGrad(grad, 3), dtypes.float32)
-  return spectral_ops.ifft3d(grad) * math_ops.complex(size, 0.)
+  size = math_ops.cast(_FFTSizeForGrad(grad, 3), grad.dtype)
+  return spectral_ops.ifft3d(grad) * size
 
 
 @ops.RegisterGradient("IFFT3D")
 def _IFFT3DGrad(_, grad):
-  rsize = 1. / math_ops.cast(_FFTSizeForGrad(grad, 3), dtypes.float32)
-  return spectral_ops.fft3d(grad) * math_ops.complex(rsize, 0.)
+  rsize = math_ops.cast(
+      1. / math_ops.cast(_FFTSizeForGrad(grad, 3), grad.dtype.real_dtype),
+      grad.dtype)
+  return spectral_ops.fft3d(grad) * rsize
 
 
 def _RFFTGradHelper(rank, irfft_fn):
-- 
GitLab


From db329cfe2dee382033ad3b3f5e1d906ff489a24d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 18:11:25 -0700
Subject: [PATCH 190/395] Automated g4 rollback of changelist 195091587

PiperOrigin-RevId: 195184798
---
 tensorflow/contrib/lite/toco/BUILD            |  1 +
 .../contrib/lite/toco/model_flags.proto       |  3 +-
 tensorflow/contrib/lite/toco/tooling_util.cc  | 79 ++++++++++++-------
 3 files changed, 53 insertions(+), 30 deletions(-)

diff --git a/tensorflow/contrib/lite/toco/BUILD b/tensorflow/contrib/lite/toco/BUILD
index f16225fd66..ce0a74724a 100644
--- a/tensorflow/contrib/lite/toco/BUILD
+++ b/tensorflow/contrib/lite/toco/BUILD
@@ -397,6 +397,7 @@ cc_library(
         ":types_proto_cc",
         "//tensorflow/core:lib",
         "@com_google_absl//absl/strings",
+        "@com_googlesource_code_re2//:re2",
         "@protobuf_archive//:protobuf_headers",
     ],
 )
diff --git a/tensorflow/contrib/lite/toco/model_flags.proto b/tensorflow/contrib/lite/toco/model_flags.proto
index d23e80c464..6c1c53658c 100644
--- a/tensorflow/contrib/lite/toco/model_flags.proto
+++ b/tensorflow/contrib/lite/toco/model_flags.proto
@@ -96,8 +96,9 @@ message RnnState {
 // model that does not already contain such MinMax information.
 message ArraysExtraInfo {
   message Entry {
-    // Next ID to use: 7.
+    // Next ID to use: 8.
     optional string name = 1;
+    optional string name_regexp = 7;
     optional double min = 2;
     optional double max = 3;
     optional IODataType data_type = 4;
diff --git a/tensorflow/contrib/lite/toco/tooling_util.cc b/tensorflow/contrib/lite/toco/tooling_util.cc
index f334c51bbb..11293a5fe5 100644
--- a/tensorflow/contrib/lite/toco/tooling_util.cc
+++ b/tensorflow/contrib/lite/toco/tooling_util.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "absl/strings/str_join.h"
 #include "absl/strings/str_replace.h"
 #include "absl/strings/str_split.h"
+#include "re2/re2.h"
 #include "tensorflow/contrib/lite/toco/dump_graphviz.h"
 #include "tensorflow/contrib/lite/toco/model_flags.pb.h"
 #include "tensorflow/contrib/lite/toco/toco_graphviz_dump_options.h"
@@ -1983,38 +1984,58 @@ void FinishBuildingRNNStates(Model* model) {
   }
 }
 
+// Returns the array names that match the ArraysExtraInfo's name and
+// name_regexp. The regexp match is for a full match.
+std::unordered_set<string> ScanArrayNames(
+    const Model& model, const toco::ArraysExtraInfo_Entry& entry) {
+  std::unordered_set<string> matches;
+  if (model.HasArray(entry.name())) {
+    matches.insert(entry.name());
+  }
+  if (!entry.name_regexp().empty()) {
+    const auto& arrays = model.GetArrayMap();
+    const RE2 name_regexp = {entry.name_regexp()};
+    for (auto it = arrays.begin(); it != arrays.end(); ++it) {
+      if (RE2::FullMatch(it->first, name_regexp)) {
+        matches.insert(it->first);
+      }
+    }
+  }
+  return matches;
+}
+
 void UseArraysExtraInfo(Model* model, bool quantize_output) {
   for (const auto& entry : model->flags.arrays_extra_info().entries()) {
-    if (!model->HasArray(entry.name())) {
-      continue;
-    }
-    auto& array = model->GetArray(entry.name());
-    if (entry.has_min() || entry.has_max()) {
-      CHECK_EQ(entry.has_min(), entry.has_max());
-      auto& minmax = array.GetOrCreateMinMax();
-      minmax.min = entry.min();
-      minmax.max = entry.max();
-    }
-    if (entry.has_data_type() && quantize_output) {
-      array.final_data_type =
-          ConvertIODataTypeToArrayDataType(entry.data_type());
-    }
-    if (entry.has_shape()) {
-      array.clear_shape();
-      // Make sure to create the shape even if there are no dims, to
-      // correctly record 0-D shapes.
-      array.mutable_shape();
-      for (int dim : entry.shape().dims()) {
-        array.mutable_shape()->mutable_dims()->push_back(dim);
+    const auto matches = ScanArrayNames(*model, entry);
+    for (const auto& matched_name : matches) {
+      auto& array = model->GetArray(matched_name);
+      if (entry.has_min() || entry.has_max()) {
+        CHECK_EQ(entry.has_min(), entry.has_max());
+        auto& minmax = array.GetOrCreateMinMax();
+        minmax.min = entry.min();
+        minmax.max = entry.max();
       }
-    }
-    if (entry.has_constant_float_value()) {
-      CHECK(array.has_shape());
-      if (array.data_type == ArrayDataType::kFloat) {
-        auto& data = array.GetMutableBuffer<ArrayDataType::kFloat>().data;
-        data.resize(RequiredBufferSizeForShape(array.shape()));
-        for (float& f : data) {
-          f = entry.constant_float_value();
+      if (entry.has_data_type() && quantize_output) {
+        array.final_data_type =
+            ConvertIODataTypeToArrayDataType(entry.data_type());
+      }
+      if (entry.has_shape()) {
+        array.clear_shape();
+        // Make sure to create the shape even if there are no dims, to
+        // correctly record 0-D shapes.
+        array.mutable_shape();
+        for (int dim : entry.shape().dims()) {
+          array.mutable_shape()->mutable_dims()->push_back(dim);
+        }
+      }
+      if (entry.has_constant_float_value()) {
+        CHECK(array.has_shape());
+        if (array.data_type == ArrayDataType::kFloat) {
+          auto& data = array.GetMutableBuffer<ArrayDataType::kFloat>().data;
+          data.resize(RequiredBufferSizeForShape(array.shape()));
+          for (float& f : data) {
+            f = entry.constant_float_value();
+          }
         }
       }
     }
-- 
GitLab


From 1a4f746ffd82376f6e9ad420d96943ff89e7013a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 18:19:16 -0700
Subject: [PATCH 191/395] Remove duplicated emplace_back floor operator.

PiperOrigin-RevId: 195185567
---
 tensorflow/contrib/lite/toco/tflite/operator.cc | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tensorflow/contrib/lite/toco/tflite/operator.cc b/tensorflow/contrib/lite/toco/tflite/operator.cc
index fce3bad326..d2e14ac5e0 100644
--- a/tensorflow/contrib/lite/toco/tflite/operator.cc
+++ b/tensorflow/contrib/lite/toco/tflite/operator.cc
@@ -901,8 +901,6 @@ std::vector<std::unique_ptr<BaseOperator>> BuildOperatorList() {
       "MINIMUM", OperatorType::kTensorFlowMinimum));
   ops.emplace_back(new SimpleOperator<TensorFlowLessOperator>(
       "LESS", OperatorType::kTensorFlowLess));
-  ops.emplace_back(
-      new SimpleOperator<FloorOperator>("FLOOR", OperatorType::kFloor));
 
   return ops;
 }
-- 
GitLab


From 2b1a03c2ad502329a1f2b1368a40913ef21e97a0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 18:35:55 -0700
Subject: [PATCH 192/395] Compute shape of segment_ids dynamically in
 _unsorted_segment_N

PiperOrigin-RevId: 195186950
---
 tensorflow/python/ops/math_ops.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py
index 7ac3bd8091..ab5997e85c 100644
--- a/tensorflow/python/ops/math_ops.py
+++ b/tensorflow/python/ops/math_ops.py
@@ -2515,7 +2515,8 @@ def _unsorted_segment_N(data, segment_ids, num_segments):
       of segment entries with 0-entries set to 1 to allow division by N.
   """
   # bincount doesn't support negative indices so we use unsorted_segment_sum
-  ones_tensor = array_ops.ones(segment_ids.shape, dtype=data.dtype)
+  segment_ids_shape = array_ops.shape_internal(segment_ids)
+  ones_tensor = array_ops.ones(segment_ids_shape, dtype=data.dtype)
   N = gen_math_ops.unsorted_segment_sum(ones_tensor, segment_ids, num_segments)
   # add dimensions for all non-reduced axes
   ndims_output = data.shape.ndims - segment_ids.shape.ndims
-- 
GitLab


From 223be4abe74592a781735a6b66e12cb0146f0830 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 18:52:02 -0700
Subject: [PATCH 193/395] Replaced calls to tensorflow::StringPiece::ToString
 with std::string conversions. That is, instances of sp.ToString() are
 replaced with std::string(sp).

This will allow tensorflow::StringPiece::ToString to be removed, which is necessary before it can be replaced with absl::string_view.

PiperOrigin-RevId: 195188185
---
 tensorflow/cc/framework/cc_op_gen.cc             |  2 +-
 tensorflow/cc/framework/scope.cc                 |  2 +-
 tensorflow/compiler/tf2xla/tf2xla_util.cc        |  2 +-
 tensorflow/compiler/tf2xla/xla_op_registry.cc    | 14 +++++++-------
 .../compiler/xla/service/llvm_ir/llvm_loop.cc    |  4 ++--
 .../compiler/xla/service/llvm_ir/llvm_loop.h     |  2 +-
 .../compiler/xla/tools/parser/hlo_lexer.cc       |  2 +-
 .../compiler/xla/tools/parser/hlo_parser.cc      |  2 +-
 tensorflow/core/debug/debug_graph_utils.cc       |  7 +++----
 tensorflow/core/debug/debug_io_utils.cc          | 10 +++++-----
 .../core/distributed_runtime/master_session.cc   |  6 +++---
 .../core/distributed_runtime/remote_device.cc    |  2 +-
 tensorflow/core/grappler/utils.h                 |  2 +-
 .../core/kernels/hexagon/graph_transferer.cc     |  2 +-
 .../kernels/hexagon/hexagon_control_wrapper.cc   |  2 +-
 tensorflow/core/lib/io/path.cc                   |  6 +++---
 tensorflow/core/lib/io/table_test.cc             |  6 +++---
 .../core/util/tensor_bundle/tensor_bundle.cc     | 16 ++++++++--------
 .../util/tensor_bundle/tensor_bundle_test.cc     |  2 +-
 tensorflow/stream_executor/kernel.cc             |  2 +-
 tensorflow/stream_executor/kernel_spec.cc        |  6 +++---
 21 files changed, 49 insertions(+), 50 deletions(-)

diff --git a/tensorflow/cc/framework/cc_op_gen.cc b/tensorflow/cc/framework/cc_op_gen.cc
index d73121c7b7..d6a4f141b6 100644
--- a/tensorflow/cc/framework/cc_op_gen.cc
+++ b/tensorflow/cc/framework/cc_op_gen.cc
@@ -440,7 +440,7 @@ string AvoidCPPKeywords(StringPiece name) {
   if (IsCPPKeyword(name)) {
     return strings::StrCat(name, "_");
   }
-  return name.ToString();
+  return std::string(name);
 }
 
 void InferArgAttributes(const OpDef::ArgDef& arg,
diff --git a/tensorflow/cc/framework/scope.cc b/tensorflow/cc/framework/scope.cc
index c143b97833..62a889181e 100644
--- a/tensorflow/cc/framework/scope.cc
+++ b/tensorflow/cc/framework/scope.cc
@@ -220,7 +220,7 @@ std::unordered_set<string> Scope::Impl::GetColocationConstraints(
     for (const string& entry : node_constraints) {
       StringPiece s(entry);
       if (str_util::ConsumePrefix(&s, kColocationGroupPrefix)) {
-        current_constraints.insert(s.ToString());
+        current_constraints.insert(std::string(s));
       }
     }
   } else {
diff --git a/tensorflow/compiler/tf2xla/tf2xla_util.cc b/tensorflow/compiler/tf2xla/tf2xla_util.cc
index 7ec85aa3cd..9203e8d9e6 100644
--- a/tensorflow/compiler/tf2xla/tf2xla_util.cc
+++ b/tensorflow/compiler/tf2xla/tf2xla_util.cc
@@ -232,7 +232,7 @@ Status PruneGraphDefInto(const tf2xla::Config& config, const GraphDef& in,
     // Push input nodes of the currently visited node to name_queue.
     for (const string& in_edge : map_entry.second->input()) {
       auto id = ParseTensorName(in_edge);
-      const string node_name = id.first.ToString();
+      const string node_name = std::string(id.first);
       if (feed_tensors.find(std::make_pair(node_name, id.second)) ==
           feed_tensors.end()) {
         name_queue.push(node_name);
diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.cc b/tensorflow/compiler/tf2xla/xla_op_registry.cc
index bbe808595d..e309cb1e34 100644
--- a/tensorflow/compiler/tf2xla/xla_op_registry.cc
+++ b/tensorflow/compiler/tf2xla/xla_op_registry.cc
@@ -311,7 +311,7 @@ XlaOpRegistry& XlaOpRegistry::Instance() {
 
 XlaOpRegistrationBuilder::XlaOpRegistrationBuilder(StringPiece name) {
   registration_.reset(new XlaOpRegistry::OpRegistration);
-  registration_->name = name.ToString();
+  registration_->name = std::string(name);
 }
 
 XlaOpRegistrationBuilder XlaOpRegistrationBuilder::Name(StringPiece name) {
@@ -323,14 +323,14 @@ XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::Device(
     gtl::ArraySlice<StringPiece> devices) {
   registration_->has_device_whitelist = true;
   for (StringPiece device : devices) {
-    registration_->device_whitelist.insert(device.ToString());
+    registration_->device_whitelist.insert(std::string(device));
   }
   return *this;
 }
 
 XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::Device(StringPiece device) {
   registration_->has_device_whitelist = true;
-  registration_->device_whitelist.insert(device.ToString());
+  registration_->device_whitelist.insert(std::string(device));
   return *this;
 }
 
@@ -347,7 +347,7 @@ XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::AllowResourceTypes() {
 XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::TypeConstraint(
     StringPiece attr_name, DataType allowed) {
   std::set<DataType>& types =
-      registration_->type_constraints[attr_name.ToString()];
+      registration_->type_constraints[std::string(attr_name)];
   types.insert(allowed);
   return *this;
 }
@@ -355,7 +355,7 @@ XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::TypeConstraint(
 XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::TypeConstraint(
     StringPiece attr_name, gtl::ArraySlice<DataType> allowed) {
   std::set<DataType>& types =
-      registration_->type_constraints[attr_name.ToString()];
+      registration_->type_constraints[std::string(attr_name)];
   for (DataType t : allowed) {
     types.insert(t);
   }
@@ -364,7 +364,7 @@ XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::TypeConstraint(
 
 XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::CompileTimeConstInput(
     StringPiece input_name) {
-  registration_->compile_time_constant_inputs.insert(input_name.ToString());
+  registration_->compile_time_constant_inputs.insert(std::string(input_name));
   return *this;
 }
 
@@ -394,7 +394,7 @@ XlaBackendRegistrar::XlaBackendRegistrar(
     StringPiece name, gtl::ArraySlice<DataType> types,
     XlaOpRegistry::BackendOpFilter op_filter) {
   XlaOpRegistry& registry = XlaOpRegistry::Instance();
-  registry.RegisterBackend(name.ToString(), types, op_filter);
+  registry.RegisterBackend(std::string(name), types, op_filter);
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc b/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc
index 7b227ce294..497b48ff22 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc
+++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc
@@ -36,8 +36,8 @@ ForLoop::ForLoop(tensorflow::StringPiece prefix, tensorflow::StringPiece suffix,
                  llvm::Value* start_index, llvm::Value* end_index,
                  llvm::Value* step, bool prevent_unrolling,
                  bool prevent_vectorization)
-    : prefix_(prefix.ToString()),
-      suffix_(suffix.ToString()),
+    : prefix_(std::string(prefix)),
+      suffix_(std::string(suffix)),
       start_index_(start_index),
       end_index_(end_index),
       step_(step),
diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h b/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h
index 20069ce5a2..d915f95db1 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h
+++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h
@@ -174,7 +174,7 @@ class ForLoopNest {
       : ForLoopNest(/*name=*/"", ir_builder) {}
 
   ForLoopNest(tensorflow::StringPiece name, llvm::IRBuilder<>* ir_builder)
-      : name_(name.ToString()),
+      : name_(std::string(name)),
         outer_loop_preheader_bb_(nullptr),
         outer_loop_exit_bb_(nullptr),
         inner_loop_body_bb_(nullptr),
diff --git a/tensorflow/compiler/xla/tools/parser/hlo_lexer.cc b/tensorflow/compiler/xla/tools/parser/hlo_lexer.cc
index fc0e444452..350db12653 100644
--- a/tensorflow/compiler/xla/tools/parser/hlo_lexer.cc
+++ b/tensorflow/compiler/xla/tools/parser/hlo_lexer.cc
@@ -230,7 +230,7 @@ TokKind HloLexer::LexIdentifier() {
     }
   }
 
-  str_val_ = identifier.ToString();
+  str_val_ = std::string(identifier);
   return TokKind::kIdent;
 }
 
diff --git a/tensorflow/compiler/xla/tools/parser/hlo_parser.cc b/tensorflow/compiler/xla/tools/parser/hlo_parser.cc
index 1bb31ddb7b..3a945fb3b1 100644
--- a/tensorflow/compiler/xla/tools/parser/hlo_parser.cc
+++ b/tensorflow/compiler/xla/tools/parser/hlo_parser.cc
@@ -242,7 +242,7 @@ bool HloParser::Error(LocTy loc, StringPiece msg) {
   std::vector<string> error_lines;
   error_lines.push_back(
       StrCat("was parsing ", line, ":", col, ": error: ", msg));
-  error_lines.push_back(lexer_.GetLine(loc).ToString());
+  error_lines.push_back(std::string(lexer_.GetLine(loc)));
   error_lines.push_back(col == 0 ? "" : StrCat(string(col - 1, ' '), "^"));
 
   error_.push_back(tensorflow::str_util::Join(error_lines, "\n"));
diff --git a/tensorflow/core/debug/debug_graph_utils.cc b/tensorflow/core/debug/debug_graph_utils.cc
index 4539ea5c0c..7641edea52 100644
--- a/tensorflow/core/debug/debug_graph_utils.cc
+++ b/tensorflow/core/debug/debug_graph_utils.cc
@@ -356,10 +356,9 @@ Status DebugNodeInserter::ParseDebugOpName(
             "Malformed attributes in debug op name \"", debug_op_name, "\"");
       }
 
-      const string key = seg.substr(0, eq_index).ToString();
-      const string value =
-          seg.substr(eq_index + 1, attribute_seg.size() - eq_index - 1)
-              .ToString();
+      const string key = std::string(seg.substr(0, eq_index));
+      const string value = std::string(
+          seg.substr(eq_index + 1, attribute_seg.size() - eq_index - 1));
       if (key.empty() || value.empty()) {
         return errors::InvalidArgument(
             "Malformed attributes in debug op name \"", debug_op_name, "\"");
diff --git a/tensorflow/core/debug/debug_io_utils.cc b/tensorflow/core/debug/debug_io_utils.cc
index baa8c08fdf..4998a7acfe 100644
--- a/tensorflow/core/debug/debug_io_utils.cc
+++ b/tensorflow/core/debug/debug_io_utils.cc
@@ -399,8 +399,8 @@ Status DebugIO::PublishDebugMetadata(
                               strings::Printf("%.14lld", session_run_index))),
           Env::Default()->NowMicros());
       status.Update(DebugFileIO::DumpEventProtoToFile(
-          event, io::Dirname(core_metadata_path).ToString(),
-          io::Basename(core_metadata_path).ToString()));
+          event, std::string(io::Dirname(core_metadata_path)),
+          std::string(io::Basename(core_metadata_path))));
     }
   }
 
@@ -632,8 +632,8 @@ Status DebugFileIO::DumpTensorToEventFile(const DebugNodeKey& debug_node_key,
   std::vector<Event> events;
   TF_RETURN_IF_ERROR(
       WrapTensorAsEvents(debug_node_key, tensor, wall_time_us, 0, &events));
-  return DumpEventProtoToFile(events[0], io::Dirname(file_path).ToString(),
-                              io::Basename(file_path).ToString());
+  return DumpEventProtoToFile(events[0], std::string(io::Dirname(file_path)),
+                              std::string(io::Basename(file_path)));
 }
 
 Status DebugFileIO::RecursiveCreateDir(Env* env, const string& dir) {
@@ -642,7 +642,7 @@ Status DebugFileIO::RecursiveCreateDir(Env* env, const string& dir) {
     return Status::OK();
   }
 
-  string parent_dir = io::Dirname(dir).ToString();
+  string parent_dir = std::string(io::Dirname(dir));
   if (!env->FileExists(parent_dir).ok()) {
     // The parent path does not exist yet, create it first.
     Status s = RecursiveCreateDir(env, parent_dir);  // Recursive call
diff --git a/tensorflow/core/distributed_runtime/master_session.cc b/tensorflow/core/distributed_runtime/master_session.cc
index 83afc5b1a4..08fbe8b144 100644
--- a/tensorflow/core/distributed_runtime/master_session.cc
+++ b/tensorflow/core/distributed_runtime/master_session.cc
@@ -606,7 +606,7 @@ Status MasterSession::ReffedClientGraph::RunPartitionsHelper(
     // inadvertently slowing down the normal run path.
     if (is_partial_) {
       for (const auto& name_index : feeds) {
-        const auto iter = part.feed_key.find(name_index.first.ToString());
+        const auto iter = part.feed_key.find(std::string(name_index.first));
         if (iter == part.feed_key.end()) {
           // The provided feed must be for a different partition.
           continue;
@@ -950,7 +950,7 @@ Status MasterSession::ReffedClientGraph::CheckFetches(
     // Skip if already fed.
     if (input.second) continue;
     TensorId id(ParseTensorName(input.first));
-    const Node* n = execution_state->get_node_by_name(id.first.ToString());
+    const Node* n = execution_state->get_node_by_name(std::string(id.first));
     if (n == nullptr) {
       return errors::NotFound("Feed ", input.first, ": not found");
     }
@@ -966,7 +966,7 @@ Status MasterSession::ReffedClientGraph::CheckFetches(
   for (size_t i = 0; i < req.num_fetches(); ++i) {
     const string& fetch = req.fetch_name(i);
     const TensorId id(ParseTensorName(fetch));
-    const Node* n = execution_state->get_node_by_name(id.first.ToString());
+    const Node* n = execution_state->get_node_by_name(std::string(id.first));
     if (n == nullptr) {
       return errors::NotFound("Fetch ", fetch, ": not found");
     }
diff --git a/tensorflow/core/distributed_runtime/remote_device.cc b/tensorflow/core/distributed_runtime/remote_device.cc
index ec26ac44b5..15e5919c54 100644
--- a/tensorflow/core/distributed_runtime/remote_device.cc
+++ b/tensorflow/core/distributed_runtime/remote_device.cc
@@ -37,7 +37,7 @@ string GetLocalDeviceName(StringPiece fullname) {
   auto pos = fullname.rfind('/');
   CHECK_NE(pos, StringPiece::npos);
   fullname.remove_prefix(pos + 1);
-  return fullname.ToString();
+  return std::string(fullname);
 }
 
 class RemoteDevice : public Device {
diff --git a/tensorflow/core/grappler/utils.h b/tensorflow/core/grappler/utils.h
index 9776e99f20..b87ae05546 100644
--- a/tensorflow/core/grappler/utils.h
+++ b/tensorflow/core/grappler/utils.h
@@ -139,7 +139,7 @@ inline StringPiece ParseNodeNameAsStringPiece(const string& name,
 
 // Returns the node name and position in a single call.
 inline string ParseNodeName(const string& name, int* position) {
-  return ParseNodeNameAsStringPiece(name, position).ToString();
+  return std::string(ParseNodeNameAsStringPiece(name, position));
 }
 
 // Add a prefix to a node name with a custom delimiter.
diff --git a/tensorflow/core/kernels/hexagon/graph_transferer.cc b/tensorflow/core/kernels/hexagon/graph_transferer.cc
index 7960cb4b05..e05de3fe8e 100644
--- a/tensorflow/core/kernels/hexagon/graph_transferer.cc
+++ b/tensorflow/core/kernels/hexagon/graph_transferer.cc
@@ -161,7 +161,7 @@ Status GraphTransferer::LoadGraphFromProto(
 
   for (const string& output_node_name : output_node_names) {
     const TensorId tid = ParseTensorName(output_node_name);
-    const string node_name = tid.first.ToString();
+    const string node_name = std::string(tid.first);
     const int port = tid.second;
     const int node_id = node_name_to_id_cache_map_.at(node_name);
     const Node* node = node_name_cache_list_.at(node_id);
diff --git a/tensorflow/core/kernels/hexagon/hexagon_control_wrapper.cc b/tensorflow/core/kernels/hexagon/hexagon_control_wrapper.cc
index 3810cbe5b5..1580b72605 100644
--- a/tensorflow/core/kernels/hexagon/hexagon_control_wrapper.cc
+++ b/tensorflow/core/kernels/hexagon/hexagon_control_wrapper.cc
@@ -168,7 +168,7 @@ bool HexagonControlWrapper::SetupGraph() {
     new_output_node_info.set_output_count(0);
 
     const TensorId tid = ParseTensorName(graph_output.name());
-    const string node_name = tid.first.ToString();
+    const string node_name = std::string(tid.first);
     const int port = tid.second;
     // Register node input for the new output node
     const GraphTransferNodeInfo* node_info =
diff --git a/tensorflow/core/lib/io/path.cc b/tensorflow/core/lib/io/path.cc
index 996fbf62e5..b62206012c 100644
--- a/tensorflow/core/lib/io/path.cc
+++ b/tensorflow/core/lib/io/path.cc
@@ -42,7 +42,7 @@ string JoinPathImpl(std::initializer_list<StringPiece> paths) {
     if (path.empty()) continue;
 
     if (result.empty()) {
-      result = path.ToString();
+      result = std::string(path);
       continue;
     }
 
@@ -124,7 +124,7 @@ StringPiece Extension(StringPiece path) {
 }
 
 string CleanPath(StringPiece unclean_path) {
-  string path = unclean_path.ToString();
+  string path = std::string(unclean_path);
   const char* src = path.c_str();
   string::iterator dst = path.begin();
 
@@ -237,7 +237,7 @@ void ParseURI(StringPiece remaining, StringPiece* scheme, StringPiece* host,
 
 string CreateURI(StringPiece scheme, StringPiece host, StringPiece path) {
   if (scheme.empty()) {
-    return path.ToString();
+    return std::string(path);
   }
   return strings::StrCat(scheme, "://", host, path);
 }
diff --git a/tensorflow/core/lib/io/table_test.cc b/tensorflow/core/lib/io/table_test.cc
index 78a3fa501c..9e3309f0a7 100644
--- a/tensorflow/core/lib/io/table_test.cc
+++ b/tensorflow/core/lib/io/table_test.cc
@@ -147,7 +147,7 @@ class Constructor {
   virtual ~Constructor() {}
 
   void Add(const string& key, const StringPiece& value) {
-    data_[key] = value.ToString();
+    data_[key] = std::string(value);
   }
 
   // Finish constructing the data structure with all the keys that have
@@ -188,7 +188,7 @@ class BlockConstructor : public Constructor {
       builder.Add(it->first, it->second);
     }
     // Open the block
-    data_ = builder.Finish().ToString();
+    data_ = std::string(builder.Finish());
     BlockContents contents;
     contents.data = data_;
     contents.cachable = false;
@@ -515,7 +515,7 @@ TEST_F(Harness, Randomized) {
       for (int e = 0; e < num_entries; e++) {
         string v;
         Add(test::RandomKey(&rnd, rnd.Skewed(4)),
-            test::RandomString(&rnd, rnd.Skewed(5), &v).ToString());
+            std::string(test::RandomString(&rnd, rnd.Skewed(5), &v)));
       }
       Test(&rnd);
     }
diff --git a/tensorflow/core/util/tensor_bundle/tensor_bundle.cc b/tensorflow/core/util/tensor_bundle/tensor_bundle.cc
index 0426fee0e2..7190614706 100644
--- a/tensorflow/core/util/tensor_bundle/tensor_bundle.cc
+++ b/tensorflow/core/util/tensor_bundle/tensor_bundle.cc
@@ -370,14 +370,14 @@ Status PadAlignment(FileOutputBuffer* out, int alignment, int64* size) {
 BundleWriter::BundleWriter(Env* env, StringPiece prefix, const Options& options)
     : env_(env),
       options_(options),
-      prefix_(prefix.ToString()),
+      prefix_(std::string(prefix)),
       tmp_metadata_path_(strings::StrCat(MetaFilename(prefix_), ".tempstate",
                                          random::New64())),
       tmp_data_path_(strings::StrCat(DataFilename(prefix_, 0, 1), ".tempstate",
                                      random::New64())),
       out_(nullptr),
       size_(0) {
-  status_ = env_->CreateDir(io::Dirname(prefix_).ToString());
+  status_ = env_->CreateDir(std::string(io::Dirname(prefix_)));
   if (!status_.ok() && !errors::IsAlreadyExists(status_)) {
     return;
   }
@@ -394,7 +394,7 @@ BundleWriter::BundleWriter(Env* env, StringPiece prefix, const Options& options)
 Status BundleWriter::Add(StringPiece key, const Tensor& val) {
   if (!status_.ok()) return status_;
   CHECK_NE(key, kHeaderEntryKey);
-  const string key_string = key.ToString();
+  const string key_string = std::string(key);
   if (entries_.find(key_string) != entries_.end()) {
     status_ = errors::InvalidArgument("Adding duplicate key: ", key);
     return status_;
@@ -445,7 +445,7 @@ Status BundleWriter::AddSlice(StringPiece full_tensor_key,
   // In the case of a sharded save, MergeBundles() is responsible for merging
   // the "slices" field of multiple metadata entries corresponding to the same
   // full tensor.
-  const string full_tensor_key_string = full_tensor_key.ToString();
+  const string full_tensor_key_string = std::string(full_tensor_key);
   BundleEntryProto* full_entry = &entries_[full_tensor_key_string];
   if (full_entry->dtype() != DT_INVALID) {
     CHECK_EQ(full_entry->dtype(), slice_tensor.dtype());
@@ -600,7 +600,7 @@ static Status MergeOneBundle(Env* env, StringPiece prefix,
   // Loops through the non-header to-merge entries.
   BundleEntryProto to_merge_entry;
   for (; iter->Valid(); iter->Next()) {
-    const string key = iter->key().ToString();
+    const string key = std::string(iter->key());
     const auto entry_iter = merge_state->entries.find(key);
 
     // Illegal: the duplicated entry is a non-slice tensor.
@@ -649,7 +649,7 @@ Status MergeBundles(Env* env, gtl::ArraySlice<string> prefixes,
   // Merges all metadata tables.
   // TODO(zhifengc): KeyValue sorter if it becomes too big.
   MergeState merge;
-  Status status = env->CreateDir(io::Dirname(merged_prefix).ToString());
+  Status status = env->CreateDir(std::string(io::Dirname(merged_prefix)));
   if (!status.ok() && !errors::IsAlreadyExists(status)) return status;
   for (int i = 0; i < prefixes.size(); ++i) {
     TF_RETURN_IF_ERROR(MergeOneBundle(env, prefixes[i], &merge));
@@ -697,7 +697,7 @@ Status MergeBundles(Env* env, gtl::ArraySlice<string> prefixes,
 
 BundleReader::BundleReader(Env* env, StringPiece prefix)
     : env_(env),
-      prefix_(prefix.ToString()),
+      prefix_(std::string(prefix)),
       metadata_(nullptr),
       table_(nullptr),
       iter_(nullptr) {
@@ -919,7 +919,7 @@ Status BundleReader::GetSliceValue(StringPiece full_tensor_key,
 
   const TensorShape full_shape(TensorShape(full_tensor_entry.shape()));
   std::vector<std::pair<TensorSlice, string>> details;
-  const string full_tensor_key_string = full_tensor_key.ToString();
+  const string full_tensor_key_string = std::string(full_tensor_key);
   const TensorSliceSet* tss =
       gtl::FindPtrOrNull(tensor_slices_, full_tensor_key_string);
 
diff --git a/tensorflow/core/util/tensor_bundle/tensor_bundle_test.cc b/tensorflow/core/util/tensor_bundle/tensor_bundle_test.cc
index 7f166f0ec0..92ce8ae00e 100644
--- a/tensorflow/core/util/tensor_bundle/tensor_bundle_test.cc
+++ b/tensorflow/core/util/tensor_bundle/tensor_bundle_test.cc
@@ -107,7 +107,7 @@ std::vector<string> AllTensorKeys(BundleReader* reader) {
   reader->Seek(kHeaderEntryKey);
   reader->Next();
   for (; reader->Valid(); reader->Next()) {
-    ret.push_back(reader->key().ToString());
+    ret.push_back(std::string(reader->key()));
   }
   return ret;
 }
diff --git a/tensorflow/stream_executor/kernel.cc b/tensorflow/stream_executor/kernel.cc
index d1aa596b73..7c1923da51 100644
--- a/tensorflow/stream_executor/kernel.cc
+++ b/tensorflow/stream_executor/kernel.cc
@@ -94,7 +94,7 @@ KernelCacheConfig KernelBase::GetPreferredCacheConfig() const {
 static const char *kStubPrefix = "__device_stub_";
 
 void KernelBase::set_name(port::StringPiece name) {
-  name_ = name.ToString();
+  name_ = std::string(name);
   port::StringPiece stubless_name = name;
   if (tensorflow::str_util::StartsWith(name, kStubPrefix)) {
     stubless_name.remove_prefix(strlen(kStubPrefix));
diff --git a/tensorflow/stream_executor/kernel_spec.cc b/tensorflow/stream_executor/kernel_spec.cc
index 6a1f0a591f..f0a5785b72 100644
--- a/tensorflow/stream_executor/kernel_spec.cc
+++ b/tensorflow/stream_executor/kernel_spec.cc
@@ -18,11 +18,11 @@ limitations under the License.
 namespace stream_executor {
 
 KernelLoaderSpec::KernelLoaderSpec(port::StringPiece kernelname)
-    : kernelname_(kernelname.ToString()) {}
+    : kernelname_(std::string(kernelname)) {}
 
 OnDiskKernelLoaderSpec::OnDiskKernelLoaderSpec(port::StringPiece filename,
                                                port::StringPiece kernelname)
-    : KernelLoaderSpec(kernelname), filename_(filename.ToString()) {}
+    : KernelLoaderSpec(kernelname), filename_(std::string(filename)) {}
 
 CudaPtxOnDisk::CudaPtxOnDisk(port::StringPiece filename,
                              port::StringPiece kernelname)
@@ -161,7 +161,7 @@ OpenCLTextOnDisk::OpenCLTextOnDisk(port::StringPiece filename,
 
 OpenCLTextInMemory::OpenCLTextInMemory(port::StringPiece text,
                                        port::StringPiece kernelname)
-    : KernelLoaderSpec(kernelname), text_(text.ToString()) {}
+    : KernelLoaderSpec(kernelname), text_(std::string(text)) {}
 
 OpenCLBinaryOnDisk::OpenCLBinaryOnDisk(port::StringPiece filename,
                                        port::StringPiece kernelname)
-- 
GitLab


From ebad5d624c6c08a3fdb4ffac6051b4888fc36790 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 2 May 2018 19:19:12 -0700
Subject: [PATCH 194/395] Update ops-related pbtxt files.

PiperOrigin-RevId: 195190335
---
 .../core/ops/compat/ops_history.v1.pbtxt      | 144 ++++++++++++++++++
 tensorflow/core/ops/ops.pbtxt                 | 102 +++++++++++--
 2 files changed, 234 insertions(+), 12 deletions(-)

diff --git a/tensorflow/core/ops/compat/ops_history.v1.pbtxt b/tensorflow/core/ops/compat/ops_history.v1.pbtxt
index cb466ef817..3db00d8180 100644
--- a/tensorflow/core/ops/compat/ops_history.v1.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history.v1.pbtxt
@@ -20998,6 +20998,30 @@ op {
     type: DT_COMPLEX64
   }
 }
+op {
+  name: "FFT"
+  input_arg {
+    name: "input"
+    type_attr: "Tcomplex"
+  }
+  output_arg {
+    name: "output"
+    type_attr: "Tcomplex"
+  }
+  attr {
+    name: "Tcomplex"
+    type: "type"
+    default_value {
+      type: DT_COMPLEX64
+    }
+    allowed_values {
+      list {
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
 op {
   name: "FFT2D"
   input_arg {
@@ -21009,6 +21033,30 @@ op {
     type: DT_COMPLEX64
   }
 }
+op {
+  name: "FFT2D"
+  input_arg {
+    name: "input"
+    type_attr: "Tcomplex"
+  }
+  output_arg {
+    name: "output"
+    type_attr: "Tcomplex"
+  }
+  attr {
+    name: "Tcomplex"
+    type: "type"
+    default_value {
+      type: DT_COMPLEX64
+    }
+    allowed_values {
+      list {
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
 op {
   name: "FFT3D"
   input_arg {
@@ -21020,6 +21068,30 @@ op {
     type: DT_COMPLEX64
   }
 }
+op {
+  name: "FFT3D"
+  input_arg {
+    name: "input"
+    type_attr: "Tcomplex"
+  }
+  output_arg {
+    name: "output"
+    type_attr: "Tcomplex"
+  }
+  attr {
+    name: "Tcomplex"
+    type: "type"
+    default_value {
+      type: DT_COMPLEX64
+    }
+    allowed_values {
+      list {
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
 op {
   name: "FIFOQueue"
   output_arg {
@@ -24711,6 +24783,30 @@ op {
     type: DT_COMPLEX64
   }
 }
+op {
+  name: "IFFT"
+  input_arg {
+    name: "input"
+    type_attr: "Tcomplex"
+  }
+  output_arg {
+    name: "output"
+    type_attr: "Tcomplex"
+  }
+  attr {
+    name: "Tcomplex"
+    type: "type"
+    default_value {
+      type: DT_COMPLEX64
+    }
+    allowed_values {
+      list {
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
 op {
   name: "IFFT2D"
   input_arg {
@@ -24722,6 +24818,30 @@ op {
     type: DT_COMPLEX64
   }
 }
+op {
+  name: "IFFT2D"
+  input_arg {
+    name: "input"
+    type_attr: "Tcomplex"
+  }
+  output_arg {
+    name: "output"
+    type_attr: "Tcomplex"
+  }
+  attr {
+    name: "Tcomplex"
+    type: "type"
+    default_value {
+      type: DT_COMPLEX64
+    }
+    allowed_values {
+      list {
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
 op {
   name: "IFFT3D"
   input_arg {
@@ -24733,6 +24853,30 @@ op {
     type: DT_COMPLEX64
   }
 }
+op {
+  name: "IFFT3D"
+  input_arg {
+    name: "input"
+    type_attr: "Tcomplex"
+  }
+  output_arg {
+    name: "output"
+    type_attr: "Tcomplex"
+  }
+  attr {
+    name: "Tcomplex"
+    type: "type"
+    default_value {
+      type: DT_COMPLEX64
+    }
+    allowed_values {
+      list {
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
+  }
+}
 op {
   name: "IRFFT"
   input_arg {
diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 207dd1c3d7..7156440b46 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -9709,33 +9709,72 @@ op {
   name: "FFT"
   input_arg {
     name: "input"
-    type: DT_COMPLEX64
+    type_attr: "Tcomplex"
   }
   output_arg {
     name: "output"
-    type: DT_COMPLEX64
+    type_attr: "Tcomplex"
+  }
+  attr {
+    name: "Tcomplex"
+    type: "type"
+    default_value {
+      type: DT_COMPLEX64
+    }
+    allowed_values {
+      list {
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
   }
 }
 op {
   name: "FFT2D"
   input_arg {
     name: "input"
-    type: DT_COMPLEX64
+    type_attr: "Tcomplex"
   }
   output_arg {
     name: "output"
-    type: DT_COMPLEX64
+    type_attr: "Tcomplex"
+  }
+  attr {
+    name: "Tcomplex"
+    type: "type"
+    default_value {
+      type: DT_COMPLEX64
+    }
+    allowed_values {
+      list {
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
   }
 }
 op {
   name: "FFT3D"
   input_arg {
     name: "input"
-    type: DT_COMPLEX64
+    type_attr: "Tcomplex"
   }
   output_arg {
     name: "output"
-    type: DT_COMPLEX64
+    type_attr: "Tcomplex"
+  }
+  attr {
+    name: "Tcomplex"
+    type: "type"
+    default_value {
+      type: DT_COMPLEX64
+    }
+    allowed_values {
+      list {
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
   }
 }
 op {
@@ -11877,33 +11916,72 @@ op {
   name: "IFFT"
   input_arg {
     name: "input"
-    type: DT_COMPLEX64
+    type_attr: "Tcomplex"
   }
   output_arg {
     name: "output"
-    type: DT_COMPLEX64
+    type_attr: "Tcomplex"
+  }
+  attr {
+    name: "Tcomplex"
+    type: "type"
+    default_value {
+      type: DT_COMPLEX64
+    }
+    allowed_values {
+      list {
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
   }
 }
 op {
   name: "IFFT2D"
   input_arg {
     name: "input"
-    type: DT_COMPLEX64
+    type_attr: "Tcomplex"
   }
   output_arg {
     name: "output"
-    type: DT_COMPLEX64
+    type_attr: "Tcomplex"
+  }
+  attr {
+    name: "Tcomplex"
+    type: "type"
+    default_value {
+      type: DT_COMPLEX64
+    }
+    allowed_values {
+      list {
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
   }
 }
 op {
   name: "IFFT3D"
   input_arg {
     name: "input"
-    type: DT_COMPLEX64
+    type_attr: "Tcomplex"
   }
   output_arg {
     name: "output"
-    type: DT_COMPLEX64
+    type_attr: "Tcomplex"
+  }
+  attr {
+    name: "Tcomplex"
+    type: "type"
+    default_value {
+      type: DT_COMPLEX64
+    }
+    allowed_values {
+      list {
+        type: DT_COMPLEX64
+        type: DT_COMPLEX128
+      }
+    }
   }
 }
 op {
-- 
GitLab


From 71f97c8cd9304a8e1cf8e309e15000d5831b212a Mon Sep 17 00:00:00 2001
From: Mostafa Alaa <mostafa.alaa.gharieb@gmail.com>
Date: Wed, 2 May 2018 19:53:18 -0700
Subject: [PATCH 195/395] Fix tf.variable_scope unique name after entering root
 scope Closes #18702.

PiperOrigin-RevId: 195192460
---
 tensorflow/python/ops/variable_scope.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/ops/variable_scope.py b/tensorflow/python/ops/variable_scope.py
index ba213ef884..adb0f59948 100644
--- a/tensorflow/python/ops/variable_scope.py
+++ b/tensorflow/python/ops/variable_scope.py
@@ -1175,7 +1175,7 @@ class _VariableScopeStore(threading.local):
 
   def close_variable_subscopes(self, scope_name):
     for k in list(self.variable_scopes_count.keys()):
-      if not scope_name or k.startswith(scope_name + "/"):
+      if scope_name is None or k.startswith(scope_name + "/"):
         self.variable_scopes_count[k] = 0
 
   def variable_scope_count(self, scope_name):
-- 
GitLab


From 0e141b75a557a646750e4af06530892af5a8da20 Mon Sep 17 00:00:00 2001
From: Taehoon Lee <me@taehoonlee.com>
Date: Thu, 3 May 2018 12:01:04 +0900
Subject: [PATCH 196/395] Fix typos (#18475)

---
 tensorflow/compiler/xla/service/copy_insertion.cc               | 2 +-
 tensorflow/compiler/xla/service/hlo_evaluator.cc                | 2 +-
 tensorflow/contrib/kfac/python/ops/optimizer.py                 | 2 +-
 .../contrib/lite/kernels/internal/optimized/optimized_ops.h     | 2 +-
 .../contrib/lite/kernels/internal/reference/reference_ops.h     | 2 +-
 tensorflow/contrib/lite/testing/generate_examples.py            | 2 +-
 tensorflow/contrib/lite/toco/toco_flags.proto                   | 2 +-
 .../contrib/opt/python/training/model_average_optimizer_test.py | 2 +-
 tensorflow/python/keras/_impl/keras/engine/training_eager.py    | 2 +-
 9 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/tensorflow/compiler/xla/service/copy_insertion.cc b/tensorflow/compiler/xla/service/copy_insertion.cc
index 40519ecc79..cbe2ba2e50 100644
--- a/tensorflow/compiler/xla/service/copy_insertion.cc
+++ b/tensorflow/compiler/xla/service/copy_insertion.cc
@@ -65,7 +65,7 @@ struct SpecialCaseCopyPolicy {
   // output tuple.
   bool copy_root_replicated_buffers = false;
   // If true, insert a copy if a buffer coming from a constant or a parameter
-  // is found wihtin the output tuple.
+  // is found within the output tuple.
   bool copy_parameters_and_constants = false;
 };
 
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.cc b/tensorflow/compiler/xla/service/hlo_evaluator.cc
index 8cf94123b7..1071f5b184 100644
--- a/tensorflow/compiler/xla/service/hlo_evaluator.cc
+++ b/tensorflow/compiler/xla/service/hlo_evaluator.cc
@@ -1193,7 +1193,7 @@ class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault {
           // specifically:
           // - For lhs, the non-contracting dimensions, including the batch
           // dimension have the same index as the `result_index`.
-          // - For rhs, the batch dimension is set seperately from other
+          // - For rhs, the batch dimension is set separately from other
           // non-contracting dimensions, since these other non-contracting
           // dimensions in rhs follow the non-contracting dimensions of lhs in
           // the resulting index.
diff --git a/tensorflow/contrib/kfac/python/ops/optimizer.py b/tensorflow/contrib/kfac/python/ops/optimizer.py
index 45a760c9f1..7203804af3 100644
--- a/tensorflow/contrib/kfac/python/ops/optimizer.py
+++ b/tensorflow/contrib/kfac/python/ops/optimizer.py
@@ -114,7 +114,7 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer):
     self._estimation_mode = estimation_mode
     self._colocate_gradients_with_ops = colocate_gradients_with_ops
 
-    # The below paramaters are required only if damping needs to be adapated.
+    # The below parameters are required only if damping needs to be adapated.
     # These parameters can be set by calling
     # set_damping_adaptation_params() explicitly.
     self._damping_adaptation_decay = 0.95
diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
index 3d6042c31f..2a70c36c95 100644
--- a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
@@ -3355,7 +3355,7 @@ inline void Concatenation(int concat_dim, const uint8* const* input_data,
                           const int32 output_zeropoint,
                           const float output_scale) {
   // The arguments input_zeropoint and input_scale are expected to be an array
-  // that have the quantization paramaters for all the inputs to the concat
+  // that have the quantization parameters for all the inputs to the concat
   // operator.
   gemmlowp::ScopedProfilingLabel label("Concatenation");
   TFLITE_DCHECK_GT(inputs_count, 1);
diff --git a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
index d41ade4c9d..445687cd15 100644
--- a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
@@ -1628,7 +1628,7 @@ inline void Concatenation(int concat_dim, const uint8* const* input_data,
                           const int32 output_zeropoint,
                           const float output_scale) {
   // The arguments input_zeropoint and input_scale are expected to be an array
-  // that have the quantization paramaters for all the inputs to the concat
+  // that have the quantization parameters for all the inputs to the concat
   // operator.
   TFLITE_DCHECK_GT(inputs_count, 1);
   int64_t concat_size = 0;
diff --git a/tensorflow/contrib/lite/testing/generate_examples.py b/tensorflow/contrib/lite/testing/generate_examples.py
index 2f8f7a1a79..e4851d6077 100644
--- a/tensorflow/contrib/lite/testing/generate_examples.py
+++ b/tensorflow/contrib/lite/testing/generate_examples.py
@@ -109,7 +109,7 @@ KNOWN_BUGS = {
 
 
 class ExtraTocoOptions(object):
-  """Additonal toco options besides input, output, shape."""
+  """Additional toco options besides input, output, shape."""
 
   def __init__(self):
     # Whether to ignore control dependency nodes.
diff --git a/tensorflow/contrib/lite/toco/toco_flags.proto b/tensorflow/contrib/lite/toco/toco_flags.proto
index a04017a6bf..802cf3e2e4 100644
--- a/tensorflow/contrib/lite/toco/toco_flags.proto
+++ b/tensorflow/contrib/lite/toco/toco_flags.proto
@@ -127,7 +127,7 @@ message TocoFlags {
   // transformations that are necessary in order to generate inference
   // code for these graphs. Such graphs should be fixed, but as a
   // temporary work-around, setting this reorder_across_fake_quant flag
-  // allows toco to perform necessary graph transformaitons on them,
+  // allows toco to perform necessary graph transformations on them,
   // at the cost of no longer faithfully matching inference and training
   // arithmetic.
   optional bool reorder_across_fake_quant = 8;
diff --git a/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py b/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py
index 6cca0a8a00..bfb3350b59 100644
--- a/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py
+++ b/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py
@@ -146,7 +146,7 @@ class ModelAverageOptimizerTest(test.TestCase):
     self.assertAllEqual(1.0, sessions[0].run(global_var_1))
     self.assertAllEqual(0, sessions[0].run(global_step))
 
-    # iteration 2, global varibale update
+    # iteration 2, global variable update
     thread_0 = self.checkedThread(
         target=self._run, args=(train_ops[0], sessions[0]))
     thread_1 = self.checkedThread(
diff --git a/tensorflow/python/keras/_impl/keras/engine/training_eager.py b/tensorflow/python/keras/_impl/keras/engine/training_eager.py
index 34adeb7599..b9c99b2222 100644
--- a/tensorflow/python/keras/_impl/keras/engine/training_eager.py
+++ b/tensorflow/python/keras/_impl/keras/engine/training_eager.py
@@ -181,7 +181,7 @@ def slice_arrays(arrays, indices, contiguous=True):
   """Slices batches out of provided arrays (workaround for eager tensors).
 
   Unfortunately eager tensors don't have the same slicing behavior as
-  Numpy arrays (they folow  the same slicing behavior as symbolic TF tensors),
+  Numpy arrays (they follow the same slicing behavior as symbolic TF tensors),
   hence we cannot use `generic_utils.slice_arrays` directly
   and we have to implement this workaround based on `concat`. This has a
   performance cost.
-- 
GitLab


From f6000468263c5db7befbf5c320e8b3af7d90b819 Mon Sep 17 00:00:00 2001
From: Andrew Selle <aselle@google.com>
Date: Wed, 2 May 2018 21:15:01 -0700
Subject: [PATCH 197/395] Expose Interpreter to tensorflow.contrib.lite

PiperOrigin-RevId: 195198645
---
 tensorflow/contrib/lite/BUILD                 |  2 +
 tensorflow/contrib/lite/python/BUILD          |  1 +
 tensorflow/contrib/lite/python/interpreter.py | 18 ++++-
 .../interpreter_wrapper.cc                    |  2 +-
 tensorflow/contrib/lite/python/lite.py        |  2 +
 .../tools/pip_package/pip_smoke_test.py       | 73 +++++++++++++------
 6 files changed, 70 insertions(+), 28 deletions(-)

diff --git a/tensorflow/contrib/lite/BUILD b/tensorflow/contrib/lite/BUILD
index 1534f97d76..10065e894c 100644
--- a/tensorflow/contrib/lite/BUILD
+++ b/tensorflow/contrib/lite/BUILD
@@ -92,6 +92,8 @@ cc_library(
     deps = [":context"],
 )
 
+exports_files(["builtin_ops.h"])
+
 cc_library(
     name = "string",
     hdrs = [
diff --git a/tensorflow/contrib/lite/python/BUILD b/tensorflow/contrib/lite/python/BUILD
index e6dcc7aa09..4920e83970 100644
--- a/tensorflow/contrib/lite/python/BUILD
+++ b/tensorflow/contrib/lite/python/BUILD
@@ -44,6 +44,7 @@ py_library(
     deps = [
         ":convert",
         ":convert_saved_model",
+        ":interpreter",
         ":op_hint",
     ],
 )
diff --git a/tensorflow/contrib/lite/python/interpreter.py b/tensorflow/contrib/lite/python/interpreter.py
index cb9c0d3121..5fbc551452 100644
--- a/tensorflow/contrib/lite/python/interpreter.py
+++ b/tensorflow/contrib/lite/python/interpreter.py
@@ -17,7 +17,19 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensorflow.contrib.lite.python.interpreter_wrapper import tensorflow_wrap_interpreter_wrapper as interpreter_wrapper
+from tensorflow.python.util.lazy_loader import LazyLoader
+
+# Lazy load since some of the performance benchmark skylark rules
+# break dependencies. Must use double quotes to match code internal rewrite
+# rule.
+# pylint: disable=g-inconsistent-quotes
+_interpreter_wrapper = LazyLoader(
+    "_interpreter_wrapper", globals(),
+    "tensorflow.contrib.lite.python.interpreter_wrapper."
+    "tensorflow_wrap_interpreter_wrapper")
+# pylint: enable=g-inconsistent-quotes
+
+del LazyLoader
 
 
 class Interpreter(object):
@@ -35,13 +47,13 @@ class Interpreter(object):
     """
     if model_path and not model_content:
       self._interpreter = (
-          interpreter_wrapper.InterpreterWrapper_CreateWrapperCPPFromFile(
+          _interpreter_wrapper.InterpreterWrapper_CreateWrapperCPPFromFile(
               model_path))
       if not self._interpreter:
         raise ValueError('Failed to open {}'.format(model_path))
     elif model_content and not model_path:
       self._interpreter = (
-          interpreter_wrapper.InterpreterWrapper_CreateWrapperCPPFromBuffer(
+          _interpreter_wrapper.InterpreterWrapper_CreateWrapperCPPFromBuffer(
               model_content, len(model_content)))
       if not self._interpreter:
         raise ValueError(
diff --git a/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc b/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc
index 04fc098129..16f4f30b94 100644
--- a/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc
+++ b/tensorflow/contrib/lite/python/interpreter_wrapper/interpreter_wrapper.cc
@@ -116,7 +116,7 @@ PyObject* PyArrayFromIntVector(const int* data, npy_intp size) {
 PyObject* PyTupleFromQuantizationParam(const TfLiteQuantizationParams& param) {
   PyObject* result = PyTuple_New(2);
   PyTuple_SET_ITEM(result, 0, PyFloat_FromDouble(param.scale));
-  PyTuple_SET_ITEM(result, 1, PyInt_FromLong(param.zero_point));
+  PyTuple_SET_ITEM(result, 1, PyLong_FromLong(param.zero_point));
   return result;
 }
 
diff --git a/tensorflow/contrib/lite/python/lite.py b/tensorflow/contrib/lite/python/lite.py
index 4ea40201f7..86b25e68ac 100644
--- a/tensorflow/contrib/lite/python/lite.py
+++ b/tensorflow/contrib/lite/python/lite.py
@@ -19,6 +19,7 @@ EXPERIMENTAL: APIs here are unstable and likely to change without notice.
 @@toco_convert
 @@toco_convert_protos
 @@tflite_from_saved_model
+@@Interpreter
 @@OpHint
 @@convert_op_hints_to_stubs
 
@@ -31,6 +32,7 @@ from __future__ import print_function
 from tensorflow.contrib.lite.python.convert import toco_convert
 from tensorflow.contrib.lite.python.convert import toco_convert_protos
 from tensorflow.contrib.lite.python.convert_saved_model import tflite_from_saved_model
+from tensorflow.contrib.lite.python.interpreter import Interpreter
 from tensorflow.contrib.lite.python.op_hint import convert_op_hints_to_stubs
 from tensorflow.contrib.lite.python.op_hint import OpHint
 # pylint: enable=unused-import
diff --git a/tensorflow/tools/pip_package/pip_smoke_test.py b/tensorflow/tools/pip_package/pip_smoke_test.py
index b23dde2019..401f833dbd 100644
--- a/tensorflow/tools/pip_package/pip_smoke_test.py
+++ b/tensorflow/tools/pip_package/pip_smoke_test.py
@@ -30,15 +30,42 @@ os.chdir(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")))
 PIP_PACKAGE_QUERY_EXPRESSION = (
     "deps(//tensorflow/tools/pip_package:build_pip_package)")
 
-# pylint: disable=g-backslash-continuation
-PY_TEST_QUERY_EXPRESSION = 'deps(\
-  filter("^((?!benchmark).)*$",\
-  kind(py_test,\
-  //tensorflow/python/... \
-  + //tensorflow/contrib/... \
-  - //tensorflow/contrib/tensorboard/... \
-  - attr(tags, "manual|no_pip", //tensorflow/...))), 1)'
-# pylint: enable=g-backslash-continuation
+
+def GetBuild(dir_base):
+  """Get the list of BUILD file all targets recursively startind at dir_base."""
+  items = []
+  for root, _, files in os.walk(dir_base):
+    for name in files:
+      if (name == "BUILD" and
+          root.find("tensorflow/contrib/lite/examples/android") == -1):
+        items.append("//" + root + ":all")
+  return items
+
+
+def BuildPyTestDependencies():
+  python_targets = GetBuild("tensorflow/python")
+  contrib_targets = GetBuild("tensorflow/contrib")
+  tensorboard_targets = GetBuild("tensorflow/contrib/tensorboard")
+  tensorflow_targets = GetBuild("tensorflow")
+  # Build list of test targets,
+  # python + contrib - tensorboard - attr(manual|pno_pip)
+  targets = " + ".join(python_targets)
+  for t in contrib_targets:
+    targets += " + " + t
+  for t in tensorboard_targets:
+    targets += " - " + t
+  targets += ' - attr(tags, "manual|no_pip", %s)' % " + ".join(
+      tensorflow_targets)
+  query_kind = "kind(py_test, %s)" % targets
+  # Skip benchmarks etc.
+  query_filter = 'filter("^((?!benchmark).)*$", %s)' % query_kind
+  # Get the dependencies
+  query_deps = "deps(%s, 1)" % query_filter
+
+  return python_targets, query_deps
+
+
+PYTHON_TARGETS, PY_TEST_QUERY_EXPRESSION = BuildPyTestDependencies()
 
 # Hard-coded blacklist of files if not included in pip package
 # TODO(amitpatankar): Clean up blacklist.
@@ -79,16 +106,6 @@ BLACKLIST = [
 ]
 
 
-def bazel_query(query_target):
-  """Run bazel query on target."""
-  try:
-    output = subprocess.check_output(
-        ["bazel", "query", "--keep_going", query_target])
-  except subprocess.CalledProcessError as e:
-    output = e.output
-  return output
-
-
 def main():
   """This script runs the pip smoke test.
 
@@ -103,14 +120,22 @@ def main():
   """
 
   # pip_package_dependencies_list is the list of included files in pip packages
-  pip_package_dependencies = bazel_query(PIP_PACKAGE_QUERY_EXPRESSION)
+  pip_package_dependencies = subprocess.check_output(
+      ["bazel", "cquery", PIP_PACKAGE_QUERY_EXPRESSION])
   pip_package_dependencies_list = pip_package_dependencies.strip().split("\n")
+  pip_package_dependencies_list = [
+      x.split()[0] for x in pip_package_dependencies_list
+  ]
   print("Pip package superset size: %d" % len(pip_package_dependencies_list))
 
   # tf_py_test_dependencies is the list of dependencies for all python
   # tests in tensorflow
-  tf_py_test_dependencies = bazel_query(PY_TEST_QUERY_EXPRESSION)
+  tf_py_test_dependencies = subprocess.check_output(
+      ["bazel", "cquery", PY_TEST_QUERY_EXPRESSION])
   tf_py_test_dependencies_list = tf_py_test_dependencies.strip().split("\n")
+  tf_py_test_dependencies_list = [
+      x.split()[0] for x in tf_py_test_dependencies.strip().split("\n")
+  ]
   print("Pytest dependency subset size: %d" % len(tf_py_test_dependencies_list))
 
   missing_dependencies = []
@@ -141,9 +166,9 @@ def main():
     for missing_dependency in missing_dependencies:
       print("\nMissing dependency: %s " % missing_dependency)
       print("Affected Tests:")
-      rdep_query = ("rdeps(kind(py_test, //tensorflow/python/...), %s)" %
-                    missing_dependency)
-      affected_tests = bazel_query(rdep_query)
+      rdep_query = ("rdeps(kind(py_test, %s), %s)" %
+                    (" + ".join(PYTHON_TARGETS), missing_dependency))
+      affected_tests = subprocess.check_output(["bazel", "cquery", rdep_query])
       affected_tests_list = affected_tests.split("\n")[:-2]
       print("\n".join(affected_tests_list))
 
-- 
GitLab


From 985351dc1ab33cedbfd7790dd9cccc36d2d4b150 Mon Sep 17 00:00:00 2001
From: Tony Wang <tonywy@google.com>
Date: Wed, 2 May 2018 21:37:04 -0700
Subject: [PATCH 198/395] Simplify getter and setter method for
 GraphOptimizationPass::name_

PiperOrigin-RevId: 195199912
---
 .../core/common_runtime/optimization_registry.cc     |  6 ++----
 .../core/common_runtime/optimization_registry.h      | 12 ++++++------
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/tensorflow/core/common_runtime/optimization_registry.cc b/tensorflow/core/common_runtime/optimization_registry.cc
index bf49a758b2..6ac047295d 100644
--- a/tensorflow/core/common_runtime/optimization_registry.cc
+++ b/tensorflow/core/common_runtime/optimization_registry.cc
@@ -36,8 +36,7 @@ Status OptimizationPassRegistry::RunGrouping(
     for (auto& phase : group->second) {
       VLOG(1) << "Running optimization phase " << phase.first;
       for (auto& pass : phase.second) {
-        VLOG(1) << "Running optimization pass: "
-                << pass->GetOptimizationPassName();
+        VLOG(1) << "Running optimization pass: " << pass->name();
         Status s = pass->Run(options);
         if (!s.ok()) return s;
       }
@@ -52,8 +51,7 @@ void OptimizationPassRegistry::LogGrouping(Grouping grouping, int vlog_level) {
     for (auto& phase : group->second) {
       for (auto& pass : phase.second) {
         VLOG(vlog_level) << "Registered optimization pass grouping " << grouping
-                         << " phase " << phase.first << ": "
-                         << pass->GetOptimizationPassName();
+                         << " phase " << phase.first << ": " << pass->name();
       }
     }
   }
diff --git a/tensorflow/core/common_runtime/optimization_registry.h b/tensorflow/core/common_runtime/optimization_registry.h
index 1b535faf19..f5d265aa24 100644
--- a/tensorflow/core/common_runtime/optimization_registry.h
+++ b/tensorflow/core/common_runtime/optimization_registry.h
@@ -65,13 +65,13 @@ class GraphOptimizationPass {
  public:
   virtual ~GraphOptimizationPass() {}
   virtual Status Run(const GraphOptimizationPassOptions& options) = 0;
-  void SetOptimizationPassName(string name) { _optimization_pass_name = name; }
-  string GetOptimizationPassName() { return _optimization_pass_name; }
+  void set_name(const string& name) { name_ = name; }
+  string name() const { return name_; }
 
  private:
-  // The name of the opitmization pass, which is the same as the inherited class
-  // name.
-  string _optimization_pass_name;
+  // The name of the opitimization pass, which is the same as the inherited
+  // class name.
+  string name_;
 };
 
 // The key is a 'phase' number. Phases are executed in increasing
@@ -118,7 +118,7 @@ class OptimizationPassRegistration {
                                int phase,
                                std::unique_ptr<GraphOptimizationPass> pass,
                                string optimization_pass_name) {
-    pass->SetOptimizationPassName(optimization_pass_name);
+    pass->set_name(optimization_pass_name);
     OptimizationPassRegistry::Global()->Register(grouping, phase,
                                                  std::move(pass));
   }
-- 
GitLab


From 17d877b21af1a3b99fd20b8ede8196040ac37486 Mon Sep 17 00:00:00 2001
From: Elson Rodriguez <elson.rodriguez@gmail.com>
Date: Wed, 2 May 2018 21:58:44 -0700
Subject: [PATCH 199/395] Enabling support for S3 and Google Storage for the
 MKL Docker image. (#19039)

---
 tensorflow/tools/docker/Dockerfile.devel-cpu-mkl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/tools/docker/Dockerfile.devel-cpu-mkl b/tensorflow/tools/docker/Dockerfile.devel-cpu-mkl
index c65e0b72bc..a6cd44ced1 100644
--- a/tensorflow/tools/docker/Dockerfile.devel-cpu-mkl
+++ b/tensorflow/tools/docker/Dockerfile.devel-cpu-mkl
@@ -35,10 +35,10 @@ ENV CI_BUILD_PYTHON=python \
     PYTHON_LIB_PATH=/usr/local/lib/python2.7/dist-packages \
     CC_OPT_FLAGS='-march=native' \
     TF_NEED_JEMALLOC=0 \
-    TF_NEED_GCP=0 \
+    TF_NEED_GCP=1 \
     TF_NEED_CUDA=0 \
     TF_NEED_HDFS=0 \
-    TF_NEED_S3=0 \
+    TF_NEED_S3=1 \
     TF_NEED_OPENCL=0 \
     TF_NEED_GDR=0 \
     TF_ENABLE_XLA=0 \
-- 
GitLab


From 283e8fe7e191f8e0e2ca6ea62b8b4553c30a6286 Mon Sep 17 00:00:00 2001
From: Suharsh Sivakumar <suharshs@google.com>
Date: Thu, 3 May 2018 00:16:09 -0700
Subject: [PATCH 200/395] Use tensorflow size to determine number of elements
 instead of the static shape, which can sometimes be missing.

PiperOrigin-RevId: 195209826
---
 tensorflow/contrib/quantize/python/fold_batch_norms.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/contrib/quantize/python/fold_batch_norms.py b/tensorflow/contrib/quantize/python/fold_batch_norms.py
index 1f286bc39a..76f695dce0 100644
--- a/tensorflow/contrib/quantize/python/fold_batch_norms.py
+++ b/tensorflow/contrib/quantize/python/fold_batch_norms.py
@@ -414,7 +414,8 @@ def _CloneWithNewOperands(layer_op, input_tensor, weight_tensor):
 def _FoldFusedBatchNormGrad(op, unused_grad_y, grad_mean, grad_var, unused_1,
                             unused_2):
   x = op.inputs[0]
-  n = x.get_shape().num_elements() / grad_mean.get_shape().num_elements()
+  n = math_ops.cast(
+      array_ops.size(x) / array_ops.size(grad_mean), dtypes.float32)
   dmean_dx = grad_mean / n
   dvar_dx = 2 * grad_var * (x - op.outputs[1]) / (n - 1)
   return (dmean_dx + dvar_dx), None, None, None, None
-- 
GitLab


From 090d21c18f16303e740136e8a4e0f62c63df4e10 Mon Sep 17 00:00:00 2001
From: wangsiyu <siyu.wsy@gmail.com>
Date: Thu, 3 May 2018 18:31:29 +0800
Subject: [PATCH 201/395] fix bug of declaring regularization loss mutiple
 times when reusing partitioned variables in tf.layers

---
 tensorflow/python/layers/base.py      | 13 ++++++++++++-
 tensorflow/python/layers/base_test.py | 15 +++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/layers/base.py b/tensorflow/python/layers/base.py
index 64db49c900..c050e6be04 100644
--- a/tensorflow/python/layers/base.py
+++ b/tensorflow/python/layers/base.py
@@ -233,7 +233,8 @@ class Layer(base_layer.Layer):
             getter=vs.get_variable)
 
         if regularizer:
-          if context.executing_eagerly() or variable not in existing_variables:
+          if context.executing_eagerly() or _should_add_regularizer(
+              variable, existing_variables):
             self._handle_weight_regularization(name, variable, regularizer)
 
         if init_graph is not None:
@@ -354,3 +355,13 @@ def _add_elements_to_collection(elements, collection_list):
       if element not in collection_set:
         collection.append(element)
 
+def _should_add_regularizer(variable, existing_variable_set):
+  result = True
+  if isinstance(variable, tf_variables.PartitionedVariable):
+    for var in variable._get_variable_list():
+      if var in existing_variable_set:
+        result = False
+        break
+  else:
+    result = variable not in existing_variable_set
+  return result
diff --git a/tensorflow/python/layers/base_test.py b/tensorflow/python/layers/base_test.py
index f08b552840..361e3de7aa 100644
--- a/tensorflow/python/layers/base_test.py
+++ b/tensorflow/python/layers/base_test.py
@@ -30,6 +30,7 @@ from tensorflow.python.layers import core as core_layers
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import init_ops
 from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import partitioned_variables
 from tensorflow.python.ops import random_ops
 from tensorflow.python.ops import state_ops
 from tensorflow.python.ops import variable_scope
@@ -95,6 +96,20 @@ class BaseLayerTest(test.TestCase):
           regularizer=regularizer)
       self.assertEqual(len(layer.losses), 1)
 
+  def testReusePartitionedVaraiblesAndRegularizers(self):
+    regularizer = lambda x: math_ops.reduce_sum(x) * 1e-3
+    partitioner = partitioned_variables.fixed_size_partitioner(3)
+    for i in xrange(2):
+      with variable_scope.variable_scope(variable_scope.get_variable_scope(),
+                                         partitioner=partitioner,
+                                         reuse=False if i == 0 else True):
+        layer = base_layers.Layer(name='my_layer')
+        variable = layer.add_variable(
+            'reg_part_var', [4, 4],
+            initializer=init_ops.zeros_initializer(),
+            regularizer=regularizer)
+    self.assertEqual(len(ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)), 3)
+
   def testNoEagerActivityRegularizer(self):
     with context.eager_mode():
       with self.assertRaisesRegexp(ValueError, 'activity_regularizer'):
-- 
GitLab


From 3e68ec6cd2ce5c9f69c83122d854dccc8ee7ff6a Mon Sep 17 00:00:00 2001
From: Letian Feng <letian.feng@hotmail.com>
Date: Thu, 3 May 2018 15:46:16 +0200
Subject: [PATCH 202/395] correct code snippets to python3 style (#19052)

---
 tensorflow/docs_src/programmers_guide/tensors.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/docs_src/programmers_guide/tensors.md b/tensorflow/docs_src/programmers_guide/tensors.md
index 58a80d5339..1248c3cabe 100644
--- a/tensorflow/docs_src/programmers_guide/tensors.md
+++ b/tensorflow/docs_src/programmers_guide/tensors.md
@@ -265,7 +265,7 @@ example:
 ```python
 constant = tf.constant([1, 2, 3])
 tensor = constant * constant
-print tensor.eval()
+print(tensor.eval())
 ```
 
 The `eval` method only works when a default `tf.Session` is active (see
@@ -306,8 +306,8 @@ Note that you rarely want to use the following pattern when printing a
 
 ``` python
 t = <<some tensorflow operation>>
-print t  # This will print the symbolic tensor when the graph is being built.
-         # This tensor does not have a value in this context.
+print(t)  # This will print the symbolic tensor when the graph is being built.
+          # This tensor does not have a value in this context.
 ```
 
 This code prints the `tf.Tensor` object (which represents deferred computation)
-- 
GitLab


From 4984a60e7147edef532ca1b15050471e81e45841 Mon Sep 17 00:00:00 2001
From: ctiijima <ctiijima@us.ibm.com>
Date: Thu, 3 May 2018 08:34:07 -0700
Subject: [PATCH 203/395] Grammar fixes on architecture.md (#19035)

---
 tensorflow/docs_src/extend/architecture.md | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensorflow/docs_src/extend/architecture.md b/tensorflow/docs_src/extend/architecture.md
index c0fc714a44..c8f522a03a 100644
--- a/tensorflow/docs_src/extend/architecture.md
+++ b/tensorflow/docs_src/extend/architecture.md
@@ -4,8 +4,8 @@ We designed TensorFlow for large-scale distributed training and inference, but
 it is also flexible enough to support experimentation with new machine
 learning models and system-level optimizations.
 
-This document describes the system architecture that makes possible this
-combination of scale and flexibility. It assumes that you have basic familiarity
+This document describes the system architecture that makes this
+combination of scale and flexibility possible. It assumes that you have basic familiarity
 with TensorFlow programming concepts such as the computation graph, operations,
 and sessions. See @{$programmers_guide/low_level_intro$this document}
 for an introduction to these topics. Some familiarity
@@ -15,8 +15,8 @@ will also be helpful.
 This document is for developers who want to extend TensorFlow in some way not
 supported by current APIs, hardware engineers who want to optimize for
 TensorFlow, implementers of machine learning systems working on scaling and
-distribution, or anyone who wants to look under Tensorflow's hood. After
-reading it you should understand TensorFlow architecture well enough to read
+distribution, or anyone who wants to look under Tensorflow's hood. By the end of this document 
+you should understand the TensorFlow architecture well enough to read
 and modify the core TensorFlow code.
 
 ## Overview
@@ -35,7 +35,7 @@ This document focuses on the following layers:
 *  **Client**:
    *  Defines the computation as a dataflow graph.
    *  Initiates graph execution using a [**session**](
-      https://www.tensorflow.org/code/tensorflow/python/client/session.py)
+      https://www.tensorflow.org/code/tensorflow/python/client/session.py).
 *  **Distributed Master**
    *  Prunes a specific subgraph from the graph, as defined by the arguments
       to Session.run().
@@ -55,7 +55,7 @@ Figure 2 illustrates the interaction of these components. "/job:worker/task:0" a
 server": a task responsible for storing and updating the model's parameters.
 Other tasks send updates to these parameters as they work on optimizing the
 parameters. This particular division of labor between tasks is not required, but
-it is common for distributed training.
+ is common for distributed training.
 
 ![TensorFlow Architecture Diagram](https://www.tensorflow.org/images/diag1.svg){: width="500"}
 
@@ -193,7 +193,7 @@ https://www.tensorflow.org/code/tensorflow/contrib/nccl/python/ops/nccl_ops.py))
 
 ## Kernel Implementations
 
-The runtime contains over 200 standard operations, including mathematical, array
+The runtime contains over 200 standard operations including mathematical, array
 manipulation, control flow, and state management operations. Each of these
 operations can have kernel implementations optimized for a variety of devices.
 Many of the operation kernels are implemented using Eigen::Tensor, which uses
-- 
GitLab


From a88a7e312581ba0c2173188019a420c888df9a10 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 3 May 2018 09:10:06 -0700
Subject: [PATCH 204/395] Post-transform pass to dedupe large constant arrays.

PiperOrigin-RevId: 195260578
---
 tensorflow/contrib/lite/toco/args.h           |   1 +
 .../contrib/lite/toco/toco_cmdline_flags.cc   |   6 +
 tensorflow/contrib/lite/toco/toco_flags.proto |   6 +-
 tensorflow/contrib/lite/toco/toco_tooling.cc  |   5 +
 tensorflow/contrib/lite/toco/tooling_util.cc  | 130 ++++++++++++++++++
 tensorflow/contrib/lite/toco/tooling_util.h   |  16 +++
 6 files changed, 163 insertions(+), 1 deletion(-)

diff --git a/tensorflow/contrib/lite/toco/args.h b/tensorflow/contrib/lite/toco/args.h
index fe30b88344..6c0311af0a 100644
--- a/tensorflow/contrib/lite/toco/args.h
+++ b/tensorflow/contrib/lite/toco/args.h
@@ -241,6 +241,7 @@ struct ParsedTocoFlags {
   Arg<bool> drop_control_dependency = Arg<bool>(false);
   Arg<bool> propagate_fake_quant_num_bits = Arg<bool>(false);
   Arg<bool> allow_nudging_weights_to_use_fast_gemm_kernel = Arg<bool>(false);
+  Arg<int64> dedupe_array_min_size_bytes = Arg<int64>(64);
 };
 
 }  // namespace toco
diff --git a/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc b/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc
index 1611c4d0c0..7786a4ada3 100644
--- a/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc
+++ b/tensorflow/contrib/lite/toco/toco_cmdline_flags.cc
@@ -148,6 +148,11 @@ bool ParseTocoFlagsFromCommandLineFlags(
            "Some fast uint8 GEMM kernels require uint8 weights to avoid the "
            "value 0. This flag allows nudging them to 1 to allow proceeding, "
            "with moderate inaccuracy."),
+      Flag("dedupe_array_min_size_bytes",
+           parsed_flags.dedupe_array_min_size_bytes.bind(),
+           parsed_flags.dedupe_array_min_size_bytes.default_value(),
+           "Minimum size of constant arrays to deduplicate; arrays smaller "
+           "will not be deduplicated."),
   };
   bool asked_for_help =
       *argc == 2 && (!strcmp(argv[1], "--help") || !strcmp(argv[1], "-help"));
@@ -239,6 +244,7 @@ void ReadTocoFlagsFromCommandLineFlags(const ParsedTocoFlags& parsed_toco_flags,
   READ_TOCO_FLAG(propagate_fake_quant_num_bits, FlagRequirement::kNone);
   READ_TOCO_FLAG(allow_nudging_weights_to_use_fast_gemm_kernel,
                  FlagRequirement::kNone);
+  READ_TOCO_FLAG(dedupe_array_min_size_bytes, FlagRequirement::kNone);
 
   // Deprecated flag handling.
   if (parsed_toco_flags.input_type.specified()) {
diff --git a/tensorflow/contrib/lite/toco/toco_flags.proto b/tensorflow/contrib/lite/toco/toco_flags.proto
index a04017a6bf..253f022e6b 100644
--- a/tensorflow/contrib/lite/toco/toco_flags.proto
+++ b/tensorflow/contrib/lite/toco/toco_flags.proto
@@ -37,7 +37,7 @@ enum FileFormat {
 // of as properties of models, instead describing how models are to be
 // processed in the context of the present tooling job.
 //
-// Next ID to use: 18.
+// Next ID to use: 19.
 message TocoFlags {
   // Input file format
   optional FileFormat input_format = 1;
@@ -161,4 +161,8 @@ message TocoFlags {
   // This flag allows nudging them to 1 to allow proceeding, with moderate
   // inaccuracy.
   optional bool allow_nudging_weights_to_use_fast_gemm_kernel = 17;
+
+  // Minimum size of constant arrays to deduplicate; arrays smaller will not be
+  // deduplicated.
+  optional int64 dedupe_array_min_size_bytes = 18 [default = 64];
 }
diff --git a/tensorflow/contrib/lite/toco/toco_tooling.cc b/tensorflow/contrib/lite/toco/toco_tooling.cc
index 7252ec2ea4..6973b22c5a 100644
--- a/tensorflow/contrib/lite/toco/toco_tooling.cc
+++ b/tensorflow/contrib/lite/toco/toco_tooling.cc
@@ -345,6 +345,11 @@ void Transform(const TocoFlags& toco_flags, Model* model) {
     EncodeConstantArraysMinMaxByWrappingThemInFakeQuantNodes(model);
   }
 
+  // Deduplicate large constant arrays.
+  if (toco_flags.has_dedupe_array_min_size_bytes()) {
+    DedupeConstantArrays(model, toco_flags.dedupe_array_min_size_bytes());
+  }
+
   LogDump(kLogLevelModelChanged, "AFTER TRANSFORMATIONS", *model);
 
   if (output_format != GRAPHVIZ_DOT && output_format != TFLITE) {
diff --git a/tensorflow/contrib/lite/toco/tooling_util.cc b/tensorflow/contrib/lite/toco/tooling_util.cc
index 11293a5fe5..86ee1f3761 100644
--- a/tensorflow/contrib/lite/toco/tooling_util.cc
+++ b/tensorflow/contrib/lite/toco/tooling_util.cc
@@ -260,6 +260,23 @@ Operator* GetFirstOpWithInput(const Model& model, const string& array_name) {
   return it == model.operators.end() ? nullptr : it->get();
 }
 
+void ReplaceArrayUsage(Model* model, const string& old_array_name,
+                       const string& new_array_name) {
+  for (auto& op_it : model->operators) {
+    Operator* op = op_it.get();
+    for (size_t i = 0; i < op->inputs.size(); ++i) {
+      if (op->inputs[i] == old_array_name) {
+        op->inputs[i] = new_array_name;
+      }
+    }
+    for (size_t i = 0; i < op->outputs.size(); ++i) {
+      if (op->outputs[i] == old_array_name) {
+        op->outputs[i] = new_array_name;
+      }
+    }
+  }
+}
+
 string FormatArraysList(const Model& model, const std::vector<string>& list) {
   if (list.empty()) {
     return "[]";
@@ -648,6 +665,65 @@ bool IsConstantParameterArray(const Model& model, const string& name) {
   return !!model.GetArray(name).buffer;
 }
 
+namespace {
+template <ArrayDataType A>
+bool CompareArrayBuffers(const Array& lhs_array, const Array& rhs_array) {
+  CHECK(lhs_array.data_type == rhs_array.data_type) << "Data types must match";
+  CHECK(lhs_array.buffer) << "LHS must be constant";
+  CHECK(rhs_array.buffer) << "RHS must be constant";
+  const auto& lhs_data = lhs_array.GetBuffer<A>().data;
+  const auto& rhs_data = rhs_array.GetBuffer<A>().data;
+  CHECK_EQ(lhs_data.size(), rhs_data.size())
+      << "Buffer sizes must match in element count";
+  for (int i = 0; i < lhs_data.size(); ++i) {
+    if (lhs_data[i] != rhs_data[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+}  // namespace
+
+bool CompareConstantArrays(const Array& lhs_array, const Array& rhs_array) {
+  bool attrs_equal =
+      lhs_array.shape() == rhs_array.shape() &&
+      lhs_array.data_type == rhs_array.data_type &&
+      lhs_array.final_data_type == rhs_array.final_data_type &&
+      lhs_array.minmax == rhs_array.minmax &&
+      lhs_array.quantization_params == rhs_array.quantization_params;
+  if (!attrs_equal) {
+    return false;
+  }
+  switch (lhs_array.data_type) {
+    case ArrayDataType::kBool:
+      return CompareArrayBuffers<ArrayDataType::kBool>(lhs_array, rhs_array);
+    case ArrayDataType::kFloat:
+      return CompareArrayBuffers<ArrayDataType::kFloat>(lhs_array, rhs_array);
+    case ArrayDataType::kInt8:
+      return CompareArrayBuffers<ArrayDataType::kInt8>(lhs_array, rhs_array);
+    case ArrayDataType::kUint8:
+      return CompareArrayBuffers<ArrayDataType::kUint8>(lhs_array, rhs_array);
+    case ArrayDataType::kInt16:
+      return CompareArrayBuffers<ArrayDataType::kInt16>(lhs_array, rhs_array);
+    case ArrayDataType::kUint16:
+      return CompareArrayBuffers<ArrayDataType::kUint16>(lhs_array, rhs_array);
+    case ArrayDataType::kInt32:
+      return CompareArrayBuffers<ArrayDataType::kInt32>(lhs_array, rhs_array);
+    case ArrayDataType::kUint32:
+      return CompareArrayBuffers<ArrayDataType::kUint32>(lhs_array, rhs_array);
+    case ArrayDataType::kInt64:
+      return CompareArrayBuffers<ArrayDataType::kInt64>(lhs_array, rhs_array);
+    case ArrayDataType::kUint64:
+      return CompareArrayBuffers<ArrayDataType::kUint64>(lhs_array, rhs_array);
+    case ArrayDataType::kString:
+      return CompareArrayBuffers<ArrayDataType::kString>(lhs_array, rhs_array);
+    default:
+      LOG(FATAL) << "Unsupported data type: "
+                 << ArrayDataTypeName(lhs_array.data_type);
+      return false;
+  }
+}
+
 namespace {
 // Take an array name, which may be something like "name:3_5" and make it
 // acceptable as a TF node name, say "name_3_5";
@@ -1072,6 +1148,60 @@ void FixEdgeArrays(Model* model) {
   }
 }
 
+void DedupeConstantArrays(Model* model, size_t min_size) {
+  // Walk all 0..N and compare with the remaining n+1..N.
+  // This lets us avoid N^2 comparisions and erase duplicate arrays while
+  // iterating.
+  const auto& array_map = model->GetArrayMap();
+  for (auto lhs_array_it = array_map.begin(); lhs_array_it != array_map.end();
+       ++lhs_array_it) {
+    const auto& lhs_array_name = lhs_array_it->first;
+    const auto& lhs_array = *lhs_array_it->second;
+    if (!IsConstantParameterArray(*model, lhs_array_name)) {
+      // Not a constant array; skip.
+      continue;
+    }
+    ArrayDataType final_data_type =
+        lhs_array.final_data_type != ArrayDataType::kNone
+            ? lhs_array.final_data_type
+            : lhs_array.data_type;
+    size_t array_byte_size =
+        lhs_array.buffer->Length() * ElementSize(final_data_type);
+    if (array_byte_size < min_size) {
+      // Too small; skip.
+      continue;
+    }
+
+    auto next_lhs_array_it = lhs_array_it;
+    ++next_lhs_array_it;
+    for (auto rhs_array_it = next_lhs_array_it;
+         rhs_array_it != array_map.end();) {
+      const auto& rhs_array_name = rhs_array_it->first;
+      const auto& rhs_array = *rhs_array_it->second;
+      ++rhs_array_it;
+      if (!IsConstantParameterArray(*model, rhs_array_name)) {
+        // Not a constant array; skip.
+        continue;
+      }
+      if (!IsDiscardableArray(*model, rhs_array_name)) {
+        // Can't remove the array as it's not discardable (such as an IO edge).
+        continue;
+      }
+      if (!CompareConstantArrays(lhs_array, rhs_array)) {
+        // Arrays aren't equal; skip.
+        continue;
+      }
+
+      // Arrays can be deduped!
+      VLOG(1) << "Deduplicating arrays; using " << lhs_array_name
+              << " in place of " << rhs_array_name;
+      ReplaceArrayUsage(model, rhs_array_name, lhs_array_name);
+      // Note: rhs_array_it above is already incremented so this is safe.
+      model->EraseArray(rhs_array_name);
+    }
+  }
+}
+
 namespace {
 void CopyArrayAttribs(const Array& source_array, Array* target_array) {
   target_array->data_type = source_array.data_type;
diff --git a/tensorflow/contrib/lite/toco/tooling_util.h b/tensorflow/contrib/lite/toco/tooling_util.h
index f5b596df0f..1f596ca8e5 100644
--- a/tensorflow/contrib/lite/toco/tooling_util.h
+++ b/tensorflow/contrib/lite/toco/tooling_util.h
@@ -88,6 +88,10 @@ std::vector<std::unique_ptr<Operator>>::iterator FindOpWithInput(
 Operator* GetOpWithInput(const Model& model, const string& array_name);
 Operator* GetFirstOpWithInput(const Model& model, const string& array_name);
 
+// Replaces all uses of the |old_array_name| with the |new_array_name|.
+void ReplaceArrayUsage(Model* model, const string& old_array_name,
+                       const string& new_array_name);
+
 std::vector<std::unique_ptr<Operator>>::const_iterator FindOp(
     const Model& model, const Operator* op);
 std::vector<std::unique_ptr<Operator>>::iterator FindOp(Model& model,
@@ -138,6 +142,9 @@ int RequiredBufferSizeForShape(const Shape& shape);
 
 bool IsConstantParameterArray(const Model& model, const string& name);
 
+// Compares two constant parameter arrays for exact equality.
+bool CompareConstantArrays(const Array& lhs_array, const Array& rhs_array);
+
 void CheckNoMissingArray(const Model& model);
 void CheckInvariants(const Model& model);
 
@@ -150,6 +157,15 @@ void FixNoOrphanedArray(Model* model);
 // Fixes input/output arrays that may have issues during export or inference.
 void FixEdgeArrays(Model* model);
 
+// Finds and deduplicates large constant arrays in the model.
+// After constant propagation runs it's possible to end up with several of the
+// same large array (whether they be zeros or otherwise).
+//
+// |min_size| is used to adjust the minimum size in bytes of an array before
+// it's considered for deduping. As deduping can make the graphs more difficult
+// to read this helps prevent small arrays from spidering out.
+void DedupeConstantArrays(Model* model, size_t min_size);
+
 // Copies the contents of an array into another.
 // Expects that the shape and data type match.
 template <ArrayDataType A>
-- 
GitLab


From 487fa7b1a48c151362ab1b16cdda6bbc78f5d6dc Mon Sep 17 00:00:00 2001
From: "Nicholas Nadeau, P.Eng., AVS" <nnadeau@users.noreply.github.com>
Date: Thu, 3 May 2018 13:47:06 -0400
Subject: [PATCH 205/395] Fixed Typos (#18806)

* fixed typos
---
 RELEASE.md                                                    | 2 +-
 tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc         | 2 +-
 tensorflow/compiler/xla/literal_util.h                        | 2 +-
 tensorflow/compiler/xla/service/conditional_simplifier.cc     | 2 +-
 tensorflow/compiler/xla/service/cpu/ir_function.h             | 4 ++--
 tensorflow/compiler/xla/service/cpu/shape_partition.h         | 2 +-
 tensorflow/compiler/xla/service/despecializer.h               | 2 +-
 tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h     | 2 +-
 tensorflow/compiler/xla/service/interpreter/README.md         | 2 +-
 tensorflow/compiler/xla/service/layout_assignment.h           | 4 ++--
 tensorflow/compiler/xla/service/reduce_precision_insertion.cc | 2 +-
 tensorflow/compiler/xla/service/source_map_util.h             | 2 +-
 tensorflow/contrib/autograph/impl/config.py                   | 2 +-
 tensorflow/contrib/autograph/operators/control_flow.py        | 2 +-
 .../boosted_trees/python/training/functions/gbdt_batch.py     | 2 +-
 .../python/ops/bijectors/cholesky_outer_product.py            | 2 +-
 tensorflow/contrib/eager/README.md                            | 2 +-
 tensorflow/contrib/ffmpeg/ffmpeg_lib.h                        | 2 +-
 .../contrib/framework/python/ops/critical_section_ops.py      | 2 +-
 .../contrib/gan/python/features/python/conditioning_utils.py  | 2 +-
 tensorflow/contrib/graph_editor/transform.py                  | 2 +-
 tensorflow/contrib/image/__init__.py                          | 2 +-
 tensorflow/contrib/kfac/examples/convnet.py                   | 2 +-
 tensorflow/contrib/kfac/python/ops/optimizer.py               | 4 ++--
 tensorflow/contrib/kfac/python/ops/placement.py               | 2 +-
 .../contrib/lite/kernels/internal/reference/reference_ops.h   | 2 +-
 tensorflow/contrib/lite/schema/schema.fbs                     | 2 +-
 tensorflow/contrib/lite/schema/schema_v0.fbs                  | 2 +-
 tensorflow/contrib/lite/schema/schema_v1.fbs                  | 2 +-
 tensorflow/contrib/lite/schema/schema_v2.fbs                  | 2 +-
 tensorflow/contrib/lite/schema/schema_v3.fbs                  | 4 ++--
 tensorflow/contrib/lite/testing/generate_examples.py          | 4 ++--
 tensorflow/contrib/lite/testing/tflite_driver.cc              | 4 ++--
 tensorflow/contrib/lite/toco/g3doc/cmdline_examples.md        | 4 ++--
 tensorflow/contrib/lite/toco/import_tensorflow.cc             | 2 +-
 tensorflow/contrib/lite/toco/tflite/operator.h                | 4 ++--
 tensorflow/contrib/lite/toco/tflite/types_test.cc             | 2 +-
 .../opt/python/training/elastic_average_optimizer_test.py     | 2 +-
 .../opt/python/training/model_average_optimizer_test.py       | 2 +-
 tensorflow/contrib/tensorrt/convert/convert_nodes.cc          | 2 +-
 tensorflow/contrib/verbs/README.md                            | 2 +-
 tensorflow/core/common_runtime/broadcaster.cc                 | 4 ++--
 tensorflow/core/common_runtime/buf_rendezvous.h               | 2 +-
 tensorflow/core/common_runtime/ring_reducer.cc                | 2 +-
 tensorflow/core/common_runtime/scoped_allocator_mgr.cc        | 2 +-
 tensorflow/core/debug/debug_io_utils.cc                       | 2 +-
 tensorflow/core/distributed_runtime/rpc/grpc_worker_cache.cc  | 2 +-
 tensorflow/core/framework/op_gen_lib.h                        | 4 ++--
 tensorflow/core/framework/op_kernel.h                         | 2 +-
 tensorflow/core/graph/while_context.h                         | 2 +-
 tensorflow/core/grappler/costs/graph_properties.cc            | 3 +--
 tensorflow/core/grappler/costs/virtual_scheduler.h            | 2 +-
 tensorflow/core/grappler/optimizers/layout_optimizer.cc       | 2 +-
 .../kernels/batching_util/adaptive_shared_batch_scheduler.h   | 2 +-
 tensorflow/core/kernels/conv_ops_gpu_3.cu.cc                  | 2 +-
 tensorflow/core/kernels/nth_element_op.cc                     | 2 +-
 tensorflow/core/kernels/roll_op.cc                            | 2 +-
 tensorflow/core/platform/cloud/gcs_file_system.cc             | 2 +-
 tensorflow/core/platform/cloud/gcs_throttle.h                 | 2 +-
 tensorflow/core/profiler/g3doc/command_line.md                | 2 +-
 tensorflow/core/protobuf/rewriter_config.proto                | 2 +-
 tensorflow/core/util/cuda_device_functions.h                  | 2 +-
 tensorflow/core/util/mkl_util.h                               | 2 +-
 tensorflow/core/util/tensor_format.h                          | 2 +-
 tensorflow/docs_src/api_guides/python/reading_data.md         | 2 +-
 tensorflow/docs_src/deploy/s3.md                              | 2 +-
 tensorflow/docs_src/mobile/mobile_intro.md                    | 2 +-
 tensorflow/python/data/util/nest.py                           | 2 +-
 tensorflow/python/estimator/estimator.py                      | 2 +-
 .../python/estimator/inputs/queues/feeding_functions.py       | 2 +-
 tensorflow/python/feature_column/feature_column.py            | 2 +-
 tensorflow/python/framework/ops.py                            | 2 +-
 tensorflow/python/framework/test_util.py                      | 2 +-
 tensorflow/python/keras/_impl/keras/engine/network.py         | 2 +-
 tensorflow/python/keras/_impl/keras/engine/saving_test.py     | 4 ++--
 tensorflow/python/keras/_impl/keras/estimator.py              | 2 +-
 tensorflow/python/kernel_tests/distributions/util_test.py     | 2 +-
 tensorflow/python/kernel_tests/manip_ops_test.py              | 2 +-
 tensorflow/python/ops/math_ops.py                             | 2 +-
 tensorflow/python/training/distribute.py                      | 2 +-
 tensorflow/python/util/util.cc                                | 2 +-
 tensorflow/python/util/util.h                                 | 2 +-
 tensorflow/stream_executor/cuda/cuda_dnn.h                    | 2 +-
 tensorflow/tensorflow.bzl                                     | 2 +-
 tensorflow/tools/graph_transforms/README.md                   | 2 +-
 third_party/examples/eager/spinn/README.md                    | 2 +-
 86 files changed, 97 insertions(+), 98 deletions(-)

diff --git a/RELEASE.md b/RELEASE.md
index 55923a2c9b..84d9d52868 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -236,7 +236,7 @@ Yoni Tsafir, yordun, Yuan (Terry) Tang, Yuxin Wu, zhengdi, Zhengsheng Wei, 田
   * Add `complex64` support to XLA compiler.
   * `bfloat` support is now added to XLA infrastructure.
   * Make `ClusterSpec` propagation work with XLA devices.
-  * Use a determinisitic executor to generate XLA graph.
+  * Use a deterministic executor to generate XLA graph.
 * `tf.contrib`:
   * `tf.contrib.distributions`:
     * Add `tf.contrib.distributions.Autoregressive`.
diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc
index f06debaf31..6d1e3325eb 100644
--- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc
+++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc
@@ -240,7 +240,7 @@ class Encapsulator {
   // Once edges between compiled and outside_compilation clusters have been
   // replaced by send/recv ops, some dependencies may no longer be apparent.
   // A clustering pass finds all the dependencies between HC nodes that are only
-  // present as a result of edges between nodes in outside_compilaton clusters.
+  // present as a result of edges between nodes in outside_compilation clusters.
   // Suppose there is a path from outside_compilation cluster C in subgraph S
   // to outside_compilation cluster D in subgraph T. If S != T then a control
   // edge is added from the call node for S to the call node for T, which
diff --git a/tensorflow/compiler/xla/literal_util.h b/tensorflow/compiler/xla/literal_util.h
index 290f388078..c6bd03bf21 100644
--- a/tensorflow/compiler/xla/literal_util.h
+++ b/tensorflow/compiler/xla/literal_util.h
@@ -286,7 +286,7 @@ class Literal {
 
   // Creates a new value that has the equivalent value as this literal, but
   // conforms to new_layout; e.g. a literal matrix that was in {0, 1}
-  // minor-to-major dimension layout can be re-layed-out as {1, 0}
+  // minor-to-major dimension layout can be re-laid-out as {1, 0}
   // minor-to-major dimension layout and the value in the cell at any given
   // logical index (i0, i1) will be the same.
   //
diff --git a/tensorflow/compiler/xla/service/conditional_simplifier.cc b/tensorflow/compiler/xla/service/conditional_simplifier.cc
index e560abc87f..e9ec796121 100644
--- a/tensorflow/compiler/xla/service/conditional_simplifier.cc
+++ b/tensorflow/compiler/xla/service/conditional_simplifier.cc
@@ -35,7 +35,7 @@ namespace xla {
 
 // Tries to replace a conditional with a call operation of the corresponding
 // computation. If the given conditional has a constant predicate, tries to
-// replace it with a call to its true/false computation as appropirate and then
+// replace it with a call to its true/false computation as appropriate and then
 // inline that computation.
 //
 // Returns true if it made a change to the graph.
diff --git a/tensorflow/compiler/xla/service/cpu/ir_function.h b/tensorflow/compiler/xla/service/cpu/ir_function.h
index 557aa4a6bf..2e55181eed 100644
--- a/tensorflow/compiler/xla/service/cpu/ir_function.h
+++ b/tensorflow/compiler/xla/service/cpu/ir_function.h
@@ -33,8 +33,8 @@ namespace cpu {
 // emitters for function and function argument access.
 // The llvm::Function is created with the standard function signature
 // used in the XLA CPU backend (see ir_function.cc for argument details).
-// In addtion IrFunction saves the callers IR insert point during contruction,
-// and restores it after desctruction.
+// In addition IrFunction saves the callers IR insert point during construction,
+// and restores it after destruction.
 //
 // Example usage:
 //
diff --git a/tensorflow/compiler/xla/service/cpu/shape_partition.h b/tensorflow/compiler/xla/service/cpu/shape_partition.h
index 33d02b70e6..db2cda2936 100644
--- a/tensorflow/compiler/xla/service/cpu/shape_partition.h
+++ b/tensorflow/compiler/xla/service/cpu/shape_partition.h
@@ -38,7 +38,7 @@ namespace cpu {
 //
 //     [0, 1), [1, 2), [2, 3), [3, 4), [4, 5) [5, 8)
 //
-//   Note that the last partition has residule because the dimension size is
+//   Note that the last partition has residual because the dimension size is
 //   not a multiple of the partition count.
 //
 //
diff --git a/tensorflow/compiler/xla/service/despecializer.h b/tensorflow/compiler/xla/service/despecializer.h
index af48f4ab6e..cc1695b7f8 100644
--- a/tensorflow/compiler/xla/service/despecializer.h
+++ b/tensorflow/compiler/xla/service/despecializer.h
@@ -25,7 +25,7 @@ namespace xla {
 
 // Creates an HloPassPipeline containing multiple HloPasses that can
 // despecialize an optimized HloModule. This is useful to run an HloModule
-// optimized for one specfic platform on a different platform (undoing platform
+// optimized for one specific platform on a different platform (undoing platform
 // specific passes) with matching numerics for comparison.
 //
 // Current despecialization passes are Defuser, ImplicitBroadcastRemover,
diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h
index b842f480c6..b41ab2162a 100644
--- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h
+++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h
@@ -38,7 +38,7 @@ namespace gpu {
 //
 // Examples of things that are not unnested computations:
 //
-//  - The reducer of a kReduce HLO.  This is emited using IrEmitterNested.
+//  - The reducer of a kReduce HLO.  This is emitted using IrEmitterNested.
 //  - The body of a fusion node.  IrEmitterUnenested emits the relevant code
 //    within a kernel function using FusedIrEmitter.  (FusedIrEmitter is not
 //    really an IrEmitter, but is more an "IR generator generator".)
diff --git a/tensorflow/compiler/xla/service/interpreter/README.md b/tensorflow/compiler/xla/service/interpreter/README.md
index 4c19a1b916..0b21b251c3 100644
--- a/tensorflow/compiler/xla/service/interpreter/README.md
+++ b/tensorflow/compiler/xla/service/interpreter/README.md
@@ -5,7 +5,7 @@ evaluating the result of the HLO graph directly with HloEvaluator, without
 lowering it further (to LLVM IR for example) before execution as other backends
 (CPU and GPU for example) do.
 
-Its key componenets are:
+Its key components are:
 
 *   [`InterpreterCompiler`] despite the inherited naming of "compiler", all
     `InterpreterCompiler` really does is the following:
diff --git a/tensorflow/compiler/xla/service/layout_assignment.h b/tensorflow/compiler/xla/service/layout_assignment.h
index c83ae0388b..9663a793fd 100644
--- a/tensorflow/compiler/xla/service/layout_assignment.h
+++ b/tensorflow/compiler/xla/service/layout_assignment.h
@@ -281,8 +281,8 @@ class LayoutAssignment : public HloPassInterface {
   // the case that no particular layout is requested.
   //
   // channel_constraints is both an input and output. Any sends or recvs that
-  // are present in channel_constraints will be layed out as constrained. Any
-  // unconstrained sends or recvs will be layed out as locally optimal and their
+  // are present in channel_constraints will be laid out as constrained. Any
+  // unconstrained sends or recvs will be laid out as locally optimal and their
   // layout will be added as a constraint to channel_constraints.
   //
   // If channel_constraints is nullptr, no kSend or kRecvs must be contained
diff --git a/tensorflow/compiler/xla/service/reduce_precision_insertion.cc b/tensorflow/compiler/xla/service/reduce_precision_insertion.cc
index e2c07e3827..688cceff0c 100644
--- a/tensorflow/compiler/xla/service/reduce_precision_insertion.cc
+++ b/tensorflow/compiler/xla/service/reduce_precision_insertion.cc
@@ -75,7 +75,7 @@ StatusOr<bool> ReducePrecisionInsertion::insert_after(
     return false;
   }
 
-  // Check that we haven't already inserted an equivalant reduce-precision
+  // Check that we haven't already inserted an equivalent reduce-precision
   // operation after this instruction.  (The zero-user case occurs when this is
   // the root instruction.)
   if (instruction->user_count() > 0) {
diff --git a/tensorflow/compiler/xla/service/source_map_util.h b/tensorflow/compiler/xla/service/source_map_util.h
index a776d745f4..18e2651abb 100644
--- a/tensorflow/compiler/xla/service/source_map_util.h
+++ b/tensorflow/compiler/xla/service/source_map_util.h
@@ -23,7 +23,7 @@ limitations under the License.
 namespace xla {
 namespace source_map_util {
 
-// Creates an INVALID_ARUGMENT status with the given format string.
+// Creates an INVALID_ARGUMENT status with the given format string.
 //
 // Also, attempts to extract the OpMetadata for parameter_number on executable
 // and append it to the status message for source mapping to user code.
diff --git a/tensorflow/contrib/autograph/impl/config.py b/tensorflow/contrib/autograph/impl/config.py
index 2600088595..878bb7e12f 100644
--- a/tensorflow/contrib/autograph/impl/config.py
+++ b/tensorflow/contrib/autograph/impl/config.py
@@ -33,7 +33,7 @@ DEFAULT_UNCOMPILED_MODULES = set((
     (utils.__name__,),
 
     # All of tensorflow's subpackages. Unlike the root tf module, they don't
-    # have well-known names. Not refering to the module directly to avoid
+    # have well-known names. Not referring to the module directly to avoid
     # circular imports.
     (
         utils.__name__[:-len('.contrib.autograph.utils')],),
diff --git a/tensorflow/contrib/autograph/operators/control_flow.py b/tensorflow/contrib/autograph/operators/control_flow.py
index 9f7202821f..671c9ccc13 100644
--- a/tensorflow/contrib/autograph/operators/control_flow.py
+++ b/tensorflow/contrib/autograph/operators/control_flow.py
@@ -174,7 +174,7 @@ def while_stmt(test, body, init_state, extra_deps, opts=None):
     Tuple containing the final state.
   """
   # TODO(mdan): Consider adding a generic mechanism for dynamic dispatch.
-  # That could be somethins as simple as a collection of dispatch rules, with
+  # That could be something as simple as a collection of dispatch rules, with
   # some prioritization.
   if any(tensor_util.is_tensor(v) for v in init_state + extra_deps):
     return _tf_while_stmt(test, body, init_state, opts)
diff --git a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py
index 08c1dcdd02..e53d86ec61 100644
--- a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py
+++ b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py
@@ -369,7 +369,7 @@ class GradientBoostedDecisionTreeModel(object):
     Returns:
       a dictionary of prediction results -
         ENSEMBLE_STAMP, PREDICTION, PARTITION_IDS,
-        NUM_LAYER_ATTEMPTED, NUM_TREES_ATTEMPED.
+        NUM_LAYER_ATTEMPTED, NUM_TREES_ATTEMPTED.
     """
     ensemble_stats = training_ops.tree_ensemble_stats(ensemble_handle,
                                                       ensemble_stamp)
diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py b/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py
index ecdb8967f4..268c8d0342 100644
--- a/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py
+++ b/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py
@@ -53,7 +53,7 @@ class CholeskyOuterProduct(bijector.Bijector):
   its spectrum), and that the product of two positive-diagonal lower-triangular
   matrices is another positive-diagonal lower-triangular matrix.
 
-  A simple inductive argument (proceding one column of L_3 at a time) shows
+  A simple inductive argument (proceeding one column of L_3 at a time) shows
   that, if `I = L_3 @ L_3.T`, with L_3 being lower-triangular with positive-
   diagonal, then `L_3 = I`. Thus, `L_1 = L_2`, proving injectivity of g.
 
diff --git a/tensorflow/contrib/eager/README.md b/tensorflow/contrib/eager/README.md
index 762685db14..4384431e7b 100644
--- a/tensorflow/contrib/eager/README.md
+++ b/tensorflow/contrib/eager/README.md
@@ -1,6 +1,6 @@
 # Eager Execution
 
-Eager execution provides an imperative interface to TensorFlow (similiar to
+Eager execution provides an imperative interface to TensorFlow (similar to
 [NumPy](http://www.numpy.org)). When you enable eager execution, TensorFlow
 operations execute immediately; you do not execute a pre-constructed graph with
 [`Session.run()`](https://www.tensorflow.org/api_docs/python/tf/Session).
diff --git a/tensorflow/contrib/ffmpeg/ffmpeg_lib.h b/tensorflow/contrib/ffmpeg/ffmpeg_lib.h
index a8d5a0dd83..bf2aa75545 100644
--- a/tensorflow/contrib/ffmpeg/ffmpeg_lib.h
+++ b/tensorflow/contrib/ffmpeg/ffmpeg_lib.h
@@ -53,7 +53,7 @@ Status CreateAudioFile(const string& audio_format_id, int32 bits_per_second,
                        int32 samples_per_second, int32 channel_count,
                        const std::vector<float>& samples, string* output_data);
 
-// Reads an video file using ffmpeg adn converts it into a RGB24 in uint8
+// Reads an video file using ffmpeg and converts it into a RGB24 in uint8
 // [frames, height, width, 3]. The w, h, and frames are obtained from ffmpeg.
 Status ReadVideoFile(const string& filename, std::vector<uint8>* output_data,
                      uint32* width, uint32* height, uint32* frames);
diff --git a/tensorflow/contrib/framework/python/ops/critical_section_ops.py b/tensorflow/contrib/framework/python/ops/critical_section_ops.py
index bd764ed57a..72835c3ad8 100644
--- a/tensorflow/contrib/framework/python/ops/critical_section_ops.py
+++ b/tensorflow/contrib/framework/python/ops/critical_section_ops.py
@@ -202,7 +202,7 @@ class CriticalSection(object):
         or lazy way that may cause a deadlock.
       ValueError: If `exclusive_resource_access` is not provided (is `True`) and
         another `CriticalSection` has an execution requesting the same
-        resources as in `*args`, `**kwargs`, and any additionaly captured
+        resources as in `*args`, `**kwargs`, and any additionally captured
         inputs in `fn`.  Note, even if `exclusive_resource_access` is `True`,
         if another execution in another `CriticalSection` was created without
         `exclusive_resource_access=True`, a `ValueError` will be raised.
diff --git a/tensorflow/contrib/gan/python/features/python/conditioning_utils.py b/tensorflow/contrib/gan/python/features/python/conditioning_utils.py
index df71187fbd..a9b8faa712 100644
--- a/tensorflow/contrib/gan/python/features/python/conditioning_utils.py
+++ b/tensorflow/contrib/gan/python/features/python/conditioning_utils.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Miscellanous utilities for TFGAN code and examples."""
+"""Miscellaneous utilities for TFGAN code and examples."""
 
 from __future__ import absolute_import
 from __future__ import division
diff --git a/tensorflow/contrib/graph_editor/transform.py b/tensorflow/contrib/graph_editor/transform.py
index a320a3f232..592d37b432 100644
--- a/tensorflow/contrib/graph_editor/transform.py
+++ b/tensorflow/contrib/graph_editor/transform.py
@@ -677,7 +677,7 @@ def copy_with_input_replacements(sgv, replacement_ts,
 
 
 def _add_control_flow_ops(ops, control_ios):
-  """Complete `ops` so that the tranformed graph is valid.
+  """Complete `ops` so that the transformed graph is valid.
 
   Partially copying a graph can lead to a malformed graph. For instance,
   copying half of a while construct is likely to result in an invalid graph.
diff --git a/tensorflow/contrib/image/__init__.py b/tensorflow/contrib/image/__init__.py
index 8f406ace1d..f230d93da4 100755
--- a/tensorflow/contrib/image/__init__.py
+++ b/tensorflow/contrib/image/__init__.py
@@ -17,7 +17,7 @@
 ### API
 
 This module provides functions for image manipulation; currently, chrominance
-transformas (including changing saturation and hue) in YIQ space and
+transforms (including changing saturation and hue) in YIQ space and
 projective transforms (including rotation) are supported.
 
 ## Image Transformation `Ops`
diff --git a/tensorflow/contrib/kfac/examples/convnet.py b/tensorflow/contrib/kfac/examples/convnet.py
index b261f41bf9..d6b1a61b71 100644
--- a/tensorflow/contrib/kfac/examples/convnet.py
+++ b/tensorflow/contrib/kfac/examples/convnet.py
@@ -325,7 +325,7 @@ def distributed_grads_only_and_ops_chief_worker(
 
   All workers perform gradient computation. Chief worker applies gradient after
   averaging the gradients obtained from all the workers. All workers block
-  execution untill the update is applied. Chief worker runs covariance and
+  execution until the update is applied. Chief worker runs covariance and
   inverse update ops. Covariance and inverse matrices are placed on parameter
   servers in a round robin manner. For further details on synchronous
   distributed optimization check `tf.train.SyncReplicasOptimizer`.
diff --git a/tensorflow/contrib/kfac/python/ops/optimizer.py b/tensorflow/contrib/kfac/python/ops/optimizer.py
index 7203804af3..b7f63d8d94 100644
--- a/tensorflow/contrib/kfac/python/ops/optimizer.py
+++ b/tensorflow/contrib/kfac/python/ops/optimizer.py
@@ -66,7 +66,7 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer):
           the local approximation with the Fisher information matrix, and to
           regularize the update direction by making it closer to the gradient.
           If damping is adapted during training then this value is used for
-          initializing damping varaible.
+          initializing damping variable.
           (Higher damping means the update looks more like a standard gradient
           update - see Tikhonov regularization.)
       layer_collection: The layer collection object, which holds the fisher
@@ -195,7 +195,7 @@ class KfacOptimizer(gradient_descent.GradientDescentOptimizer):
       min_damping: `float`(Optional), Minimum value the damping parameter
         can take. Default value 1e-5.
       damping_adaptation_decay: `float`(Optional), The `damping` parameter is
-        multipled by the `damping_adaptation_decay` every
+        multiplied by the `damping_adaptation_decay` every
         `damping_adaptation_interval` number of iterations. Default value 0.99.
       damping_adaptation_interval: `int`(Optional), Number of steps in between
         updating the `damping` parameter. Default value 5.
diff --git a/tensorflow/contrib/kfac/python/ops/placement.py b/tensorflow/contrib/kfac/python/ops/placement.py
index 8a20ebe198..c4454325ae 100644
--- a/tensorflow/contrib/kfac/python/ops/placement.py
+++ b/tensorflow/contrib/kfac/python/ops/placement.py
@@ -51,7 +51,7 @@ class RoundRobinPlacementMixin(object):
     self._inv_devices = inv_devices
 
   def make_vars_and_create_op_thunks(self, scope=None):
-    """Make vars and create op thunks w/ a round-robin device placement strat.
+    """Make vars and create op thunks w/ a round-robin device placement start.
 
     For each factor, all of that factor's cov variables and their associated
     update ops will be placed on a particular device.  A new device is chosen
diff --git a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
index 445687cd15..e2e1cf4478 100644
--- a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
@@ -1814,7 +1814,7 @@ inline void LstmCell(const float* input_data, const Dims<4>& input_dims,
 // requiring a power-of-two representation interval. Thus, we should right
 // away quantize this array to a power-of-two interval; otherwise,
 // implementation will need to rescale that, losing any benefit that a tighter
-// representation interval might otherwise yield, while introducting some
+// representation interval might otherwise yield, while introducing some
 // numerical error and computational overhead.
 //
 // Now, Logistic and Tanh
diff --git a/tensorflow/contrib/lite/schema/schema.fbs b/tensorflow/contrib/lite/schema/schema.fbs
index b16baf02dc..ff56c31720 100644
--- a/tensorflow/contrib/lite/schema/schema.fbs
+++ b/tensorflow/contrib/lite/schema/schema.fbs
@@ -65,7 +65,7 @@ table Tensor {
   quantization:QuantizationParameters;  // Optional.
 }
 
-// A list of builtin operators. Builtin operators a slighlty faster than custom
+// A list of builtin operators. Builtin operators are slightly faster than custom
 // ones, but not by much. Moreover, while custom operators accept an opaque
 // object containing configuration parameters, builtins have a predetermined
 // set of acceptable options.
diff --git a/tensorflow/contrib/lite/schema/schema_v0.fbs b/tensorflow/contrib/lite/schema/schema_v0.fbs
index 852ea988f3..891d8366cc 100644
--- a/tensorflow/contrib/lite/schema/schema_v0.fbs
+++ b/tensorflow/contrib/lite/schema/schema_v0.fbs
@@ -48,7 +48,7 @@ table Tensor {
   quantization:QuantizationParameters;  // Optional.
 }
 
-// A list of builtin operators. Builtin operators a slighlty faster than custom
+// A list of builtin operators. Builtin operators are slightly faster than custom
 // ones, but not by much. Moreover, while custom operators accept an opaque
 // object containing configuration parameters, builtins have a predetermined
 // set of acceptable options.
diff --git a/tensorflow/contrib/lite/schema/schema_v1.fbs b/tensorflow/contrib/lite/schema/schema_v1.fbs
index 06cd9408ed..b438b569e6 100644
--- a/tensorflow/contrib/lite/schema/schema_v1.fbs
+++ b/tensorflow/contrib/lite/schema/schema_v1.fbs
@@ -53,7 +53,7 @@ table Tensor {
   quantization:QuantizationParameters;  // Optional.
 }
 
-// A list of builtin operators. Builtin operators a slighlty faster than custom
+// A list of builtin operators. Builtin operators are slightly faster than custom
 // ones, but not by much. Moreover, while custom operators accept an opaque
 // object containing configuration parameters, builtins have a predetermined
 // set of acceptable options.
diff --git a/tensorflow/contrib/lite/schema/schema_v2.fbs b/tensorflow/contrib/lite/schema/schema_v2.fbs
index 96731c8aae..b90408ff6d 100644
--- a/tensorflow/contrib/lite/schema/schema_v2.fbs
+++ b/tensorflow/contrib/lite/schema/schema_v2.fbs
@@ -54,7 +54,7 @@ table Tensor {
   quantization:QuantizationParameters;  // Optional.
 }
 
-// A list of builtin operators. Builtin operators a slighlty faster than custom
+// A list of builtin operators. Builtin operators are slightly faster than custom
 // ones, but not by much. Moreover, while custom operators accept an opaque
 // object containing configuration parameters, builtins have a predetermined
 // set of acceptable options.
diff --git a/tensorflow/contrib/lite/schema/schema_v3.fbs b/tensorflow/contrib/lite/schema/schema_v3.fbs
index cedefe08f3..020da38493 100644
--- a/tensorflow/contrib/lite/schema/schema_v3.fbs
+++ b/tensorflow/contrib/lite/schema/schema_v3.fbs
@@ -53,7 +53,7 @@ table Tensor {
   type:TensorType;
   // An index that refers to the buffers table at the root of the model. Or,
   // if there is no data buffer associated (i.e. intermediate results), then
-  // this is 0 (which refers to an always existant empty buffer).
+  // this is 0 (which refers to an always existent empty buffer).
   //
   // The data_buffer itself is an opaque container, with the assumption that the
   // target device is little-endian. In addition, all builtin operators assume
@@ -64,7 +64,7 @@ table Tensor {
   quantization:QuantizationParameters;  // Optional.
 }
 
-// A list of builtin operators. Builtin operators a slighlty faster than custom
+// A list of builtin operators. Builtin operators are slightly faster than custom
 // ones, but not by much. Moreover, while custom operators accept an opaque
 // object containing configuration parameters, builtins have a predetermined
 // set of acceptable options.
diff --git a/tensorflow/contrib/lite/testing/generate_examples.py b/tensorflow/contrib/lite/testing/generate_examples.py
index e4851d6077..fd09332165 100644
--- a/tensorflow/contrib/lite/testing/generate_examples.py
+++ b/tensorflow/contrib/lite/testing/generate_examples.py
@@ -1758,7 +1758,7 @@ def make_strided_slice_tests(zip_path):
           "shrink_axis_mask": [None, 1, 8, 11, 15, -1],
           "constant_indices": [False, True],
       },
-      # TODO(b/73170889) Restore test paramaters removed in cl/191608113.
+      # TODO(b/73170889) Restore test parameters removed in cl/191608113.
       # 2-D
       {
           "dtype": [tf.float32, tf.int32, tf.int64],
@@ -1899,7 +1899,7 @@ def make_lstm_tests(zip_path):
     return inputs_after_split, [out]
 
   def build_inputs(parameters, sess, inputs, outputs):
-    """Feed inputs, assign vairables, and freeze graph."""
+    """Feed inputs, assign variables, and freeze graph."""
 
     with tf.variable_scope("", reuse=True):
       kernel = tf.get_variable("rnn/basic_lstm_cell/kernel")
diff --git a/tensorflow/contrib/lite/testing/tflite_driver.cc b/tensorflow/contrib/lite/testing/tflite_driver.cc
index 58fe5bd6e4..75ac24719a 100644
--- a/tensorflow/contrib/lite/testing/tflite_driver.cc
+++ b/tensorflow/contrib/lite/testing/tflite_driver.cc
@@ -226,8 +226,8 @@ void TfLiteDriver::SetExpectation(int id, const string& csv_values) {
   if (!IsValid()) return;
   auto* tensor = interpreter_->tensor(id);
   if (expected_output_.count(id) != 0) {
-    fprintf(stderr, "Overriden expectation for tensor %d\n", id);
-    Invalidate("Overriden expectation");
+    fprintf(stderr, "Overridden expectation for tensor %d\n", id);
+    Invalidate("Overridden expectation");
   }
   expected_output_[id].reset(new Expectation);
   switch (tensor->type) {
diff --git a/tensorflow/contrib/lite/toco/g3doc/cmdline_examples.md b/tensorflow/contrib/lite/toco/g3doc/cmdline_examples.md
index 495014c6fc..7680cdd344 100644
--- a/tensorflow/contrib/lite/toco/g3doc/cmdline_examples.md
+++ b/tensorflow/contrib/lite/toco/g3doc/cmdline_examples.md
@@ -115,7 +115,7 @@ bazel run --config=opt \
 
 In order to evaluate the possible benefit of generating a quantized graph, TOCO
 allows "dummy-quantization" on float graphs. The flags `--default_ranges_min`
-and `--default_ranges_max` accept plausable values for the min-max ranges of the
+and `--default_ranges_max` accept plausible values for the min-max ranges of the
 values in all arrays that do not have min-max information. "Dummy-quantization"
 will produce lower accuracy but will emulate the performance of a correctly
 quantized model.
@@ -338,7 +338,7 @@ below outline the use cases for each.
 ### Using `--output_format=GRAPHVIZ_DOT`
 
 The first way to get a graphviz rendering is to pass `GRAPHVIZ_DOT` into
-`--output_format`. This results in a plausable visualization of the graph. This
+`--output_format`. This results in a plausible visualization of the graph. This
 reduces the requirements that normally exist during conversion between other
 input and output formats. For example, this may be useful if conversion from
 TENSORFLOW_GRAPHDEF to TFLITE is failing.
diff --git a/tensorflow/contrib/lite/toco/import_tensorflow.cc b/tensorflow/contrib/lite/toco/import_tensorflow.cc
index 453ff29b0d..8efe6ab7b9 100644
--- a/tensorflow/contrib/lite/toco/import_tensorflow.cc
+++ b/tensorflow/contrib/lite/toco/import_tensorflow.cc
@@ -144,7 +144,7 @@ ArrayDataType ConvertDataType(tensorflow::DataType dtype) {
   else if (dtype == DT_STRING)
     return ArrayDataType::kString;
   else
-    LOG(INFO) << "Unsupported data type in placehoder op: " << dtype;
+    LOG(INFO) << "Unsupported data type in placeholder op: " << dtype;
   return ArrayDataType::kNone;
 }
 
diff --git a/tensorflow/contrib/lite/toco/tflite/operator.h b/tensorflow/contrib/lite/toco/tflite/operator.h
index 88af3d6ab6..85f7bdafe0 100644
--- a/tensorflow/contrib/lite/toco/tflite/operator.h
+++ b/tensorflow/contrib/lite/toco/tflite/operator.h
@@ -25,10 +25,10 @@ namespace tflite {
 
 class BaseOperator;
 
-// Return a map contained all knwo TF Lite Operators, keyed by their names.
+// Return a map contained all know TF Lite Operators, keyed by their names.
 std::map<string, std::unique_ptr<BaseOperator>> BuildOperatorByNameMap();
 
-// Return a map contained all knwo TF Lite Operators, keyed by the type of
+// Return a map contained all know TF Lite Operators, keyed by the type of
 // their tf.mini counterparts.
 std::map<OperatorType, std::unique_ptr<BaseOperator>> BuildOperatorByTypeMap();
 
diff --git a/tensorflow/contrib/lite/toco/tflite/types_test.cc b/tensorflow/contrib/lite/toco/tflite/types_test.cc
index 29fb0b2af2..efb849f422 100644
--- a/tensorflow/contrib/lite/toco/tflite/types_test.cc
+++ b/tensorflow/contrib/lite/toco/tflite/types_test.cc
@@ -44,7 +44,7 @@ template <ArrayDataType T>
 Array ToFlatBufferAndBack(std::initializer_list<::toco::DataType<T>> items) {
   // NOTE: This test does not construct the full buffers list. Since
   // Deserialize normally takes a buffer, we need to synthesize one and provide
-  // an index that is non-zero so the buffer is not assumed to be emtpy.
+  // an index that is non-zero so the buffer is not assumed to be empty.
   Array src;
   src.data_type = T;
   src.GetMutableBuffer<T>().data = items;
diff --git a/tensorflow/contrib/opt/python/training/elastic_average_optimizer_test.py b/tensorflow/contrib/opt/python/training/elastic_average_optimizer_test.py
index 37539b9599..5ed8057b86 100644
--- a/tensorflow/contrib/opt/python/training/elastic_average_optimizer_test.py
+++ b/tensorflow/contrib/opt/python/training/elastic_average_optimizer_test.py
@@ -58,7 +58,7 @@ def create_local_cluster(num_workers, num_ps, protocol="grpc"):
 
 
 # Creates the workers and return their sessions, graphs, train_ops.
-# Cheif worker will update at last
+# Chief worker will update at last
 def _get_workers(num_workers, period, workers, moving_rate):
   sessions = []
   graphs = []
diff --git a/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py b/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py
index bfb3350b59..3acd940268 100644
--- a/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py
+++ b/tensorflow/contrib/opt/python/training/model_average_optimizer_test.py
@@ -57,7 +57,7 @@ def create_local_cluster(num_workers, num_ps, protocol="grpc"):
 
 
 # Creates the workers and return their sessions, graphs, train_ops.
-# Cheif worker will update at last
+# Chief worker will update at last
 def _get_workers(num_workers, steps, workers):
   sessions = []
   graphs = []
diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
index 4d3710a514..3767596f8c 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
@@ -2145,7 +2145,7 @@ tensorflow::Status ConvertCalibrationNodeToEngineNode(
   if (!status.ok() || !calib_res->calibrator_) {
     return tensorflow::errors::FailedPrecondition(
         "You must run calibration"
-        " and inference conversion in the same proces");
+        " and inference conversion in the same process");
   }
 
   calib_res->calibrator_->setDone();
diff --git a/tensorflow/contrib/verbs/README.md b/tensorflow/contrib/verbs/README.md
index 4b6104a8b4..3137bfd03e 100644
--- a/tensorflow/contrib/verbs/README.md
+++ b/tensorflow/contrib/verbs/README.md
@@ -159,7 +159,7 @@ When the receiver receives the RDMA write, it will locate the relevant **RdmaTen
 	* step_id - Step ID.
 	* request_index - Request index.
 	* remote_addr/rkey - Address/rkey of the reallocated result/proxy tensor.
-* **RDMA_MESSAGE_ERROR_STATUS** - (sender ==> receiver) Notify the receiver that an error had occured on the sender side, so it can propagate it to the upper levels.
+* **RDMA_MESSAGE_ERROR_STATUS** - (sender ==> receiver) Notify the receiver that an error had occurred on the sender side, so it can propagate it to the upper levels.
 	* type - The message type.
 	* name (name_size) - Name of the requested tensor.
 	* step_id - Step ID.
diff --git a/tensorflow/core/common_runtime/broadcaster.cc b/tensorflow/core/common_runtime/broadcaster.cc
index 5e8af8653d..e42d3f6b92 100644
--- a/tensorflow/core/common_runtime/broadcaster.cc
+++ b/tensorflow/core/common_runtime/broadcaster.cc
@@ -80,7 +80,7 @@ void Broadcaster::Run(StatusCallback done) {
 // continuing to occupy its current position.  Hence we calculate as
 // though each device's rank is actually r+1, then subtract 1 again to
 // get the descendent ranks.  If the source is not rank 0 then its
-// decendents include both {0,1} and the descendents of its current
+// descendants include both {0,1} and the descendents of its current
 // position.  Where a non-0-rank source is a descendent of another
 // device, no send to it is necessary.
 
@@ -115,7 +115,7 @@ void Broadcaster::TreeSendTo(const CollectiveParams& cp,
   DCHECK_NE(successor_rank, my_rank);
   if (cp.is_source && source_rank != 0) {
     // The source sends to rank 0,1 in addition to its positional
-    // decendents.
+    // descendants.
     if (cp.group.group_size > 1) {
       targets->push_back(0);
     }
diff --git a/tensorflow/core/common_runtime/buf_rendezvous.h b/tensorflow/core/common_runtime/buf_rendezvous.h
index e94e88b323..9eb9f060f6 100644
--- a/tensorflow/core/common_runtime/buf_rendezvous.h
+++ b/tensorflow/core/common_runtime/buf_rendezvous.h
@@ -79,7 +79,7 @@ class BufRendezvous {
                   const ProducerCallback& done);
 
   // Called to request access to a Tensor value corresponding to key.
-  // Consumer is provide with a Hook as soon as availble.
+  // Consumer is provide with a Hook as soon as available.
   void ConsumeBuf(const string& key, const ConsumerCallback& done);
 
   // Consumer must call this function when it's done reading the Hook provided
diff --git a/tensorflow/core/common_runtime/ring_reducer.cc b/tensorflow/core/common_runtime/ring_reducer.cc
index a17281835e..a74c502a92 100644
--- a/tensorflow/core/common_runtime/ring_reducer.cc
+++ b/tensorflow/core/common_runtime/ring_reducer.cc
@@ -275,7 +275,7 @@ void RingReducer::InitRingField(RingField* rf, int chunk_idx, int subdiv_idx,
   // Note on field indexing: There are group_size_ devices in the
   // instance, implying the same number of chunks per tensor, where a
   // chunk is the unit of data transferred in a time step.  However, if
-  // a device can simultaenously send data by 2 or more independent
+  // a device can simultaneously send data by 2 or more independent
   // channels we can speed up the transfer by subdividing chunks and
   // processing multiple subdivisions at once.  So the actual number
   // of RingFields is group_size_ * num_subdivs_.
diff --git a/tensorflow/core/common_runtime/scoped_allocator_mgr.cc b/tensorflow/core/common_runtime/scoped_allocator_mgr.cc
index be79cc4507..c045596a69 100644
--- a/tensorflow/core/common_runtime/scoped_allocator_mgr.cc
+++ b/tensorflow/core/common_runtime/scoped_allocator_mgr.cc
@@ -104,7 +104,7 @@ ScopedAllocatorContainer::~ScopedAllocatorContainer() {
   // contents deleted via Drop.  When when a step ends early
   // (e.g. through abnormal termination) we need to clean up
   // explicitly.  So long as graph execution of the associated step has
-  // completey terminated this should be safe.
+  // completely terminated this should be safe.
   for (auto& it : allocators_) {
     if (it.second.field_index == ScopedAllocator::kBackingIndex) {
       delete it.second.scoped_allocator;
diff --git a/tensorflow/core/debug/debug_io_utils.cc b/tensorflow/core/debug/debug_io_utils.cc
index baa8c08fdf..63be4ef53a 100644
--- a/tensorflow/core/debug/debug_io_utils.cc
+++ b/tensorflow/core/debug/debug_io_utils.cc
@@ -52,7 +52,7 @@ namespace {
 
 // Creates an Event proto representing a chunk of a Tensor. This method only
 // populates the field of the Event proto that represent the envelope
-// informaion (e.g., timestmap, device_name, num_chunks, chunk_index, dtype,
+// information (e.g., timestamp, device_name, num_chunks, chunk_index, dtype,
 // shape). It does not set the value.tensor field, which should be set by the
 // caller separately.
 Event PrepareChunkEventProto(const DebugNodeKey& debug_node_key,
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_worker_cache.cc b/tensorflow/core/distributed_runtime/rpc/grpc_worker_cache.cc
index 18998bbccb..b9f21ea211 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_worker_cache.cc
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_worker_cache.cc
@@ -115,7 +115,7 @@ class GrpcWorkerCache : public WorkerCachePartial {
 
   size_t AssignWorkerToThread(const string& target) {
     // Round-robin target assignment, but keeps the same target on the same
-    // polling thread always, as this is important for gRPC performace
+    // polling thread always, as this is important for gRPC performance
     mutex_lock lock(assignment_mu_);
     auto it = target_assignments_.find(target);
     if (it == target_assignments_.end()) {
diff --git a/tensorflow/core/framework/op_gen_lib.h b/tensorflow/core/framework/op_gen_lib.h
index ff38e4b221..533dd64805 100644
--- a/tensorflow/core/framework/op_gen_lib.h
+++ b/tensorflow/core/framework/op_gen_lib.h
@@ -59,14 +59,14 @@ class ApiDefMap {
   // You can call this method multiple times to load multiple
   // sets of files. Api definitions are merged if the same
   // op definition is loaded multiple times. Later-loaded
-  // definitions take precedense.
+  // definitions take precedence.
   // ApiDefs loaded from files must contain a subset of ops defined
   // in the OpList passed to the constructor.
   Status LoadFileList(Env* env, const std::vector<string>& filenames);
 
   // Load a single file. Api definitions are merged if the same
   // op definition is loaded multiple times. Later-loaded
-  // definitions take precedense.
+  // definitions take precedence.
   // ApiDefs loaded from file must contain a subset of ops defined
   // in the OpList passed to the constructor.
   Status LoadFile(Env* env, const string& filename);
diff --git a/tensorflow/core/framework/op_kernel.h b/tensorflow/core/framework/op_kernel.h
index 67943377b9..f577664709 100644
--- a/tensorflow/core/framework/op_kernel.h
+++ b/tensorflow/core/framework/op_kernel.h
@@ -534,7 +534,7 @@ class OpKernelContext {
     Rendezvous* rendezvous = nullptr;
 
     // Mechanism for executing a collective op that needs to coordinate
-    // with parallel instances runing on other devices.
+    // with parallel instances running on other devices.
     CollectiveExecutor* collective_executor = nullptr;
 
     // The session state for this op.
diff --git a/tensorflow/core/graph/while_context.h b/tensorflow/core/graph/while_context.h
index 5944e36897..2a83eb7bd8 100644
--- a/tensorflow/core/graph/while_context.h
+++ b/tensorflow/core/graph/while_context.h
@@ -31,7 +31,7 @@ namespace tensorflow {
 // future to support these features.
 //
 // TODO(skyewm): de/serialize in MetaGraphDef so imported while loops will be
-// differentiable. Figure out backwards compatability story.
+// differentiable. Figure out backwards compatibility story.
 class WhileContext {
  public:
   WhileContext(StringPiece frame_name, std::vector<Node*> enter_nodes,
diff --git a/tensorflow/core/grappler/costs/graph_properties.cc b/tensorflow/core/grappler/costs/graph_properties.cc
index 2c7b57971a..fd0547cf86 100644
--- a/tensorflow/core/grappler/costs/graph_properties.cc
+++ b/tensorflow/core/grappler/costs/graph_properties.cc
@@ -574,7 +574,6 @@ class SymbolicShapeRefiner {
     }
   };
 
-  // Compute the shape of the tensors outputed by node 'node' at output port
   // 'port_index' as the union of shape1 and shape2.
   ShapeHandle OutputAsUnion(const NodeDef* node, int port_index,
                             ShapeHandle shape1, ShapeHandle shape2) {
@@ -968,7 +967,7 @@ Status GraphProperties::PropagateShapes(
     const std::unordered_map<const NodeDef*, const NodeDef*>& resource_handles,
     int num_loops) const {
   // Limit the number of iterations to prevent infinite loops in the presence of
-  // incorrect shape functions. The algoritm should converge in at most
+  // incorrect shape functions. The algorithm should converge in at most
   // num_nested_loops^2 * max_rank. We approximate max_rank with the constant 4.
   // The same applies to resources.
   VLOG(1) << "Propagating " << new_shapes->size() << " new shapes through "
diff --git a/tensorflow/core/grappler/costs/virtual_scheduler.h b/tensorflow/core/grappler/costs/virtual_scheduler.h
index 67bf1e6980..34d48819ac 100644
--- a/tensorflow/core/grappler/costs/virtual_scheduler.h
+++ b/tensorflow/core/grappler/costs/virtual_scheduler.h
@@ -328,7 +328,7 @@ class VirtualScheduler {
   Costs graph_costs_;                   // Graph cost.
   std::map<string, Costs> op_to_cost_;  // Per-op cost.
 
-  // Auxilliary data structures for constructing NodeState and DeviceState.
+  // Auxiliary data structures for constructing NodeState and DeviceState.
   GraphProperties graph_properties_;
   Cluster* cluster_;  // Not owned.
 
diff --git a/tensorflow/core/grappler/optimizers/layout_optimizer.cc b/tensorflow/core/grappler/optimizers/layout_optimizer.cc
index 87ab460862..e08ab1eb67 100644
--- a/tensorflow/core/grappler/optimizers/layout_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/layout_optimizer.cc
@@ -2183,7 +2183,7 @@ Status LayoutOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
 
   TuningConfig config;
   config.no_gemm = true;
-  // TODO(yaozhang): Enable tuning with various TuningConfig choices wtih
+  // TODO(yaozhang): Enable tuning with various TuningConfig choices with
   // the measurement-based estimator.
   status = Tune(item, graph_properties, config, output);
   if (!status.ok()) {
diff --git a/tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler.h b/tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler.h
index f5ced95feb..ae652961db 100644
--- a/tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler.h
+++ b/tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler.h
@@ -76,7 +76,7 @@ class AdaptiveSharedBatchScheduler
           AdaptiveSharedBatchScheduler<TaskType>> {
  public:
   ~AdaptiveSharedBatchScheduler() {
-    // Finish processing batches before destorying other class members.
+    // Finish processing batches before destroying other class members.
     batch_thread_pool_.reset();
   }
 
diff --git a/tensorflow/core/kernels/conv_ops_gpu_3.cu.cc b/tensorflow/core/kernels/conv_ops_gpu_3.cu.cc
index 2503b475dc..8e426ddf2b 100644
--- a/tensorflow/core/kernels/conv_ops_gpu_3.cu.cc
+++ b/tensorflow/core/kernels/conv_ops_gpu_3.cu.cc
@@ -595,7 +595,7 @@ constexpr bool TileSizeOnNonLongSideFrontier(int TileLongSide,
   // For a tile size combination (longside, shortside), lying on the frontier
   // implies that (longside, shortside) is on or within the frontier but
   // (longside*2, shortside) or (longside, shortside+1) is not. With the above
-  // critereon, we simply need to use !TileSizeOnLongSideFrontier to ensure that
+  // criterion, we simply need to use !TileSizeOnLongSideFrontier to ensure that
   // it is not on the long side frontier.
   return !TileSizeOutsideFrontier(TileLongSide, TileShortSide, size_of_t) &&
          (TileSizeOutsideFrontier(TileLongSide * 2, TileShortSide, size_of_t) ||
diff --git a/tensorflow/core/kernels/nth_element_op.cc b/tensorflow/core/kernels/nth_element_op.cc
index 7f12eb953a..0e43cc19aa 100644
--- a/tensorflow/core/kernels/nth_element_op.cc
+++ b/tensorflow/core/kernels/nth_element_op.cc
@@ -114,7 +114,7 @@ struct NthElementFunctor<CPUDevice, T> {
 
     auto worker_threads = *(context->device()->tensorflow_cpu_worker_threads());
     // The average time complexity of partition-based nth_element (BFPRT) is
-    // O(n), althought the worst time complexity could be O(n^2). Here, 20 is a
+    // O(n), although the worst time complexity could be O(n^2). Here, 20 is a
     // empirical factor of cost_per_unit.
     Shard(worker_threads.num_threads, worker_threads.workers, num_rows,
           20 * last_dim, SubNthElement);
diff --git a/tensorflow/core/kernels/roll_op.cc b/tensorflow/core/kernels/roll_op.cc
index 4b630809c5..f5ebf0ea2e 100644
--- a/tensorflow/core/kernels/roll_op.cc
+++ b/tensorflow/core/kernels/roll_op.cc
@@ -84,7 +84,7 @@ void DoRoll(OpKernelContext* context, const int64 num_elements,
   // Shard
   auto worker_threads = context->device()->tensorflow_cpu_worker_threads();
   // 15 - expiramentally determined with float and bool types
-  const int cost_per_element = 15 * sizeof(T);  // rough esitmate
+  const int cost_per_element = 15 * sizeof(T);  // rough estimate
   Shard(worker_threads->num_threads, worker_threads->workers, num_elements,
         cost_per_element, std::move(work));
 }
diff --git a/tensorflow/core/platform/cloud/gcs_file_system.cc b/tensorflow/core/platform/cloud/gcs_file_system.cc
index 2d9c99c124..4da7510c01 100644
--- a/tensorflow/core/platform/cloud/gcs_file_system.cc
+++ b/tensorflow/core/platform/cloud/gcs_file_system.cc
@@ -103,7 +103,7 @@ constexpr char kResolveCacheSecs[] = "GCS_RESOLVE_REFRESH_SECS";
 // The environment variable to configure the http request's connection timeout.
 constexpr char kRequestConnectionTimeout[] =
     "GCS_REQUEST_CONNECTION_TIMEOUT_SECS";
-// The environment varaible to configure the http request's idle timeout.
+// The environment variable to configure the http request's idle timeout.
 constexpr char kRequestIdleTimeout[] = "GCS_REQUEST_IDLE_TIMEOUT_SECS";
 // The environment variable to configure the overall request timeout for
 // metadata requests.
diff --git a/tensorflow/core/platform/cloud/gcs_throttle.h b/tensorflow/core/platform/cloud/gcs_throttle.h
index 97a858e3fe..8c9e2e074c 100644
--- a/tensorflow/core/platform/cloud/gcs_throttle.h
+++ b/tensorflow/core/platform/cloud/gcs_throttle.h
@@ -132,7 +132,7 @@ class GcsThrottle {
    * UpdateState updates the available_tokens_ and last_updated_secs_ variables.
    *
    * UpdateState should be called in order to mark the passage of time, and
-   * therefore add tokens to the availble_tokens_ pool.
+   * therefore add tokens to the available_tokens_ pool.
    */
   void UpdateState() EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
diff --git a/tensorflow/core/profiler/g3doc/command_line.md b/tensorflow/core/profiler/g3doc/command_line.md
index bbaf55e613..cc6d9def47 100644
--- a/tensorflow/core/profiler/g3doc/command_line.md
+++ b/tensorflow/core/profiler/g3doc/command_line.md
@@ -82,7 +82,7 @@ bazel-bin/tensorflow/core/profiler/profiler \
 #
 # Alternatively, user can pass separate files.
 #
-# --graph_path contains the model architecutre and tensor shapes.
+# --graph_path contains the model architecture and tensor shapes.
 # --run_meta_path contains the memory and time information.
 # --op_log_path contains float operation and code traces.
 # --checkpoint_path contains the model checkpoint data.
diff --git a/tensorflow/core/protobuf/rewriter_config.proto b/tensorflow/core/protobuf/rewriter_config.proto
index a15ccdfd87..5372ef24b8 100644
--- a/tensorflow/core/protobuf/rewriter_config.proto
+++ b/tensorflow/core/protobuf/rewriter_config.proto
@@ -32,7 +32,7 @@ message RewriterConfig {
     AGGRESSIVE = 3;
   }
 
-  // Enum controling the number of times to run optimizers. The default is to
+  // Enum controlling the number of times to run optimizers. The default is to
   // run them once.
   enum NumIterationsType {
     DEFAULT_NUM_ITERS = 0;
diff --git a/tensorflow/core/util/cuda_device_functions.h b/tensorflow/core/util/cuda_device_functions.h
index f2d4e470c8..b91f8bb8ef 100644
--- a/tensorflow/core/util/cuda_device_functions.h
+++ b/tensorflow/core/util/cuda_device_functions.h
@@ -537,7 +537,7 @@ __device__ detail::ToTypeIfConvertible<U, T> CudaAtomicSub(T* ptr, U value) {
   return atomicSub(ptr, value);
 }
 
-// Specializations of substraction which add the negative value.
+// Specializations of subtraction which add the negative value.
 __device__ inline float CudaAtomicSub(float* ptr, float value) {
   return CudaAtomicAdd(ptr, -value);
 }
diff --git a/tensorflow/core/util/mkl_util.h b/tensorflow/core/util/mkl_util.h
index 50a8e30574..8105121e7c 100644
--- a/tensorflow/core/util/mkl_util.h
+++ b/tensorflow/core/util/mkl_util.h
@@ -1359,7 +1359,7 @@ inline memory::dims MklDnnDimsInNCHW(const memory::dims& in_dims,
 /// Map MklDnn memory::dims object into TensorShape object.
 ///
 /// This function will simply map input shape in MKL-DNN memory::dims format
-/// in Tensorflow's TensorShape object by perserving dimension order.
+/// in Tensorflow's TensorShape object by preserving dimension order.
 ///
 /// @input MKL-DNN memory::dims object
 /// @output TensorShape corresponding to memory::dims
diff --git a/tensorflow/core/util/tensor_format.h b/tensorflow/core/util/tensor_format.h
index 646673512c..517b85a5ba 100644
--- a/tensorflow/core/util/tensor_format.h
+++ b/tensorflow/core/util/tensor_format.h
@@ -61,7 +61,7 @@ enum FilterTensorFormat {
   FORMAT_OIHW = 1,
 
   // OIHW_VECT_I is the most performant tensor format for cudnn6's quantized
-  // int8 convolution and fused convolution. It is analagous to the NCHW_VECT_C
+  // int8 convolution and fused convolution. It is analogous to the NCHW_VECT_C
   // data format. It is laid out in the same order as OIHW, except that the size
   // of the Input Channels dimension is divided by 4, and a new dimension of
   // size 4 is appended, which packs 4 adjacent input channel weights into an
diff --git a/tensorflow/docs_src/api_guides/python/reading_data.md b/tensorflow/docs_src/api_guides/python/reading_data.md
index b3ca958370..5bbbfd3216 100644
--- a/tensorflow/docs_src/api_guides/python/reading_data.md
+++ b/tensorflow/docs_src/api_guides/python/reading_data.md
@@ -184,7 +184,7 @@ The recommended way to read a TFRecord file is with a @{tf.data.TFRecordDataset}
     dataset = dataset.map(decode)
 ```
 
-To acomplish the same task with a queue based input pipeline requires the following code 
+To accomplish the same task with a queue based input pipeline requires the following code
 (using the same `decode` function from the above example): 
 
 ``` python
diff --git a/tensorflow/docs_src/deploy/s3.md b/tensorflow/docs_src/deploy/s3.md
index ef3b030e32..9ef9674338 100644
--- a/tensorflow/docs_src/deploy/s3.md
+++ b/tensorflow/docs_src/deploy/s3.md
@@ -1,6 +1,6 @@
 # How to run TensorFlow on S3
 
-Tensorflow supports reading and writing data to S3. S3 is an object storage API which is nearly ubiquitious, and can help in situations where data must accessed by multiple actors, such as in distributed training.
+Tensorflow supports reading and writing data to S3. S3 is an object storage API which is nearly ubiquitous, and can help in situations where data must accessed by multiple actors, such as in distributed training.
 
 This document guides you through the required setup, and provides examples on usage.
 
diff --git a/tensorflow/docs_src/mobile/mobile_intro.md b/tensorflow/docs_src/mobile/mobile_intro.md
index 69b63ae7d2..39dda0b45f 100644
--- a/tensorflow/docs_src/mobile/mobile_intro.md
+++ b/tensorflow/docs_src/mobile/mobile_intro.md
@@ -212,7 +212,7 @@ handle the task then it will be difficult to train a computer to do better.
 
 After you’ve solved any fundamental issues with your use case, you need to
 create a labeled dataset to define what problem you’re trying to solve. This
-step is extremely important, moreso than picking which model to use. You want it
+step is extremely important, more than picking which model to use. You want it
 to be as representative as possible of your actual use case, since the model
 will only be effective at the task you teach it. It’s also worth investing in
 tools to make labeling the data as efficient and accurate as possible. For
diff --git a/tensorflow/python/data/util/nest.py b/tensorflow/python/data/util/nest.py
index eff6e02c14..7ee3d92cad 100644
--- a/tensorflow/python/data/util/nest.py
+++ b/tensorflow/python/data/util/nest.py
@@ -114,7 +114,7 @@ def is_sequence(seq):
   NOTE(mrry): This differs from `tensorflow.python.util.nest.is_sequence()`,
   which *does* treat a Python list as a sequence. For ergonomic
   reasons, `tf.data` users would prefer to treat lists as
-  implict `tf.Tensor` objects, and dicts as (nested) sequences.
+  implicit `tf.Tensor` objects, and dicts as (nested) sequences.
 
   Args:
     seq: an input sequence.
diff --git a/tensorflow/python/estimator/estimator.py b/tensorflow/python/estimator/estimator.py
index 3691c99dda..bd16e33262 100644
--- a/tensorflow/python/estimator/estimator.py
+++ b/tensorflow/python/estimator/estimator.py
@@ -883,7 +883,7 @@ class Estimator(object):
             model_fn_lib.ModeKeys.TRAIN,
             self.config)
 
-        # TODO(anjalisridhar): Figure out how to resolve the folowing scaffold
+        # TODO(anjalisridhar): Figure out how to resolve the following scaffold
         # parameters: init_feed_dict, init_fn.
         scaffold_list = self._distribution.unwrap(
             grouped_estimator_spec.scaffold)
diff --git a/tensorflow/python/estimator/inputs/queues/feeding_functions.py b/tensorflow/python/estimator/inputs/queues/feeding_functions.py
index 8e5d8141a1..8e2ec83020 100644
--- a/tensorflow/python/estimator/inputs/queues/feeding_functions.py
+++ b/tensorflow/python/estimator/inputs/queues/feeding_functions.py
@@ -52,7 +52,7 @@ def _fill_array(arr, seq, fillvalue=0):
   If length of seq is less than arr padded length, fillvalue used.
   Args:
     arr: Padded tensor of shape [batch_size, ..., max_padded_dim_len].
-    seq: Non-padded list of data sampels of shape
+    seq: Non-padded list of data samples of shape
       [batch_size, ..., padded_dim(None)]
     fillvalue: Default fillvalue to use.
   """
diff --git a/tensorflow/python/feature_column/feature_column.py b/tensorflow/python/feature_column/feature_column.py
index c16c3cda48..9e6429e59e 100644
--- a/tensorflow/python/feature_column/feature_column.py
+++ b/tensorflow/python/feature_column/feature_column.py
@@ -48,7 +48,7 @@ should choose depends on (1) the feature type and (2) the model type.
 
       embedded_dept_column = embedding_column(
           categorical_column_with_vocabulary_list(
-              "department", ["math", "philosphy", ...]), dimension=10)
+              "department", ["math", "philosophy", ...]), dimension=10)
 
   * Wide (aka linear) models (`LinearClassifier`, `LinearRegressor`).
 
diff --git a/tensorflow/python/framework/ops.py b/tensorflow/python/framework/ops.py
index dd9acdd9eb..908d0da35e 100644
--- a/tensorflow/python/framework/ops.py
+++ b/tensorflow/python/framework/ops.py
@@ -2573,7 +2573,7 @@ def set_shape_and_handle_data_for_outputs(op):
 
   When _USE_C_API = True, this is lazily called when a tensor's shape is first
   requested. Usually this should work automatically, but some edge cases may
-  require manaully calling this first to make sure Tensor._shape_val and
+  require manually calling this first to make sure Tensor._shape_val and
   Tensor._handle_data are set (e.g. manually overriding _handle_data, copying a
   Tensor).
   """
diff --git a/tensorflow/python/framework/test_util.py b/tensorflow/python/framework/test_util.py
index dc56d88066..5e02e7e3ec 100644
--- a/tensorflow/python/framework/test_util.py
+++ b/tensorflow/python/framework/test_util.py
@@ -674,7 +674,7 @@ def run_in_graph_and_eager_modes(__unused__=None,
 
 
   Args:
-    __unused__: Prevents sliently skipping tests.
+    __unused__: Prevents silently skipping tests.
     config: An optional config_pb2.ConfigProto to use to configure the
       session when executing graphs.
     use_gpu: If True, attempt to run as many operations as possible on GPU.
diff --git a/tensorflow/python/keras/_impl/keras/engine/network.py b/tensorflow/python/keras/_impl/keras/engine/network.py
index a0229be346..3197d49fce 100644
--- a/tensorflow/python/keras/_impl/keras/engine/network.py
+++ b/tensorflow/python/keras/_impl/keras/engine/network.py
@@ -115,7 +115,7 @@ class Network(base_layer.Layer):
     # Entries are unique. Includes input and output layers.
     self._layers = []
 
-    # Used in symbolic mode only, only in conjonction with graph-networks
+    # Used in symbolic mode only, only in conjunction with graph-networks
     self._outbound_nodes = []
     self._inbound_nodes = []
 
diff --git a/tensorflow/python/keras/_impl/keras/engine/saving_test.py b/tensorflow/python/keras/_impl/keras/engine/saving_test.py
index 709a8e9fb1..c0b16b6bf5 100644
--- a/tensorflow/python/keras/_impl/keras/engine/saving_test.py
+++ b/tensorflow/python/keras/_impl/keras/engine/saving_test.py
@@ -457,7 +457,7 @@ class TestWholeModelSaving(test.TestCase):
       with h5py.File(fname, 'r') as h5file:
         num_names_arrays = len([attr for attr in h5file['model_weights'].attrs
                                 if attr.startswith('layer_names')])
-      # The chunking of layer names array should have happend.
+      # The chunking of layer names array should have happened.
       self.assertGreater(num_names_arrays, 0)
       out2 = model.predict(x)
       self.assertAllClose(out, out2, atol=1e-05)
@@ -502,7 +502,7 @@ class TestWholeModelSaving(test.TestCase):
         num_weight_arrays = len(
             [attr for attr in h5file['model_weights']['nested_model'].attrs
              if attr.startswith('weight_names')])
-      # The chunking of layer names array should have happend.
+      # The chunking of layer names array should have happened.
       self.assertGreater(num_weight_arrays, 0)
       out2 = model.predict(x)
       self.assertAllClose(out, out2, atol=1e-05)
diff --git a/tensorflow/python/keras/_impl/keras/estimator.py b/tensorflow/python/keras/_impl/keras/estimator.py
index c3c3fceb45..5c79c964c8 100644
--- a/tensorflow/python/keras/_impl/keras/estimator.py
+++ b/tensorflow/python/keras/_impl/keras/estimator.py
@@ -72,7 +72,7 @@ def _any_variable_initalized():
   """Check if any variable has been initialized in the Keras model.
 
   Returns:
-    boolean, True if at least one variable has been initalized, else False.
+    boolean, True if at least one variable has been initialized, else False.
   """
   variables = variables_module.global_variables()
   for v in variables:
diff --git a/tensorflow/python/kernel_tests/distributions/util_test.py b/tensorflow/python/kernel_tests/distributions/util_test.py
index f54f146e0a..5ec80b95ee 100644
--- a/tensorflow/python/kernel_tests/distributions/util_test.py
+++ b/tensorflow/python/kernel_tests/distributions/util_test.py
@@ -703,7 +703,7 @@ class FillTriangularTest(test.TestCase):
       raise ValueError("Invalid shape.")
     n = np.int32(n)
     # We can't do: `x[..., -(n**2-m):]` because this doesn't correctly handle
-    # `m == n == 1`. Hence, we do absoulte indexing.
+    # `m == n == 1`. Hence, we do absolute indexing.
     x_tail = x[..., (m - (n * n - m)):]
     y = np.concatenate(
         [x, x_tail[..., ::-1]] if upper else [x_tail, x[..., ::-1]],
diff --git a/tensorflow/python/kernel_tests/manip_ops_test.py b/tensorflow/python/kernel_tests/manip_ops_test.py
index f31426713c..dc3ea38671 100644
--- a/tensorflow/python/kernel_tests/manip_ops_test.py
+++ b/tensorflow/python/kernel_tests/manip_ops_test.py
@@ -93,7 +93,7 @@ class RollTest(test_util.TensorFlowTestCase):
   def testNegativeAxis(self):
     self._testAll(np.random.randint(-100, 100, (5)).astype(np.int32), 3, -1)
     self._testAll(np.random.randint(-100, 100, (4, 4)).astype(np.int32), 3, -2)
-    # Make sure negative axis shoudl be 0 <= axis + dims < dims
+    # Make sure negative axis should be 0 <= axis + dims < dims
     with self.test_session():
       with self.assertRaisesRegexp(errors_impl.InvalidArgumentError,
                                    "is out of range"):
diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py
index 7ac3bd8091..477f870060 100644
--- a/tensorflow/python/ops/math_ops.py
+++ b/tensorflow/python/ops/math_ops.py
@@ -1285,7 +1285,7 @@ def reduce_sum(input_tensor,
     The reduced tensor, of the same dtype as the input_tensor.
 
   @compatibility(numpy)
-  Equivalent to np.sum appart the fact that numpy upcast uint8 and int32 to
+  Equivalent to np.sum apart the fact that numpy upcast uint8 and int32 to
   int64 while tensorflow returns the same dtype as the input.
   @end_compatibility
   """
diff --git a/tensorflow/python/training/distribute.py b/tensorflow/python/training/distribute.py
index c16b05102e..30c857fd18 100644
--- a/tensorflow/python/training/distribute.py
+++ b/tensorflow/python/training/distribute.py
@@ -734,7 +734,7 @@ class DistributionStrategy(object):
     `fn` may call `tf.get_tower_context()` to access methods such as
     `tower_id()` and `merge_call()`.
 
-    `merge_call()` is used to communicate betwen the towers and
+    `merge_call()` is used to communicate between the towers and
     re-enter the cross-tower context. All towers pause their execution
     having encountered a `merge_call()` call. After that the
     `merge_fn`-function is executed. Its results are then unwrapped and
diff --git a/tensorflow/python/util/util.cc b/tensorflow/python/util/util.cc
index 70aee4a3f6..9c8d50da73 100644
--- a/tensorflow/python/util/util.cc
+++ b/tensorflow/python/util/util.cc
@@ -234,7 +234,7 @@ void SetDifferentKeysError(PyObject* dict1, PyObject* dict2, string* error_msg,
 
 // Returns true iff there were no "internal" errors. In other words,
 // errors that has nothing to do with structure checking.
-// If an "internal" error occured, the appropriate Python error will be
+// If an "internal" error occurred, the appropriate Python error will be
 // set and the caller can propage it directly to the user.
 //
 // Both `error_msg` and `is_type_error` must be non-null. `error_msg` must
diff --git a/tensorflow/python/util/util.h b/tensorflow/python/util/util.h
index c325baa5f8..4bb80d8289 100644
--- a/tensorflow/python/util/util.h
+++ b/tensorflow/python/util/util.h
@@ -97,7 +97,7 @@ PyObject* AssertSameStructure(PyObject* o1, PyObject* o2, bool check_types);
 // used instead. The same convention is followed in `pack_sequence_as`. This
 // correctly repacks dicts and `OrderedDict`s after they have been flattened,
 // and also allows flattening an `OrderedDict` and then repacking it back using
-// a correponding plain dict, or vice-versa.
+// a corresponding plain dict, or vice-versa.
 // Dictionaries with non-sortable keys cannot be flattened.
 //
 // Args:
diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.h b/tensorflow/stream_executor/cuda/cuda_dnn.h
index 7d53dbe4a5..dfe2779949 100644
--- a/tensorflow/stream_executor/cuda/cuda_dnn.h
+++ b/tensorflow/stream_executor/cuda/cuda_dnn.h
@@ -639,7 +639,7 @@ class CudnnSupport : public dnn::DnnSupport {
   // Guards the enqueueing of DNN operations via the dnn_handle_ below, and
   // access to current_dnn_stream_.
   //
-  // This is a public member because we need to add thread safty annotations in
+  // This is a public member because we need to add thread safety annotations in
   // the cudnn wrapper functions in the cc file, which need to access this
   // mutex (the annotations require C++ permission checks).
   mutex dnn_handle_mutex_;
diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl
index e5cc886b32..f8f83add6a 100644
--- a/tensorflow/tensorflow.bzl
+++ b/tensorflow/tensorflow.bzl
@@ -1492,7 +1492,7 @@ def tf_py_wrap_cc(name,
 # This macro is for running python tests against system installed pip package
 # on Windows.
 #
-# py_test is built as an exectuable python zip file on Windows, which contains all
+# py_test is built as an executable python zip file on Windows, which contains all
 # dependencies of the target. Because of the C++ extensions, it would be very
 # inefficient if the py_test zips all runfiles, plus we don't need them when running
 # tests against system installed pip package. So we'd like to get rid of the deps
diff --git a/tensorflow/tools/graph_transforms/README.md b/tensorflow/tools/graph_transforms/README.md
index 67badb4869..9f6f553ba1 100644
--- a/tensorflow/tools/graph_transforms/README.md
+++ b/tensorflow/tools/graph_transforms/README.md
@@ -388,7 +388,7 @@ input is collapsed down into a simple constant.
 Args:
 
 *   clear_output_shapes: Clears tensor shape information saved as attributes.
-    Some older graphs containes out-of-date information and may cause import
+    Some older graphs contains out-of-date information and may cause import
     errors. Defaults to true.
 
 Prerequisites: None
diff --git a/third_party/examples/eager/spinn/README.md b/third_party/examples/eager/spinn/README.md
index 7f477d1920..fbb1fde837 100644
--- a/third_party/examples/eager/spinn/README.md
+++ b/third_party/examples/eager/spinn/README.md
@@ -70,7 +70,7 @@ Other eager execution examples can be found under [tensorflow/contrib/eager/pyth
 - After training, you may use the model to perform inference on input data in
   the SNLI data format. The premise and hypotheses sentences are specified with
   the command-line flags `--inference_premise` and `--inference_hypothesis`,
-  respecitvely. Each sentence should include the words, as well as parentheses
+  respectively. Each sentence should include the words, as well as parentheses
   representing a binary parsing of the sentence. The words and parentheses
   should all be separated by spaces. For instance,
 
-- 
GitLab


From ebcde41d721ec554a7840cb18e4e8a7a489e424a Mon Sep 17 00:00:00 2001
From: Aditya Yogi <yogiadi@users.noreply.github.com>
Date: Thu, 3 May 2018 23:43:40 +0530
Subject: [PATCH 206/395] Update learning.py (#19064)

---
 tensorflow/contrib/slim/python/slim/learning.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/slim/python/slim/learning.py b/tensorflow/contrib/slim/python/slim/learning.py
index 8a2c74742a..6e55b9407b 100644
--- a/tensorflow/contrib/slim/python/slim/learning.py
+++ b/tensorflow/contrib/slim/python/slim/learning.py
@@ -571,7 +571,7 @@ def train(train_op,
       default, two `Boolean`, scalar ops called "should_stop" and "should_log"
       are provided.
     log_every_n_steps: The frequency, in terms of global steps, that the loss
-      and global step and logged.
+      and global step are logged.
     graph: The graph to pass to the supervisor. If no graph is supplied the
       default graph is used.
     master: The address of the tensorflow master.
-- 
GitLab


From 85a47596caf89705aae8ffcb57fcdaecb22fe356 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 3 May 2018 11:16:06 -0700
Subject: [PATCH 207/395] [XLA] Redesign: add ExecuteGraph to grpc service.

PiperOrigin-RevId: 195281004
---
 tensorflow/compiler/xla/rpc/BUILD               | 2 +-
 tensorflow/compiler/xla/rpc/grpc_client_test.cc | 4 ++--
 tensorflow/compiler/xla/rpc/grpc_service.cc     | 7 +++++++
 tensorflow/compiler/xla/rpc/grpc_service.h      | 4 ++++
 third_party/libxsmm.BUILD                       | 2 +-
 5 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/tensorflow/compiler/xla/rpc/BUILD b/tensorflow/compiler/xla/rpc/BUILD
index 977f863787..0d56a9a477 100644
--- a/tensorflow/compiler/xla/rpc/BUILD
+++ b/tensorflow/compiler/xla/rpc/BUILD
@@ -55,7 +55,7 @@ tf_cc_test(
     deps = [
         ":grpc_stub",
         "//tensorflow/compiler/xla/client",
-        "//tensorflow/compiler/xla/client:computation_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
         "//tensorflow/compiler/xla/tests:literal_test_util",
         "//tensorflow/core:framework_internal",
         "//tensorflow/core:lib",
diff --git a/tensorflow/compiler/xla/rpc/grpc_client_test.cc b/tensorflow/compiler/xla/rpc/grpc_client_test.cc
index b559ee4b5a..10997c0719 100644
--- a/tensorflow/compiler/xla/rpc/grpc_client_test.cc
+++ b/tensorflow/compiler/xla/rpc/grpc_client_test.cc
@@ -24,7 +24,7 @@ limitations under the License.
 #include "grpc++/security/credentials.h"
 
 #include "tensorflow/compiler/xla/client/client.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/rpc/grpc_stub.h"
 #include "tensorflow/compiler/xla/tests/literal_test_util.h"
 #include "tensorflow/core/lib/io/path.h"
@@ -84,7 +84,7 @@ TEST_F(GRPCClientTestBase, ItsAlive) {
 }
 
 TEST_F(GRPCClientTestBase, AxpyTenValues) {
-  ComputationBuilder builder(client_.get(), "axpy_10");
+  XlaBuilder builder("axpy_10");
   auto alpha = builder.ConstantR0<float>(3.1415926535);
   auto x = builder.ConstantR1<float>(
       {-1.0, 1.0, 2.0, -2.0, -3.0, 3.0, 4.0, -4.0, -5.0, 5.0});
diff --git a/tensorflow/compiler/xla/rpc/grpc_service.cc b/tensorflow/compiler/xla/rpc/grpc_service.cc
index 0b100bd108..ffb72fc73c 100644
--- a/tensorflow/compiler/xla/rpc/grpc_service.cc
+++ b/tensorflow/compiler/xla/rpc/grpc_service.cc
@@ -75,6 +75,13 @@ namespace xla {
       [this, arg, result]() { return service_->Execute(arg, result); });
 }
 
+::grpc::Status GRPCService::ExecuteGraph(::grpc::ServerContext* /*context*/,
+                                         const ExecuteGraphRequest* arg,
+                                         ExecuteResponse* result) {
+  return DelegateRPC(
+      [this, arg, result]() { return service_->ExecuteGraph(arg, result); });
+}
+
 ::grpc::Status GRPCService::ExecuteAsync(::grpc::ServerContext* context,
                                          const ExecuteAsyncRequest* arg,
                                          ExecuteAsyncResponse* result) {
diff --git a/tensorflow/compiler/xla/rpc/grpc_service.h b/tensorflow/compiler/xla/rpc/grpc_service.h
index fad74375bd..50f02796f2 100644
--- a/tensorflow/compiler/xla/rpc/grpc_service.h
+++ b/tensorflow/compiler/xla/rpc/grpc_service.h
@@ -54,6 +54,10 @@ class GRPCService : public grpc::XlaService::Service {
                          const ExecuteRequest* arg,
                          ExecuteResponse* result) override;
 
+  ::grpc::Status ExecuteGraph(::grpc::ServerContext* context,
+                              const ExecuteGraphRequest* arg,
+                              ExecuteResponse* result) override;
+
   ::grpc::Status ExecuteAsync(::grpc::ServerContext* context,
                               const ExecuteAsyncRequest* arg,
                               ExecuteAsyncResponse* result) override;
diff --git a/third_party/libxsmm.BUILD b/third_party/libxsmm.BUILD
index 4124f2db63..78ed1f4e16 100644
--- a/third_party/libxsmm.BUILD
+++ b/third_party/libxsmm.BUILD
@@ -38,8 +38,8 @@ genrule(
         ":libxsmm_interface",
     ],
     visibility = [
-        "//tensorflow/core/kernels:__pkg__",
         "//third_party/eigen3:__pkg__",
+        "//tensorflow/core/kernels:__pkg__",
     ],
 )
 
-- 
GitLab


From a16ba4fc0d3faec077c689f3f361264978a2d3cb Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 3 May 2018 12:00:57 -0700
Subject: [PATCH 208/395] Do not delegate temporary tensors to NNAPI.   - also
 added delegation for MUL, and set the default scale to be 0.0f.

PiperOrigin-RevId: 195288948
---
 tensorflow/contrib/lite/nnapi_delegate.cc | 39 +++++++++++++++++++----
 1 file changed, 33 insertions(+), 6 deletions(-)

diff --git a/tensorflow/contrib/lite/nnapi_delegate.cc b/tensorflow/contrib/lite/nnapi_delegate.cc
index 6a78f30fd1..e1895dd38e 100644
--- a/tensorflow/contrib/lite/nnapi_delegate.cc
+++ b/tensorflow/contrib/lite/nnapi_delegate.cc
@@ -72,11 +72,23 @@ NNAPIDelegate::~NNAPIDelegate() {
 // Adds the tensors of the interpreter to the NN API model.
 // Returns the number of operands added.
 uint32_t addTensorOperands(tflite::Interpreter* interpreter,
-                           ANeuralNetworksModel* nn_model) {
+                           ANeuralNetworksModel* nn_model,
+                           const std::vector<uint32_t>& skip_list) {
   uint32_t next_id = 0;
   for (size_t i = 0; i < interpreter->tensors_size(); i++) {
+    // skip temporaries tensors.
+    bool shouldSkip = false;
+    for (auto skip_idx : skip_list) {
+      if (i == skip_idx) {
+        shouldSkip = true;
+        break;
+      }
+    }
+    if (shouldSkip) continue;
+
     int32_t nn_type = 0;
-    float scale = 1.0f;
+    // NNAPI requires 32-bit float scale to be zero, tflite doesn't care
+    float scale = 0.0f;
     int32_t zeroPoint = 0;
     TfLiteTensor* tensor = interpreter->tensor(i);
     switch (tensor->type) {
@@ -116,11 +128,11 @@ uint32_t addTensorOperands(tflite::Interpreter* interpreter,
       if (const NNAPIAllocation* alloc = dynamic_cast<const NNAPIAllocation*>(
               static_cast<const Allocation*>(tensor->allocation))) {
         CHECK_NN(ANeuralNetworksModel_setOperandValueFromMemory(
-            nn_model, i, alloc->memory(), alloc->offset(tensor->data.raw),
+            nn_model, next_id, alloc->memory(), alloc->offset(tensor->data.raw),
             tensor->bytes));
       } else {
         CHECK_NN(ANeuralNetworksModel_setOperandValue(
-            nn_model, i, tensor->data.raw, tensor->bytes));
+            nn_model, next_id, tensor->data.raw, tensor->bytes));
       }
     }
     ++next_id;
@@ -253,6 +265,10 @@ void AddOpsAndParams(tflite::Interpreter* interpreter,
         nn_op_type = ANEURALNETWORKS_ADD;
         add_add_params();
         break;
+      case tflite::BuiltinOperator_MUL:
+        nn_op_type = ANEURALNETWORKS_MUL;
+        add_add_params();
+        break;
       case tflite::BuiltinOperator_AVERAGE_POOL_2D:
         add_pooling_params(node.builtin_data);
         nn_op_type = ANEURALNETWORKS_AVERAGE_POOL_2D;
@@ -330,7 +346,6 @@ void AddOpsAndParams(tflite::Interpreter* interpreter,
       case tflite::BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM:
       case tflite::BuiltinOperator_L2_NORMALIZATION:
       case tflite::BuiltinOperator_LOCAL_RESPONSE_NORMALIZATION:
-      case tflite::BuiltinOperator_MUL:
       case tflite::BuiltinOperator_PAD:
       case tflite::BuiltinOperator_RESIZE_BILINEAR:
       case tflite::BuiltinOperator_CALL:
@@ -381,7 +396,19 @@ TfLiteStatus NNAPIDelegate::BuildGraph(Interpreter* interpreter) {
   if (!nn_model_) {
     CHECK_NN(ANeuralNetworksModel_create(&nn_model_));
 
-    uint32_t next_id = addTensorOperands(interpreter, nn_model_);
+    // Find all the temporary tensors and put them in a skip_list.
+    std::vector<uint32_t> skip_list;
+    for (size_t i = 0; i < interpreter->nodes_size(); i++) {
+      const auto* node_and_registration = interpreter->node_and_registration(i);
+      const TfLiteNode& node = node_and_registration->first;
+      if (node.temporaries != nullptr) {
+        for (int j = 0; j < node.temporaries->size; j++) {
+          skip_list.push_back(static_cast<uint32_t>(node.temporaries->data[j]));
+        }
+      }
+    }
+
+    uint32_t next_id = addTensorOperands(interpreter, nn_model_, skip_list);
     AddOpsAndParams(interpreter, nn_model_, next_id);
     CHECK_NN(ANeuralNetworksModel_identifyInputsAndOutputs(
         nn_model_, static_cast<uint32_t>(interpreter->inputs().size()),
-- 
GitLab


From e5854637cc3f8099586f18ed144fd6d4f90a6fc7 Mon Sep 17 00:00:00 2001
From: Yao Zhang <yaozhang@google.com>
Date: Thu, 3 May 2018 12:19:01 -0700
Subject: [PATCH 209/395] Simplify file reading and support SavedModel.

PiperOrigin-RevId: 195291836
---
 .../python/grappler/cost_analyzer_tool.py     | 75 ++++++++++---------
 1 file changed, 41 insertions(+), 34 deletions(-)

diff --git a/tensorflow/python/grappler/cost_analyzer_tool.py b/tensorflow/python/grappler/cost_analyzer_tool.py
index 0853db2524..e6229e1856 100644
--- a/tensorflow/python/grappler/cost_analyzer_tool.py
+++ b/tensorflow/python/grappler/cost_analyzer_tool.py
@@ -21,11 +21,13 @@ from __future__ import print_function
 import argparse
 import sys
 
+from google.protobuf import message
 from google.protobuf import text_format
 from tensorflow.contrib.fused_conv.ops import gen_fused_conv2d_bias_activation_op  # pylint: disable=unused-import
 from tensorflow.core.framework import graph_pb2
 from tensorflow.core.protobuf import meta_graph_pb2
 from tensorflow.core.protobuf import rewriter_config_pb2
+from tensorflow.core.protobuf import saved_model_pb2
 from tensorflow.python.framework import importer
 from tensorflow.python.framework import ops
 from tensorflow.python.grappler import cost_analyzer
@@ -37,33 +39,42 @@ from tensorflow.python.training import saver
 
 def get_metagraph():
   """Constructs and returns a MetaGraphDef from the input file."""
-  if FLAGS.metagraphdef:
-    with gfile.GFile(FLAGS.metagraphdef) as meta_file:
-      metagraph = meta_graph_pb2.MetaGraphDef()
-      if FLAGS.metagraphdef.endswith(".pbtxt"):
-        text_format.Merge(meta_file.read(), metagraph)
-      else:
-        metagraph.ParseFromString(meta_file.read())
-    if FLAGS.fetch is not None:
-      fetch_collection = meta_graph_pb2.CollectionDef()
-      for fetch in FLAGS.fetch.split(","):
-        fetch_collection.node_list.value.append(fetch)
-      metagraph.collection_def["train_op"].CopyFrom(fetch_collection)
-  else:
-    with gfile.GFile(FLAGS.graphdef) as graph_file:
-      graph_def = graph_pb2.GraphDef()
-      if FLAGS.graphdef.endswith(".pbtxt"):
-        text_format.Merge(graph_file.read(), graph_def)
-      else:
-        graph_def.ParseFromString(graph_file.read())
-      importer.import_graph_def(graph_def, name="")
-      graph = ops.get_default_graph()
-      for fetch in FLAGS.fetch.split(","):
-        fetch_op = graph.get_operation_by_name(fetch)
-        graph.add_to_collection("train_op", fetch_op)
-      metagraph = saver.export_meta_graph(
-          graph_def=graph.as_graph_def(), graph=graph)
-  return metagraph
+  with gfile.GFile(FLAGS.input) as input_file:
+    input_data = input_file.read()
+    try:
+      saved_model = saved_model_pb2.SavedModel()
+      text_format.Merge(input_data, saved_model)
+      meta_graph = saved_model.meta_graphs[0]
+    except text_format.ParseError:
+      try:
+        saved_model.ParseFromString(input_data)
+        meta_graph = saved_model.meta_graphs[0]
+      except message.DecodeError:
+        try:
+          meta_graph = meta_graph_pb2.MetaGraphDef()
+          text_format.Merge(input_data, meta_graph)
+        except text_format.ParseError:
+          try:
+            meta_graph.ParseFromString(input_data)
+          except message.DecodeError:
+            try:
+              graph_def = graph_pb2.GraphDef()
+              text_format.Merge(input_data, graph_def)
+            except text_format.ParseError:
+              try:
+                graph_def.ParseFromString(input_data)
+              except message.DecodeError:
+                raise ValueError("Invalid input file.")
+            importer.import_graph_def(graph_def, name="")
+            graph = ops.get_default_graph()
+            meta_graph = saver.export_meta_graph(
+                graph_def=graph.as_graph_def(), graph=graph)
+  if FLAGS.fetch is not None:
+    fetch_collection = meta_graph_pb2.CollectionDef()
+    for fetch in FLAGS.fetch.split(","):
+      fetch_collection.node_list.value.append(fetch)
+    meta_graph.collection_def["train_op"].CopyFrom(fetch_collection)
+  return meta_graph
 
 
 def main(_):
@@ -85,15 +96,11 @@ def main(_):
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument(
-      "--metagraphdef",
+      "--input",
       type=str,
       default=None,
-      help="Input .meta MetaGraphDef file path.")
-  parser.add_argument(
-      "--graphdef",
-      type=str,
-      default=None,
-      help="Input .pb GraphDef file path.")
+      help="Input file path. Accept SavedModel, MetaGraphDef, and GraphDef in "
+      "either binary or text format.")
   parser.add_argument(
       "--fetch",
       type=str,
-- 
GitLab


From 4b767a835b61061ef4d167dc1ee935f2f85a3e87 Mon Sep 17 00:00:00 2001
From: Allen Lavoie <allenl@google.com>
Date: Thu, 3 May 2018 12:53:47 -0700
Subject: [PATCH 210/395] Small fix for an eager colab notebook.

PiperOrigin-RevId: 195296384
---
 .../contrib/eager/python/examples/notebooks/1_basics.ipynb      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/eager/python/examples/notebooks/1_basics.ipynb b/tensorflow/contrib/eager/python/examples/notebooks/1_basics.ipynb
index 0279db80fa..9fd2d8d125 100644
--- a/tensorflow/contrib/eager/python/examples/notebooks/1_basics.ipynb
+++ b/tensorflow/contrib/eager/python/examples/notebooks/1_basics.ipynb
@@ -478,7 +478,7 @@
       "source": [
         "# Time GPU-based matrix multiplications.\n",
         "\n",
-        "if is_gpu_available:\n",
+        "if tf.test.is_gpu_available():\n",
         "  # First use of the GPU will be slow:\n",
         "  print(\"Time to conduct first matmul on GPU:\")\n",
         "  %time tf.matmul(gpu_tensor, gpu_tensor)\n",
-- 
GitLab


From 775d1c03c1772c0c2e10e5884af8d9363cfdf314 Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Thu, 3 May 2018 12:59:33 -0700
Subject: [PATCH 211/395] [TF:XLA] Bump open source llvm revision to r331442

PiperOrigin-RevId: 195297133
---
 tensorflow/workspace.bzl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 94cac4f8fa..8b6ad0a138 100644
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -452,11 +452,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""):
   tf_http_archive(
       name = "llvm",
       urls = [
-          "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/a5108a08ceab35886a7df07c86f96aedd3d94bb7.tar.gz",
-          "https://github.com/llvm-mirror/llvm/archive/a5108a08ceab35886a7df07c86f96aedd3d94bb7.tar.gz",
+          "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/b3f6a6a61625296bb532a65c0bf51b91b05b3361.tar.gz",
+          "https://github.com/llvm-mirror/llvm/archive/b3f6a6a61625296bb532a65c0bf51b91b05b3361.tar.gz",
       ],
-      sha256 = "79cae03ebbdfd812bb69c460e1325ca069b5c576f7c7071f8216cf2b0975e36f",
-      strip_prefix = "llvm-a5108a08ceab35886a7df07c86f96aedd3d94bb7",
+      sha256 = "93895b289a78a47a1e75652e12a1b9a6c119f086a509b00e0084cf2bb944b709",
+      strip_prefix = "llvm-b3f6a6a61625296bb532a65c0bf51b91b05b3361",
       build_file = clean_dep("//third_party/llvm:llvm.BUILD"),
   )
 
-- 
GitLab


From ceda30408f66a7eea86dc359164deb662d5a32d0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 3 May 2018 13:00:56 -0700
Subject: [PATCH 212/395] Enable unary chain hoisting optimization for
 concat/split/splitv by default.

PiperOrigin-RevId: 195297330
---
 tensorflow/core/grappler/op_types.cc          | 38 ++++++++++++-------
 tensorflow/core/grappler/op_types.h           |  4 ++
 .../optimizers/arithmetic_optimizer.cc        | 18 ++++++---
 .../optimizers/arithmetic_optimizer.h         |  2 +-
 .../optimizers/arithmetic_optimizer_test.cc   | 16 ++++----
 5 files changed, 51 insertions(+), 27 deletions(-)

diff --git a/tensorflow/core/grappler/op_types.cc b/tensorflow/core/grappler/op_types.cc
index 7c936dfca1..c48dc00941 100644
--- a/tensorflow/core/grappler/op_types.cc
+++ b/tensorflow/core/grappler/op_types.cc
@@ -476,28 +476,40 @@ bool IsInvolution(const NodeDef& node) {
   return involution_ops->count(node.op()) > 0;
 }
 
+bool IsValueAndOrderAndShapePreserving(const NodeDef& node) {
+  if (NumNonControlInputs(node) == 1 && IsAggregate(node)) {
+    return true;
+  }
+  static const std::unordered_set<string>*
+      value_and_order_and_shape_preserving_ops =
+          CHECK_NOTNULL((new const std::unordered_set<string>{
+              "CheckNumerics",
+              "DebugGradientIdentity",
+              "DeepCopy"
+              "Enter",
+              "Exit",
+              "Identity",
+              "IdentityN",
+              "PreventGradient",
+              "Print",
+              "Snapshot",
+              "StopGradient",
+          }));
+  return value_and_order_and_shape_preserving_ops->count(node.op()) > 0;
+}
+
 bool IsValueAndOrderPreserving(const NodeDef& node) {
   if (NumNonControlInputs(node) == 1 && IsAggregate(node)) {
     return true;
   }
   static const std::unordered_set<string>* value_and_order_preserving_ops =
       CHECK_NOTNULL((new const std::unordered_set<string>{
-          "CheckNumerics",
-          "DebugGradientIdentity",
-          "DeepCopy"
-          "Enter",
-          "Exit",
           "ExpandDims",
-          "Identity",
-          "IdentityN",
-          "PreventGradient",
-          "Print",
-          "Reshape",
           "Snapshot",
           "Squeeze",
-          "StopGradient",
       }));
-  return value_and_order_preserving_ops->count(node.op()) > 0;
+  return value_and_order_preserving_ops->count(node.op()) > 0 ||
+         IsValueAndOrderAndShapePreserving(node);
 }
 
 bool IsValuePreserving(const NodeDef& node) {
@@ -564,7 +576,7 @@ bool IsUnaryElementWise(const NodeDef& node) {
           "Tanh",
       }));
   return element_wise_ops->count(node.op()) > 0 ||
-         (!IsIdentityN(node) && IsValueAndOrderPreserving(node));
+         (!IsIdentityN(node) && IsValueAndOrderAndShapePreserving(node));
 }
 
 bool HasOpDef(const NodeDef& node) {
diff --git a/tensorflow/core/grappler/op_types.h b/tensorflow/core/grappler/op_types.h
index 7a1b438768..e33dd21538 100644
--- a/tensorflow/core/grappler/op_types.h
+++ b/tensorflow/core/grappler/op_types.h
@@ -174,6 +174,10 @@ bool ModifiesInputsInPlace(const NodeDef& node);
 // own inverse such that f(f(x)) == x.
 bool IsInvolution(const NodeDef& node);
 
+// Returns true if the op preserves the order and value of elements
+// and shape of its first input tensor.
+bool IsValueAndOrderAndShapePreserving(const NodeDef& node);
+
 // Returns true if the op preserves the order and value of elements in its
 // first input tensor and possible changes its shape.
 bool IsValueAndOrderPreserving(const NodeDef& node);
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
index d6510ba681..2a5654f752 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
@@ -1400,6 +1400,11 @@ class HoistCWiseUnaryChainsStage : public ArithmeticOptimizerStage {
       return n > 1;
     } else if (IsSplit(*node) || IsSplitV(*node)) {
       const int num_split = node->attr().at("num_split").i();
+      if (NumNonControlOutputs(*node, *ctx().node_map) > num_split) {
+        // TODO(rmlarsen): Remove this constraint when we have optimizations
+        // in place for merging slices into splits.
+        return false;
+      }
       return num_split > 1 && !IsAlreadyOptimized(*node);
     }
     return false;
@@ -1458,13 +1463,13 @@ class HoistCWiseUnaryChainsStage : public ArithmeticOptimizerStage {
     if (tails.empty()) {
       return Status::OK();
     }
-    AddControlInputs(ctrl_inputs, root_node);
     AddToOptimizationQueue(root_node);
     optimized_nodes_.insert(root_node->name());
     if (node_is_concat_) {
+      AddControlInputs(ctrl_inputs, root_node);
       return HoistChainForConcat(prefix_length, tails, root_node);
     } else {
-      return HoistChainForSplit(prefix_length, tails, root_node);
+      return HoistChainForSplit(prefix_length, tails, ctrl_inputs, root_node);
     }
   }
 
@@ -1542,9 +1547,8 @@ class HoistCWiseUnaryChainsStage : public ArithmeticOptimizerStage {
           IsInPreserveSet(*op)) {
         return false;
       }
-      if (node_is_concat_ &&
-          ctx().node_map->GetOutputs(op->name()).size() > 1) {
-        // TODO(rmlarsen): Allow and hoist outgoing control edges.
+      if (ctx().node_map->GetOutputs(op->name()).size() > 1) {
+        // TODO(rmlarsen): Allow outgoing control edges.
         return false;
       }
     }
@@ -1612,6 +1616,7 @@ class HoistCWiseUnaryChainsStage : public ArithmeticOptimizerStage {
   }
 
   Status HoistChainForSplit(const int prefix_length, const ChainLinkSet& tails,
+                            std::set<string>* ctrl_inputs,
                             NodeDef* split_node) {
     // Create a new chain before the split node to process the input tensor.
     const string& split_name = split_node->name();
@@ -1646,6 +1651,9 @@ class HoistCWiseUnaryChainsStage : public ArithmeticOptimizerStage {
     cur_copy->add_input(orig_input);
     ctx().node_map->UpdateOutput(NodeName(orig_input), split_name,
                                  cur_copy->name());
+    // Make sure all the control inputs are satisfied before running the first
+    // node in the new chain.
+    AddControlInputs(ctrl_inputs, cur_copy);
 
     // Connect all consumers of the tail nodes directly to the
     // output port of Split from which the chain started.
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
index 3b297ec0aa..6309dc1a33 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
@@ -65,7 +65,7 @@ class ArithmeticOptimizer : public GraphOptimizer {
     bool remove_redundant_bitcast = true;
     bool remove_redundant_cast = true;
     bool remove_negation = true;
-    bool hoist_cwise_unary_chains = false;
+    bool hoist_cwise_unary_chains = true;
     bool convert_sqrt_div_to_rsqrt_mul = false;
 
     // Choose which arithmetic optimizer stages will be enabled for a given
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
index f903f53a35..d32743f3f2 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
@@ -2320,16 +2320,16 @@ TEST_F(ArithmeticOptimizerTest, HoistCWiseUnaryIntoSplit) {
     EXPECT_NE(node.name(), "cos_exp_b2");
 
     if (node.name() == "split1") {
-      EXPECT_EQ(3, node.input_size());
+      EXPECT_EQ(2, node.input_size());
       EXPECT_EQ("axis", node.input(0));
       EXPECT_EQ("ArithmeticOptimizer/_sin_a_split1", node.input(1));
-      EXPECT_EQ("^ctrl1", node.input(2));
       found++;
     }
     if (node.name() == "ArithmeticOptimizer/_sin_a_split1") {
       EXPECT_EQ("Sin", node.op());
-      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ(2, node.input_size());
       EXPECT_EQ("x", node.input(0));
+      EXPECT_EQ("^ctrl1", node.input(1));
       found++;
     }
     if (node.name() == "id_a") {
@@ -2349,8 +2349,11 @@ TEST_F(ArithmeticOptimizerTest, HoistCWiseUnaryIntoSplit) {
     }
     if (node.name() == "ArithmeticOptimizer/_exp_a2_split2") {
       EXPECT_EQ("Exp", node.op());
-      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ(4, node.input_size());
       EXPECT_EQ("x", node.input(0));
+      EXPECT_EQ("^ctrl1", node.input(1));
+      EXPECT_EQ("^ctrl2", node.input(2));
+      EXPECT_EQ("^ctrl3", node.input(3));
       found++;
     }
     if (node.name() == "ArithmeticOptimizer/_cos_exp_a2_split2") {
@@ -2360,13 +2363,10 @@ TEST_F(ArithmeticOptimizerTest, HoistCWiseUnaryIntoSplit) {
       found++;
     }
     if (node.name() == "split2") {
-      EXPECT_EQ(6, node.input_size());
+      EXPECT_EQ(3, node.input_size());
       EXPECT_EQ("ArithmeticOptimizer/_cos_exp_a2_split2", node.input(0));
       EXPECT_EQ("size_splits2", node.input(1));
       EXPECT_EQ("axis", node.input(2));
-      EXPECT_EQ("^ctrl1", node.input(3));
-      EXPECT_EQ("^ctrl2", node.input(4));
-      EXPECT_EQ("^ctrl3", node.input(5));
       found++;
     }
     if (node.name() == "id_a2") {
-- 
GitLab


From fded0f901c99087b100191273e28692f9b4569ee Mon Sep 17 00:00:00 2001
From: Ruoxin Sang <rxsang@google.com>
Date: Thu, 3 May 2018 13:03:48 -0700
Subject: [PATCH 213/395] Change all std::bind usages in GCS to lambdas. Fix
 the wrong #define Guard name in retrying_file_system.h.

PiperOrigin-RevId: 195297877
---
 .../core/platform/cloud/gcs_dns_cache.cc      |  4 +-
 .../core/platform/cloud/gcs_file_system.cc    |  5 +-
 .../platform/cloud/retrying_file_system.h     | 81 ++++++++++---------
 .../core/platform/cloud/retrying_utils.cc     |  6 +-
 4 files changed, 52 insertions(+), 44 deletions(-)

diff --git a/tensorflow/core/platform/cloud/gcs_dns_cache.cc b/tensorflow/core/platform/cloud/gcs_dns_cache.cc
index 4d9aff4d24..f2e64662a9 100644
--- a/tensorflow/core/platform/cloud/gcs_dns_cache.cc
+++ b/tensorflow/core/platform/cloud/gcs_dns_cache.cc
@@ -71,8 +71,8 @@ void GcsDnsCache::AnnotateRequest(HttpRequest* request) {
     addresses_ = ResolveNames(kCachedDomainNames);
 
     // Note: we opt to use a thread instead of a delayed closure.
-    worker_.reset(env_->StartThread(
-        {}, "gcs_dns_worker", std::bind(&GcsDnsCache::WorkerThread, this)));
+    worker_.reset(env_->StartThread({}, "gcs_dns_worker",
+                                    [this]() { return WorkerThread(); }));
     started_ = true;
   }
 
diff --git a/tensorflow/core/platform/cloud/gcs_file_system.cc b/tensorflow/core/platform/cloud/gcs_file_system.cc
index f1e18403ec..488f9cc75d 100644
--- a/tensorflow/core/platform/cloud/gcs_file_system.cc
+++ b/tensorflow/core/platform/cloud/gcs_file_system.cc
@@ -1397,8 +1397,7 @@ Status GcsFileSystem::RenameObject(const string& src, const string& target) {
   // on the server side, we can't just retry the whole RenameFile operation
   // because the source object is already gone.
   return RetryingUtils::DeleteWithRetries(
-      std::bind(&GcsFileSystem::DeleteFile, this, src),
-      initial_retry_delay_usec_);
+      [this, &src]() { return DeleteFile(src); }, initial_retry_delay_usec_);
 }
 
 Status GcsFileSystem::IsDirectory(const string& fname) {
@@ -1454,7 +1453,7 @@ Status GcsFileSystem::DeleteRecursively(const string& dirname,
     // and therefore RetryingFileSystem won't pay attention to the failures,
     // we need to make sure these failures are properly retried.
     const auto& delete_file_status = RetryingUtils::DeleteWithRetries(
-        std::bind(&GcsFileSystem::DeleteFile, this, full_path),
+        [this, &full_path]() { return DeleteFile(full_path); },
         initial_retry_delay_usec_);
     if (!delete_file_status.ok()) {
       if (IsDirectory(full_path).ok()) {
diff --git a/tensorflow/core/platform/cloud/retrying_file_system.h b/tensorflow/core/platform/cloud/retrying_file_system.h
index 399a21617e..92aa72be89 100644
--- a/tensorflow/core/platform/cloud/retrying_file_system.h
+++ b/tensorflow/core/platform/cloud/retrying_file_system.h
@@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#ifndef TENSORFLOW_CORE_PLATFORM_RETRYING_FILE_SYSTEM_H_
-#define TENSORFLOW_CORE_PLATFORM_RETRYING_FILE_SYSTEM_H_
+#ifndef TENSORFLOW_CORE_PLATFORM_CLOUD_RETRYING_FILE_SYSTEM_H_
+#define TENSORFLOW_CORE_PLATFORM_CLOUD_RETRYING_FILE_SYSTEM_H_
 
 #include <functional>
 #include <string>
@@ -54,74 +54,80 @@ class RetryingFileSystem : public FileSystem {
 
   Status FileExists(const string& fname) override {
     return RetryingUtils::CallWithRetries(
-        std::bind(&FileSystem::FileExists, base_file_system_.get(), fname),
+        [this, &fname]() { return base_file_system_->FileExists(fname); },
         initial_delay_microseconds_);
   }
 
   Status GetChildren(const string& dir, std::vector<string>* result) override {
     return RetryingUtils::CallWithRetries(
-        std::bind(&FileSystem::GetChildren, base_file_system_.get(), dir,
-                  result),
+        [this, &dir, result]() {
+          return base_file_system_->GetChildren(dir, result);
+        },
         initial_delay_microseconds_);
   }
 
   Status GetMatchingPaths(const string& pattern,
                           std::vector<string>* result) override {
     return RetryingUtils::CallWithRetries(
-        std::bind(&FileSystem::GetMatchingPaths, base_file_system_.get(),
-                  pattern, result),
+        [this, &pattern, result]() {
+          return base_file_system_->GetMatchingPaths(pattern, result);
+        },
         initial_delay_microseconds_);
   }
 
   Status Stat(const string& fname, FileStatistics* stat) override {
     return RetryingUtils::CallWithRetries(
-        std::bind(&FileSystem::Stat, base_file_system_.get(), fname, stat),
+        [this, &fname, stat]() { return base_file_system_->Stat(fname, stat); },
         initial_delay_microseconds_);
   }
 
   Status DeleteFile(const string& fname) override {
     return RetryingUtils::DeleteWithRetries(
-        std::bind(&FileSystem::DeleteFile, base_file_system_.get(), fname),
+        [this, &fname]() { return base_file_system_->DeleteFile(fname); },
         initial_delay_microseconds_);
   }
 
   Status CreateDir(const string& dirname) override {
     return RetryingUtils::CallWithRetries(
-        std::bind(&FileSystem::CreateDir, base_file_system_.get(), dirname),
+        [this, &dirname]() { return base_file_system_->CreateDir(dirname); },
         initial_delay_microseconds_);
   }
 
   Status DeleteDir(const string& dirname) override {
     return RetryingUtils::DeleteWithRetries(
-        std::bind(&FileSystem::DeleteDir, base_file_system_.get(), dirname),
+        [this, &dirname]() { return base_file_system_->DeleteDir(dirname); },
         initial_delay_microseconds_);
   }
 
   Status GetFileSize(const string& fname, uint64* file_size) override {
     return RetryingUtils::CallWithRetries(
-        std::bind(&FileSystem::GetFileSize, base_file_system_.get(), fname,
-                  file_size),
+        [this, &fname, file_size]() {
+          return base_file_system_->GetFileSize(fname, file_size);
+        },
         initial_delay_microseconds_);
   }
 
   Status RenameFile(const string& src, const string& target) override {
     return RetryingUtils::CallWithRetries(
-        std::bind(&FileSystem::RenameFile, base_file_system_.get(), src,
-                  target),
+        [this, &src, &target]() {
+          return base_file_system_->RenameFile(src, target);
+        },
         initial_delay_microseconds_);
   }
 
   Status IsDirectory(const string& dirname) override {
     return RetryingUtils::CallWithRetries(
-        std::bind(&FileSystem::IsDirectory, base_file_system_.get(), dirname),
+        [this, &dirname]() { return base_file_system_->IsDirectory(dirname); },
         initial_delay_microseconds_);
   }
 
   Status DeleteRecursively(const string& dirname, int64* undeleted_files,
                            int64* undeleted_dirs) override {
     return RetryingUtils::DeleteWithRetries(
-        std::bind(&FileSystem::DeleteRecursively, base_file_system_.get(),
-                  dirname, undeleted_files, undeleted_dirs),
+        [this, &dirname, undeleted_files, undeleted_dirs]() {
+          return base_file_system_->DeleteRecursively(dirname, undeleted_files,
+                                                      undeleted_dirs);
+        },
         initial_delay_microseconds_);
   }
 
@@ -148,8 +154,9 @@ class RetryingRandomAccessFile : public RandomAccessFile {
   Status Read(uint64 offset, size_t n, StringPiece* result,
               char* scratch) const override {
     return RetryingUtils::CallWithRetries(
-        std::bind(&RandomAccessFile::Read, base_file_.get(), offset, n, result,
-                  scratch),
+        [this, offset, n, result, scratch]() {
+          return base_file_->Read(offset, n, result, scratch);
+        },
         initial_delay_microseconds_);
   }
 
@@ -172,23 +179,20 @@ class RetryingWritableFile : public WritableFile {
 
   Status Append(const StringPiece& data) override {
     return RetryingUtils::CallWithRetries(
-        std::bind(&WritableFile::Append, base_file_.get(), data),
+        [this, &data]() { return base_file_->Append(data); },
         initial_delay_microseconds_);
   }
   Status Close() override {
     return RetryingUtils::CallWithRetries(
-        std::bind(&WritableFile::Close, base_file_.get()),
-        initial_delay_microseconds_);
+        [this]() { return base_file_->Close(); }, initial_delay_microseconds_);
   }
   Status Flush() override {
     return RetryingUtils::CallWithRetries(
-        std::bind(&WritableFile::Flush, base_file_.get()),
-        initial_delay_microseconds_);
+        [this]() { return base_file_->Flush(); }, initial_delay_microseconds_);
   }
   Status Sync() override {
     return RetryingUtils::CallWithRetries(
-        std::bind(&WritableFile::Sync, base_file_.get()),
-        initial_delay_microseconds_);
+        [this]() { return base_file_->Sync(); }, initial_delay_microseconds_);
   }
 
  private:
@@ -203,8 +207,9 @@ Status RetryingFileSystem<Underlying>::NewRandomAccessFile(
     const string& filename, std::unique_ptr<RandomAccessFile>* result) {
   std::unique_ptr<RandomAccessFile> base_file;
   TF_RETURN_IF_ERROR(RetryingUtils::CallWithRetries(
-      std::bind(&FileSystem::NewRandomAccessFile, base_file_system_.get(),
-                filename, &base_file),
+      [this, &filename, &base_file]() {
+        return base_file_system_->NewRandomAccessFile(filename, &base_file);
+      },
       initial_delay_microseconds_));
   result->reset(new retrying_internals::RetryingRandomAccessFile(
       std::move(base_file), initial_delay_microseconds_));
@@ -216,8 +221,9 @@ Status RetryingFileSystem<Underlying>::NewWritableFile(
     const string& filename, std::unique_ptr<WritableFile>* result) {
   std::unique_ptr<WritableFile> base_file;
   TF_RETURN_IF_ERROR(RetryingUtils::CallWithRetries(
-      std::bind(&FileSystem::NewWritableFile, base_file_system_.get(), filename,
-                &base_file),
+      [this, &filename, &base_file]() {
+        return base_file_system_->NewWritableFile(filename, &base_file);
+      },
       initial_delay_microseconds_));
   result->reset(new retrying_internals::RetryingWritableFile(
       std::move(base_file), initial_delay_microseconds_));
@@ -229,8 +235,9 @@ Status RetryingFileSystem<Underlying>::NewAppendableFile(
     const string& filename, std::unique_ptr<WritableFile>* result) {
   std::unique_ptr<WritableFile> base_file;
   TF_RETURN_IF_ERROR(RetryingUtils::CallWithRetries(
-      std::bind(&FileSystem::NewAppendableFile, base_file_system_.get(),
-                filename, &base_file),
+      [this, &filename, &base_file]() {
+        return base_file_system_->NewAppendableFile(filename, &base_file);
+      },
       initial_delay_microseconds_));
   result->reset(new retrying_internals::RetryingWritableFile(
       std::move(base_file), initial_delay_microseconds_));
@@ -241,11 +248,13 @@ template <typename Underlying>
 Status RetryingFileSystem<Underlying>::NewReadOnlyMemoryRegionFromFile(
     const string& filename, std::unique_ptr<ReadOnlyMemoryRegion>* result) {
   return RetryingUtils::CallWithRetries(
-      std::bind(&FileSystem::NewReadOnlyMemoryRegionFromFile,
-                base_file_system_.get(), filename, result),
+      [this, &filename, result]() {
+        return base_file_system_->NewReadOnlyMemoryRegionFromFile(filename,
+                                                                  result);
+      },
       initial_delay_microseconds_);
 }
 
 }  // namespace tensorflow
 
-#endif  // TENSORFLOW_CORE_PLATFORM_RETRYING_FILE_SYSTEM_H_
+#endif  // TENSORFLOW_CORE_PLATFORM_CLOUD_RETRYING_FILE_SYSTEM_H_
diff --git a/tensorflow/core/platform/cloud/retrying_utils.cc b/tensorflow/core/platform/cloud/retrying_utils.cc
index 99691ecfb9..d2df422024 100644
--- a/tensorflow/core/platform/cloud/retrying_utils.cc
+++ b/tensorflow/core/platform/cloud/retrying_utils.cc
@@ -44,9 +44,9 @@ bool IsRetriable(error::Code code) {
 
 Status RetryingUtils::CallWithRetries(const std::function<Status()>& f,
                                       const int64 initial_delay_microseconds) {
-  return CallWithRetries(f, initial_delay_microseconds,
-                         std::bind(&Env::SleepForMicroseconds, Env::Default(),
-                                   std::placeholders::_1));
+  return CallWithRetries(f, initial_delay_microseconds, [](int64 micros) {
+    return Env::Default()->SleepForMicroseconds(micros);
+  });
 }
 
 Status RetryingUtils::CallWithRetries(
-- 
GitLab


From 278e68cedbb80c6f3342856bfccf688a808e461a Mon Sep 17 00:00:00 2001
From: Benoit Steiner <bsteiner@google.com>
Date: Thu, 3 May 2018 13:09:28 -0700
Subject: [PATCH 214/395] Simplified the implementation of shape_n since the
 optimized code path isn't needed anymore and can be incorrect in some rare
 cases.

PiperOrigin-RevId: 195298813
---
 tensorflow/python/ops/array_ops.py                | 10 +---------
 tensorflow/python/profiler/model_analyzer_test.py |  2 +-
 2 files changed, 2 insertions(+), 10 deletions(-)

diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py
index e235047aff..96df15684b 100644
--- a/tensorflow/python/ops/array_ops.py
+++ b/tensorflow/python/ops/array_ops.py
@@ -263,15 +263,7 @@ def shape_n(input, out_type=dtypes.int32, name=None):
       type `out_type`.
   """
 
-  output = gen_array_ops.shape_n(input, out_type=out_type, name=name)
-  if not context.executing_eagerly():
-    for i, input_tensor in enumerate(input):
-      input_tensor = ops.convert_to_tensor(input_tensor)
-      input_shape = input_tensor.get_shape()
-      if input_shape.is_fully_defined():
-        output[i] = constant(
-            input_shape.as_list(), dtype=out_type, name=name)
-  return output
+  return gen_array_ops.shape_n(input, out_type=out_type, name=name)
 
 
 @tf_export("size")
diff --git a/tensorflow/python/profiler/model_analyzer_test.py b/tensorflow/python/profiler/model_analyzer_test.py
index 04ba28c219..75580fc630 100644
--- a/tensorflow/python/profiler/model_analyzer_test.py
+++ b/tensorflow/python/profiler/model_analyzer_test.py
@@ -232,7 +232,7 @@ class PrintModelAnalysisTest(test.TestCase):
 
         self.assertLess(0, tfprof_node.total_exec_micros)
         self.assertEqual(2844, tfprof_node.total_parameters)
-        self.assertLess(168800, tfprof_node.total_float_ops)
+        self.assertLess(145660, tfprof_node.total_float_ops)
         self.assertEqual(8, len(tfprof_node.children))
         self.assertEqual('_TFProfRoot', tfprof_node.name)
         self.assertEqual(
-- 
GitLab


From 41dcb67efd272e9ce0e5071433f42a9d540ec6dc Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 3 May 2018 13:09:30 -0700
Subject: [PATCH 215/395] Fix bugs in model pruner.

PiperOrigin-RevId: 195298816
---
 tensorflow/core/grappler/optimizers/model_pruner.cc | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensorflow/core/grappler/optimizers/model_pruner.cc b/tensorflow/core/grappler/optimizers/model_pruner.cc
index 3311e97010..36eab4999d 100644
--- a/tensorflow/core/grappler/optimizers/model_pruner.cc
+++ b/tensorflow/core/grappler/optimizers/model_pruner.cc
@@ -70,6 +70,7 @@ Status ModelPruner::Optimize(Cluster* cluster, const GrapplerItem& item,
     }
     // Try to keep the nodes ordered somewhat topologically since this helps
     // further optimizations perform better.
+    runnable_item.graph.mutable_node()->Reserve(keep.size());
     for (int i = keep.size() - 1; i >= 0; --i) {
       *runnable_item.graph.add_node() = *keep[i];
     }
@@ -113,6 +114,7 @@ Status ModelPruner::Optimize(Cluster* cluster, const GrapplerItem& item,
     }
   }
 
+  pruned_graph->Clear();
   *pruned_graph->mutable_library() = item.graph.library();
   *pruned_graph->mutable_versions() = item.graph.versions();
 
@@ -122,6 +124,7 @@ Status ModelPruner::Optimize(Cluster* cluster, const GrapplerItem& item,
   }
 
   const bool fetches_are_known = !item.fetch.empty();
+  pruned_graph->mutable_node()->Reserve(runnable_item.graph.node_size());
   for (auto& node : runnable_item.graph.node()) {
     if (!fetches_are_known ||
         nodes_to_delete.find(&node) == nodes_to_delete.end()) {
@@ -134,6 +137,7 @@ Status ModelPruner::Optimize(Cluster* cluster, const GrapplerItem& item,
   VLOG(1) << "Pruned " << nodes_to_delete.size()
           << " nodes from the graph. The graph now contains "
           << pruned_graph->node_size() << " nodes.";
+  CHECK_LE(pruned_graph->node_size(), item.graph.node_size());
 
   return Status::OK();
 }
-- 
GitLab


From 5a64e609d0eb94244067f5d7514605863c9f37c3 Mon Sep 17 00:00:00 2001
From: Allen Lavoie <allenl@google.com>
Date: Thu, 3 May 2018 13:22:33 -0700
Subject: [PATCH 216/395] Checkpointable: Utilities to read object metadata

Useful for inspecting checkpoints programatically (e.g. in unit tests).

PiperOrigin-RevId: 195300780
---
 tensorflow/contrib/checkpoint/__init__.py     |  4 ++
 .../contrib/checkpoint/python/visualize.py    | 16 +-------
 .../eager/python/examples/spinn/spinn_test.py | 10 ++---
 .../python/training/checkpointable_utils.py   | 38 +++++++++++++++++++
 .../training/checkpointable_utils_test.py     | 16 ++++++++
 5 files changed, 64 insertions(+), 20 deletions(-)

diff --git a/tensorflow/contrib/checkpoint/__init__.py b/tensorflow/contrib/checkpoint/__init__.py
index 1192cc44a1..d2c30f1215 100644
--- a/tensorflow/contrib/checkpoint/__init__.py
+++ b/tensorflow/contrib/checkpoint/__init__.py
@@ -16,7 +16,9 @@
 
 
 For creating and managing dependencies:
+@@CheckpointableObjectGraph
 @@dot_graph_from_checkpoint
+@@object_metadata
 @@split_dependency
 """
 
@@ -26,6 +28,8 @@ from __future__ import print_function
 
 from tensorflow.contrib.checkpoint.python.split_dependency import split_dependency
 from tensorflow.contrib.checkpoint.python.visualize import dot_graph_from_checkpoint
+from tensorflow.core.protobuf.checkpointable_object_graph_pb2 import CheckpointableObjectGraph
+from tensorflow.python.training.checkpointable_utils import object_metadata
 
 from tensorflow.python.util.all_util import remove_undocumented
 
diff --git a/tensorflow/contrib/checkpoint/python/visualize.py b/tensorflow/contrib/checkpoint/python/visualize.py
index 86fbdb41d2..9a3b23bb2c 100644
--- a/tensorflow/contrib/checkpoint/python/visualize.py
+++ b/tensorflow/contrib/checkpoint/python/visualize.py
@@ -17,10 +17,9 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensorflow.core.protobuf import checkpointable_object_graph_pb2
 from tensorflow.python import pywrap_tensorflow
-from tensorflow.python.framework import errors_impl
 from tensorflow.python.training import checkpointable
+from tensorflow.python.training import checkpointable_utils
 
 
 def dot_graph_from_checkpoint(save_path):
@@ -52,20 +51,9 @@ def dot_graph_from_checkpoint(save_path):
     A graph in DOT format as a string.
   """
   reader = pywrap_tensorflow.NewCheckpointReader(save_path)
-  try:
-    object_graph_string = reader.get_tensor(
-        checkpointable.OBJECT_GRAPH_PROTO_KEY)
-  except errors_impl.NotFoundError:
-    raise ValueError(
-        ('The specified checkpoint "%s" does not appear to be object-based (it '
-         'is missing the key "%s"). Likely it was created with a name-based '
-         'saver and does not contain an object dependency graph.') % (
-             save_path, checkpointable.OBJECT_GRAPH_PROTO_KEY))
+  object_graph = checkpointable_utils.object_metadata(save_path)
   shape_map = reader.get_variable_to_shape_map()
   dtype_map = reader.get_variable_to_dtype_map()
-  object_graph = (
-      checkpointable_object_graph_pb2.CheckpointableObjectGraph())
-  object_graph.ParseFromString(object_graph_string)
   graph = 'digraph {\n'
   def _escape(name):
     return name.replace('"', '\\"')
diff --git a/tensorflow/contrib/eager/python/examples/spinn/spinn_test.py b/tensorflow/contrib/eager/python/examples/spinn/spinn_test.py
index f825a2a736..1e4746d01c 100644
--- a/tensorflow/contrib/eager/python/examples/spinn/spinn_test.py
+++ b/tensorflow/contrib/eager/python/examples/spinn/spinn_test.py
@@ -34,10 +34,10 @@ import tensorflow.contrib.eager as tfe
 from tensorflow.contrib.eager.python.examples.spinn import data
 from third_party.examples.eager.spinn import spinn
 from tensorflow.contrib.summary import summary_test_util
-from tensorflow.core.protobuf import checkpointable_object_graph_pb2
 from tensorflow.python.eager import test
 from tensorflow.python.framework import test_util
-from tensorflow.python.training import checkpoint_utils
+from tensorflow.python.training import checkpointable_utils
+from tensorflow.python.training import saver
 # pylint: enable=g-bad-import-order
 
 
@@ -421,10 +421,8 @@ class SpinnTest(test_util.TensorFlowTestCase):
 
     # 5. Verify that checkpoints exist and contains all the expected variables.
     self.assertTrue(glob.glob(os.path.join(config.logdir, "ckpt*")))
-    object_graph_string = checkpoint_utils.load_variable(
-        config.logdir, name="_CHECKPOINTABLE_OBJECT_GRAPH")
-    object_graph = checkpointable_object_graph_pb2.CheckpointableObjectGraph()
-    object_graph.ParseFromString(object_graph_string)
+    object_graph = checkpointable_utils.object_metadata(
+        saver.latest_checkpoint(config.logdir))
     ckpt_variable_names = set()
     for node in object_graph.nodes:
       for attribute in node.attributes:
diff --git a/tensorflow/python/training/checkpointable_utils.py b/tensorflow/python/training/checkpointable_utils.py
index 9cdd53cbf9..cf4112ff99 100644
--- a/tensorflow/python/training/checkpointable_utils.py
+++ b/tensorflow/python/training/checkpointable_utils.py
@@ -159,6 +159,44 @@ def add_variable(checkpointable, name, shape=None, dtype=dtypes.float32,
       initializer=initializer, getter=_default_getter)
 
 
+def object_metadata(save_path):
+  """Retrieves information about the objects in a checkpoint.
+
+  Example usage:
+
+  ```python
+  object_graph = tf.contrib.checkpoint.object_metadata(
+      tf.train.latest_checkpoint(checkpoint_directory))
+  ckpt_variable_names = set()
+  for node in object_graph.nodes:
+    for attribute in node.attributes:
+      ckpt_variable_names.add(attribute.full_name)
+  ```
+
+  Args:
+    save_path: The path to the checkpoint, as returned by `save` or
+      `tf.train.latest_checkpoint`.
+  Returns:
+    A parsed `tf.contrib.checkpoint.CheckpointableObjectGraph` protocol buffer.
+  Raises:
+    ValueError: If an object graph was not found in the checkpoint.
+  """
+  reader = pywrap_tensorflow.NewCheckpointReader(save_path)
+  try:
+    object_graph_string = reader.get_tensor(
+        checkpointable_lib.OBJECT_GRAPH_PROTO_KEY)
+  except errors_impl.NotFoundError:
+    raise ValueError(
+        ('The specified checkpoint "%s" does not appear to be object-based (it '
+         'is missing the key "%s"). Likely it was created with a name-based '
+         'saver and does not contain an object dependency graph.') % (
+             save_path, checkpointable_lib.OBJECT_GRAPH_PROTO_KEY))
+  object_graph_proto = (
+      checkpointable_object_graph_pb2.CheckpointableObjectGraph())
+  object_graph_proto.ParseFromString(object_graph_string)
+  return object_graph_proto
+
+
 def _breadth_first_checkpointable_traversal(root_checkpointable):
   """Find shortest paths to all variables owned by dependencies of root."""
   bfs_sorted = []
diff --git a/tensorflow/python/training/checkpointable_utils_test.py b/tensorflow/python/training/checkpointable_utils_test.py
index 40dfeb28d5..3b8166bf37 100644
--- a/tensorflow/python/training/checkpointable_utils_test.py
+++ b/tensorflow/python/training/checkpointable_utils_test.py
@@ -155,6 +155,22 @@ class InterfaceTests(test.TestCase):
     self.assertEqual(dtypes.float64, v2.dtype)
     self.assertAllEqual([1., 1., 1.], self.evaluate(v2))
 
+  def testObjectMetadata(self):
+    with context.eager_mode():
+      checkpoint_directory = self.get_temp_dir()
+      checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
+      dense = core.Dense(1)
+      checkpoint = checkpointable_utils.Checkpoint(dense=dense)
+      dense(constant_op.constant([[1.]]))
+      save_path = checkpoint.save(checkpoint_prefix)
+
+    objects = checkpointable_utils.object_metadata(save_path)
+    all_variable_names = []
+    for obj in objects.nodes:
+      for attribute in obj.attributes:
+        all_variable_names.append(attribute.full_name)
+    self.assertIn("dense/kernel", all_variable_names)
+
 
 class _MirroringSaveable(saver_lib.BaseSaverBuilder.SaveableObject):
 
-- 
GitLab


From 7529268d692c1c888f93924e6ca5e10fd3183b80 Mon Sep 17 00:00:00 2001
From: Shanqing Cai <cais@google.com>
Date: Thu, 3 May 2018 13:30:12 -0700
Subject: [PATCH 217/395] tfdbg + tflearn: replace deprecated classes and
 methods in example & docs

* `tf.contrib.learn.Experiment` is deprecated. Remove it from debug_tflearn_iris.py.
* Use `tf.estimator.DNNClassifier`, instead of the older one from `tf.contrib.learn`.
* Use `train()`, instead of `fit()` of Estimators.
* `Estimator.predict()` supports hooks. Add example lines for that.

PiperOrigin-RevId: 195301913
---
 .../docs_src/programmers_guide/debugger.md    | 89 ++++++-------------
 tensorflow/python/debug/BUILD                 |  1 -
 .../debug/examples/debug_tflearn_iris.py      | 83 ++++++++---------
 3 files changed, 65 insertions(+), 108 deletions(-)

diff --git a/tensorflow/docs_src/programmers_guide/debugger.md b/tensorflow/docs_src/programmers_guide/debugger.md
index f7817b06d4..6bd941886d 100644
--- a/tensorflow/docs_src/programmers_guide/debugger.md
+++ b/tensorflow/docs_src/programmers_guide/debugger.md
@@ -34,7 +34,7 @@ type of bug in TensorFlow model development.
 The following example is for users who use the low-level
 [`Session`](https://www.tensorflow.org/api_docs/python/tf/Session) API of
 TensorFlow. A later section of this document describes how to use **tfdbg**
-with a higher-level API, namely tf-learn `Estimator`s and `Experiment`s.
+with a higher-level API, namely `Estimator`s.
 To *observe* such an issue, run the following command without the debugger (the
 source code can be found
 [here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/debug/examples/debug_mnist.py)):
@@ -418,21 +418,20 @@ run -f has_inf_or_nan`
 Confirm that no tensors are flagged as containing `nan` or `inf` values, and
 accuracy now continues to rise rather than getting stuck. Success!
 
-## Debugging tf-learn Estimators and Experiments
+## Debugging TensorFlow Estimators
 
 This section explains how to debug TensorFlow programs that use the `Estimator`
-and `Experiment` APIs. Part of the convenience provided by these APIs is that
+APIs. Part of the convenience provided by these APIs is that
 they manage `Session`s internally. This makes the `LocalCLIDebugWrapperSession`
 described in the preceding sections inapplicable. Fortunately, you can still
 debug them by using special `hook`s provided by `tfdbg`.
 
-### Debugging tf.contrib.learn Estimators
-
-Currently, `tfdbg` can debug the
-@{tf.contrib.learn.BaseEstimator.fit$`fit()`}
-@{tf.contrib.learn.BaseEstimator.evaluate$`evaluate()`}
-methods of tf-learn `Estimator`s. To debug `Estimator.fit()`,
-create a `LocalCLIDebugHook` and supply it in the `monitors` argument. For example:
+`tfdbg` can debug the
+@{tf.estimator.Estimator.train$`train()`},
+@{tf.estimator.Estimator.evaluate$`evaluate()`} and
+@{tf.estimator.Estimator.predict$`predict()`}
+methods of tf-learn `Estimator`s. To debug `Estimator.train()`,
+create a `LocalCLIDebugHook` and supply it in the `hooks` argument. For example:
 
 ```python
 # First, let your BUILD target depend on "//tensorflow/python/debug:debug_py"
@@ -443,67 +442,33 @@ from tensorflow.python import debug as tf_debug
 # Create a LocalCLIDebugHook and use it as a monitor when calling fit().
 hooks = [tf_debug.LocalCLIDebugHook()]
 
-classifier.fit(x=training_set.data,
-               y=training_set.target,
-               steps=1000,
-               monitors=hooks)
+# To debug `train`:
+classifier.train(input_fn,
+                 steps=1000,
+                 hooks=hooks)
 ```
 
-To debug `Estimator.evaluate()`, assign hooks to the `hooks` parameter, as in
-the following example:
+Similarly, to debug `Estimator.evaluate()` and `Estimator.predict()`, assign
+hooks to the `hooks` parameter, as in the following example:
 
 ```python
-accuracy_score = classifier.evaluate(x=test_set.data,
-                                     y=test_set.target,
+# To debug `evaluate`:
+accuracy_score = classifier.evaluate(eval_input_fn,
                                      hooks=hooks)["accuracy"]
-```
 
+# To debug `predict`:
+predict_results = classifier.predict(predict_input_fn, hooks=hooks)
+```
 
 [debug_tflearn_iris.py](https://www.tensorflow.org/code/tensorflow/python/debug/examples/debug_tflearn_iris.py),
-based on [tf-learn's iris tutorial](https://www.tensorflow.org/versions/r1.2/get_started/tflearn), contains a full example of how to
-use the tfdbg with `Estimator`s. To run this example, do:
+based on [tf-learn's iris tutorial](https://www.tensorflow.org/versions/r1.8/get_started/tflearn),
+contains a full example of how to use the tfdbg with `Estimator`s.
+To run this example, do:
 
 ```none
 python -m tensorflow.python.debug.examples.debug_tflearn_iris --debug
 ```
 
-### Debugging tf.contrib.learn Experiments
-
-`Experiment` is a construct in `tf.contrib.learn` at a higher level than
-`Estimator`.
-It provides a single interface for training and evaluating a model. To debug
-the `train()` and `evaluate()` calls to an `Experiment` object, you can
-use the keyword arguments `train_monitors` and `eval_hooks`, respectively, when
-calling its constructor. For example:
-
-```python
-# First, let your BUILD target depend on "//tensorflow/python/debug:debug_py"
-# (You don't need to worry about the BUILD dependency if you are using a pip
-#  install of open-source TensorFlow.)
-from tensorflow.python import debug as tf_debug
-
-hooks = [tf_debug.LocalCLIDebugHook()]
-
-ex = experiment.Experiment(classifier,
-                           train_input_fn=iris_input_fn,
-                           eval_input_fn=iris_input_fn,
-                           train_steps=FLAGS.train_steps,
-                           eval_delay_secs=0,
-                           eval_steps=1,
-                           train_monitors=hooks,
-                           eval_hooks=hooks)
-
-ex.train()
-accuracy_score = ex.evaluate()["accuracy"]
-```
-
-To build and run the `debug_tflearn_iris` example in the `Experiment` mode, do:
-
-```none
-python -m tensorflow.python.debug.examples.debug_tflearn_iris \
-    --use_experiment --debug
-```
-
 The `LocalCLIDebugHook` also allows you to configure a `watch_fn` that can be
 used to flexibly specify what `Tensor`s to watch on different `Session.run()`
 calls, as a function of the `fetches` and `feed_dict` and other states. See
@@ -573,7 +538,7 @@ Often, your model is running on a remote machine or a process that you don't
 have terminal access to. To perform model debugging in such cases, you can use
 the `offline_analyzer` binary of `tfdbg` (described below). It operates on
 dumped data directories. This can be done to both the lower-level `Session` API
-and the higher-level `Estimator` and `Experiment` APIs.
+and the higher-level `Estimator` API.
 
 ### Debugging Remote tf.Sessions
 
@@ -636,7 +601,7 @@ can be inspected offline. See
 [the proto definition](https://www.tensorflow.org/code/tensorflow/core/protobuf/debug.proto)
 for more details.
 
-### Debugging Remotely-Running tf-learn Estimators and Experiments
+### Debugging Remotely-Running Estimators
 
 If your remote TensorFlow server runs `Estimator`s,
 you can use the non-interactive `DumpingDebugHook`. For example:
@@ -652,8 +617,8 @@ hooks = [tf_debug.DumpingDebugHook("/shared/storage/location/tfdbg_dumps_1")]
 
 Then this `hook` can be used in the same way as the `LocalCLIDebugHook` examples
 described earlier in this document.
-As the training and/or evalution of `Estimator` or `Experiment`
-happens, tfdbg creates directories having the following name pattern:
+As the training, evalution or prediction happens with `Estimator`,
+tfdbg creates directories having the following name pattern:
 `/shared/storage/location/tfdbg_dumps_1/run_<epoch_timestamp_microsec>_<uuid>`.
 Each directory corresponds to a `Session.run()` call that underlies
 the `fit()` or `evaluate()` call. You can load these directories and inspect
diff --git a/tensorflow/python/debug/BUILD b/tensorflow/python/debug/BUILD
index b5760df1ed..183994ddaa 100644
--- a/tensorflow/python/debug/BUILD
+++ b/tensorflow/python/debug/BUILD
@@ -449,7 +449,6 @@ py_binary(
     deps = [
         ":debug_py",
         "//tensorflow:tensorflow_py",
-        "//third_party/py/numpy",
         "@six_archive//:six",
     ],
 )
diff --git a/tensorflow/python/debug/examples/debug_tflearn_iris.py b/tensorflow/python/debug/examples/debug_tflearn_iris.py
index 4f4666ee4f..00090b21fe 100644
--- a/tensorflow/python/debug/examples/debug_tflearn_iris.py
+++ b/tensorflow/python/debug/examples/debug_tflearn_iris.py
@@ -22,11 +22,9 @@ import os
 import sys
 import tempfile
 
-import numpy as np
 from six.moves import urllib
 import tensorflow as tf
 
-from tensorflow.contrib.learn.python.learn import experiment
 from tensorflow.contrib.learn.python.learn.datasets import base
 from tensorflow.python import debug as tf_debug
 
@@ -82,28 +80,34 @@ def iris_input_fn():
 def main(_):
   # Load datasets.
   if FLAGS.fake_data:
-    training_set = tf.contrib.learn.datasets.base.Dataset(
-        np.random.random([120, 4]),
-        np.random.random_integers(3, size=[120]) - 1)
-    test_set = tf.contrib.learn.datasets.base.Dataset(
-        np.random.random([30, 4]),
-        np.random.random_integers(3, size=[30]) - 1)
+    def training_input_fn():
+      return ({"features": tf.random_normal([128, 4])},
+              tf.random_uniform([128], minval=0, maxval=3, dtype=tf.int32))
+    def test_input_fn():
+      return ({"features": tf.random_normal([32, 4])},
+              tf.random_uniform([32], minval=0, maxval=3, dtype=tf.int32))
+    feature_columns = [
+        tf.feature_column.numeric_column("features", shape=(4,))]
   else:
     training_data_path, test_data_path = maybe_download_data(FLAGS.data_dir)
-    training_set = tf.contrib.learn.datasets.base.load_csv_with_header(
-        filename=training_data_path,
-        target_dtype=np.int,
-        features_dtype=np.float32)
-    test_set = tf.contrib.learn.datasets.base.load_csv_with_header(
-        filename=test_data_path, target_dtype=np.int, features_dtype=np.float32)
-
-  # Specify that all features have real-value data
-  feature_columns = [tf.contrib.layers.real_valued_column("", dimension=4)]
+    column_names = [
+        "sepal_length", "sepal_width", "petal_length", "petal_width", "label"]
+    batch_size = 32
+    def training_input_fn():
+      return tf.contrib.data.make_csv_dataset(
+          [training_data_path], batch_size,
+          column_names=column_names, label_name="label")
+    def test_input_fn():
+      return tf.contrib.data.make_csv_dataset(
+          [test_data_path], batch_size,
+          column_names=column_names, label_name="label")
+    feature_columns = [tf.feature_column.numeric_column(feature)
+                       for feature in column_names[:-1]]
 
   # Build 3 layer DNN with 10, 20, 10 units respectively.
   model_dir = FLAGS.model_dir or tempfile.mkdtemp(prefix="debug_tflearn_iris_")
 
-  classifier = tf.contrib.learn.DNNClassifier(
+  classifier = tf.estimator.DNNClassifier(
       feature_columns=feature_columns,
       hidden_units=[10, 20, 10],
       n_classes=3,
@@ -121,32 +125,23 @@ def main(_):
     debug_hook = tf_debug.TensorBoardDebugHook(FLAGS.tensorboard_debug_address)
   hooks = [debug_hook]
 
-  if not FLAGS.use_experiment:
-    # Fit model.
-    classifier.fit(x=training_set.data,
-                   y=training_set.target,
+  # Train model, using tfdbg hook.
+  classifier.train(training_input_fn,
                    steps=FLAGS.train_steps,
-                   monitors=hooks)
+                   hooks=hooks)
 
-    # Evaluate accuracy.
-    accuracy_score = classifier.evaluate(x=test_set.data,
-                                         y=test_set.target,
-                                         hooks=hooks)["accuracy"]
-  else:
-    ex = experiment.Experiment(classifier,
-                               train_input_fn=iris_input_fn,
-                               eval_input_fn=iris_input_fn,
-                               train_steps=FLAGS.train_steps,
-                               eval_delay_secs=0,
-                               eval_steps=1,
-                               train_monitors=hooks,
-                               eval_hooks=hooks)
-    ex.train()
-    accuracy_score = ex.evaluate()["accuracy"]
+  # Evaluate accuracy, using tfdbg hook.
+  accuracy_score = classifier.evaluate(test_input_fn,
+                                       steps=FLAGS.eval_steps,
+                                       hooks=hooks)["accuracy"]
 
   print("After training %d steps, Accuracy = %f" %
         (FLAGS.train_steps, accuracy_score))
 
+  # Make predictions, using tfdbg hook.
+  predict_results = classifier.predict(test_input_fn, hooks=hooks)
+  print("A prediction result: %s" % predict_results.next())
+
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
@@ -165,14 +160,12 @@ if __name__ == "__main__":
       "--train_steps",
       type=int,
       default=10,
-      help="Number of steps to run trainer.")
+      help="Number of steps to run training for.")
   parser.add_argument(
-      "--use_experiment",
-      type="bool",
-      nargs="?",
-      const=True,
-      default=False,
-      help="Use tf.contrib.learn Experiment to run training and evaluation")
+      "--eval_steps",
+      type=int,
+      default=1,
+      help="Number of steps to run evaluation foir.")
   parser.add_argument(
       "--ui_type",
       type=str,
-- 
GitLab


From 2dc7575123ffa0e6413fc3d2700968ef25f049de Mon Sep 17 00:00:00 2001
From: Sergii Khomenko <x-sam@brainscode.com>
Date: Fri, 4 May 2018 04:22:09 +0200
Subject: [PATCH 218/395] Fix minor typos (#19070)

---
 tensorflow/python/estimator/training.py            | 2 +-
 tensorflow/python/feature_column/feature_column.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/estimator/training.py b/tensorflow/python/estimator/training.py
index 534c357067..95366132d9 100644
--- a/tensorflow/python/estimator/training.py
+++ b/tensorflow/python/estimator/training.py
@@ -588,7 +588,7 @@ class _TrainingExecutor(object):
     # max_steps, the evaluator will send the final export signal. There is a
     # small chance that the Estimator.train stopping logic sees a different
     # global_step value (due to global step race condition and the fact the
-    # saver sees a larger value for checkpoing saving), which does not end
+    # saver sees a larger value for checkpoint saving), which does not end
     # the training. When the training ends, a new checkpoint is generated, which
     # triggers the listener again. So, it could be the case the final export is
     # triggered twice.
diff --git a/tensorflow/python/feature_column/feature_column.py b/tensorflow/python/feature_column/feature_column.py
index 9e6429e59e..40386ae7aa 100644
--- a/tensorflow/python/feature_column/feature_column.py
+++ b/tensorflow/python/feature_column/feature_column.py
@@ -280,7 +280,7 @@ def input_layer(features,
 # TODO(akshayka): InputLayer should be a subclass of Layer, and it
 # should implement the logic in input_layer using Layer's build-and-call
 # paradigm; input_layer should create an instance of InputLayer and
-# return the result of inovking its apply method, just as functional layers do.
+# return the result of invoking its apply method, just as functional layers do.
 class InputLayer(object):
   """An object-oriented version of `input_layer` that reuses variables."""
 
@@ -834,7 +834,7 @@ def shared_embedding_columns(
     tensor_name_in_ckpt=None, max_norm=None, trainable=True):
   """List of dense columns that convert from sparse, categorical input.
 
-  This is similar to `embedding_column`, except that that it produces a list of
+  This is similar to `embedding_column`, except that it produces a list of
   embedding columns that share the same embedding weights.
 
   Use this when your inputs are sparse and of the same type (e.g. watched and
-- 
GitLab


From fe9b2637cfe39cf11eb3d0494948a733b7fc1d7d Mon Sep 17 00:00:00 2001
From: Karl Lessard <karl@kubx.ca>
Date: Thu, 29 Mar 2018 05:28:16 +0800
Subject: [PATCH 219/395] Parse op definition and generate a Java Op class.

---
 tensorflow/java/BUILD                         |   4 +
 tensorflow/java/src/gen/cc/java_defs.h        |  76 ++--
 tensorflow/java/src/gen/cc/op_gen_main.cc     |  22 +-
 tensorflow/java/src/gen/cc/op_generator.cc    | 406 +++++++++++++++--
 tensorflow/java/src/gen/cc/op_generator.h     |  42 +-
 tensorflow/java/src/gen/cc/op_parser.cc       | 417 ++++++++++++++++++
 tensorflow/java/src/gen/cc/op_parser.h        | 137 ++++++
 tensorflow/java/src/gen/cc/source_writer.cc   | 127 +++---
 tensorflow/java/src/gen/cc/source_writer.h    |  55 ++-
 .../java/src/gen/cc/source_writer_test.cc     |  82 ++--
 tensorflow/java/src/gen/gen_ops.bzl           |  29 +-
 .../src/gen/resources/license.snippet.java    |  14 +
 12 files changed, 1201 insertions(+), 210 deletions(-)
 create mode 100644 tensorflow/java/src/gen/cc/op_parser.cc
 create mode 100644 tensorflow/java/src/gen/cc/op_parser.h
 create mode 100644 tensorflow/java/src/gen/resources/license.snippet.java

diff --git a/tensorflow/java/BUILD b/tensorflow/java/BUILD
index ab7d698a45..635a4e807d 100644
--- a/tensorflow/java/BUILD
+++ b/tensorflow/java/BUILD
@@ -70,6 +70,7 @@ filegroup(
 
 tf_java_op_gen_srcjar(
     name = "java_op_gen_sources",
+    api_def_srcs = ["//tensorflow/core/api_def:base_api_def"],
     gen_base_package = "org.tensorflow.op",
     gen_tool = "java_op_gen_tool",
     ops_libs = [
@@ -111,11 +112,13 @@ cc_library(
     name = "java_op_gen_lib",
     srcs = [
         "src/gen/cc/op_generator.cc",
+        "src/gen/cc/op_parser.cc",
         "src/gen/cc/source_writer.cc",
     ],
     hdrs = [
         "src/gen/cc/java_defs.h",
         "src/gen/cc/op_generator.h",
+        "src/gen/cc/op_parser.h",
         "src/gen/cc/source_writer.h",
     ],
     copts = tf_copts(),
@@ -124,6 +127,7 @@ cc_library(
         "//tensorflow/core:framework_internal",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
+        "//tensorflow/core:op_gen_lib",
     ],
 )
 
diff --git a/tensorflow/java/src/gen/cc/java_defs.h b/tensorflow/java/src/gen/cc/java_defs.h
index 59f8beaee7..2065477f58 100644
--- a/tensorflow/java/src/gen/cc/java_defs.h
+++ b/tensorflow/java/src/gen/cc/java_defs.h
@@ -18,12 +18,15 @@ limitations under the License.
 
 #include <string>
 #include <list>
+#include <map>
+#include <utility>
 
 namespace tensorflow {
 namespace java {
 
 // An enumeration of different modifiers commonly used in Java
 enum Modifier {
+  PACKAGE   = 0,
   PUBLIC    = (1 << 0),
   PROTECTED = (1 << 1),
   PRIVATE   = (1 << 2),
@@ -72,6 +75,12 @@ class Type {
     // Reflection API does
     return Type(Type::PRIMITIVE, "void");
   }
+  static Type Generic(const string& name) {
+    return Type(Type::GENERIC, name);
+  }
+  static Type Wildcard() {
+    return Type(Type::GENERIC, "");
+  }
   static Type Class(const string& name, const string& package = "") {
     return Type(Type::CLASS, name, package);
   }
@@ -81,9 +90,6 @@ class Type {
   static Type Enum(const string& name, const string& package = "") {
     return Type(Type::ENUM, name, package);
   }
-  static Type Generic(const string& name = "") {
-    return Type(Type::GENERIC, name);
-  }
   static Type ClassOf(const Type& type) {
     return Class("Class").add_parameter(type);
   }
@@ -96,11 +102,10 @@ class Type {
   const Kind& kind() const { return kind_; }
   const string& name() const { return name_; }
   const string& package() const { return package_; }
-  const string& description() const { return description_; }
-  Type& description(const string& description) {
-    description_ = description;
-    return *this;
+  const string full_name() const {
+    return package_.empty() ? name_ : package_ + "." + name_;
   }
+  bool unknown() const { return name_.empty(); }  // only wildcards has no name
   const std::list<Type>& parameters() const { return parameters_; }
   Type& add_parameter(const Type& parameter) {
     parameters_.push_back(parameter);
@@ -120,14 +125,6 @@ class Type {
     }
     return *this;
   }
-  // Returns true if "type" is of a known collection type (only a few for now)
-  bool IsCollection() const {
-    return name_ == "List" || name_ == "Iterable";
-  }
-  // Returns true if this instance is a wildcard (<?>)
-  bool IsWildcard() const {
-    return kind_ == GENERIC && name_.empty();
-  }
 
  protected:
   Type(Kind kind, const string& name, const string& package = "")
@@ -137,7 +134,6 @@ class Type {
   Kind kind_;
   string name_;
   string package_;
-  string description_;
   std::list<Type> parameters_;
   std::list<Annotation> annotations_;
   std::list<Type> supertypes_;
@@ -180,16 +176,11 @@ class Variable {
   const string& name() const { return name_; }
   const Type& type() const { return type_; }
   bool variadic() const { return variadic_; }
-  const string& description() const { return description_; }
-  Variable& description(const string& description) {
-    description_ = description;
-    return *this;
-  }
+
  private:
   string name_;
   Type type_;
   bool variadic_;
-  string description_;
 
   Variable(const string& name, const Type& type, bool variadic)
     : name_(name), type_(type), variadic_(variadic) {}
@@ -210,16 +201,6 @@ class Method {
   bool constructor() const { return constructor_; }
   const string& name() const { return name_; }
   const Type& return_type() const { return return_type_; }
-  const string& description() const { return description_; }
-  Method& description(const string& description) {
-    description_ = description;
-    return *this;
-  }
-  const string& return_description() const { return return_description_; }
-  Method& return_description(const string& description) {
-    return_description_ = description;
-    return *this;
-  }
   const std::list<Variable>& arguments() const { return arguments_; }
   Method& add_argument(const Variable& var) {
     arguments_.push_back(var);
@@ -235,8 +216,6 @@ class Method {
   string name_;
   Type return_type_;
   bool constructor_;
-  string description_;
-  string return_description_;
   std::list<Variable> arguments_;
   std::list<Annotation> annotations_;
 
@@ -244,6 +223,35 @@ class Method {
     : name_(name), return_type_(return_type), constructor_(constructor) {}
 };
 
+// A definition of a documentation bloc for a Java element (JavaDoc)
+class Javadoc {
+ public:
+  static Javadoc Create(const string& brief = "") {
+    return Javadoc(brief);
+  }
+  const string& brief() const { return brief_; }
+  const string& details() const { return description_; }
+  Javadoc& details(const string description) {
+    description_ = description;
+    return *this;
+  }
+  const std::list<std::pair<string, string>> tags() const { return tags_; }
+  Javadoc& add_tag(const string& tag, const string& text) {
+    tags_.push_back(std::make_pair(tag, text));
+    return *this;
+  }
+  Javadoc& add_param_tag(const string& name, const string& text) {
+    return add_tag("param", name + " " + text);
+  }
+
+ private:
+  string brief_;
+  string description_;
+  std::list<std::pair<string, string>> tags_;
+
+  explicit Javadoc(const string& brief) : brief_(brief) {}
+};
+
 }  // namespace java
 }  // namespace tensorflow
 
diff --git a/tensorflow/java/src/gen/cc/op_gen_main.cc b/tensorflow/java/src/gen/cc/op_gen_main.cc
index bea99f3d7f..015200023f 100644
--- a/tensorflow/java/src/gen/cc/op_gen_main.cc
+++ b/tensorflow/java/src/gen/cc/op_gen_main.cc
@@ -48,8 +48,11 @@ const char kUsageHeader[] =
     "through\n"
     "the 'org.tensorflow.op.Ops' API as a group until the generated classes "
     "are compiled using an appropriate annotation processor.\n\n"
-    "Finally, the '--base_package' overrides the default parent package "
-    "under which the generated subpackage and classes are to be located.\n\n";
+    "The '--base_package' overrides the default parent package under which "
+    "the generated subpackage and classes are to be located.\n\n"
+    "Finally, a list of directories of API proto definitions can be provided "
+    "to override default values found in the ops definitions, ordered by\n"
+    "priority (the last having precedence over the first).\n\n";
 
 }  // namespace java
 }  // namespace tensorflow
@@ -60,7 +63,7 @@ int main(int argc, char* argv[]) {
   tensorflow::string base_package = "org.tensorflow.op";
   std::vector<tensorflow::Flag> flag_list = {
       tensorflow::Flag("output_dir", &output_dir,
-                       "Root directory into which output files are generated"),
+          "Root directory into which output files are generated"),
       tensorflow::Flag(
           "lib_name", &lib_name,
           "A name, in snake_case, used to classify this set of operations"),
@@ -72,12 +75,15 @@ int main(int argc, char* argv[]) {
   bool parsed_flags_ok = tensorflow::Flags::Parse(&argc, argv, flag_list);
   tensorflow::port::InitMain(usage.c_str(), &argc, &argv);
   QCHECK(parsed_flags_ok && !lib_name.empty() && !output_dir.empty()) << usage;
-
-  tensorflow::java::OpGenerator generator;
+  std::vector<tensorflow::string> api_dirs;
+  if (argc > 1) {
+    api_dirs = tensorflow::str_util::Split(argv[1], ",",
+        tensorflow::str_util::SkipEmpty());
+  }
+  tensorflow::java::OpGenerator generator(base_package, output_dir, api_dirs);
   tensorflow::OpList ops;
-  tensorflow::OpRegistry::Global()->Export(true, &ops);
-  tensorflow::Status status =
-      generator.Run(ops, lib_name, base_package, output_dir);
+  tensorflow::OpRegistry::Global()->Export(false, &ops);
+  tensorflow::Status status = generator.Run(ops, lib_name);
   TF_QCHECK_OK(status);
 
   return 0;
diff --git a/tensorflow/java/src/gen/cc/op_generator.cc b/tensorflow/java/src/gen/cc/op_generator.cc
index def06baf2d..c9b57f5706 100644
--- a/tensorflow/java/src/gen/cc/op_generator.cc
+++ b/tensorflow/java/src/gen/cc/op_generator.cc
@@ -14,53 +14,409 @@ limitations under the License.
 ==============================================================================*/
 
 #include <string>
+#include <map>
+#include <vector>
+#include <list>
+#include <memory>
 
+#include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/io/path.h"
 #include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/framework/op_gen_lib.h"
+#include "tensorflow/java/src/gen/cc/java_defs.h"
+#include "tensorflow/java/src/gen/cc/source_writer.h"
+#include "tensorflow/java/src/gen/cc/op_parser.h"
 #include "tensorflow/java/src/gen/cc/op_generator.h"
 
 namespace tensorflow {
 namespace java {
 namespace {
 
-string CamelCase(const string& str, char delimiter, bool upper) {
-  string result;
-  bool cap = upper;
-  for (string::const_iterator it = str.begin(); it != str.end(); ++it) {
-    const char c = *it;
-    if (c == delimiter) {
-      cap = true;
-    } else if (cap) {
-      result += toupper(c);
-      cap = false;
+const char* kLicenseSnippet =
+    "tensorflow/java/src/gen/resources/license.snippet.java";
+
+const std::map<string, Type> kPrimitiveAttrTypes = {
+  { "Boolean", Type::Boolean() },
+  { "Byte", Type::Byte() },
+  { "Character", Type::Byte() },
+  { "Float", Type::Float() },
+  { "Integer", Type::Long() },
+  { "Long", Type::Long() },
+  { "Short", Type::Long() },
+  { "Double", Type::Float() },
+};
+
+enum RenderMode {
+  DEFAULT,
+  SINGLE_OUTPUT,
+  SINGLE_LIST_OUTPUT
+};
+
+void CollectOpDependencies(const OpSpec& op, RenderMode mode,
+    std::list<Type>* out) {
+  out->push_back(Type::Class("Operation", "org.tensorflow"));
+  out->push_back(Type::Class("OperationBuilder", "org.tensorflow"));
+  out->push_back(Type::Class("Scope", "org.tensorflow.op"));
+  if (mode == SINGLE_OUTPUT) {
+    out->push_back(Type::Class("Output", "org.tensorflow"));
+  } else if (mode == SINGLE_LIST_OUTPUT) {
+    out->push_back(Type::Interface("Iterator", "java.util"));
+  }
+  // Don't pay attention to duplicate types in the dependency list, they will
+  // be filtered out by the SourceWriter.
+  for (const OpSpec::Operand& input : op.inputs()) {
+    out->push_back(input.var().type());
+    if (input.iterable()) {
+      out->push_back(Type::Class("Operands", "org.tensorflow.op"));
+    }
+  }
+  for (const OpSpec::Operand& output : op.outputs()) {
+    out->push_back(output.var().type());
+    if (output.iterable()) {
+      out->push_back(Type::Class("Arrays", "java.util"));
+    }
+  }
+  for (const OpSpec::Operand& attribute : op.attributes()) {
+    out->push_back(attribute.var().type());
+    if (attribute.var().type().name() == "Class") {
+      out->push_back(Type::Enum("DataType", "org.tensorflow"));
+    }
+  }
+  for (const OpSpec::Operand& option : op.options()) {
+    out->push_back(option.var().type());
+  }
+}
+
+void WriteSetAttrDirective(const OpSpec::Operand& attr, bool optional,
+    SourceWriter* writer) {
+  string var = optional ? "opts." + attr.var().name() : attr.var().name();
+  if (attr.iterable()) {
+    const Type& type = attr.data_type();
+    std::map<string, Type>::const_iterator it =
+      kPrimitiveAttrTypes.find(type.name());
+    if (it != kPrimitiveAttrTypes.end()) {
+      string array = attr.var().name() + "Array";
+      writer->AppendType(it->second)
+          .Append("[] " + array + " = new ")
+          .AppendType(it->second)
+          .Append("[" + var + ".size()];")
+          .EndLine();
+      writer->BeginBlock("for (int i = 0; i < " + array + ".length; ++i)")
+          .Append(array + "[i] = " + var + ".get(i);")
+          .EndLine()
+          .EndBlock()
+          .Append("opBuilder.setAttr(\"" + attr.graph_name() + "\", " + array)
+          .Append(");")
+          .EndLine();
     } else {
-      result += c;
+      writer->Append("opBuilder.setAttr(\"" + attr.graph_name() + "\", " + var)
+          .Append(".toArray(new ")
+          .AppendType(type)
+          .Append("[" + var + ".size()]));")
+          .EndLine();
     }
+  } else {
+    Type type = attr.var().type();
+    writer->Append("opBuilder.setAttr(\"" + attr.graph_name() + "\", ");
+    if (type.name() == "Class") {
+      writer->Append("DataType.fromClass(" + attr.var().name() + "));");
+    } else {
+      writer->Append(var + ");");
+    }
+    writer->EndLine();
   }
-  return result;
 }
 
-}  // namespace
+void RenderFactoryMethod(const OpSpec& op, const Type& op_class,
+    SourceWriter* writer) {
+  Method factory = Method::Create("create", op_class);
+  Javadoc factory_doc = Javadoc::Create(
+      "Factory method to create a class to wrap a new " + op_class.name()
+      + " operation to the graph.");
+  Variable scope =
+      Variable::Create("scope", Type::Class("Scope", "org.tensorflow.op"));
+  factory.add_argument(scope);
+  factory_doc.add_param_tag(scope.name(), "Current graph scope");
+  for (const OpSpec::Operand& input : op.inputs()) {
+    factory.add_argument(input.var());
+    factory_doc.add_param_tag(input.var().name(), input.description());
+  }
+  for (const OpSpec::Operand& attribute : op.attributes()) {
+    factory.add_argument(attribute.var());
+    factory_doc.add_param_tag(attribute.var().name(), attribute.description());
+  }
+  if (!op.options().empty()) {
+    factory.add_argument(Variable::Varargs("options", Type::Class("Options")));
+    factory_doc.add_param_tag("options", "carries optional attributes values");
+  }
+  factory_doc.add_tag("return", "a new instance of " + op_class.name());
+  writer->BeginMethod(factory, PUBLIC|STATIC, &factory_doc);
+  writer->Append("OperationBuilder opBuilder = scope.graph().opBuilder(\""
+      + op.graph_name() + "\", scope.makeOpName(\""
+      + op_class.name() + "\"));");
+  writer->EndLine();
 
-OpGenerator::OpGenerator() : env(Env::Default()) {}
+  for (const OpSpec::Operand& input : op.inputs()) {
+    if (input.iterable()) {
+      writer->Append("opBuilder.addInputList(Operands.asOutputs("
+          + input.var().name() + "));");
+      writer->EndLine();
+    } else {
+      writer->Append("opBuilder.addInput(" + input.var().name()
+          + ".asOutput());");
+      writer->EndLine();
+    }
+  }
+  for (const OpSpec::Operand& attribute : op.attributes()) {
+    WriteSetAttrDirective(attribute, false, writer);
+  }
+  if (!op.options().empty()) {
+    writer->BeginBlock("if (options != null)")
+        .BeginBlock("for (Options opts : options)");
+    for (const OpSpec::Operand& option : op.options()) {
+      writer->BeginBlock("if (opts." + option.var().name() + " != null)");
+      WriteSetAttrDirective(option, true, writer);
+      writer->EndBlock();
+    }
+    writer->EndBlock().EndBlock();
+  }
+  writer->Append("return new ")
+      .AppendType(op_class)
+      .Append("(opBuilder.build());")
+      .EndLine();
+  writer->EndMethod();
+}
 
-OpGenerator::~OpGenerator() {}
+void RenderConstructor(const OpSpec& op, const Type& op_class,
+    SourceWriter* writer) {
+  Method constructor = Method::ConstructorFor(op_class)
+    .add_argument(
+        Variable::Create("operation",
+            Type::Class("Operation", "org.tensorflow")));
+  for (const OpSpec::Operand& output : op.outputs()) {
+    if (output.iterable() && !output.data_type().unknown()) {
+      constructor.add_annotation(
+          Annotation::Create("SuppressWarnings").attributes("\"unchecked\""));
+      break;
+    }
+  }
+  writer->BeginMethod(constructor, PRIVATE)
+      .Append("super(operation);")
+      .EndLine();
+  if (op.outputs().size() > 0) {
+    writer->Append("int outputIdx = 0;")
+        .EndLine();
+    for (const OpSpec::Operand& output : op.outputs()) {
+      if (output.iterable()) {
+        string var_length = output.var().name() + "Length";
+        writer->Append("int " + var_length)
+            .Append(" = operation.outputListLength(\"" + output.graph_name()
+                + "\");")
+            .EndLine()
+            .Append(output.var().name() + " = Arrays.asList(");
+        if (!output.data_type().unknown()) {
+          writer->Append("(")
+              .AppendType(output.var().type().parameters().front())
+              .Append("[])");
+        }
+        writer->Append("operation.outputList(outputIdx, " + var_length + "));")
+            .EndLine()
+            .Append("outputIdx += " + var_length + ";")
+            .EndLine();
+      } else {
+        writer->Append(output.var().name()
+                + " = operation.output(outputIdx++);")
+            .EndLine();
+      }
+    }
+  }
+  writer->EndMethod();
+}
 
-Status OpGenerator::Run(const OpList& ops, const string& lib_name,
-                        const string& base_package, const string& output_dir) {
-  const string package =
-      base_package + '.' + str_util::StringReplace(lib_name, "_", "", true);
-  const string package_path =
-      output_dir + '/' + str_util::StringReplace(package, ".", "/", true);
-  const string group = CamelCase(lib_name, '_', false);
+void RenderGettersAndSetters(const OpSpec& op, SourceWriter* writer) {
+  for (const OpSpec::Operand& option : op.options()) {
+    Method setter = Method::Create(option.var().name(), Type::Class("Options"))
+        .add_argument(option.var());
+    Javadoc setter_doc = Javadoc::Create()
+        .add_param_tag(option.var().name(), option.description());
+    writer->BeginMethod(setter, PUBLIC|STATIC, &setter_doc)
+        .Append("return new Options()." + option.var().name() + "("
+            + option.var().name() + ");")
+        .EndLine()
+        .EndMethod();
+  }
+  for (const OpSpec::Operand& output : op.outputs()) {
+    Method getter = Method::Create(output.var().name(), output.var().type());
+    Javadoc getter_doc = Javadoc::Create(output.description());
+    writer->BeginMethod(getter, PUBLIC, &getter_doc)
+        .Append("return " + output.var().name() + ";")
+        .EndLine()
+        .EndMethod();
+  }
+}
+
+void RenderInterfaceImpl(const OpSpec& op, RenderMode mode,
+    SourceWriter* writer) {
+  OpSpec::Operand output = op.outputs().front();
+
+  if (mode == SINGLE_OUTPUT) {
+    bool cast2obj = output.data_type().unknown();
+    Type return_type = Type::Class("Output", "org.tensorflow")
+        .add_parameter(cast2obj ? Type::Class("Object") : output.data_type());
+    Method as_output = Method::Create("asOutput", return_type)
+        .add_annotation(Annotation::Create("Override"));
+    if (cast2obj) {
+      as_output.add_annotation(
+          Annotation::Create("SuppressWarnings").attributes("\"unchecked\""));
+    }
+    writer->BeginMethod(as_output, PUBLIC);
+    if (cast2obj) {
+      writer->Append("return (").AppendType(return_type).Append(") ");
+    } else {
+      writer->Append("return ");
+    }
+    writer->Append(output.var().name() + ";")
+        .EndLine()
+        .EndMethod();
+
+  } else if (mode == SINGLE_LIST_OUTPUT) {
+    Type operand = Type::Interface("Operand", "org.tensorflow");
+    if (output.data_type().unknown()) {
+      operand.add_parameter(Type::Class("Object"));
+    } else {
+      operand.add_parameter(output.data_type());
+    }
+    Type return_type = Type::Interface("Iterator", "java.util")
+        .add_parameter(operand);
+    Method iterator = Method::Create("iterator", return_type)
+        .add_annotation(Annotation::Create("Override"))
+        .add_annotation(Annotation::Create("SuppressWarnings")
+            .attributes("{\"rawtypes\", \"unchecked\"}"));
+    // cast the output list using a raw List
+    writer->BeginMethod(iterator, PUBLIC)
+        .Append("return (" + return_type.name() + ") ")
+        .Append(output.var().name() + ".iterator();")
+        .EndLine()
+        .EndMethod();
+  }
+}
+
+void RenderOptionsClass(const OpSpec& op, SourceWriter* writer) {
+  Type options_class = Type::Class("Options");
+  Javadoc options_doc = Javadoc::Create(
+      "Class holding optional attributes of this operation");
+  writer->BeginInnerType(options_class, PUBLIC | STATIC, &options_doc);
+  for (const OpSpec::Operand& option : op.options()) {
+    Method setter = Method::Create(option.var().name(), options_class)
+        .add_argument(option.var());
+    Javadoc setter_doc = Javadoc::Create()
+        .add_param_tag(option.var().name(), option.description());
+    writer->BeginMethod(setter, PUBLIC, &setter_doc)
+        .Append("this." + option.var().name() + " = " + option.var().name()
+            + ";")
+        .EndLine()
+        .Append("return this;")
+        .EndLine()
+        .EndMethod();
+  }
+  writer->EndLine();
+  for (const OpSpec::Operand& option : op.options()) {
+    writer->WriteField(option.var(), PRIVATE);
+  }
+  Method constructor = Method::ConstructorFor(options_class);
+  writer->BeginMethod(constructor, PRIVATE).EndMethod();
+  writer->EndType();
+}
 
-  if (!env->FileExists(package_path).ok()) {
-    TF_CHECK_OK(env->RecursivelyCreateDir(package_path));
+void RenderEndpoint(const OpSpec& op, const OpSpec::Endpoint& endpoint,
+    SourceWriter* writer) {
+  RenderMode mode = DEFAULT;
+  if (op.outputs().size() == 1) {
+    mode = op.outputs().front().iterable() ? SINGLE_LIST_OUTPUT : SINGLE_OUTPUT;
+  }
+  std::list<Type> dependencies;
+  CollectOpDependencies(op, mode, &dependencies);
+  const Type& op_class = endpoint.type();
+  writer->WriteFromFile(kLicenseSnippet)
+      .EndLine()
+      .Append("// This file is machine generated, DO NOT EDIT!")
+      .EndLine()
+      .EndLine()
+      .BeginType(op_class, PUBLIC|FINAL, &dependencies, &endpoint.javadoc());
+  if (!op.options().empty()) {
+    RenderOptionsClass(op, writer);
   }
+  RenderFactoryMethod(op, op_class, writer);
+  RenderGettersAndSetters(op, writer);
+  if (mode != DEFAULT) {
+    RenderInterfaceImpl(op, mode, writer);
+  }
+  writer->EndLine();
+  for (const OpSpec::Operand& output : op.outputs()) {
+    writer->WriteField(output.var(), PRIVATE);
+  }
+  RenderConstructor(op, op_class, writer);
+  writer->EndType();
+}
+
+}  // namespace
+
+OpGenerator::OpGenerator(const string& base_package, const string& output_dir,
+    const std::vector<string>& api_dirs, Env* env)
+  : base_package_(base_package), output_dir_(output_dir), api_dirs_(api_dirs),
+    env_(env) {
+}
 
+Status OpGenerator::Run(const OpList& op_list, const string& lib_name) {
   LOG(INFO) << "Generating Java wrappers for '" << lib_name << "' operations";
-  // TODO(karllessard) generate wrappers from list of ops
+  ApiDefMap api_map(op_list);
+  if (!api_dirs_.empty()) {
+    // Only load api files that correspond to the requested "op_list"
+    for (const auto& op : op_list.op()) {
+      for (const auto& api_def_dir : api_dirs_) {
+        const std::string api_def_file_pattern =
+            io::JoinPath(api_def_dir, "api_def_" + op.name() + ".pbtxt");
+        if (env_->FileExists(api_def_file_pattern).ok()) {
+          TF_CHECK_OK(api_map.LoadFile(env_, api_def_file_pattern));
+        }
+      }
+    }
+  }
+  api_map.UpdateDocs();
+  for (const auto& op_def : op_list.op()) {
+    const ApiDef* api_def = api_map.GetApiDef(op_def.name());
+    if (api_def->visibility() != ApiDef::SKIP) {
+      Status status = GenerateOp(op_def, *api_def, lib_name);
+      if (status != Status::OK()) {
+        LOG(ERROR) << "Fail to generate Java wrapper for operation \""
+            << op_def.name() << "\"";
+      }
+    }
+  }
+  return Status::OK();
+}
+
+Status OpGenerator::GenerateOp(const OpDef& op_def, const ApiDef& api_def,
+    const string& lib_name) {
+  std::unique_ptr<OpSpec> op;
+  OpParser op_parser(op_def, api_def, lib_name, base_package_);
+  op_parser.Parse(&op);
+  for (const OpSpec::Endpoint& endpoint : op->endpoints()) {
+    string package_path = io::JoinPath(output_dir_,
+        str_util::StringReplace(endpoint.type().package(), ".", "/", true));
+    if (!env_->FileExists(package_path).ok()) {
+      TF_CHECK_OK(Env::Default()->RecursivelyCreateDir(package_path));
+    }
+    string file_path =
+        io::JoinPath(package_path, endpoint.type().name() + ".java");
+    std::unique_ptr<tensorflow::WritableFile> file;
+    TF_CHECK_OK(env_->NewWritableFile(file_path, &file));
 
+    SourceFileWriter writer(file.get());
+    RenderEndpoint(*op, endpoint, &writer);
+  }
   return Status::OK();
 }
 
diff --git a/tensorflow/java/src/gen/cc/op_generator.h b/tensorflow/java/src/gen/cc/op_generator.h
index 4b55ed3ed9..19d8db95fb 100644
--- a/tensorflow/java/src/gen/cc/op_generator.h
+++ b/tensorflow/java/src/gen/cc/op_generator.h
@@ -17,34 +17,42 @@ limitations under the License.
 #define TENSORFLOW_JAVA_SRC_GEN_CC_OP_GENERATOR_H_
 
 #include <string>
+#include <vector>
 
-#include "tensorflow/core/framework/op.h"
-#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/api_def.pb.h"
 #include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/lib/core/status.h"
 
 namespace tensorflow {
 namespace java {
 
-/// \brief A generator of Java operation wrappers.
-///
-/// Such generator is normally ran only once per executable, outputting
-/// wrappers for the all registered operations it has been compiled with.
-/// Nonetheless, it is designed to support multiple runs, giving a different
-/// list of operations on each cycle.
+// A generator of Java operation wrappers.
+//
+// Such generator is normally ran only once per executable, outputting
+// wrappers for the all registered operations it has been compiled with.
+// Nonetheless, it is designed to support multiple runs, giving a different
+// list of operations on each cycle.
 class OpGenerator {
  public:
-  OpGenerator();
-  virtual ~OpGenerator();
+  OpGenerator(const string& base_package, const string& output_dir,
+      const std::vector<string>& api_dirs, Env* env = Env::Default());
+  virtual ~OpGenerator() = default;
 
-  /// \brief Generates wrappers for the given list of 'ops'.
-  ///
-  /// Output files are generated in <output_dir>/<base_package>/<lib_package>,
-  /// where 'lib_package' is derived from 'lib_name'.
-  Status Run(const OpList& ops, const string& lib_name,
-             const string& base_package, const string& output_dir);
+  // Generates wrappers for the given list of 'ops'.
+  //
+  // Output files are generated in <output_dir>/<base_package>/<lib_package>,
+  // where 'lib_package' is derived from 'lib_name'.
+  Status Run(const OpList& op_list, const string& lib_name);
 
  private:
-  Env* env;
+  string base_package_;
+  string output_dir_;
+  std::vector<string> api_dirs_;
+  Env* env_;
+
+  Status GenerateOp(const OpDef& op_def, const ApiDef& api_def,
+    const string& lib_name);
 };
 
 }  // namespace java
diff --git a/tensorflow/java/src/gen/cc/op_parser.cc b/tensorflow/java/src/gen/cc/op_parser.cc
new file mode 100644
index 0000000000..0541e343d8
--- /dev/null
+++ b/tensorflow/java/src/gen/cc/op_parser.cc
@@ -0,0 +1,417 @@
+/* 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.
+==============================================================================*/
+
+#include <map>
+#include <vector>
+#include <utility>
+#include <string>
+
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/java/src/gen/cc/op_parser.h"
+
+namespace tensorflow {
+namespace java {
+namespace {
+
+string SnakeToCamelCase(const string& str, bool upper = false) {
+  string result;
+  bool cap = upper;
+  for (string::const_iterator it = str.begin(); it != str.end(); ++it) {
+    const char c = *it;
+    if (c == '_') {
+      cap = true;
+    } else if (cap) {
+      result += toupper(c);
+      cap = false;
+    } else {
+      result += c;
+    }
+  }
+  return result;
+}
+
+bool IsRealNumber(DataType type) {
+  for (DataType dt : RealNumberTypes()) {
+    if (type == dt) {
+      return true;
+    }
+  }
+  return false;
+}
+
+bool IsRealNumbers(const AttrValue& values) {
+  if (values.has_list()) {
+    for (int i = 0; i < values.list().type_size(); ++i) {
+      if (!IsRealNumber(values.list().type(i))) {
+        return false;
+      }
+    }
+    return true;
+  }
+  return IsRealNumber(values.type());
+}
+
+string ParseDocumentation(const string& text) {
+  std::stringstream javadoc_text;
+  string::const_iterator c_iter = text.cbegin();
+  bool code = false;
+  bool emphasis = false;
+  bool list = false;
+  while (c_iter != text.cend()) {
+    char c = *c_iter++;
+    int count = 1;
+    switch (c) {
+    case '\n':
+      if (!code) {
+        // consumes all subsequent newlines, if there are more than one,
+        // then there are two choices:
+        // - if the next line starts with an asterisk, we are enumerating
+        //   a list of items
+        // - otherwise, we are starting a new paragraph
+        for (; c_iter != text.cend() && *c_iter == '\n'; ++count, ++c_iter) {}
+        if (c_iter != text.cend()) {
+          if (count > 1) {
+            if (*c_iter != '*' && list) {
+              javadoc_text << "</li>\n</ul>\n";
+              list = false;
+            } else if (*c_iter == '*' && !list) {
+              javadoc_text << "\n<ul>\n<li>";
+              list = true;
+              c_iter++;
+            } else {
+              javadoc_text << "\n<p>\n";
+            }
+          } else if (list && *c_iter == '*') {
+            javadoc_text << "</li>\n<li>";
+            c_iter++;
+          } else {
+            javadoc_text << '\n';
+          }
+        }
+      }
+      break;
+    case '`':
+      // consumes all subsequent backquotes, those are use enclose code.
+      // if there are more than 3, we are dealing with a pre-formatted block,
+      // otherwise it is a single-line code snippet
+      for (; c_iter != text.cend() && *c_iter == '`'; ++count, ++c_iter) {}
+      if (count >= 3) {
+        javadoc_text << (code ? "\n}</pre>" : "<pre>{@code\n");
+      } else {
+        javadoc_text << (code ? "}" : "{@code ");
+      }
+      code = !code;
+      break;
+    case '*':
+      if (!code) {
+        // consumes all subsequent asterisks, if there are more than one, then
+        // we put the text in bold, otherwise in italic
+        for (; c_iter != text.cend() && *c_iter == '*'; ++count, ++c_iter) {}
+        if (count > 1) {
+          javadoc_text << (emphasis ? "</b>" : "<b>");
+        } else {
+          javadoc_text << (emphasis ? "</i>" : "<i>");
+        }
+        emphasis = !emphasis;
+      } else {
+        javadoc_text << '*';
+      }
+      break;
+    default:
+      javadoc_text << c;
+      break;
+    }
+  }
+  return javadoc_text.str();
+}
+
+}  // namespace
+
+OpParser::OpParser(const OpDef& op_def, const ApiDef& api_def,
+    const string& lib_name, const string& base_package)
+  : op_def_(op_def), op_api_(api_def), lib_name_(lib_name),
+    base_package_(base_package) {
+}
+
+void OpParser::Parse(std::unique_ptr<OpSpec>* op_ptr) {
+  visited_attrs_.clear();
+  next_generic_ = 'T';
+  op_ptr->reset(new OpSpec(op_api_.graph_op_name()));
+  for (const string& next_input_name : op_api_.arg_order()) {
+    for (int i = 0; i < op_def_.input_arg().size(); ++i) {
+      if (op_def_.input_arg(i).name() == next_input_name) {
+        ParseInput(op_def_.input_arg(i), op_api_.in_arg(i), op_ptr->get());
+        break;
+      }
+    }
+  }
+  for (int i = 0; i < op_def_.attr().size(); ++i) {
+    ParseAttribute(op_def_.attr(i), op_api_.attr(i), op_ptr->get());
+  }
+  for (int i = 0; i < op_def_.output_arg().size(); ++i) {
+    ParseOutput(op_def_.output_arg(i), op_api_.out_arg(i), op_ptr->get());
+  }
+  BuildEndpoints(op_ptr->get());
+}
+
+void OpParser::BuildEndpoints(OpSpec* op) {
+  Javadoc op_doc = Javadoc::Create(ParseDocumentation(op_api_.summary()))
+    .details(ParseDocumentation(op_api_.description()));
+  std::vector<Type> op_supertypes;
+  op_supertypes.push_back(Type::Class("PrimitiveOp", "org.tensorflow.op"));
+  std::map<string, const Type*> op_generics;
+  for (const OpSpec::Operand& output : op->outputs()) {
+    // declare generic output parameters at the Op class level
+    const Type& data_type = output.data_type();
+    if (data_type.kind() == Type::GENERIC && !data_type.unknown()
+        && op_generics.find(data_type.name()) == op_generics.end()) {
+      op_generics.insert(std::make_pair(data_type.name(), &data_type));
+      op_doc.add_param_tag("<" + data_type.name() + ">",
+          "data type of output '" + output.var().name() + "'");
+    }
+    // implement the Op as an (iteration of) Operand if it has only one output
+    if (op->outputs().size() == 1) {
+      Type operand_inf(Type::Interface("Operand", "org.tensorflow"));
+      operand_inf.add_parameter(data_type.unknown() ?
+          Type::Class("Object") : data_type);
+      op_supertypes.push_back(output.iterable() ?
+          Type::IterableOf(operand_inf) : operand_inf);
+    }
+  }
+  for (const auto& endpoint_def : op_api_.endpoint()) {
+    std::vector<string> name_tokens = str_util::Split(endpoint_def.name(), ".");
+    // if the endpoint specifies a package, use it, otherwise derive it from the
+    // op library name.
+    string name;
+    string package;
+    if (name_tokens.size() > 1) {
+      package = str_util::Lowercase(name_tokens.at(0));
+      name = name_tokens.at(1);
+    } else {
+      package = str_util::StringReplace(lib_name_, "_", "", true);
+      name = name_tokens.at(0);
+    }
+    Type endpoint(Type::Class(name, base_package_ + "." + package));
+    Javadoc endpoint_doc(op_doc);
+    for (const auto& parameter : op_generics) {
+      endpoint.add_parameter(*parameter.second);
+    }
+    for (const Type& supertype : op_supertypes) {
+      endpoint.add_supertype(supertype);
+    }
+    if (endpoint_def.deprecation_version() > 0) {
+      string explanation;
+      if (op_api_.endpoint(0).deprecation_version() == 0) {
+        explanation = ", use {@link "
+            + op->endpoints().at(0).type().full_name()
+            + "} instead";
+      } else {
+        explanation = op_def_.deprecation().explanation();
+      }
+      endpoint_doc.add_tag("deprecated", explanation);
+      endpoint.add_annotation(Annotation::Create("Deprecated"));
+    }
+    // only visible ops should be annotated for exposure in the Ops Graph API
+    if (op_api_.visibility() != ApiDef::HIDDEN) {
+      string group_name = SnakeToCamelCase(lib_name_);
+      endpoint.add_annotation(
+          Annotation::Create("Operator", "org.tensorflow.op.annotation")
+            .attributes("group = \"" + group_name + "\""));
+    }
+    op->add_endpoint(endpoint, endpoint_doc);
+  }
+}
+
+void OpParser::ParseInput(const OpDef_ArgDef& input_def,
+    const ApiDef::Arg& input_api, OpSpec* op) {
+  bool iterable = false;
+  Type data_type = DataTypeOf(input_def, &iterable);
+  Type type = Type::Interface("Operand", "org.tensorflow")
+    .add_parameter(data_type);
+  if (iterable) {
+    type = Type::IterableOf(type);
+  }
+  op->add_input(OpSpec::Operand(input_api.name(),
+      Variable::Create(SnakeToCamelCase(input_api.rename_to()), type),
+      data_type,
+      ParseDocumentation(input_api.description()),
+      iterable));
+}
+
+void OpParser::ParseOutput(const OpDef_ArgDef& output_def,
+    const ApiDef::Arg& output_api, OpSpec* op) {
+  bool iterable = false;
+  Type data_type = DataTypeOf(output_def, &iterable);
+  Type type = Type::Class("Output", "org.tensorflow")
+    .add_parameter(data_type);
+  if (iterable) {
+    type = Type::ListOf(type);
+  }
+  op->add_output(OpSpec::Operand(output_api.name(),
+      Variable::Create(SnakeToCamelCase(output_api.rename_to()), type),
+      data_type,
+      ParseDocumentation(output_api.description()),
+      iterable));
+}
+
+void OpParser::ParseAttribute(const OpDef_AttrDef& attr_def,
+    const ApiDef::Attr& attr_api, OpSpec* op) {
+  // do not parse attributes already visited, they have probably been inferred
+  // before as an input argument type
+  if (visited_attrs_.find(attr_def.name()) != visited_attrs_.cend()) {
+    return;
+  }
+  bool iterable = false;
+  Type data_type = DataTypeOf(attr_def, &iterable);
+  // generic attributes should be passed as an explicit type
+  bool explicit_type = data_type.kind() == Type::GENERIC && !iterable;
+  Type type = explicit_type ?
+      Type::Class("Class").add_parameter(data_type) : data_type;
+  if (iterable) {
+    type = Type::ListOf(data_type);
+  }
+  OpSpec::Operand attr(attr_api.name(),
+      Variable::Create(SnakeToCamelCase(attr_api.rename_to()), type),
+      data_type,
+      ParseDocumentation(attr_api.description()),
+      iterable);
+  // attributes with a default value are optional
+  if (attr_api.has_default_value() && !explicit_type) {
+    op->add_option(attr);
+  } else {
+    op->add_attribute(attr);
+  }
+  visited_attrs_.insert(std::make_pair(attr_api.name(), data_type));
+}
+
+Type OpParser::DataTypeOf(const OpDef_ArgDef& arg, bool* iterable_out) {
+  if (!arg.number_attr().empty()) {
+    visited_attrs_.insert(std::make_pair(arg.number_attr(), Type::Int()));
+    *iterable_out = true;
+  }
+  if (arg.type() != DataType::DT_INVALID) {
+    // resolve type from DataType
+    switch (arg.type()) {
+      case DataType::DT_BOOL:
+        return Type::Class("Boolean");
+
+      case DataType::DT_STRING:
+        return Type::Class("String");
+
+      case DataType::DT_FLOAT:
+        return Type::Class("Float");
+
+      case DataType::DT_DOUBLE:
+        return Type::Class("Double");
+
+      case DataType::DT_UINT8:
+        return Type::Class("UInt8", "org.tensorflow.types");
+
+      case DataType::DT_INT32:
+        return Type::Class("Integer");
+
+      case DataType::DT_INT64:
+        return Type::Class("Long");
+
+      case DataType::DT_RESOURCE:
+        // TODO(karllessard) create a Resource utility class that could be
+        // used to store a resource and its type (passed in a second argument).
+        // For now, we need to force a wildcard and we will unfortunately lose
+        // track of the resource type.
+        return Type::Wildcard();
+
+      default:
+        break;
+    }
+  } else {
+    // resolve type from type attribute
+    string attr_name = arg.type_attr();
+    if (attr_name.empty()) {
+      attr_name = arg.type_list_attr();
+      if (!attr_name.empty()) {
+        *iterable_out = true;
+        Type type = Type::Wildcard();
+        visited_attrs_.insert(std::make_pair(attr_name, type));
+        return type;
+      }
+    }
+    for (const auto& attr : op_def_.attr()) {
+      if (attr.name() == attr_name) {
+        Type type = DataTypeOf(attr, iterable_out);
+        visited_attrs_.insert(std::make_pair(attr_name, type));
+        return type;
+      }
+    }
+  }
+  LOG(WARNING) << "Data type for arg \"" << arg.name() << "\" is unknown";
+  return Type::Wildcard();
+}
+
+Type OpParser::DataTypeOf(const OpDef_AttrDef& attr, bool* iterable_out) {
+  std::map<string, Type>::const_iterator it = visited_attrs_.find(attr.name());
+  if (it != visited_attrs_.cend()) {
+    return it->second;
+  }
+  string attr_type = attr.type();
+  if (attr.type().compare(0, 5, "list(") == 0) {
+    attr_type = attr_type.substr(5, attr.type().find_last_of(')') - 5);
+    *iterable_out = true;
+  }
+  if (attr_type == "type") {
+    if (*iterable_out) {
+      return Type::Enum("DataType", "org.tensorflow");
+    }
+    return GetNextGenericTensorType(attr.allowed_values());
+  }
+  if (attr_type == "string") {
+    return Type::Class("String");
+  }
+  if (attr_type == "int") {
+    return Type::Class("Integer");
+  }
+  if (attr_type == "float") {
+    return Type::Class("Float");
+  }
+  if (attr_type == "bool") {
+    return Type::Class("Boolean");
+  }
+  if (attr_type == "shape") {
+    return Type::Class("Shape", "org.tensorflow");
+  }
+  if (attr_type == "tensor") {
+    return Type::Class("Tensor", "org.tensorflow")
+      .add_parameter(Type::Wildcard());
+  }
+  LOG(WARNING) << "Data type for attribute \"" << attr_type << "\" is unknown";
+  return *iterable_out ? Type::Wildcard() : Type::Class("Object");
+}
+
+Type OpParser::GetNextGenericTensorType(const AttrValue& allowed_values)  {
+  Type generic = Type::Generic(string(1, next_generic_));
+  next_generic_ = (next_generic_ == 'Z') ? 'A' : next_generic_ + 1;
+
+  // when only real numbers are allowed, enforce that restriction in the Java by
+  // extending the generic from java.lang.Number
+  if (IsRealNumbers(allowed_values)) {
+    generic.add_supertype(Type::Class("Number"));
+  }
+  return generic;
+}
+
+}  // namespace java
+}  // namespace tensorflow
diff --git a/tensorflow/java/src/gen/cc/op_parser.h b/tensorflow/java/src/gen/cc/op_parser.h
new file mode 100644
index 0000000000..42855127cc
--- /dev/null
+++ b/tensorflow/java/src/gen/cc/op_parser.h
@@ -0,0 +1,137 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_JAVA_SRC_GEN_CC_OP_PARSER_H_
+#define TENSORFLOW_JAVA_SRC_GEN_CC_OP_PARSER_H_
+
+#include <map>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/api_def.pb.h"
+#include "tensorflow/java/src/gen/cc/java_defs.h"
+
+namespace tensorflow {
+namespace java {
+
+// Specification of a TensorFlow operation to generate.
+//
+// This is the result of an operation definition parsing, see OpParser::Parse().
+class OpSpec {
+ public:
+  class Endpoint {
+   public:
+    Endpoint(const Type& type, const Javadoc& javadoc)
+      : type_(type), javadoc_(javadoc) {}
+    const Type& type() const { return type_; }
+    const Javadoc& javadoc() const { return javadoc_; }
+
+   private:
+    Type type_;
+    Javadoc javadoc_;
+  };
+
+  class Operand {
+   public:
+    Operand(const string& graph_name, const Variable& var,
+        const Type& data_type, const string& description, bool iterable)
+     : graph_name_(graph_name), var_(var), data_type_(data_type),
+       description_(description), iterable_(iterable) {}
+    const string& graph_name() const { return graph_name_; }
+    const Variable& var() const { return var_; }
+    Variable* var_ptr() { return &var_; }
+    const Type& data_type() const { return data_type_; }
+    const string& description() const { return description_; }
+    bool iterable() const { return iterable_; }
+
+   private:
+    string graph_name_;
+    Variable var_;
+    Type data_type_;
+    string description_;
+    bool iterable_;
+  };
+
+  explicit OpSpec(const string& graph_name) : graph_name_(graph_name) {}
+  const string& graph_name() const { return graph_name_; }
+  const std::vector<Endpoint> endpoints() const { return endpoints_; }
+  void add_endpoint(const Type& type, const Javadoc& javadoc) {
+    endpoints_.push_back(Endpoint(type, javadoc));
+  }
+  const std::vector<Operand>& inputs() const { return inputs_; }
+  void add_input(const Operand& input) {
+    inputs_.push_back(input);
+  }
+  const std::vector<Operand>& outputs() const { return outputs_; }
+  void add_output(const Operand& output) {
+    outputs_.push_back(output);
+  }
+  const std::vector<Operand>& attributes() const { return attributes_; }
+  void add_attribute(const Operand& attribute) {
+    attributes_.push_back(attribute);
+  }
+  const std::vector<Operand>& options() const { return options_; }
+  void add_option(const Operand& option) {
+    options_.push_back(option);
+  }
+
+ private:
+  string graph_name_;
+  std::vector<Endpoint> endpoints_;
+  std::vector<Operand> inputs_;
+  std::vector<Operand> outputs_;
+  std::vector<Operand> attributes_;
+  std::vector<Operand> options_;
+};
+
+// A parser of ops proto definitions.
+//
+// This object parses the definition and the api of an TensorFlow operation to
+// produce a specification that can be used for Java source code rendering.
+class OpParser {
+ public:
+  OpParser(const OpDef& op_def, const ApiDef& api_def, const string& lib_name,
+      const string& base_package);
+  virtual ~OpParser() = default;
+
+  // Produces an operation specification from its proto definitions.
+  void Parse(std::unique_ptr<OpSpec>* op_ptr);
+
+ private:
+  OpDef op_def_;
+  ApiDef op_api_;
+  string lib_name_;
+  string base_package_;
+  std::map<string, Type> visited_attrs_;
+  char next_generic_ = 0;
+
+  void BuildEndpoints(OpSpec* op);
+  void ParseInput(const OpDef_ArgDef& input_def,
+      const ApiDef::Arg& input_api, OpSpec* op);
+  void ParseOutput(const OpDef_ArgDef& output_def,
+      const ApiDef::Arg& output_api, OpSpec* op);
+  void ParseAttribute(const OpDef_AttrDef& attr_def,
+      const ApiDef::Attr& attr_api, OpSpec* op);
+  Type DataTypeOf(const OpDef_ArgDef& arg_def, bool *iterable_out);
+  Type DataTypeOf(const OpDef_AttrDef& attr_def, bool *iterable_out);
+  Type GetNextGenericTensorType(const AttrValue& allowed_values);
+};
+
+}  // namespace java
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_JAVA_SRC_GEN_CC_OP_PARSER_H_
diff --git a/tensorflow/java/src/gen/cc/source_writer.cc b/tensorflow/java/src/gen/cc/source_writer.cc
index a02f75ad6e..b1de5af6ba 100644
--- a/tensorflow/java/src/gen/cc/source_writer.cc
+++ b/tensorflow/java/src/gen/cc/source_writer.cc
@@ -15,7 +15,7 @@ limitations under the License.
 
 #include <string>
 #include <algorithm>
-#include <deque>
+#include <list>
 
 #include "tensorflow/java/src/gen/cc/source_writer.h"
 
@@ -83,20 +83,20 @@ SourceWriter& SourceWriter::Append(const StringPiece& str) {
 }
 
 SourceWriter& SourceWriter::AppendType(const Type& type) {
-  if (type.kind() == Type::Kind::GENERIC && type.name().empty()) {
+  if (type.unknown()) {
     Append("?");
   } else {
     Append(type.name());
-  }
-  if (!type.parameters().empty()) {
-    Append("<");
-    for (const Type& t : type.parameters()) {
-      if (&t != &type.parameters().front()) {
-        Append(", ");
+    if (!type.parameters().empty()) {
+      Append("<");
+      for (const Type& t : type.parameters()) {
+        if (&t != &type.parameters().front()) {
+          Append(", ");
+        }
+        AppendType(t);
       }
-      AppendType(t);
+      Append(">");
     }
-    Append(">");
   }
   return *this;
 }
@@ -107,7 +107,21 @@ SourceWriter& SourceWriter::EndLine() {
   return *this;
 }
 
-SourceWriter& SourceWriter::BeginMethod(const Method& method, int modifiers) {
+SourceWriter& SourceWriter::BeginBlock(const string& expression) {
+  if (!expression.empty()) {
+    Append(expression + " {");
+  } else {
+    Append(newline_ ? "{" : " {");
+  }
+  return EndLine().Indent(2);
+}
+
+SourceWriter& SourceWriter::EndBlock() {
+  return Indent(-2).Append("}").EndLine();
+}
+
+SourceWriter& SourceWriter::BeginMethod(const Method& method, int modifiers,
+    const Javadoc* javadoc) {
   GenericNamespace* generic_namespace = PushGenericNamespace(modifiers);
   if (!method.constructor()) {
     generic_namespace->Visit(method.return_type());
@@ -116,8 +130,9 @@ SourceWriter& SourceWriter::BeginMethod(const Method& method, int modifiers) {
     generic_namespace->Visit(v.type());
   }
   EndLine();
-  WriteDoc(method.description(), method.return_description(),
-      &method.arguments());
+  if (javadoc != nullptr) {
+    WriteJavadoc(*javadoc);
+  }
   if (!method.annotations().empty()) {
     WriteAnnotations(method.annotations());
   }
@@ -145,29 +160,35 @@ SourceWriter& SourceWriter::EndMethod() {
   return *this;
 }
 
-SourceWriter& SourceWriter::BeginType(const Type& type,
-    const std::list<Type>* dependencies, int modifiers) {
+SourceWriter& SourceWriter::BeginType(const Type& type, int modifiers,
+    const std::list<Type>* extra_dependencies, const Javadoc* javadoc) {
   if (!type.package().empty()) {
     Append("package ").Append(type.package()).Append(";").EndLine();
   }
-  if (dependencies != nullptr && !dependencies->empty()) {
-    TypeImporter type_importer(type.package());
-    for (const Type& t : *dependencies) {
+  TypeImporter type_importer(type.package());
+  type_importer.Visit(type);
+  if (extra_dependencies != nullptr) {
+    for (const Type& t : *extra_dependencies) {
       type_importer.Visit(t);
     }
+  }
+  if (!type_importer.imports().empty()) {
     EndLine();
     for (const string& s : type_importer.imports()) {
       Append("import ").Append(s).Append(";").EndLine();
     }
   }
-  return BeginInnerType(type, modifiers);
+  return BeginInnerType(type, modifiers, javadoc);
 }
 
-SourceWriter& SourceWriter::BeginInnerType(const Type& type, int modifiers) {
+SourceWriter& SourceWriter::BeginInnerType(const Type& type, int modifiers,
+    const Javadoc* javadoc) {
   GenericNamespace* generic_namespace = PushGenericNamespace(modifiers);
   generic_namespace->Visit(type);
   EndLine();
-  WriteDoc(type.description());
+  if (javadoc != nullptr) {
+    WriteJavadoc(*javadoc);
+  }
   if (!type.annotations().empty()) {
     WriteAnnotations(type.annotations());
   }
@@ -200,14 +221,15 @@ SourceWriter& SourceWriter::EndType() {
   return *this;
 }
 
-SourceWriter& SourceWriter::WriteFields(const std::list<Variable>& fields,
-    int modifiers) {
-  EndLine();
-  for (const Variable& v : fields) {
-    WriteModifiers(modifiers);
-    AppendType(v.type()).Append(" ").Append(v.name()).Append(";");
-    EndLine();
+SourceWriter& SourceWriter::WriteField(const Variable& field, int modifiers,
+    const Javadoc* javadoc) {
+  // If present, write field javadoc only as one brief line
+  if (javadoc != nullptr && !javadoc->brief().empty()) {
+    Append("/** ").Append(javadoc->brief()).Append(" */").EndLine();
   }
+  WriteModifiers(modifiers);
+  AppendType(field.type()).Append(" ").Append(field.name()).Append(";");
+  EndLine();
   return *this;
 }
 
@@ -228,39 +250,33 @@ SourceWriter& SourceWriter::WriteModifiers(int modifiers) {
   return *this;
 }
 
-SourceWriter& SourceWriter::WriteDoc(const string& description,
-    const string& return_description, const std::list<Variable>* parameters) {
-  if (description.empty() && return_description.empty()
-      && (parameters == nullptr || parameters->empty())) {
-    return *this;  // no doc to write
-  }
+SourceWriter& SourceWriter::WriteJavadoc(const Javadoc& javadoc) {
+  Append("/**").Prefix(" * ").EndLine();
   bool do_line_break = false;
-  Append("/**").EndLine().Prefix(" * ");
-  if (!description.empty()) {
-    Write(description).EndLine();
+  if (!javadoc.brief().empty()) {
+    Write(javadoc.brief()).EndLine();
     do_line_break = true;
   }
-  if (parameters != nullptr && !parameters->empty()) {
+  if (!javadoc.details().empty()) {
     if (do_line_break) {
-      EndLine();
-      do_line_break = false;
-    }
-    for (const Variable& v : *parameters) {
-      Append("@param ").Append(v.name());
-      if (!v.description().empty()) {
-        Append(" ").Write(v.description());
-      }
-      EndLine();
+      Append("<p>").EndLine();
     }
+    Write(javadoc.details()).EndLine();
+    do_line_break = true;
   }
-  if (!return_description.empty()) {
+  if (!javadoc.tags().empty()) {
     if (do_line_break) {
       EndLine();
-      do_line_break = false;
     }
-    Append("@return ").Write(return_description).EndLine();
+    for (const auto& p : javadoc.tags()) {
+      Append("@" + p.first);
+      if (!p.second.empty()) {
+        Append(" ").Write(p.second);
+      }
+      EndLine();
+    }
   }
-  return Prefix("").Append(" **/").EndLine();
+  return Prefix("").Append(" */").EndLine();
 }
 
 SourceWriter& SourceWriter::WriteAnnotations(
@@ -311,20 +327,19 @@ void SourceWriter::PopGenericNamespace() {
 void SourceWriter::TypeVisitor::Visit(const Type& type) {
   DoVisit(type);
   for (const Type& t : type.parameters()) {
-    DoVisit(t);
+    Visit(t);
   }
   for (const Annotation& t : type.annotations()) {
     DoVisit(t);
   }
   for (const Type& t : type.supertypes()) {
-    DoVisit(t);
+    Visit(t);
   }
 }
 
 void SourceWriter::GenericNamespace::DoVisit(const Type& type) {
   // ignore non-generic parameters, wildcards and generics already declared
-  if (type.kind() == Type::GENERIC
-      && !type.IsWildcard()
+  if (type.kind() == Type::GENERIC && !type.unknown()
       && generic_names_.find(type.name()) == generic_names_.end()) {
     declared_types_.push_back(&type);
     generic_names_.insert(type.name());
@@ -333,7 +348,7 @@ void SourceWriter::GenericNamespace::DoVisit(const Type& type) {
 
 void SourceWriter::TypeImporter::DoVisit(const Type& type) {
   if (!type.package().empty() && type.package() != current_package_) {
-    imports_.insert(type.package() + '.' + type.name());
+    imports_.insert(type.full_name());
   }
 }
 
diff --git a/tensorflow/java/src/gen/cc/source_writer.h b/tensorflow/java/src/gen/cc/source_writer.h
index f011acd30a..1f0febe9a3 100644
--- a/tensorflow/java/src/gen/cc/source_writer.h
+++ b/tensorflow/java/src/gen/cc/source_writer.h
@@ -93,25 +93,22 @@ class SourceWriter {
   // This method appends a new opening brace to the current data and indent the
   // next lines according to Google Java Style Guide. The block can optionally
   // be preceded by an expression (e.g. Append("if(true)").BeginBlock();)
-  SourceWriter& BeginBlock() {
-    return Append(newline_ ? "{" : " {").EndLine().Indent(2);
-  }
+  SourceWriter& BeginBlock(const string& expr = "");
 
   // Ends the current block of source code.
   //
   // This method appends a new closing brace to the current data and outdent the
   // next lines back to the margin used before BeginBlock() was invoked.
-  SourceWriter& EndBlock() {
-    return Indent(-2).Append("}").EndLine();
-  }
+  SourceWriter& EndBlock();
 
   // Begins to write a method.
   //
   // This method outputs the signature of the Java method from the data passed
-  // in the 'method' parameter and starts a new block. Additionnal modifiers can
-  // also be passed in parameter to define the accesses and the scope of this
-  // method.
-  SourceWriter& BeginMethod(const Method& method, int modifiers = 0);
+  // in the 'method' parameter and starts a new block. Modifiers are also passed
+  // in parameter to define the access scope of this method and, optionally,
+  // a Javadoc.
+  SourceWriter& BeginMethod(const Method& method, int modifiers,
+      const Javadoc* javadoc = nullptr);
 
   // Ends the current method.
   //
@@ -122,22 +119,24 @@ class SourceWriter {
   // Begins to write the main type of a source file.
   //
   // This method outputs the declaration of the Java type from the data passed
-  // in the 'type' parameter and starts a new block. Additionnal modifiers can
-  // also be passed in parameter to define the accesses and the scope of this
-  // type.
+  // in the 'type' parameter and starts a new block. Modifiers are also passed
+  // in parameter to define the access scope of this type and, optionally,
+  // a Javadoc.
   //
-  // If not null, all types found in the 'dependencies' list will be imported
-  // before declaring the new type.
-  SourceWriter& BeginType(const Type& clazz,
-      const std::list<Type>* dependencies, int modifiers = 0);
+  // If not null, all types found in the 'extra_dependencies' list will be
+  // imported before declaring the new type.
+  SourceWriter& BeginType(const Type& clazz, int modifiers,
+      const std::list<Type>* extra_dependencies = nullptr,
+      const Javadoc* javadoc = nullptr);
 
   // Begins to write a new inner type.
   //
   // This method outputs the declaration of the Java type from the data passed
-  // in the 'type' parameter and starts a new block. Additionnal modifiers can
-  // also be passed in parameter to define the accesses and the scope of this
-  // type.
-  SourceWriter& BeginInnerType(const Type& type, int modifiers = 0);
+  // in the 'type' parameter and starts a new block. Modifiers are also passed
+  // in parameter to define the accesses and the scope of this type and,
+  // optionally, a Javadoc.
+  SourceWriter& BeginInnerType(const Type& type, int modifiers,
+      const Javadoc* javadoc = nullptr);
 
   // Ends the current type.
   //
@@ -145,13 +144,13 @@ class SourceWriter {
   // BeginType() or BeginInnerType() prior to this.
   SourceWriter& EndType();
 
-  // Writes a list of variables as fields of a type.
+  // Writes a variable as fields of a type.
   //
   // This method must be called within the definition of a type (see BeginType()
-  // or BeginInnerType()). Additional modifiers can also be passed in parameter
-  // to define the accesses and the scope of those fields.
-  SourceWriter& WriteFields(const std::list<Variable>& fields,
-      int modifiers = 0);
+  // or BeginInnerType()). Modifiers are also be passed in parameter to define
+  // the accesses and the scope of this field and, optionally, a Javadoc.
+  SourceWriter& WriteField(const Variable& field, int modifiers,
+      const Javadoc* javadoc = nullptr);
 
  protected:
   virtual void DoAppend(const StringPiece& str) = 0;
@@ -207,9 +206,7 @@ class SourceWriter {
   std::stack<GenericNamespace*> generic_namespaces_;
 
   SourceWriter& WriteModifiers(int modifiers);
-  SourceWriter& WriteDoc(const string& description,
-    const string& return_description = "",
-    const std::list<Variable>* parameters = nullptr);
+  SourceWriter& WriteJavadoc(const Javadoc& javadoc);
   SourceWriter& WriteAnnotations(const std::list<Annotation>& annotations);
   SourceWriter& WriteGenerics(const std::list<const Type*>& generics);
   GenericNamespace* PushGenericNamespace(int modifiers);
diff --git a/tensorflow/java/src/gen/cc/source_writer_test.cc b/tensorflow/java/src/gen/cc/source_writer_test.cc
index 4bce2fea70..8bd42d9d0e 100644
--- a/tensorflow/java/src/gen/cc/source_writer_test.cc
+++ b/tensorflow/java/src/gen/cc/source_writer_test.cc
@@ -250,7 +250,7 @@ TEST(StreamTest, Types) {
         .AppendType(generic).Append(", ")
         .AppendType(Type::ListOf(generic)).Append(", ")
         .AppendType(Type::ListOf(Type::IterableOf(generic))).Append(", ")
-        .AppendType(Type::ListOf(Type::Generic()));
+        .AppendType(Type::ListOf(Type::Wildcard()));
 
   const char* expected =
       "int, String, T, List<T>, List<Iterable<T>>, List<?>";
@@ -282,7 +282,7 @@ TEST(WriteType, SimpleClass) {
   SourceBufferWriter writer;
   Type clazz = Type::Class("Test", "org.tensorflow");
 
-  writer.BeginType(clazz, nullptr, PUBLIC).EndType();
+  writer.BeginType(clazz, PUBLIC).EndType();
 
   const char* expected =
       "package org.tensorflow;\n\n"
@@ -300,7 +300,7 @@ TEST(WriteType, SimpleClassWithDependencies) {
   deps.push_back(Type::Class("SamePackageType", "org.tensorflow"));
   deps.push_back(Type::Class("NoPackageType"));
 
-  writer.BeginType(clazz, &deps, PUBLIC).EndType();
+  writer.BeginType(clazz, PUBLIC, &deps).EndType();
 
   const char* expected =
       "package org.tensorflow;\n\n"
@@ -313,18 +313,21 @@ TEST(WriteType, SimpleClassWithDependencies) {
 TEST(WriteType, AnnotatedAndDocumentedClass) {
   SourceBufferWriter writer;
   Type clazz = Type::Class("Test", "org.tensorflow");
-  clazz.description("This class has a\n<p>\nmultiline description.");
+  Javadoc clazz_doc;
+  clazz_doc.brief("Javadoc test")
+      .details("This is a\nmultiline description.");
   clazz.add_annotation(Annotation::Create("Bean"));
   clazz.add_annotation(Annotation::Create("SuppressWarnings")
       .attributes("\"rawtypes\""));
 
-  writer.BeginType(clazz, nullptr, PUBLIC).EndType();
+  writer.BeginType(clazz, PUBLIC, nullptr, &clazz_doc).EndType();
 
   const char* expected =
       "package org.tensorflow;\n\n"
       "/**\n"
-      " * This class has a\n"
+      " * Javadoc test\n"
       " * <p>\n"
+      " * This is a\n"
       " * multiline description.\n"
       " **/\n"
       "@Bean\n"
@@ -339,7 +342,7 @@ TEST(WriteType, ParameterizedClass) {
   clazz.add_parameter(Type::Generic("T"));
   clazz.add_parameter(Type::Generic("U").add_supertype(Type::Class("Number")));
 
-  writer.BeginType(clazz, nullptr, PUBLIC).EndType();
+  writer.BeginType(clazz, PUBLIC).EndType();
 
   const char* expected =
       "package org.tensorflow;\n\n"
@@ -358,7 +361,7 @@ TEST(WriteType, ParameterizedClassAndSupertypes) {
   clazz.add_supertype(Type::Interface("Runnable"));
   clazz.add_supertype(Type::Class("SuperTest").add_parameter(type_t));
 
-  writer.BeginType(clazz, nullptr, PUBLIC).EndType();
+  writer.BeginType(clazz, PUBLIC).EndType();
 
   const char* expected =
       "package org.tensorflow;\n\n"
@@ -372,24 +375,24 @@ TEST(WriteType, ParameterizedClassFields) {
   Type clazz = Type::Class("Test", "org.tensorflow");
   Type type_t = Type::Generic("T").add_supertype(Type::Class("Number"));
   clazz.add_parameter(type_t);
-  std::list<Variable> static_fields;
-  static_fields.push_back(Variable::Create("field1", Type::Class("String")));
-  std::list<Variable> member_fields;
-  member_fields.push_back(Variable::Create("field2", Type::Class("String")));
-  member_fields.push_back(Variable::Create("field3", type_t));
-
-  writer.BeginType(clazz, nullptr, PUBLIC)
-          .WriteFields(static_fields, STATIC | PUBLIC | FINAL)
-          .WriteFields(member_fields, PRIVATE)
+  Variable field1 = Variable::Create("field1", Type::Class("String"));
+  Variable field2 = Variable::Create("field2", Type::Class("String"));
+  Variable field3 = Variable::Create("field3", type_t);
+  Javadoc field3_doc;
+  field3_doc.brief("This variable is documented");
+
+  writer.BeginType(clazz, PUBLIC)
+          .WriteField(field1, STATIC | PUBLIC | FINAL)
+          .WriteField(field2, PRIVATE)
+          .WriteField(field3, PRIVATE, &field3_doc)
         .EndType();
 
   const char* expected =
       "package org.tensorflow;\n\n"
       "public class Test<T extends Number> {\n"
-      "  \n"
       "  public static final String field1;\n"
-      "  \n"
       "  private String field2;\n"
+      "  /** This variable is documented */\n"
       "  private T field3;\n"
       "}\n";
   ASSERT_STREQ(expected, writer.str().data());
@@ -400,7 +403,7 @@ TEST(WriteType, SimpleInnerClass) {
   Type clazz = Type::Class("Test", "org.tensorflow");
   Type inner_class = Type::Class("InnerTest");
 
-  writer.BeginType(clazz, nullptr, PUBLIC)
+  writer.BeginType(clazz, PUBLIC)
           .BeginInnerType(inner_class, PUBLIC)
           .EndType()
         .EndType();
@@ -423,7 +426,7 @@ TEST(WriteType, StaticParameterizedInnerClass) {
   Type inner_class = Type::Class("InnerTest");
   inner_class.add_parameter(type_t);
 
-  writer.BeginType(clazz, nullptr, PUBLIC)
+  writer.BeginType(clazz, PUBLIC)
           .BeginInnerType(inner_class, PUBLIC | STATIC)
           .EndType()
         .EndType();
@@ -443,7 +446,7 @@ TEST(WriteMethod, SimpleMethod) {
   Type clazz = Type::Class("Test", "org.tensorflow");
   Method method = Method::Create("doNothing", Type::Void());
 
-  writer.BeginType(clazz, nullptr, PUBLIC)
+  writer.BeginType(clazz, PUBLIC)
           .BeginMethod(method, PUBLIC).EndMethod()
         .EndType();
 
@@ -461,13 +464,15 @@ TEST(WriteMethod, AnnotatedAndDocumentedMethod) {
   SourceBufferWriter writer;
   Type clazz = Type::Class("Test", "org.tensorflow");
   Method method = Method::Create("doNothing", Type::Void());
-  method.description("This method has a\n<p>\nmultiline description.");
+  Javadoc method_doc;
+  method_doc.brief("Javadoc test")
+      .details("This method has a\nmultiline description.");
   method.add_annotation(Annotation::Create("Override"));
   method.add_annotation(Annotation::Create("SuppressWarnings")
       .attributes("\"rawtypes\""));
 
-  writer.BeginType(clazz, nullptr, PUBLIC)
-          .BeginMethod(method, PUBLIC).EndMethod()
+  writer.BeginType(clazz, PUBLIC)
+          .BeginMethod(method, PUBLIC, &method_doc).EndMethod()
         .EndType();
 
   const char* expected =
@@ -475,8 +480,9 @@ TEST(WriteMethod, AnnotatedAndDocumentedMethod) {
       "public class Test {\n"
       "  \n"
       "  /**\n"
-      "   * This method has a\n"
+      "   * Javadoc test\n"
       "   * <p>\n"
+      "   * This method has a\n"
       "   * multiline description.\n"
       "   **/\n"
       "  @Override\n"
@@ -490,16 +496,18 @@ TEST(WriteMethod, AnnotatedAndDocumentedMethod) {
 TEST(WriteMethod, DocumentedMethodWithArguments) {
   SourceBufferWriter writer;
   Type clazz = Type::Class("Test", "org.tensorflow");
+  Variable reverse = Variable::Create("reverse", Type::Boolean());
   Method method = Method::Create("boolToInt", Type::Int());
-  method.description("Converts a boolean to an int");
-  method.return_description("int value for this boolean");
   method.add_argument(Variable::Create("b", Type::Boolean()));
-  Variable reverse = Variable::Create("reverse", Type::Boolean());
-  reverse.description("if true, value is reversed");
   method.add_argument(reverse);
-
-  writer.BeginType(clazz, nullptr, PUBLIC)
-          .BeginMethod(method, PUBLIC)
+  Javadoc method_doc;
+  method_doc.brief("Converts a boolean to an int")
+      .details("This method will convert\na boolean to an int")
+      .add_param_tag(reverse.name(), "if true, value is reversed")
+      .add_tag("return", "int value for this boolean");
+
+  writer.BeginType(clazz, PUBLIC)
+          .BeginMethod(method, PUBLIC, &method_doc)
             .Append("if (b && !reverse)")
             .BeginBlock()
               .Append("return 1;").EndLine()
@@ -514,8 +522,10 @@ TEST(WriteMethod, DocumentedMethodWithArguments) {
       "  \n"
       "  /**\n"
       "   * Converts a boolean to an int\n"
+      "   * <p>\n"
+      "   * This method will convert\n"
+      "   * a boolean to an int\n"
       "   * \n"
-      "   * @param b\n"
       "   * @param reverse if true, value is reversed\n"
       "   * @return int value for this boolean\n"
       "   **/\n"
@@ -536,7 +546,7 @@ TEST(WriteMethod, ParameterizedMethod) {
   clazz.add_parameter(type_t);
   Method method = Method::Create("doNothing", type_t);
 
-  writer.BeginType(clazz, nullptr, PUBLIC)
+  writer.BeginType(clazz, PUBLIC)
           .BeginMethod(method, PUBLIC)
             .Append("return null;").EndLine()
           .EndMethod()
@@ -560,7 +570,7 @@ TEST(WriteMethod, StaticParameterizedMethod) {
   clazz.add_parameter(type_t);
   Method method = Method::Create("doNothing", type_t);
 
-  writer.BeginType(clazz, nullptr, PUBLIC)
+  writer.BeginType(clazz, PUBLIC)
           .BeginMethod(method, PUBLIC | STATIC)
             .Append("return null;").EndLine()
           .EndMethod()
diff --git a/tensorflow/java/src/gen/gen_ops.bzl b/tensorflow/java/src/gen/gen_ops.bzl
index a6650fc4ea..1e7899cf7a 100644
--- a/tensorflow/java/src/gen/gen_ops.bzl
+++ b/tensorflow/java/src/gen/gen_ops.bzl
@@ -1,9 +1,11 @@
 # -*- Python -*-
 
-load("//tensorflow:tensorflow.bzl",
-     "tf_binary_additional_srcs",
-     "tf_cc_binary",
-     "tf_copts")
+load(
+    "//tensorflow:tensorflow.bzl",
+    "tf_binary_additional_srcs",
+    "tf_cc_binary",
+    "tf_copts",
+)
 
 # Given a list of "ops_libs" (a list of files in the core/ops directory
 # without their .cc extensions), generate Java wrapper code for all operations
@@ -27,16 +29,31 @@ def tf_java_op_gen_srcjar(name,
                           ops_libs_pkg="//tensorflow/core",
                           out_dir="ops/",
                           out_src_dir="src/main/java/",
+                          api_def_srcs=[],
                           visibility=["//tensorflow/java:__pkg__"]):
 
   gen_tools = []
   gen_cmds = ["rm -rf $(@D)"]  # Always start from fresh when generating source files
+  srcs = api_def_srcs[:]
 
   # Construct an op generator binary for each ops library.
   for ops_lib in ops_libs:
     gen_lib = ops_lib[:ops_lib.rfind("_")]
     out_gen_tool = out_dir + ops_lib + "_gen_tool"
 
+    if not api_def_srcs:
+      api_def_args_str = ","
+    else:
+      api_def_args = []
+      for api_def_src in api_def_srcs:
+        # Add directory of the first ApiDef source to args.
+        # We are assuming all ApiDefs in a single api_def_src are in the
+        # same directory.
+        api_def_args.append(
+            " $$(dirname $$(echo $(locations " + api_def_src +
+            ") | cut -d\" \" -f1))")
+      api_def_args_str = ",".join(api_def_args)
+
     tf_cc_binary(
         name=out_gen_tool,
         copts=tf_copts(),
@@ -48,7 +65,8 @@ def tf_java_op_gen_srcjar(name,
     gen_cmds += ["$(location :" + out_gen_tool + ")" +
                  " --output_dir=$(@D)/" + out_src_dir +
                  " --lib_name=" + gen_lib +
-                 " --base_package=" + gen_base_package]
+                 " --base_package=" + gen_base_package +
+                 " " + api_def_args_str]
 
   # Generate a source archive containing generated code for these ops.
   gen_srcjar = out_dir + name + ".srcjar"
@@ -57,6 +75,7 @@ def tf_java_op_gen_srcjar(name,
   gen_tools += tf_binary_additional_srcs()
   native.genrule(
       name=name,
+      srcs=srcs,
       outs=[gen_srcjar],
       tools=gen_tools,
       cmd="&&".join(gen_cmds))
diff --git a/tensorflow/java/src/gen/resources/license.snippet.java b/tensorflow/java/src/gen/resources/license.snippet.java
new file mode 100644
index 0000000000..90285ec669
--- /dev/null
+++ b/tensorflow/java/src/gen/resources/license.snippet.java
@@ -0,0 +1,14 @@
+/* Copyright 2018 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.
+==============================================================================*/
-- 
GitLab


From 7e80197f020895fea41eda36b08135b747a9a4f1 Mon Sep 17 00:00:00 2001
From: "karl@kubx.ca" <karl@kubx.ca>
Date: Fri, 6 Apr 2018 08:56:54 -0400
Subject: [PATCH 220/395] Improve Javadoc and include first code review

---
 tensorflow/java/BUILD                         |  23 +-
 tensorflow/java/src/gen/cc/java_defs.h        |  12 +-
 tensorflow/java/src/gen/cc/op_gen_main.cc     |  48 +-
 tensorflow/java/src/gen/cc/op_generator.cc    | 224 ++++++----
 tensorflow/java/src/gen/cc/op_generator.h     |  25 +-
 tensorflow/java/src/gen/cc/op_parser.cc       | 417 ------------------
 tensorflow/java/src/gen/cc/op_parser.h        | 137 ------
 tensorflow/java/src/gen/cc/op_specs.cc        | 390 ++++++++++++++++
 tensorflow/java/src/gen/cc/op_specs.h         | 152 +++++++
 tensorflow/java/src/gen/cc/source_writer.cc   |   2 +-
 tensorflow/java/src/gen/cc/source_writer.h    |   2 +-
 .../java/src/gen/cc/source_writer_test.cc     |  20 +-
 tensorflow/java/src/gen/gen_ops.bzl           |  68 +--
 ...ense.snippet.java => license.java.snippet} |   0
 14 files changed, 760 insertions(+), 760 deletions(-)
 delete mode 100644 tensorflow/java/src/gen/cc/op_parser.cc
 delete mode 100644 tensorflow/java/src/gen/cc/op_parser.h
 create mode 100644 tensorflow/java/src/gen/cc/op_specs.cc
 create mode 100644 tensorflow/java/src/gen/cc/op_specs.h
 rename tensorflow/java/src/gen/resources/{license.snippet.java => license.java.snippet} (100%)

diff --git a/tensorflow/java/BUILD b/tensorflow/java/BUILD
index 635a4e807d..17566e1a9c 100644
--- a/tensorflow/java/BUILD
+++ b/tensorflow/java/BUILD
@@ -68,9 +68,13 @@ filegroup(
     ],
 )
 
+# Build the gen tool as a library, as it will be linked to a core/ops binary
+# files before making it an executable.
 tf_java_op_gen_srcjar(
     name = "java_op_gen_sources",
-    api_def_srcs = ["//tensorflow/core/api_def:base_api_def"],
+    api_def_srcs = [
+        "//tensorflow/core/api_def:base_api_def",
+    ],
     gen_base_package = "org.tensorflow.op",
     gen_tool = "java_op_gen_tool",
     ops_libs = [
@@ -95,30 +99,17 @@ tf_java_op_gen_srcjar(
     ],
 )
 
-# Build the gen tool as a library, as it will be linked to a core/ops binary
-# file before making it an executable. See tf_java_op_gen_srcjar().
-cc_library(
-    name = "java_op_gen_tool",
-    srcs = [
-        "src/gen/cc/op_gen_main.cc",
-    ],
-    copts = tf_copts(),
-    deps = [
-        ":java_op_gen_lib",
-    ],
-)
-
 cc_library(
     name = "java_op_gen_lib",
     srcs = [
         "src/gen/cc/op_generator.cc",
-        "src/gen/cc/op_parser.cc",
+        "src/gen/cc/op_specs.cc",
         "src/gen/cc/source_writer.cc",
     ],
     hdrs = [
         "src/gen/cc/java_defs.h",
         "src/gen/cc/op_generator.h",
-        "src/gen/cc/op_parser.h",
+        "src/gen/cc/op_specs.h",
         "src/gen/cc/source_writer.h",
     ],
     copts = tf_copts(),
diff --git a/tensorflow/java/src/gen/cc/java_defs.h b/tensorflow/java/src/gen/cc/java_defs.h
index 2065477f58..81ac67eb2f 100644
--- a/tensorflow/java/src/gen/cc/java_defs.h
+++ b/tensorflow/java/src/gen/cc/java_defs.h
@@ -1,4 +1,4 @@
-/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2018 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.
@@ -230,12 +230,12 @@ class Javadoc {
     return Javadoc(brief);
   }
   const string& brief() const { return brief_; }
-  const string& details() const { return description_; }
-  Javadoc& details(const string description) {
-    description_ = description;
+  const string& details() const { return details_; }
+  Javadoc& details(const string& details) {
+    details_ = details;
     return *this;
   }
-  const std::list<std::pair<string, string>> tags() const { return tags_; }
+  const std::list<std::pair<string, string>>& tags() const { return tags_; }
   Javadoc& add_tag(const string& tag, const string& text) {
     tags_.push_back(std::make_pair(tag, text));
     return *this;
@@ -246,7 +246,7 @@ class Javadoc {
 
  private:
   string brief_;
-  string description_;
+  string details_;
   std::list<std::pair<string, string>> tags_;
 
   explicit Javadoc(const string& brief) : brief_(brief) {}
diff --git a/tensorflow/java/src/gen/cc/op_gen_main.cc b/tensorflow/java/src/gen/cc/op_gen_main.cc
index 015200023f..458141b877 100644
--- a/tensorflow/java/src/gen/cc/op_gen_main.cc
+++ b/tensorflow/java/src/gen/cc/op_gen_main.cc
@@ -1,4 +1,4 @@
-/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2018 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.
@@ -36,55 +36,41 @@ const char kUsageHeader[] =
     "Operation wrappers are generated under the path specified by the "
     "'--output_dir' argument. This path can be absolute or relative to the\n"
     "current working directory and will be created if it does not exists.\n\n"
-    "The '--lib_name' argument is used to classify the set of operations. If "
-    "the chosen name contains more than one word, it must be provided in \n"
-    "snake_case. This value is declined into other meaningful names, such as "
-    "the group and package of the generated operations. For example,\n"
-    "'--lib_name=my_lib' generates the operations under the "
-    "'org.tensorflow.op.mylib' package and add them to the 'myLib()' operator\n"
-    "group.\n\n"
-    "Note that the operator group assigned to the generated wrappers is just "
-    "an annotation tag at this stage. Operations will not be available "
-    "through\n"
-    "the 'org.tensorflow.op.Ops' API as a group until the generated classes "
-    "are compiled using an appropriate annotation processor.\n\n"
+    "Note that the operations will not be available through the "
+    "'org.tensorflow.op.Ops' API until the generated classes are compiled\n"
+    "using an appropriate annotation processor.\n\n"
     "The '--base_package' overrides the default parent package under which "
     "the generated subpackage and classes are to be located.\n\n"
-    "Finally, a list of directories of API proto definitions can be provided "
-    "to override default values found in the ops definitions, ordered by\n"
-    "priority (the last having precedence over the first).\n\n";
+    "Finally, the `--api_dirs` argument takes a list of comma-seperated "
+    "directories of API definitions can be provided to override default\n"
+    "values found in the ops definitions. Directories are ordered by priority "
+    "(the last having precedence over the first).\n\n";
 
 }  // namespace java
 }  // namespace tensorflow
 
 int main(int argc, char* argv[]) {
-  tensorflow::string lib_name;
   tensorflow::string output_dir;
   tensorflow::string base_package = "org.tensorflow.op";
+  tensorflow::string api_dirs_str;
   std::vector<tensorflow::Flag> flag_list = {
       tensorflow::Flag("output_dir", &output_dir,
           "Root directory into which output files are generated"),
-      tensorflow::Flag(
-          "lib_name", &lib_name,
-          "A name, in snake_case, used to classify this set of operations"),
-      tensorflow::Flag(
-          "base_package", &base_package,
-          "Package parent to the generated subpackage and classes")};
+      tensorflow::Flag("base_package", &base_package,
+          "Package parent to the generated subpackage and classes"),
+      tensorflow::Flag("api_dirs", &api_dirs_str,
+          "List of directories that contains the ops api definitions")};
   tensorflow::string usage = tensorflow::java::kUsageHeader;
   usage += tensorflow::Flags::Usage(argv[0], flag_list);
   bool parsed_flags_ok = tensorflow::Flags::Parse(&argc, argv, flag_list);
   tensorflow::port::InitMain(usage.c_str(), &argc, &argv);
-  QCHECK(parsed_flags_ok && !lib_name.empty() && !output_dir.empty()) << usage;
-  std::vector<tensorflow::string> api_dirs;
-  if (argc > 1) {
-    api_dirs = tensorflow::str_util::Split(argv[1], ",",
-        tensorflow::str_util::SkipEmpty());
-  }
+  QCHECK(parsed_flags_ok && !output_dir.empty()) << usage;
+  std::vector<tensorflow::string> api_dirs = tensorflow::str_util::Split(
+      api_dirs_str, ",", tensorflow::str_util::SkipEmpty());
   tensorflow::java::OpGenerator generator(base_package, output_dir, api_dirs);
   tensorflow::OpList ops;
   tensorflow::OpRegistry::Global()->Export(false, &ops);
-  tensorflow::Status status = generator.Run(ops, lib_name);
-  TF_QCHECK_OK(status);
+  TF_CHECK_OK(generator.Run(ops));
 
   return 0;
 }
diff --git a/tensorflow/java/src/gen/cc/op_generator.cc b/tensorflow/java/src/gen/cc/op_generator.cc
index c9b57f5706..c32ad3b109 100644
--- a/tensorflow/java/src/gen/cc/op_generator.cc
+++ b/tensorflow/java/src/gen/cc/op_generator.cc
@@ -1,4 +1,4 @@
-/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2018 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.
@@ -18,6 +18,7 @@ limitations under the License.
 #include <vector>
 #include <list>
 #include <memory>
+#include <set>
 
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/str_util.h"
@@ -27,15 +28,15 @@ limitations under the License.
 #include "tensorflow/core/framework/op_gen_lib.h"
 #include "tensorflow/java/src/gen/cc/java_defs.h"
 #include "tensorflow/java/src/gen/cc/source_writer.h"
-#include "tensorflow/java/src/gen/cc/op_parser.h"
 #include "tensorflow/java/src/gen/cc/op_generator.h"
+#include "tensorflow/java/src/gen/cc/op_specs.h"
 
 namespace tensorflow {
 namespace java {
 namespace {
 
 const char* kLicenseSnippet =
-    "tensorflow/java/src/gen/resources/license.snippet.java";
+    "tensorflow/java/src/gen/resources/license.java.snippet";
 
 const std::map<string, Type> kPrimitiveAttrTypes = {
   { "Boolean", Type::Boolean() },
@@ -66,34 +67,34 @@ void CollectOpDependencies(const OpSpec& op, RenderMode mode,
   }
   // Don't pay attention to duplicate types in the dependency list, they will
   // be filtered out by the SourceWriter.
-  for (const OpSpec::Operand& input : op.inputs()) {
+  for (const ArgumentSpec& input : op.inputs()) {
     out->push_back(input.var().type());
     if (input.iterable()) {
       out->push_back(Type::Class("Operands", "org.tensorflow.op"));
     }
   }
-  for (const OpSpec::Operand& output : op.outputs()) {
+  for (const ArgumentSpec& output : op.outputs()) {
     out->push_back(output.var().type());
     if (output.iterable()) {
       out->push_back(Type::Class("Arrays", "java.util"));
     }
   }
-  for (const OpSpec::Operand& attribute : op.attributes()) {
+  for (const AttributeSpec& attribute : op.attributes()) {
     out->push_back(attribute.var().type());
     if (attribute.var().type().name() == "Class") {
       out->push_back(Type::Enum("DataType", "org.tensorflow"));
     }
   }
-  for (const OpSpec::Operand& option : op.options()) {
-    out->push_back(option.var().type());
+  for (const AttributeSpec& optional_attribute : op.optional_attributes()) {
+    out->push_back(optional_attribute.var().type());
   }
 }
 
-void WriteSetAttrDirective(const OpSpec::Operand& attr, bool optional,
+void WriteSetAttrDirective(const AttributeSpec& attr, bool optional,
     SourceWriter* writer) {
   string var = optional ? "opts." + attr.var().name() : attr.var().name();
   if (attr.iterable()) {
-    const Type& type = attr.data_type();
+    const Type& type = attr.type();
     std::map<string, Type>::const_iterator it =
       kPrimitiveAttrTypes.find(type.name());
     if (it != kPrimitiveAttrTypes.end()) {
@@ -107,11 +108,11 @@ void WriteSetAttrDirective(const OpSpec::Operand& attr, bool optional,
           .Append(array + "[i] = " + var + ".get(i);")
           .EndLine()
           .EndBlock()
-          .Append("opBuilder.setAttr(\"" + attr.graph_name() + "\", " + array)
+          .Append("opBuilder.setAttr(\"" + attr.op_def_name() + "\", " + array)
           .Append(");")
           .EndLine();
     } else {
-      writer->Append("opBuilder.setAttr(\"" + attr.graph_name() + "\", " + var)
+      writer->Append("opBuilder.setAttr(\"" + attr.op_def_name() + "\", " + var)
           .Append(".toArray(new ")
           .AppendType(type)
           .Append("[" + var + ".size()]));")
@@ -119,7 +120,7 @@ void WriteSetAttrDirective(const OpSpec::Operand& attr, bool optional,
     }
   } else {
     Type type = attr.var().type();
-    writer->Append("opBuilder.setAttr(\"" + attr.graph_name() + "\", ");
+    writer->Append("opBuilder.setAttr(\"" + attr.op_def_name() + "\", ");
     if (type.name() == "Class") {
       writer->Append("DataType.fromClass(" + attr.var().name() + "));");
     } else {
@@ -139,26 +140,26 @@ void RenderFactoryMethod(const OpSpec& op, const Type& op_class,
       Variable::Create("scope", Type::Class("Scope", "org.tensorflow.op"));
   factory.add_argument(scope);
   factory_doc.add_param_tag(scope.name(), "Current graph scope");
-  for (const OpSpec::Operand& input : op.inputs()) {
+  for (const ArgumentSpec& input : op.inputs()) {
     factory.add_argument(input.var());
     factory_doc.add_param_tag(input.var().name(), input.description());
   }
-  for (const OpSpec::Operand& attribute : op.attributes()) {
+  for (const AttributeSpec& attribute : op.attributes()) {
     factory.add_argument(attribute.var());
     factory_doc.add_param_tag(attribute.var().name(), attribute.description());
   }
-  if (!op.options().empty()) {
+  if (!op.optional_attributes().empty()) {
     factory.add_argument(Variable::Varargs("options", Type::Class("Options")));
     factory_doc.add_param_tag("options", "carries optional attributes values");
   }
   factory_doc.add_tag("return", "a new instance of " + op_class.name());
   writer->BeginMethod(factory, PUBLIC|STATIC, &factory_doc);
   writer->Append("OperationBuilder opBuilder = scope.graph().opBuilder(\""
-      + op.graph_name() + "\", scope.makeOpName(\""
+      + op.graph_op_name() + "\", scope.makeOpName(\""
       + op_class.name() + "\"));");
   writer->EndLine();
 
-  for (const OpSpec::Operand& input : op.inputs()) {
+  for (const ArgumentSpec& input : op.inputs()) {
     if (input.iterable()) {
       writer->Append("opBuilder.addInputList(Operands.asOutputs("
           + input.var().name() + "));");
@@ -169,15 +170,15 @@ void RenderFactoryMethod(const OpSpec& op, const Type& op_class,
       writer->EndLine();
     }
   }
-  for (const OpSpec::Operand& attribute : op.attributes()) {
+  for (const AttributeSpec& attribute : op.attributes()) {
     WriteSetAttrDirective(attribute, false, writer);
   }
-  if (!op.options().empty()) {
+  if (!op.optional_attributes().empty()) {
     writer->BeginBlock("if (options != null)")
         .BeginBlock("for (Options opts : options)");
-    for (const OpSpec::Operand& option : op.options()) {
-      writer->BeginBlock("if (opts." + option.var().name() + " != null)");
-      WriteSetAttrDirective(option, true, writer);
+    for (const AttributeSpec& attribute : op.optional_attributes()) {
+      writer->BeginBlock("if (opts." + attribute.var().name() + " != null)");
+      WriteSetAttrDirective(attribute, true, writer);
       writer->EndBlock();
     }
     writer->EndBlock().EndBlock();
@@ -195,8 +196,8 @@ void RenderConstructor(const OpSpec& op, const Type& op_class,
     .add_argument(
         Variable::Create("operation",
             Type::Class("Operation", "org.tensorflow")));
-  for (const OpSpec::Operand& output : op.outputs()) {
-    if (output.iterable() && !output.data_type().unknown()) {
+  for (const ArgumentSpec& output : op.outputs()) {
+    if (output.iterable() && !output.type().unknown()) {
       constructor.add_annotation(
           Annotation::Create("SuppressWarnings").attributes("\"unchecked\""));
       break;
@@ -208,15 +209,15 @@ void RenderConstructor(const OpSpec& op, const Type& op_class,
   if (op.outputs().size() > 0) {
     writer->Append("int outputIdx = 0;")
         .EndLine();
-    for (const OpSpec::Operand& output : op.outputs()) {
+    for (const ArgumentSpec& output : op.outputs()) {
       if (output.iterable()) {
         string var_length = output.var().name() + "Length";
         writer->Append("int " + var_length)
-            .Append(" = operation.outputListLength(\"" + output.graph_name()
+            .Append(" = operation.outputListLength(\"" + output.op_def_name()
                 + "\");")
             .EndLine()
             .Append(output.var().name() + " = Arrays.asList(");
-        if (!output.data_type().unknown()) {
+        if (!output.type().unknown()) {
           writer->Append("(")
               .AppendType(output.var().type().parameters().front())
               .Append("[])");
@@ -236,18 +237,19 @@ void RenderConstructor(const OpSpec& op, const Type& op_class,
 }
 
 void RenderGettersAndSetters(const OpSpec& op, SourceWriter* writer) {
-  for (const OpSpec::Operand& option : op.options()) {
-    Method setter = Method::Create(option.var().name(), Type::Class("Options"))
-        .add_argument(option.var());
+  for (const AttributeSpec& attribute : op.optional_attributes()) {
+    Method setter =
+        Method::Create(attribute.var().name(), Type::Class("Options"))
+            .add_argument(attribute.var());
     Javadoc setter_doc = Javadoc::Create()
-        .add_param_tag(option.var().name(), option.description());
+        .add_param_tag(attribute.var().name(), attribute.description());
     writer->BeginMethod(setter, PUBLIC|STATIC, &setter_doc)
-        .Append("return new Options()." + option.var().name() + "("
-            + option.var().name() + ");")
+        .Append("return new Options()." + attribute.var().name() + "("
+            + attribute.var().name() + ");")
         .EndLine()
         .EndMethod();
   }
-  for (const OpSpec::Operand& output : op.outputs()) {
+  for (const ArgumentSpec& output : op.outputs()) {
     Method getter = Method::Create(output.var().name(), output.var().type());
     Javadoc getter_doc = Javadoc::Create(output.description());
     writer->BeginMethod(getter, PUBLIC, &getter_doc)
@@ -259,12 +261,12 @@ void RenderGettersAndSetters(const OpSpec& op, SourceWriter* writer) {
 
 void RenderInterfaceImpl(const OpSpec& op, RenderMode mode,
     SourceWriter* writer) {
-  OpSpec::Operand output = op.outputs().front();
+  ArgumentSpec output = op.outputs().front();
 
   if (mode == SINGLE_OUTPUT) {
-    bool cast2obj = output.data_type().unknown();
+    bool cast2obj = output.type().unknown();
     Type return_type = Type::Class("Output", "org.tensorflow")
-        .add_parameter(cast2obj ? Type::Class("Object") : output.data_type());
+        .add_parameter(cast2obj ? Type::Class("Object") : output.type());
     Method as_output = Method::Create("asOutput", return_type)
         .add_annotation(Annotation::Create("Override"));
     if (cast2obj) {
@@ -283,10 +285,10 @@ void RenderInterfaceImpl(const OpSpec& op, RenderMode mode,
 
   } else if (mode == SINGLE_LIST_OUTPUT) {
     Type operand = Type::Interface("Operand", "org.tensorflow");
-    if (output.data_type().unknown()) {
+    if (output.type().unknown()) {
       operand.add_parameter(Type::Class("Object"));
     } else {
-      operand.add_parameter(output.data_type());
+      operand.add_parameter(output.type());
     }
     Type return_type = Type::Interface("Iterator", "java.util")
         .add_parameter(operand);
@@ -308,57 +310,119 @@ void RenderOptionsClass(const OpSpec& op, SourceWriter* writer) {
   Javadoc options_doc = Javadoc::Create(
       "Class holding optional attributes of this operation");
   writer->BeginInnerType(options_class, PUBLIC | STATIC, &options_doc);
-  for (const OpSpec::Operand& option : op.options()) {
-    Method setter = Method::Create(option.var().name(), options_class)
-        .add_argument(option.var());
+  for (const AttributeSpec& attribute : op.optional_attributes()) {
+    Method setter = Method::Create(attribute.var().name(), options_class)
+        .add_argument(attribute.var());
     Javadoc setter_doc = Javadoc::Create()
-        .add_param_tag(option.var().name(), option.description());
+        .add_param_tag(attribute.var().name(), attribute.description());
     writer->BeginMethod(setter, PUBLIC, &setter_doc)
-        .Append("this." + option.var().name() + " = " + option.var().name()
-            + ";")
+        .Append("this." + attribute.var().name() + " = "
+            + attribute.var().name() + ";")
         .EndLine()
         .Append("return this;")
         .EndLine()
         .EndMethod();
   }
   writer->EndLine();
-  for (const OpSpec::Operand& option : op.options()) {
-    writer->WriteField(option.var(), PRIVATE);
+  for (const AttributeSpec& optional_attribute : op.optional_attributes()) {
+    writer->WriteField(optional_attribute.var(), PRIVATE);
   }
   Method constructor = Method::ConstructorFor(options_class);
   writer->BeginMethod(constructor, PRIVATE).EndMethod();
   writer->EndType();
 }
 
-void RenderEndpoint(const OpSpec& op, const OpSpec::Endpoint& endpoint,
-    SourceWriter* writer) {
+inline Type ClassOf(const EndpointSpec& endpoint, const string& base_package) {
+  return Type::Class(endpoint.name(),
+      base_package + "." + str_util::Lowercase(endpoint.package()));
+}
+
+void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
+    const string& base_package, const string& output_dir, Env* env) {
+  Type op_class(ClassOf(endpoint, base_package)
+      .add_supertype(Type::Class("PrimitiveOp", "org.tensorflow.op")));
+  Javadoc op_javadoc(endpoint.javadoc());
+
+  // implement Operand (or Iterable<Operand>) if the op has only one output
   RenderMode mode = DEFAULT;
   if (op.outputs().size() == 1) {
-    mode = op.outputs().front().iterable() ? SINGLE_LIST_OUTPUT : SINGLE_OUTPUT;
+    const ArgumentSpec& output = op.outputs().front();
+    Type operand_type(output.type().unknown() ?
+        Type::Class("Object") : output.type());
+    Type operand_inf(Type::Interface("Operand", "org.tensorflow")
+        .add_parameter(operand_type));
+    if (output.iterable()) {
+      mode = SINGLE_LIST_OUTPUT;
+      op_class.add_supertype(Type::IterableOf(operand_inf));
+    } else {
+      mode = SINGLE_OUTPUT;
+      op_class.add_supertype(operand_inf);
+    }
+  }
+  // declare all outputs generics at the op class level
+  std::set<string> generics;
+  for (const ArgumentSpec& output : op.outputs()) {
+    if (output.type().kind() == Type::GENERIC && !output.type().unknown()
+        && generics.find(output.type().name()) == generics.end()) {
+      op_class.add_parameter(output.type());
+      op_javadoc.add_param_tag("<" + output.type().name() + ">",
+          "data type of output {@code " + output.var().name() + "}");
+      generics.insert(output.type().name());
+    }
+  }
+  // handle endpoint deprecation
+  if (endpoint.deprecated()) {
+    op_class.add_annotation(Annotation::Create("Deprecated"));
+    string explanation;
+    if (!op.endpoints().front().deprecated()) {
+      explanation = "use {@link " +
+          ClassOf(op.endpoints().front(), base_package).full_name()
+          + "} instead";
+    } else {
+      explanation = op.deprecation_explanation();
+    }
+    op_javadoc.add_tag("deprecated", explanation);
   }
+  // expose the op in the Ops Graph API only if it is visible
+  if (!op.hidden()) {
+    op_class.add_annotation(
+        Annotation::Create("Operator", "org.tensorflow.op.annotation")
+          .attributes("group = \"" + endpoint.package() + "\""));
+  }
+  // create op class file
+  string op_dir = io::JoinPath(output_dir,
+      str_util::StringReplace(op_class.package(), ".", "/", true));
+  if (!env->FileExists(op_dir).ok()) {
+    TF_CHECK_OK(Env::Default()->RecursivelyCreateDir(op_dir));
+  }
+  std::unique_ptr<tensorflow::WritableFile> op_file;
+  TF_CHECK_OK(env->NewWritableFile(
+      io::JoinPath(op_dir, op_class.name() + ".java"), &op_file));
+
+  // render endpoint source code
+  SourceFileWriter writer(op_file.get());
   std::list<Type> dependencies;
   CollectOpDependencies(op, mode, &dependencies);
-  const Type& op_class = endpoint.type();
-  writer->WriteFromFile(kLicenseSnippet)
+  writer.WriteFromFile(kLicenseSnippet)
       .EndLine()
       .Append("// This file is machine generated, DO NOT EDIT!")
       .EndLine()
       .EndLine()
-      .BeginType(op_class, PUBLIC|FINAL, &dependencies, &endpoint.javadoc());
-  if (!op.options().empty()) {
-    RenderOptionsClass(op, writer);
+      .BeginType(op_class, PUBLIC|FINAL, &dependencies, &op_javadoc);
+  if (!op.optional_attributes().empty()) {
+    RenderOptionsClass(op, &writer);
   }
-  RenderFactoryMethod(op, op_class, writer);
-  RenderGettersAndSetters(op, writer);
+  RenderFactoryMethod(op, op_class, &writer);
+  RenderGettersAndSetters(op, &writer);
   if (mode != DEFAULT) {
-    RenderInterfaceImpl(op, mode, writer);
+    RenderInterfaceImpl(op, mode, &writer);
   }
-  writer->EndLine();
-  for (const OpSpec::Operand& output : op.outputs()) {
-    writer->WriteField(output.var(), PRIVATE);
+  writer.EndLine();
+  for (const ArgumentSpec& output : op.outputs()) {
+    writer.WriteField(output.var(), PRIVATE);
   }
-  RenderConstructor(op, op_class, writer);
-  writer->EndType();
+  RenderConstructor(op, op_class, &writer);
+  writer.EndType();
 }
 
 }  // namespace
@@ -369,8 +433,7 @@ OpGenerator::OpGenerator(const string& base_package, const string& output_dir,
     env_(env) {
 }
 
-Status OpGenerator::Run(const OpList& op_list, const string& lib_name) {
-  LOG(INFO) << "Generating Java wrappers for '" << lib_name << "' operations";
+Status OpGenerator::Run(const OpList& op_list) {
   ApiDefMap api_map(op_list);
   if (!api_dirs_.empty()) {
     // Only load api files that correspond to the requested "op_list"
@@ -388,37 +451,14 @@ Status OpGenerator::Run(const OpList& op_list, const string& lib_name) {
   for (const auto& op_def : op_list.op()) {
     const ApiDef* api_def = api_map.GetApiDef(op_def.name());
     if (api_def->visibility() != ApiDef::SKIP) {
-      Status status = GenerateOp(op_def, *api_def, lib_name);
-      if (status != Status::OK()) {
-        LOG(ERROR) << "Fail to generate Java wrapper for operation \""
-            << op_def.name() << "\"";
+      OpSpec op(OpSpec::Create(op_def, *api_def));
+      for (const EndpointSpec& endpoint : op.endpoints()) {
+        GenerateOp(op, endpoint, base_package_, output_dir_, env_);
       }
     }
   }
   return Status::OK();
 }
 
-Status OpGenerator::GenerateOp(const OpDef& op_def, const ApiDef& api_def,
-    const string& lib_name) {
-  std::unique_ptr<OpSpec> op;
-  OpParser op_parser(op_def, api_def, lib_name, base_package_);
-  op_parser.Parse(&op);
-  for (const OpSpec::Endpoint& endpoint : op->endpoints()) {
-    string package_path = io::JoinPath(output_dir_,
-        str_util::StringReplace(endpoint.type().package(), ".", "/", true));
-    if (!env_->FileExists(package_path).ok()) {
-      TF_CHECK_OK(Env::Default()->RecursivelyCreateDir(package_path));
-    }
-    string file_path =
-        io::JoinPath(package_path, endpoint.type().name() + ".java");
-    std::unique_ptr<tensorflow::WritableFile> file;
-    TF_CHECK_OK(env_->NewWritableFile(file_path, &file));
-
-    SourceFileWriter writer(file.get());
-    RenderEndpoint(*op, endpoint, &writer);
-  }
-  return Status::OK();
-}
-
 }  // namespace java
 }  // namespace tensorflow
diff --git a/tensorflow/java/src/gen/cc/op_generator.h b/tensorflow/java/src/gen/cc/op_generator.h
index 19d8db95fb..06b08e852a 100644
--- a/tensorflow/java/src/gen/cc/op_generator.h
+++ b/tensorflow/java/src/gen/cc/op_generator.h
@@ -1,4 +1,4 @@
-/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2018 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.
@@ -23,36 +23,33 @@ limitations under the License.
 #include "tensorflow/core/framework/api_def.pb.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/java/src/gen/cc/op_specs.h"
 
 namespace tensorflow {
 namespace java {
 
 // A generator of Java operation wrappers.
 //
-// Such generator is normally ran only once per executable, outputting
-// wrappers for the all registered operations it has been compiled with.
-// Nonetheless, it is designed to support multiple runs, giving a different
-// list of operations on each cycle.
+// This generator takes a list of ops definitions in input and outputs
+// a Java Op wrapper for each of them in the provided directory. The same
+// generator instance can be invoked multiple times with a different list of
+// ops definitions.
 class OpGenerator {
  public:
   OpGenerator(const string& base_package, const string& output_dir,
       const std::vector<string>& api_dirs, Env* env = Env::Default());
-  virtual ~OpGenerator() = default;
 
   // Generates wrappers for the given list of 'ops'.
   //
   // Output files are generated in <output_dir>/<base_package>/<lib_package>,
-  // where 'lib_package' is derived from 'lib_name'.
-  Status Run(const OpList& op_list, const string& lib_name);
+  // where 'lib_package' is derived from ops endpoints.
+  Status Run(const OpList& op_list);
 
  private:
-  string base_package_;
-  string output_dir_;
-  std::vector<string> api_dirs_;
+  const string base_package_;
+  const string output_dir_;
+  const std::vector<string> api_dirs_;
   Env* env_;
-
-  Status GenerateOp(const OpDef& op_def, const ApiDef& api_def,
-    const string& lib_name);
 };
 
 }  // namespace java
diff --git a/tensorflow/java/src/gen/cc/op_parser.cc b/tensorflow/java/src/gen/cc/op_parser.cc
deleted file mode 100644
index 0541e343d8..0000000000
--- a/tensorflow/java/src/gen/cc/op_parser.cc
+++ /dev/null
@@ -1,417 +0,0 @@
-/* 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.
-==============================================================================*/
-
-#include <map>
-#include <vector>
-#include <utility>
-#include <string>
-
-#include "tensorflow/core/framework/op.h"
-#include "tensorflow/core/framework/types.h"
-#include "tensorflow/core/lib/strings/str_util.h"
-#include "tensorflow/core/platform/logging.h"
-#include "tensorflow/java/src/gen/cc/op_parser.h"
-
-namespace tensorflow {
-namespace java {
-namespace {
-
-string SnakeToCamelCase(const string& str, bool upper = false) {
-  string result;
-  bool cap = upper;
-  for (string::const_iterator it = str.begin(); it != str.end(); ++it) {
-    const char c = *it;
-    if (c == '_') {
-      cap = true;
-    } else if (cap) {
-      result += toupper(c);
-      cap = false;
-    } else {
-      result += c;
-    }
-  }
-  return result;
-}
-
-bool IsRealNumber(DataType type) {
-  for (DataType dt : RealNumberTypes()) {
-    if (type == dt) {
-      return true;
-    }
-  }
-  return false;
-}
-
-bool IsRealNumbers(const AttrValue& values) {
-  if (values.has_list()) {
-    for (int i = 0; i < values.list().type_size(); ++i) {
-      if (!IsRealNumber(values.list().type(i))) {
-        return false;
-      }
-    }
-    return true;
-  }
-  return IsRealNumber(values.type());
-}
-
-string ParseDocumentation(const string& text) {
-  std::stringstream javadoc_text;
-  string::const_iterator c_iter = text.cbegin();
-  bool code = false;
-  bool emphasis = false;
-  bool list = false;
-  while (c_iter != text.cend()) {
-    char c = *c_iter++;
-    int count = 1;
-    switch (c) {
-    case '\n':
-      if (!code) {
-        // consumes all subsequent newlines, if there are more than one,
-        // then there are two choices:
-        // - if the next line starts with an asterisk, we are enumerating
-        //   a list of items
-        // - otherwise, we are starting a new paragraph
-        for (; c_iter != text.cend() && *c_iter == '\n'; ++count, ++c_iter) {}
-        if (c_iter != text.cend()) {
-          if (count > 1) {
-            if (*c_iter != '*' && list) {
-              javadoc_text << "</li>\n</ul>\n";
-              list = false;
-            } else if (*c_iter == '*' && !list) {
-              javadoc_text << "\n<ul>\n<li>";
-              list = true;
-              c_iter++;
-            } else {
-              javadoc_text << "\n<p>\n";
-            }
-          } else if (list && *c_iter == '*') {
-            javadoc_text << "</li>\n<li>";
-            c_iter++;
-          } else {
-            javadoc_text << '\n';
-          }
-        }
-      }
-      break;
-    case '`':
-      // consumes all subsequent backquotes, those are use enclose code.
-      // if there are more than 3, we are dealing with a pre-formatted block,
-      // otherwise it is a single-line code snippet
-      for (; c_iter != text.cend() && *c_iter == '`'; ++count, ++c_iter) {}
-      if (count >= 3) {
-        javadoc_text << (code ? "\n}</pre>" : "<pre>{@code\n");
-      } else {
-        javadoc_text << (code ? "}" : "{@code ");
-      }
-      code = !code;
-      break;
-    case '*':
-      if (!code) {
-        // consumes all subsequent asterisks, if there are more than one, then
-        // we put the text in bold, otherwise in italic
-        for (; c_iter != text.cend() && *c_iter == '*'; ++count, ++c_iter) {}
-        if (count > 1) {
-          javadoc_text << (emphasis ? "</b>" : "<b>");
-        } else {
-          javadoc_text << (emphasis ? "</i>" : "<i>");
-        }
-        emphasis = !emphasis;
-      } else {
-        javadoc_text << '*';
-      }
-      break;
-    default:
-      javadoc_text << c;
-      break;
-    }
-  }
-  return javadoc_text.str();
-}
-
-}  // namespace
-
-OpParser::OpParser(const OpDef& op_def, const ApiDef& api_def,
-    const string& lib_name, const string& base_package)
-  : op_def_(op_def), op_api_(api_def), lib_name_(lib_name),
-    base_package_(base_package) {
-}
-
-void OpParser::Parse(std::unique_ptr<OpSpec>* op_ptr) {
-  visited_attrs_.clear();
-  next_generic_ = 'T';
-  op_ptr->reset(new OpSpec(op_api_.graph_op_name()));
-  for (const string& next_input_name : op_api_.arg_order()) {
-    for (int i = 0; i < op_def_.input_arg().size(); ++i) {
-      if (op_def_.input_arg(i).name() == next_input_name) {
-        ParseInput(op_def_.input_arg(i), op_api_.in_arg(i), op_ptr->get());
-        break;
-      }
-    }
-  }
-  for (int i = 0; i < op_def_.attr().size(); ++i) {
-    ParseAttribute(op_def_.attr(i), op_api_.attr(i), op_ptr->get());
-  }
-  for (int i = 0; i < op_def_.output_arg().size(); ++i) {
-    ParseOutput(op_def_.output_arg(i), op_api_.out_arg(i), op_ptr->get());
-  }
-  BuildEndpoints(op_ptr->get());
-}
-
-void OpParser::BuildEndpoints(OpSpec* op) {
-  Javadoc op_doc = Javadoc::Create(ParseDocumentation(op_api_.summary()))
-    .details(ParseDocumentation(op_api_.description()));
-  std::vector<Type> op_supertypes;
-  op_supertypes.push_back(Type::Class("PrimitiveOp", "org.tensorflow.op"));
-  std::map<string, const Type*> op_generics;
-  for (const OpSpec::Operand& output : op->outputs()) {
-    // declare generic output parameters at the Op class level
-    const Type& data_type = output.data_type();
-    if (data_type.kind() == Type::GENERIC && !data_type.unknown()
-        && op_generics.find(data_type.name()) == op_generics.end()) {
-      op_generics.insert(std::make_pair(data_type.name(), &data_type));
-      op_doc.add_param_tag("<" + data_type.name() + ">",
-          "data type of output '" + output.var().name() + "'");
-    }
-    // implement the Op as an (iteration of) Operand if it has only one output
-    if (op->outputs().size() == 1) {
-      Type operand_inf(Type::Interface("Operand", "org.tensorflow"));
-      operand_inf.add_parameter(data_type.unknown() ?
-          Type::Class("Object") : data_type);
-      op_supertypes.push_back(output.iterable() ?
-          Type::IterableOf(operand_inf) : operand_inf);
-    }
-  }
-  for (const auto& endpoint_def : op_api_.endpoint()) {
-    std::vector<string> name_tokens = str_util::Split(endpoint_def.name(), ".");
-    // if the endpoint specifies a package, use it, otherwise derive it from the
-    // op library name.
-    string name;
-    string package;
-    if (name_tokens.size() > 1) {
-      package = str_util::Lowercase(name_tokens.at(0));
-      name = name_tokens.at(1);
-    } else {
-      package = str_util::StringReplace(lib_name_, "_", "", true);
-      name = name_tokens.at(0);
-    }
-    Type endpoint(Type::Class(name, base_package_ + "." + package));
-    Javadoc endpoint_doc(op_doc);
-    for (const auto& parameter : op_generics) {
-      endpoint.add_parameter(*parameter.second);
-    }
-    for (const Type& supertype : op_supertypes) {
-      endpoint.add_supertype(supertype);
-    }
-    if (endpoint_def.deprecation_version() > 0) {
-      string explanation;
-      if (op_api_.endpoint(0).deprecation_version() == 0) {
-        explanation = ", use {@link "
-            + op->endpoints().at(0).type().full_name()
-            + "} instead";
-      } else {
-        explanation = op_def_.deprecation().explanation();
-      }
-      endpoint_doc.add_tag("deprecated", explanation);
-      endpoint.add_annotation(Annotation::Create("Deprecated"));
-    }
-    // only visible ops should be annotated for exposure in the Ops Graph API
-    if (op_api_.visibility() != ApiDef::HIDDEN) {
-      string group_name = SnakeToCamelCase(lib_name_);
-      endpoint.add_annotation(
-          Annotation::Create("Operator", "org.tensorflow.op.annotation")
-            .attributes("group = \"" + group_name + "\""));
-    }
-    op->add_endpoint(endpoint, endpoint_doc);
-  }
-}
-
-void OpParser::ParseInput(const OpDef_ArgDef& input_def,
-    const ApiDef::Arg& input_api, OpSpec* op) {
-  bool iterable = false;
-  Type data_type = DataTypeOf(input_def, &iterable);
-  Type type = Type::Interface("Operand", "org.tensorflow")
-    .add_parameter(data_type);
-  if (iterable) {
-    type = Type::IterableOf(type);
-  }
-  op->add_input(OpSpec::Operand(input_api.name(),
-      Variable::Create(SnakeToCamelCase(input_api.rename_to()), type),
-      data_type,
-      ParseDocumentation(input_api.description()),
-      iterable));
-}
-
-void OpParser::ParseOutput(const OpDef_ArgDef& output_def,
-    const ApiDef::Arg& output_api, OpSpec* op) {
-  bool iterable = false;
-  Type data_type = DataTypeOf(output_def, &iterable);
-  Type type = Type::Class("Output", "org.tensorflow")
-    .add_parameter(data_type);
-  if (iterable) {
-    type = Type::ListOf(type);
-  }
-  op->add_output(OpSpec::Operand(output_api.name(),
-      Variable::Create(SnakeToCamelCase(output_api.rename_to()), type),
-      data_type,
-      ParseDocumentation(output_api.description()),
-      iterable));
-}
-
-void OpParser::ParseAttribute(const OpDef_AttrDef& attr_def,
-    const ApiDef::Attr& attr_api, OpSpec* op) {
-  // do not parse attributes already visited, they have probably been inferred
-  // before as an input argument type
-  if (visited_attrs_.find(attr_def.name()) != visited_attrs_.cend()) {
-    return;
-  }
-  bool iterable = false;
-  Type data_type = DataTypeOf(attr_def, &iterable);
-  // generic attributes should be passed as an explicit type
-  bool explicit_type = data_type.kind() == Type::GENERIC && !iterable;
-  Type type = explicit_type ?
-      Type::Class("Class").add_parameter(data_type) : data_type;
-  if (iterable) {
-    type = Type::ListOf(data_type);
-  }
-  OpSpec::Operand attr(attr_api.name(),
-      Variable::Create(SnakeToCamelCase(attr_api.rename_to()), type),
-      data_type,
-      ParseDocumentation(attr_api.description()),
-      iterable);
-  // attributes with a default value are optional
-  if (attr_api.has_default_value() && !explicit_type) {
-    op->add_option(attr);
-  } else {
-    op->add_attribute(attr);
-  }
-  visited_attrs_.insert(std::make_pair(attr_api.name(), data_type));
-}
-
-Type OpParser::DataTypeOf(const OpDef_ArgDef& arg, bool* iterable_out) {
-  if (!arg.number_attr().empty()) {
-    visited_attrs_.insert(std::make_pair(arg.number_attr(), Type::Int()));
-    *iterable_out = true;
-  }
-  if (arg.type() != DataType::DT_INVALID) {
-    // resolve type from DataType
-    switch (arg.type()) {
-      case DataType::DT_BOOL:
-        return Type::Class("Boolean");
-
-      case DataType::DT_STRING:
-        return Type::Class("String");
-
-      case DataType::DT_FLOAT:
-        return Type::Class("Float");
-
-      case DataType::DT_DOUBLE:
-        return Type::Class("Double");
-
-      case DataType::DT_UINT8:
-        return Type::Class("UInt8", "org.tensorflow.types");
-
-      case DataType::DT_INT32:
-        return Type::Class("Integer");
-
-      case DataType::DT_INT64:
-        return Type::Class("Long");
-
-      case DataType::DT_RESOURCE:
-        // TODO(karllessard) create a Resource utility class that could be
-        // used to store a resource and its type (passed in a second argument).
-        // For now, we need to force a wildcard and we will unfortunately lose
-        // track of the resource type.
-        return Type::Wildcard();
-
-      default:
-        break;
-    }
-  } else {
-    // resolve type from type attribute
-    string attr_name = arg.type_attr();
-    if (attr_name.empty()) {
-      attr_name = arg.type_list_attr();
-      if (!attr_name.empty()) {
-        *iterable_out = true;
-        Type type = Type::Wildcard();
-        visited_attrs_.insert(std::make_pair(attr_name, type));
-        return type;
-      }
-    }
-    for (const auto& attr : op_def_.attr()) {
-      if (attr.name() == attr_name) {
-        Type type = DataTypeOf(attr, iterable_out);
-        visited_attrs_.insert(std::make_pair(attr_name, type));
-        return type;
-      }
-    }
-  }
-  LOG(WARNING) << "Data type for arg \"" << arg.name() << "\" is unknown";
-  return Type::Wildcard();
-}
-
-Type OpParser::DataTypeOf(const OpDef_AttrDef& attr, bool* iterable_out) {
-  std::map<string, Type>::const_iterator it = visited_attrs_.find(attr.name());
-  if (it != visited_attrs_.cend()) {
-    return it->second;
-  }
-  string attr_type = attr.type();
-  if (attr.type().compare(0, 5, "list(") == 0) {
-    attr_type = attr_type.substr(5, attr.type().find_last_of(')') - 5);
-    *iterable_out = true;
-  }
-  if (attr_type == "type") {
-    if (*iterable_out) {
-      return Type::Enum("DataType", "org.tensorflow");
-    }
-    return GetNextGenericTensorType(attr.allowed_values());
-  }
-  if (attr_type == "string") {
-    return Type::Class("String");
-  }
-  if (attr_type == "int") {
-    return Type::Class("Integer");
-  }
-  if (attr_type == "float") {
-    return Type::Class("Float");
-  }
-  if (attr_type == "bool") {
-    return Type::Class("Boolean");
-  }
-  if (attr_type == "shape") {
-    return Type::Class("Shape", "org.tensorflow");
-  }
-  if (attr_type == "tensor") {
-    return Type::Class("Tensor", "org.tensorflow")
-      .add_parameter(Type::Wildcard());
-  }
-  LOG(WARNING) << "Data type for attribute \"" << attr_type << "\" is unknown";
-  return *iterable_out ? Type::Wildcard() : Type::Class("Object");
-}
-
-Type OpParser::GetNextGenericTensorType(const AttrValue& allowed_values)  {
-  Type generic = Type::Generic(string(1, next_generic_));
-  next_generic_ = (next_generic_ == 'Z') ? 'A' : next_generic_ + 1;
-
-  // when only real numbers are allowed, enforce that restriction in the Java by
-  // extending the generic from java.lang.Number
-  if (IsRealNumbers(allowed_values)) {
-    generic.add_supertype(Type::Class("Number"));
-  }
-  return generic;
-}
-
-}  // namespace java
-}  // namespace tensorflow
diff --git a/tensorflow/java/src/gen/cc/op_parser.h b/tensorflow/java/src/gen/cc/op_parser.h
deleted file mode 100644
index 42855127cc..0000000000
--- a/tensorflow/java/src/gen/cc/op_parser.h
+++ /dev/null
@@ -1,137 +0,0 @@
-/* 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.
-==============================================================================*/
-
-#ifndef TENSORFLOW_JAVA_SRC_GEN_CC_OP_PARSER_H_
-#define TENSORFLOW_JAVA_SRC_GEN_CC_OP_PARSER_H_
-
-#include <map>
-#include <memory>
-#include <string>
-#include <vector>
-
-#include "tensorflow/core/framework/op_def.pb.h"
-#include "tensorflow/core/framework/api_def.pb.h"
-#include "tensorflow/java/src/gen/cc/java_defs.h"
-
-namespace tensorflow {
-namespace java {
-
-// Specification of a TensorFlow operation to generate.
-//
-// This is the result of an operation definition parsing, see OpParser::Parse().
-class OpSpec {
- public:
-  class Endpoint {
-   public:
-    Endpoint(const Type& type, const Javadoc& javadoc)
-      : type_(type), javadoc_(javadoc) {}
-    const Type& type() const { return type_; }
-    const Javadoc& javadoc() const { return javadoc_; }
-
-   private:
-    Type type_;
-    Javadoc javadoc_;
-  };
-
-  class Operand {
-   public:
-    Operand(const string& graph_name, const Variable& var,
-        const Type& data_type, const string& description, bool iterable)
-     : graph_name_(graph_name), var_(var), data_type_(data_type),
-       description_(description), iterable_(iterable) {}
-    const string& graph_name() const { return graph_name_; }
-    const Variable& var() const { return var_; }
-    Variable* var_ptr() { return &var_; }
-    const Type& data_type() const { return data_type_; }
-    const string& description() const { return description_; }
-    bool iterable() const { return iterable_; }
-
-   private:
-    string graph_name_;
-    Variable var_;
-    Type data_type_;
-    string description_;
-    bool iterable_;
-  };
-
-  explicit OpSpec(const string& graph_name) : graph_name_(graph_name) {}
-  const string& graph_name() const { return graph_name_; }
-  const std::vector<Endpoint> endpoints() const { return endpoints_; }
-  void add_endpoint(const Type& type, const Javadoc& javadoc) {
-    endpoints_.push_back(Endpoint(type, javadoc));
-  }
-  const std::vector<Operand>& inputs() const { return inputs_; }
-  void add_input(const Operand& input) {
-    inputs_.push_back(input);
-  }
-  const std::vector<Operand>& outputs() const { return outputs_; }
-  void add_output(const Operand& output) {
-    outputs_.push_back(output);
-  }
-  const std::vector<Operand>& attributes() const { return attributes_; }
-  void add_attribute(const Operand& attribute) {
-    attributes_.push_back(attribute);
-  }
-  const std::vector<Operand>& options() const { return options_; }
-  void add_option(const Operand& option) {
-    options_.push_back(option);
-  }
-
- private:
-  string graph_name_;
-  std::vector<Endpoint> endpoints_;
-  std::vector<Operand> inputs_;
-  std::vector<Operand> outputs_;
-  std::vector<Operand> attributes_;
-  std::vector<Operand> options_;
-};
-
-// A parser of ops proto definitions.
-//
-// This object parses the definition and the api of an TensorFlow operation to
-// produce a specification that can be used for Java source code rendering.
-class OpParser {
- public:
-  OpParser(const OpDef& op_def, const ApiDef& api_def, const string& lib_name,
-      const string& base_package);
-  virtual ~OpParser() = default;
-
-  // Produces an operation specification from its proto definitions.
-  void Parse(std::unique_ptr<OpSpec>* op_ptr);
-
- private:
-  OpDef op_def_;
-  ApiDef op_api_;
-  string lib_name_;
-  string base_package_;
-  std::map<string, Type> visited_attrs_;
-  char next_generic_ = 0;
-
-  void BuildEndpoints(OpSpec* op);
-  void ParseInput(const OpDef_ArgDef& input_def,
-      const ApiDef::Arg& input_api, OpSpec* op);
-  void ParseOutput(const OpDef_ArgDef& output_def,
-      const ApiDef::Arg& output_api, OpSpec* op);
-  void ParseAttribute(const OpDef_AttrDef& attr_def,
-      const ApiDef::Attr& attr_api, OpSpec* op);
-  Type DataTypeOf(const OpDef_ArgDef& arg_def, bool *iterable_out);
-  Type DataTypeOf(const OpDef_AttrDef& attr_def, bool *iterable_out);
-  Type GetNextGenericTensorType(const AttrValue& allowed_values);
-};
-
-}  // namespace java
-}  // namespace tensorflow
-
-#endif  // TENSORFLOW_JAVA_SRC_GEN_CC_OP_PARSER_H_
diff --git a/tensorflow/java/src/gen/cc/op_specs.cc b/tensorflow/java/src/gen/cc/op_specs.cc
new file mode 100644
index 0000000000..a727f7ae90
--- /dev/null
+++ b/tensorflow/java/src/gen/cc/op_specs.cc
@@ -0,0 +1,390 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include <map>
+#include <vector>
+#include <string>
+#include <utility>
+
+#include "re2/re2.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/java/src/gen/cc/op_specs.h"
+
+namespace tensorflow {
+namespace java {
+namespace {
+
+inline bool IsRealNumbers(const AttrValue& values) {
+  if (!values.has_list()) {
+    return RealNumberTypes().Contains(values.type());
+  }
+  for (int i = 0; i < values.list().type_size(); ++i) {
+    if (!RealNumberTypes().Contains(values.list().type(i))) {
+      return false;
+    }
+  }
+  return true;
+}
+
+class TypeResolver {
+ public:
+  explicit TypeResolver(const OpDef& op_def) : op_def_(op_def) {}
+
+  Type TypeOf(const OpDef_ArgDef& arg_def, bool *iterable_out);
+  Type TypeOf(const OpDef_AttrDef& attr_def, bool *iterable_out);
+  bool IsAttributeVisited(const string& attr_name) {
+    return visited_attrs_.find(attr_name) != visited_attrs_.cend();
+  }
+ private:
+  const OpDef op_def_;
+  std::map<std::string, Type> visited_attrs_;
+  char next_generic_ = 'T';
+};
+
+Type TypeResolver::TypeOf(const OpDef_ArgDef& arg_def,
+    bool* iterable_out) {
+  *iterable_out = false;
+  if (!arg_def.number_attr().empty()) {
+    // when number_attr is set, argument has to be a list of tensors
+    *iterable_out = true;
+    visited_attrs_.insert(std::make_pair(arg_def.number_attr(), Type::Int()));
+  }
+  Type type = Type::Wildcard();
+  if (arg_def.type() != DataType::DT_INVALID) {
+    // resolve type from DataType
+    switch (arg_def.type()) {
+      case DataType::DT_BOOL:
+        type = Type::Class("Boolean");
+        break;
+      case DataType::DT_STRING:
+        type = Type::Class("String");
+        break;
+      case DataType::DT_FLOAT:
+        type = Type::Class("Float");
+        break;
+      case DataType::DT_DOUBLE:
+        type = Type::Class("Double");
+        break;
+      case DataType::DT_UINT8:
+        type = Type::Class("UInt8", "org.tensorflow.types");
+        break;
+      case DataType::DT_INT32:
+        type = Type::Class("Integer");
+        break;
+      case DataType::DT_INT64:
+        type = Type::Class("Long");
+        break;
+      case DataType::DT_RESOURCE:
+        // TODO(karllessard) create a Resource utility class that could be
+        // used to store a resource and its type (passed in a second argument).
+        // For now, we need to force a wildcard and we will unfortunately lose
+        // track of the resource type.
+        break;
+      default:
+        // Any other datatypes does not have a equivalent in Java and must
+        // remain a wildcard (e.g. DT_COMPLEX64, DT_QINT8, ...)
+        break;
+    }
+  } else if (!arg_def.type_attr().empty()) {
+    // resolve type from attribute (if already visited, retrieve its type)
+    if (IsAttributeVisited(arg_def.type_attr())) {
+      type = visited_attrs_.at(arg_def.type_attr());
+    } else {
+      for (const auto& attr_def : op_def_.attr()) {
+        if (attr_def.name() == arg_def.type_attr()) {
+          type = TypeOf(attr_def, iterable_out);
+          break;
+        }
+      }
+    }
+  } else if (!arg_def.type_list_attr().empty()) {
+    // type is a list of tensors that can be of different data types, so leave
+    // it as a list of wildcards
+    *iterable_out = true;
+    visited_attrs_.insert(std::make_pair(arg_def.type_list_attr(), type));
+
+  } else {
+    LOG(FATAL) << "Cannot resolve data type of argument \"" << arg_def.name()
+        << "\" in operation \"" << op_def_.name() << "\"";
+  }
+  return type;
+}
+
+Type TypeResolver::TypeOf(const OpDef_AttrDef& attr_def,
+    bool* iterable_out) {
+  *iterable_out = false;
+  StringPiece attr_type = attr_def.type();
+  if (str_util::ConsumePrefix(&attr_type, "list(")) {
+    attr_type.remove_suffix(1);  // remove closing brace
+    *iterable_out = true;
+  }
+  Type type = *iterable_out ? Type::Wildcard() : Type::Class("Object");
+  if (attr_type == "type") {
+    if (*iterable_out) {
+      type = Type::Enum("DataType", "org.tensorflow");
+    } else {
+      type = Type::Generic(string(1, next_generic_));
+      next_generic_ = (next_generic_ == 'Z') ? 'A' : next_generic_ + 1;
+      if (IsRealNumbers(attr_def.allowed_values())) {
+        // enforce real numbers datasets by extending java.lang.Number
+        type.add_supertype(Type::Class("Number"));
+      }
+    }
+  } else if (attr_type == "string") {
+    type = Type::Class("String");
+
+  } else if (attr_type == "int") {
+    type = Type::Class("Integer");
+
+  } else if (attr_type == "float") {
+    type = Type::Class("Float");
+
+  } else if (attr_type == "bool") {
+    type = Type::Class("Boolean");
+
+  } else if (attr_type == "shape") {
+    type = Type::Class("Shape", "org.tensorflow");
+
+  } else if (attr_type == "tensor") {
+    type = Type::Class("Tensor", "org.tensorflow")
+        .add_parameter(Type::Wildcard());
+
+  } else {
+    LOG(FATAL) << "Cannot resolve data type for attribute \"" << attr_type
+        << "\" in operation \"" << op_def_.name() << "\"";
+  }
+  visited_attrs_.insert(std::make_pair(attr_def.name(), type));
+  return type;
+}
+
+string SnakeToCamelCase(const string& str, bool upper = false) {
+  string result;
+  bool cap = upper;
+  for (string::const_iterator it = str.begin(); it != str.end(); ++it) {
+    const char c = *it;
+    if (c == '_') {
+      cap = true;
+    } else if (cap) {
+      result += toupper(c);
+      cap = false;
+    } else {
+      result += c;
+    }
+  }
+  return result;
+}
+
+bool FindAndCut(re2::StringPiece* input, const RE2& expr,
+    re2::StringPiece* before_match, re2::StringPiece* ret_match = nullptr) {
+  re2::StringPiece match;
+  bool matches =
+      expr.Match(*input, 0, input->size(), RE2::UNANCHORED, &match, 1);
+  if (matches) {
+    before_match->set(input->data(), match.begin() - input->begin());
+    input->remove_prefix(match.end() - before_match->begin());
+    if (ret_match != nullptr) {
+      *ret_match = match;
+    }
+  } else {
+    *before_match = *input;
+    if (ret_match != nullptr) {
+      ret_match->set(nullptr, 0);
+    }
+  }
+  return matches;
+}
+
+string ParseDocumentation(const string& mdtext) {
+  std::stringstream javadoc_text;
+  re2::StringPiece input(mdtext);
+  re2::StringPiece text;
+  bool in_list = false;
+  do {
+    re2::StringPiece markup;
+    FindAndCut(&input,
+        "\n+\\*[[:blank:]]+|\n{2,}|`{3,}|`{1,2}|\\*{1,2}\\b|\\[",
+        &text, &markup);
+    javadoc_text << text;
+    if (markup.empty()) {
+      break;  // we are done parsing
+    }
+    if (markup.starts_with("\n")) {
+      javadoc_text << "\n";
+      if (markup.contains("* ")) {
+        javadoc_text << (in_list ? "</li>\n" : "<ul>\n") << "<li>\n";
+        in_list = true;
+      } else if (markup.starts_with("\n\n")) {
+        if (in_list) {
+          javadoc_text << "</li>\n</ul>\n";
+          in_list = false;
+        } else if (!input.starts_with("```")) {
+          javadoc_text << "<p>\n";
+        }
+      }
+    } else if (markup.starts_with("```") && text.empty()) {
+      re2::StringPiece language;
+      RE2::Consume(&input, "[\\w\\+]+", &language);
+      if (FindAndCut(&input, markup.ToString() + "\n*", &text)) {
+        javadoc_text << "<pre>\n{@code" << text << "}\n</pre>\n";
+      } else {
+        javadoc_text << markup << language;
+      }
+    } else if (markup.starts_with("`")) {
+      if (FindAndCut(&input, markup, &text)) {
+        javadoc_text << "{@code " << text << "}";
+      } else {
+        javadoc_text << markup;
+      }
+    } else if (markup == "**") {
+      if (FindAndCut(&input, "\\b\\*{2}", &text)) {
+        javadoc_text << "<b>" << text << "</b>";
+      } else {
+        javadoc_text << markup;
+      }
+    } else if (markup == "*") {
+      if (FindAndCut(&input, "\\b\\*{1}", &text)) {
+        javadoc_text << "<i>" << text << "</i>";
+      } else {
+        javadoc_text << markup;
+      }
+    } else if (markup == "[") {
+      string label;
+      string link;
+      if (RE2::Consume(&input, "([^\\[]+)\\]\\((http.+)\\)", &label, &link)) {
+        javadoc_text << "<a href=\"" << link << "\">" << label << "</a>";
+      } else {
+        javadoc_text << markup;
+      }
+    }
+  } while (!input.empty());
+
+  return javadoc_text.str();
+}
+
+ArgumentSpec CreateInput(const OpDef_ArgDef& input_def,
+    const ApiDef::Arg& input_api_def, TypeResolver* type_resolver) {
+  bool iterable = false;
+  Type type = type_resolver->TypeOf(input_def, &iterable);
+  Type var_type = Type::Interface("Operand", "org.tensorflow")
+    .add_parameter(type);
+  if (iterable) {
+    var_type = Type::IterableOf(var_type);
+  }
+  return ArgumentSpec(input_api_def.name(),
+      Variable::Create(SnakeToCamelCase(input_api_def.rename_to()), var_type),
+      type,
+      ParseDocumentation(input_api_def.description()),
+      iterable);
+}
+
+AttributeSpec CreateAttribute(const OpDef_AttrDef& attr_def,
+    const ApiDef::Attr& attr_api_def, TypeResolver* type_resolver) {
+  bool iterable = false;
+  Type type = type_resolver->TypeOf(attr_def, &iterable);
+  // type attributes must be passed explicitly in methods as a Class<> parameter
+  bool is_explicit = type.kind() == Type::GENERIC && !iterable;
+  Type var_type = is_explicit ? Type::Class("Class").add_parameter(type) : type;
+  if (iterable) {
+    var_type = Type::ListOf(type);
+  }
+  return AttributeSpec(attr_api_def.name(),
+      Variable::Create(SnakeToCamelCase(attr_api_def.rename_to()), var_type),
+      type,
+      ParseDocumentation(attr_api_def.description()),
+      iterable,
+      attr_api_def.has_default_value() && !is_explicit);
+}
+
+ArgumentSpec CreateOutput(const OpDef_ArgDef& output_def,
+    const ApiDef::Arg& output_api, TypeResolver* type_resolver) {
+  bool iterable = false;
+  Type type = type_resolver->TypeOf(output_def, &iterable);
+  Type var_type = Type::Class("Output", "org.tensorflow")
+    .add_parameter(type);
+  if (iterable) {
+    var_type = Type::ListOf(var_type);
+  }
+  return ArgumentSpec(output_api.name(),
+      Variable::Create(SnakeToCamelCase(output_api.rename_to()), var_type),
+      type,
+      ParseDocumentation(output_api.description()),
+      iterable);
+}
+
+EndpointSpec CreateEndpoint(const OpDef& op_def, const ApiDef& api_def,
+    const ApiDef_Endpoint& endpoint_def) {
+
+  std::vector<string> name_tokens = str_util::Split(endpoint_def.name(), ".");
+  string package;
+  string name;
+  if (name_tokens.size() > 1) {
+    package = name_tokens.at(0);
+    name = name_tokens.at(1);
+  } else {
+    package = "core";  // generate unclassified ops in the 'core' package
+    name = name_tokens.at(0);
+  }
+  return EndpointSpec(package,
+      name,
+      Javadoc::Create(ParseDocumentation(api_def.summary()))
+          .details(ParseDocumentation(api_def.description())),
+      endpoint_def.deprecation_version() > 0);
+}
+
+}  // namespace
+
+OpSpec OpSpec::Create(const OpDef& op_def, const ApiDef& api_def) {
+  OpSpec op(api_def.graph_op_name(),
+      api_def.visibility() == ApiDef::HIDDEN,
+      op_def.deprecation().explanation());
+  TypeResolver type_resolver(op_def);
+  for (const string& next_input_name : api_def.arg_order()) {
+    for (int i = 0; i < op_def.input_arg().size(); ++i) {
+      if (op_def.input_arg(i).name() == next_input_name) {
+        op.inputs_.push_back(CreateInput(op_def.input_arg(i), api_def.in_arg(i),
+            &type_resolver));
+        break;
+      }
+    }
+  }
+  for (int i = 0; i < op_def.attr().size(); ++i) {
+    // do not parse attributes already visited, they have probably been inferred
+    // before as an input argument type
+    if (!type_resolver.IsAttributeVisited(op_def.attr(i).name())) {
+      AttributeSpec attr = CreateAttribute(op_def.attr(i), api_def.attr(i),
+          &type_resolver);
+      // attributes with a default value are optional
+      if (attr.optional()) {
+        op.optional_attributes_.push_back(attr);
+      } else {
+        op.attributes_.push_back(attr);
+      }
+    }
+  }
+  for (int i = 0; i < op_def.output_arg().size(); ++i) {
+    op.outputs_.push_back(CreateOutput(op_def.output_arg(i), api_def.out_arg(i),
+        &type_resolver));
+  }
+  for (const auto& endpoint_def : api_def.endpoint()) {
+    op.endpoints_.push_back(CreateEndpoint(op_def, api_def, endpoint_def));
+  }
+  return op;
+}
+
+}  // namespace java
+}  // namespace tensorflow
diff --git a/tensorflow/java/src/gen/cc/op_specs.h b/tensorflow/java/src/gen/cc/op_specs.h
new file mode 100644
index 0000000000..55c2c3f307
--- /dev/null
+++ b/tensorflow/java/src/gen/cc/op_specs.h
@@ -0,0 +1,152 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_JAVA_SRC_GEN_CC_OP_SPECS_H_
+#define TENSORFLOW_JAVA_SRC_GEN_CC_OP_SPECS_H_
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/api_def.pb.h"
+#include "tensorflow/java/src/gen/cc/java_defs.h"
+
+namespace tensorflow {
+namespace java {
+
+class EndpointSpec {
+ public:
+  // A specification for an operation endpoint
+  //
+  // package: package of this endpoint (from which also derives its package)
+  // name: name of this endpoint class
+  // javadoc: the endpoint class documentation
+  // deprecated: true if this endpoint is now deprecated
+  EndpointSpec(const string& package, const string& name,
+      const Javadoc& javadoc, bool deprecated)
+    : package_(package), name_(name), javadoc_(javadoc),
+      deprecated_(deprecated) {}
+
+  const string& package() const { return package_; }
+  const string& name() const { return name_; }
+  const Javadoc& javadoc() const { return javadoc_; }
+  bool deprecated() const { return deprecated_; }
+
+ private:
+  const string package_;
+  const string name_;
+  const Javadoc javadoc_;
+  const bool deprecated_;
+};
+
+class ArgumentSpec {
+ public:
+  // A specification for an operation argument
+  //
+  // op_def_name: argument name, as known by TensorFlow core
+  // var: a variable to represent this argument in Java
+  // type: the tensor type of this argument
+  // description: a description of this argument, in javadoc
+  // iterable: true if this argument is a list
+  ArgumentSpec(const string& op_def_name, const Variable& var,
+      const Type& type, const string& description, bool iterable)
+    : op_def_name_(op_def_name), var_(var), type_(type),
+      description_(description), iterable_(iterable) {}
+  virtual ~ArgumentSpec() = default;
+
+  const string& op_def_name() const { return op_def_name_; }
+  const Variable& var() const { return var_; }
+  const Type& type() const { return type_; }
+  const string& description() const { return description_; }
+  bool iterable() const { return iterable_; }
+
+ private:
+  const string op_def_name_;
+  const Variable var_;
+  const Type type_;
+  const string description_;
+  const bool iterable_;
+};
+
+class AttributeSpec : public ArgumentSpec {
+ public:
+  // A specification for an operation attribute
+  //
+  // op_def_name: attribute name, as known by TensorFlow core
+  // var: a variable to represent this attribute in Java
+  // type: the type of this attribute
+  // description: a description of this attribute, in javadoc
+  // iterable: true if this attribute is a list
+  // optional: true if this attribute does not require to be set explicitly
+  AttributeSpec(const string& op_def_name, const Variable& var,
+      const Type& type, const string& description, bool iterable,
+      bool optional)
+    : ArgumentSpec(op_def_name, var, type, description, iterable),
+      optional_(optional) {}
+  virtual ~AttributeSpec() = default;
+
+  bool optional() const { return optional_; }
+
+ private:
+  const bool optional_;
+};
+
+class OpSpec {
+ public:
+  // Parses an op definition and its API to produce a specification used for
+  // rendering its Java wrapper
+  //
+  // op_def: Op definition
+  // api_def: Op API definition
+  static OpSpec Create(const OpDef& op_def, const ApiDef& api_def);
+
+  const string& graph_op_name() const { return graph_op_name_; }
+  bool hidden() const { return hidden_; }
+  const string& deprecation_explanation() const {
+    return deprecation_explanation_;
+  }
+  const std::vector<EndpointSpec> endpoints() const { return endpoints_; }
+  const std::vector<ArgumentSpec>& inputs() const { return inputs_; }
+  const std::vector<ArgumentSpec>& outputs() const { return outputs_; }
+  const std::vector<AttributeSpec>& attributes() const { return attributes_; }
+  const std::vector<AttributeSpec>& optional_attributes() const {
+    return optional_attributes_;
+  }
+
+ private:
+  // A specification for an operation
+  //
+  // graph_op_name: name of this op, as known by TensorFlow core engine
+  // hidden: true if this op should not be visible through the Graph Ops API
+  // deprecation_explanation: message to show if all endpoints are deprecated
+  explicit OpSpec(const string& graph_op_name, bool hidden,
+      const string& deprecation_explanation)
+    : graph_op_name_(graph_op_name), hidden_(hidden),
+      deprecation_explanation_(deprecation_explanation) {}
+
+  const string graph_op_name_;
+  const bool hidden_;
+  const string deprecation_explanation_;
+  std::vector<EndpointSpec> endpoints_;
+  std::vector<ArgumentSpec> inputs_;
+  std::vector<ArgumentSpec> outputs_;
+  std::vector<AttributeSpec> attributes_;
+  std::vector<AttributeSpec> optional_attributes_;
+};
+
+}  // namespace java
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_JAVA_SRC_GEN_CC_OP_SPECS_H_
diff --git a/tensorflow/java/src/gen/cc/source_writer.cc b/tensorflow/java/src/gen/cc/source_writer.cc
index b1de5af6ba..7e427787f9 100644
--- a/tensorflow/java/src/gen/cc/source_writer.cc
+++ b/tensorflow/java/src/gen/cc/source_writer.cc
@@ -1,4 +1,4 @@
-/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2018 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.
diff --git a/tensorflow/java/src/gen/cc/source_writer.h b/tensorflow/java/src/gen/cc/source_writer.h
index 1f0febe9a3..bcae33ccce 100644
--- a/tensorflow/java/src/gen/cc/source_writer.h
+++ b/tensorflow/java/src/gen/cc/source_writer.h
@@ -1,4 +1,4 @@
-/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2018 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.
diff --git a/tensorflow/java/src/gen/cc/source_writer_test.cc b/tensorflow/java/src/gen/cc/source_writer_test.cc
index 8bd42d9d0e..875ad99ae2 100644
--- a/tensorflow/java/src/gen/cc/source_writer_test.cc
+++ b/tensorflow/java/src/gen/cc/source_writer_test.cc
@@ -1,4 +1,4 @@
-/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2018 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.
@@ -313,8 +313,7 @@ TEST(WriteType, SimpleClassWithDependencies) {
 TEST(WriteType, AnnotatedAndDocumentedClass) {
   SourceBufferWriter writer;
   Type clazz = Type::Class("Test", "org.tensorflow");
-  Javadoc clazz_doc;
-  clazz_doc.brief("Javadoc test")
+  Javadoc clazz_doc = Javadoc::Create("Javadoc test")
       .details("This is a\nmultiline description.");
   clazz.add_annotation(Annotation::Create("Bean"));
   clazz.add_annotation(Annotation::Create("SuppressWarnings")
@@ -329,7 +328,7 @@ TEST(WriteType, AnnotatedAndDocumentedClass) {
       " * <p>\n"
       " * This is a\n"
       " * multiline description.\n"
-      " **/\n"
+      " */\n"
       "@Bean\n"
       "@SuppressWarnings(\"rawtypes\")\n"
       "public class Test {\n}\n";
@@ -378,8 +377,7 @@ TEST(WriteType, ParameterizedClassFields) {
   Variable field1 = Variable::Create("field1", Type::Class("String"));
   Variable field2 = Variable::Create("field2", Type::Class("String"));
   Variable field3 = Variable::Create("field3", type_t);
-  Javadoc field3_doc;
-  field3_doc.brief("This variable is documented");
+  Javadoc field3_doc = Javadoc::Create("This variable is documented");
 
   writer.BeginType(clazz, PUBLIC)
           .WriteField(field1, STATIC | PUBLIC | FINAL)
@@ -464,8 +462,7 @@ TEST(WriteMethod, AnnotatedAndDocumentedMethod) {
   SourceBufferWriter writer;
   Type clazz = Type::Class("Test", "org.tensorflow");
   Method method = Method::Create("doNothing", Type::Void());
-  Javadoc method_doc;
-  method_doc.brief("Javadoc test")
+  Javadoc method_doc = Javadoc::Create("Javadoc test")
       .details("This method has a\nmultiline description.");
   method.add_annotation(Annotation::Create("Override"));
   method.add_annotation(Annotation::Create("SuppressWarnings")
@@ -484,7 +481,7 @@ TEST(WriteMethod, AnnotatedAndDocumentedMethod) {
       "   * <p>\n"
       "   * This method has a\n"
       "   * multiline description.\n"
-      "   **/\n"
+      "   */\n"
       "  @Override\n"
       "  @SuppressWarnings(\"rawtypes\")\n"
       "  public void doNothing() {\n"
@@ -500,8 +497,7 @@ TEST(WriteMethod, DocumentedMethodWithArguments) {
   Method method = Method::Create("boolToInt", Type::Int());
   method.add_argument(Variable::Create("b", Type::Boolean()));
   method.add_argument(reverse);
-  Javadoc method_doc;
-  method_doc.brief("Converts a boolean to an int")
+  Javadoc method_doc = Javadoc::Create("Converts a boolean to an int")
       .details("This method will convert\na boolean to an int")
       .add_param_tag(reverse.name(), "if true, value is reversed")
       .add_tag("return", "int value for this boolean");
@@ -528,7 +524,7 @@ TEST(WriteMethod, DocumentedMethodWithArguments) {
       "   * \n"
       "   * @param reverse if true, value is reversed\n"
       "   * @return int value for this boolean\n"
-      "   **/\n"
+      "   */\n"
       "  public int boolToInt(boolean b, boolean reverse) {\n"
       "    if (b && !reverse) {\n"
       "      return 1;\n"
diff --git a/tensorflow/java/src/gen/gen_ops.bzl b/tensorflow/java/src/gen/gen_ops.bzl
index 1e7899cf7a..7017b52649 100644
--- a/tensorflow/java/src/gen/gen_ops.bzl
+++ b/tensorflow/java/src/gen/gen_ops.bzl
@@ -32,50 +32,52 @@ def tf_java_op_gen_srcjar(name,
                           api_def_srcs=[],
                           visibility=["//tensorflow/java:__pkg__"]):
 
-  gen_tools = []
   gen_cmds = ["rm -rf $(@D)"]  # Always start from fresh when generating source files
   srcs = api_def_srcs[:]
 
-  # Construct an op generator binary for each ops library.
-  for ops_lib in ops_libs:
-    gen_lib = ops_lib[:ops_lib.rfind("_")]
-    out_gen_tool = out_dir + ops_lib + "_gen_tool"
+  if not api_def_srcs:
+    api_def_args_str = ","
+  else:
+    api_def_args = []
+    for api_def_src in api_def_srcs:
+      # Add directory of the first ApiDef source to args.
+      # We are assuming all ApiDefs in a single api_def_src are in the
+      # same directory.
+      api_def_args.append(
+          "$$(dirname $$(echo $(locations " + api_def_src +
+          ") | cut -d\" \" -f1))")
+    api_def_args_str = ",".join(api_def_args)
 
-    if not api_def_srcs:
-      api_def_args_str = ","
-    else:
-      api_def_args = []
-      for api_def_src in api_def_srcs:
-        # Add directory of the first ApiDef source to args.
-        # We are assuming all ApiDefs in a single api_def_src are in the
-        # same directory.
-        api_def_args.append(
-            " $$(dirname $$(echo $(locations " + api_def_src +
-            ") | cut -d\" \" -f1))")
-      api_def_args_str = ",".join(api_def_args)
+  gen_tool_deps = [":java_op_gen_lib"]
+  for ops_lib in ops_libs:
+    gen_tool_deps.append(ops_libs_pkg + ":" + ops_lib + "_op_lib")
 
-    tf_cc_binary(
-        name=out_gen_tool,
-        copts=tf_copts(),
-        linkopts=["-lm"],
-        linkstatic=1,  # Faster to link this one-time-use binary dynamically
-        deps=[gen_tool, ops_libs_pkg + ":" + ops_lib + "_op_lib"])
+  tf_cc_binary(
+      name=gen_tool,
+      srcs=[
+          "src/gen/cc/op_gen_main.cc",
+      ],
+      copts=tf_copts(),
+      linkopts=["-lm"],
+      linkstatic=1,  # Faster to link this one-time-use binary dynamically
+      deps = gen_tool_deps)
 
-    gen_tools += [":" + out_gen_tool]
-    gen_cmds += ["$(location :" + out_gen_tool + ")" +
-                 " --output_dir=$(@D)/" + out_src_dir +
-                 " --lib_name=" + gen_lib +
-                 " --base_package=" + gen_base_package +
-                 " " + api_def_args_str]
+  gen_cmds += ["$(location :" + gen_tool + ")" +
+               " --output_dir=$(@D)/" + out_src_dir +
+               " --base_package=" + gen_base_package +
+               " --api_dirs=" + api_def_args_str]
 
   # Generate a source archive containing generated code for these ops.
   gen_srcjar = out_dir + name + ".srcjar"
   gen_cmds += ["$(location @local_jdk//:jar) cMf $(location :" + gen_srcjar + ") -C $(@D) src"]
-  gen_tools += ["@local_jdk//:jar"] + ["@local_jdk//:jdk"]
-  gen_tools += tf_binary_additional_srcs()
+
   native.genrule(
       name=name,
       srcs=srcs,
       outs=[gen_srcjar],
-      tools=gen_tools,
-      cmd="&&".join(gen_cmds))
+      tools=[
+          "@local_jdk//:jar",
+          "@local_jdk//:jdk",
+          gen_tool
+      ] + tf_binary_additional_srcs(),
+      cmd=" && ".join(gen_cmds))
diff --git a/tensorflow/java/src/gen/resources/license.snippet.java b/tensorflow/java/src/gen/resources/license.java.snippet
similarity index 100%
rename from tensorflow/java/src/gen/resources/license.snippet.java
rename to tensorflow/java/src/gen/resources/license.java.snippet
-- 
GitLab


From 6fee70dd4c82502fefa8259f0d8dbefcece58c60 Mon Sep 17 00:00:00 2001
From: "karl@kubx.ca" <karl@kubx.ca>
Date: Tue, 24 Apr 2018 09:09:01 -0400
Subject: [PATCH 221/395] Comments and little improvements to documentation
 parser

---
 tensorflow/java/src/gen/cc/op_specs.cc | 65 ++++++++++++++------------
 1 file changed, 36 insertions(+), 29 deletions(-)

diff --git a/tensorflow/java/src/gen/cc/op_specs.cc b/tensorflow/java/src/gen/cc/op_specs.cc
index a727f7ae90..3645fcf836 100644
--- a/tensorflow/java/src/gen/cc/op_specs.cc
+++ b/tensorflow/java/src/gen/cc/op_specs.cc
@@ -192,51 +192,51 @@ string SnakeToCamelCase(const string& str, bool upper = false) {
 bool FindAndCut(re2::StringPiece* input, const RE2& expr,
     re2::StringPiece* before_match, re2::StringPiece* ret_match = nullptr) {
   re2::StringPiece match;
-  bool matches =
-      expr.Match(*input, 0, input->size(), RE2::UNANCHORED, &match, 1);
-  if (matches) {
-    before_match->set(input->data(), match.begin() - input->begin());
-    input->remove_prefix(match.end() - before_match->begin());
-    if (ret_match != nullptr) {
-      *ret_match = match;
-    }
-  } else {
-    *before_match = *input;
-    if (ret_match != nullptr) {
-      ret_match->set(nullptr, 0);
-    }
+  if (!expr.Match(*input, 0, input->size(), RE2::UNANCHORED, &match, 1)) {
+    return false;
   }
-  return matches;
+  before_match->set(input->data(), match.begin() - input->begin());
+  input->remove_prefix(match.end() - before_match->begin());
+  if (ret_match != nullptr) {
+    *ret_match = match;
+  }
+  return true;
 }
 
-string ParseDocumentation(const string& mdtext) {
+string ParseDocumentation(re2::StringPiece input) {
+  // TODO(karllessard) This is a very minimalist utility method for converting
+  // markdown syntax, as found in ops descriptions, to Javadoc/html tags. Check
+  // for alternatives to increase the level of support for markups.
   std::stringstream javadoc_text;
-  re2::StringPiece input(mdtext);
-  re2::StringPiece text;
   bool in_list = false;
-  do {
+  const RE2 markup_expr(
+      "\n+\\*[[:blank:]]+|\n{2,}|`{3,}|`{1,2}|\\*{1,2}\\b|\\[");
+  while (true) {
+    re2::StringPiece text;
     re2::StringPiece markup;
-    FindAndCut(&input,
-        "\n+\\*[[:blank:]]+|\n{2,}|`{3,}|`{1,2}|\\*{1,2}\\b|\\[",
-        &text, &markup);
-    javadoc_text << text;
-    if (markup.empty()) {
-      break;  // we are done parsing
+    if (!FindAndCut(&input, markup_expr, &text, &markup)) {
+      javadoc_text << input;
+      break;  // end of loop
     }
+    javadoc_text << text;
     if (markup.starts_with("\n")) {
       javadoc_text << "\n";
       if (markup.contains("* ")) {
+        // starts a list item
         javadoc_text << (in_list ? "</li>\n" : "<ul>\n") << "<li>\n";
         in_list = true;
       } else if (markup.starts_with("\n\n")) {
         if (in_list) {
+          // ends the current list
           javadoc_text << "</li>\n</ul>\n";
           in_list = false;
         } else if (!input.starts_with("```")) {
+          // starts new paragraph (not required if a <pre> block follows)
           javadoc_text << "<p>\n";
         }
       }
     } else if (markup.starts_with("```") && text.empty()) {
+      // create a multiline code block
       re2::StringPiece language;
       RE2::Consume(&input, "[\\w\\+]+", &language);
       if (FindAndCut(&input, markup.ToString() + "\n*", &text)) {
@@ -245,34 +245,41 @@ string ParseDocumentation(const string& mdtext) {
         javadoc_text << markup << language;
       }
     } else if (markup.starts_with("`")) {
+      // write inlined code
       if (FindAndCut(&input, markup, &text)) {
         javadoc_text << "{@code " << text << "}";
       } else {
         javadoc_text << markup;
       }
     } else if (markup == "**") {
+      // emphase text (strong)
       if (FindAndCut(&input, "\\b\\*{2}", &text)) {
-        javadoc_text << "<b>" << text << "</b>";
+        javadoc_text << "<b>" << ParseDocumentation(text) << "</b>";
       } else {
         javadoc_text << markup;
       }
     } else if (markup == "*") {
+      // emphase text (light)
       if (FindAndCut(&input, "\\b\\*{1}", &text)) {
-        javadoc_text << "<i>" << text << "</i>";
+        javadoc_text << "<i>" << ParseDocumentation(text) << "</i>";
       } else {
         javadoc_text << markup;
       }
     } else if (markup == "[") {
+      // add an external link
       string label;
       string link;
       if (RE2::Consume(&input, "([^\\[]+)\\]\\((http.+)\\)", &label, &link)) {
-        javadoc_text << "<a href=\"" << link << "\">" << label << "</a>";
+        javadoc_text << "<a href=\"" << link << "\">"
+            << ParseDocumentation(label)
+            << "</a>";
       } else {
         javadoc_text << markup;
       }
+    } else {
+      javadoc_text << markup;
     }
-  } while (!input.empty());
-
+  }
   return javadoc_text.str();
 }
 
-- 
GitLab


From 4bfedb4f2edc4bd71984d79145ab6b0293fe8096 Mon Sep 17 00:00:00 2001
From: "karl@kubx.ca" <karl@kubx.ca>
Date: Wed, 25 Apr 2018 00:39:37 -0400
Subject: [PATCH 222/395] Improve again javadoc readability and quality

---
 tensorflow/java/src/gen/cc/op_generator.cc |  7 +--
 tensorflow/java/src/gen/cc/op_specs.cc     | 57 ++++++++++++----------
 2 files changed, 35 insertions(+), 29 deletions(-)

diff --git a/tensorflow/java/src/gen/cc/op_generator.cc b/tensorflow/java/src/gen/cc/op_generator.cc
index c32ad3b109..00f84bc9cd 100644
--- a/tensorflow/java/src/gen/cc/op_generator.cc
+++ b/tensorflow/java/src/gen/cc/op_generator.cc
@@ -305,10 +305,11 @@ void RenderInterfaceImpl(const OpSpec& op, RenderMode mode,
   }
 }
 
-void RenderOptionsClass(const OpSpec& op, SourceWriter* writer) {
+void RenderOptionsClass(const OpSpec& op, const Type& op_class,
+    SourceWriter* writer) {
   Type options_class = Type::Class("Options");
   Javadoc options_doc = Javadoc::Create(
-      "Class holding optional attributes of this operation");
+      "Optional attributes for {@link " + op_class.full_name() + "}");
   writer->BeginInnerType(options_class, PUBLIC | STATIC, &options_doc);
   for (const AttributeSpec& attribute : op.optional_attributes()) {
     Method setter = Method::Create(attribute.var().name(), options_class)
@@ -410,7 +411,7 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
       .EndLine()
       .BeginType(op_class, PUBLIC|FINAL, &dependencies, &op_javadoc);
   if (!op.optional_attributes().empty()) {
-    RenderOptionsClass(op, &writer);
+    RenderOptionsClass(op, op_class, &writer);
   }
   RenderFactoryMethod(op, op_class, &writer);
   RenderGettersAndSetters(op, &writer);
diff --git a/tensorflow/java/src/gen/cc/op_specs.cc b/tensorflow/java/src/gen/cc/op_specs.cc
index 3645fcf836..a0e7a180f2 100644
--- a/tensorflow/java/src/gen/cc/op_specs.cc
+++ b/tensorflow/java/src/gen/cc/op_specs.cc
@@ -204,13 +204,21 @@ bool FindAndCut(re2::StringPiece* input, const RE2& expr,
 }
 
 string ParseDocumentation(re2::StringPiece input) {
+  std::stringstream javadoc_text;
+
   // TODO(karllessard) This is a very minimalist utility method for converting
   // markdown syntax, as found in ops descriptions, to Javadoc/html tags. Check
   // for alternatives to increase the level of support for markups.
-  std::stringstream javadoc_text;
+  std::vector<string> markups_subexpr;
+  markups_subexpr.push_back("\n+\\*\\s+");  // lists
+  markups_subexpr.push_back("\n{2,}");  // paragraphs
+  markups_subexpr.push_back("`{3,}\\s*[^\\s\n]*\\s*\n");  // code blocks
+  markups_subexpr.push_back("`+");  // inlined code and code blocks
+  markups_subexpr.push_back("\\*{1,2}\\b");  // text emphasis
+  markups_subexpr.push_back("\\[");  // hyperlinks
+  const RE2 markup_expr(str_util::Join(markups_subexpr, "|"));
+
   bool in_list = false;
-  const RE2 markup_expr(
-      "\n+\\*[[:blank:]]+|\n{2,}|`{3,}|`{1,2}|\\*{1,2}\\b|\\[");
   while (true) {
     re2::StringPiece text;
     re2::StringPiece markup;
@@ -221,52 +229,48 @@ string ParseDocumentation(re2::StringPiece input) {
     javadoc_text << text;
     if (markup.starts_with("\n")) {
       javadoc_text << "\n";
-      if (markup.contains("* ")) {
-        // starts a list item
+      if (markup.contains("*")) {
+        // new list item
         javadoc_text << (in_list ? "</li>\n" : "<ul>\n") << "<li>\n";
         in_list = true;
-      } else if (markup.starts_with("\n\n")) {
-        if (in_list) {
-          // ends the current list
-          javadoc_text << "</li>\n</ul>\n";
-          in_list = false;
-        } else if (!input.starts_with("```")) {
-          // starts new paragraph (not required if a <pre> block follows)
-          javadoc_text << "<p>\n";
-        }
+      } else if (in_list) {
+        // end of list
+        javadoc_text << "</li>\n</ul>\n";
+        in_list = false;
+      } else if (!input.starts_with("```")) {
+        // new paragraph (not required if a <pre> block follows)
+        javadoc_text << "<p>\n";
       }
-    } else if (markup.starts_with("```") && text.empty()) {
-      // create a multiline code block
-      re2::StringPiece language;
-      RE2::Consume(&input, "[\\w\\+]+", &language);
-      if (FindAndCut(&input, markup.ToString() + "\n*", &text)) {
-        javadoc_text << "<pre>\n{@code" << text << "}\n</pre>\n";
+    } else if (markup.starts_with("```")) {
+      // code blocks
+      if (FindAndCut(&input, "```\\s*\n*", &text)) {
+        javadoc_text << "<pre>{@code\n" << text << "}</pre>\n";
       } else {
-        javadoc_text << markup << language;
+        javadoc_text << markup;
       }
     } else if (markup.starts_with("`")) {
-      // write inlined code
+      // inlined code
       if (FindAndCut(&input, markup, &text)) {
         javadoc_text << "{@code " << text << "}";
       } else {
         javadoc_text << markup;
       }
     } else if (markup == "**") {
-      // emphase text (strong)
+      // text emphasis (strong)
       if (FindAndCut(&input, "\\b\\*{2}", &text)) {
         javadoc_text << "<b>" << ParseDocumentation(text) << "</b>";
       } else {
         javadoc_text << markup;
       }
     } else if (markup == "*") {
-      // emphase text (light)
+      // text emphasis (normal)
       if (FindAndCut(&input, "\\b\\*{1}", &text)) {
         javadoc_text << "<i>" << ParseDocumentation(text) << "</i>";
       } else {
         javadoc_text << markup;
       }
-    } else if (markup == "[") {
-      // add an external link
+    } else if (markup.starts_with("[")) {
+      // hyperlinks
       string label;
       string link;
       if (RE2::Consume(&input, "([^\\[]+)\\]\\((http.+)\\)", &label, &link)) {
@@ -277,6 +281,7 @@ string ParseDocumentation(re2::StringPiece input) {
         javadoc_text << markup;
       }
     } else {
+      // safe fallback
       javadoc_text << markup;
     }
   }
-- 
GitLab


From eac1479f04181fb107c85af29a709eb369831972 Mon Sep 17 00:00:00 2001
From: "karl@kubx.ca" <karl@kubx.ca>
Date: Mon, 30 Apr 2018 07:38:48 -0400
Subject: [PATCH 223/395] Simplify and improve generics handling in generator

---
 tensorflow/java/build_defs.bzl             |   1 +
 tensorflow/java/src/gen/cc/op_gen_main.cc  |   4 +-
 tensorflow/java/src/gen/cc/op_generator.cc | 155 +++++++++------------
 tensorflow/java/src/gen/cc/op_generator.h  |  13 +-
 tensorflow/java/src/gen/cc/op_specs.cc     |  81 ++++++-----
 tensorflow/java/src/gen/cc/op_specs.h      |  16 ++-
 6 files changed, 132 insertions(+), 138 deletions(-)

diff --git a/tensorflow/java/build_defs.bzl b/tensorflow/java/build_defs.bzl
index ab7f60d03d..e1916ca4d9 100644
--- a/tensorflow/java/build_defs.bzl
+++ b/tensorflow/java/build_defs.bzl
@@ -15,6 +15,7 @@ JAVA_VERSION_OPTS = [
 XLINT_OPTS = [
     "-Werror",
     "-Xlint:all",
+    "-Xlint:-processing",
     "-Xlint:-serial",
     "-Xlint:-try",
     "-Xlint:-classfile", # see b/32750402, go/javac-warnings#classfile
diff --git a/tensorflow/java/src/gen/cc/op_gen_main.cc b/tensorflow/java/src/gen/cc/op_gen_main.cc
index 458141b877..a508c96516 100644
--- a/tensorflow/java/src/gen/cc/op_gen_main.cc
+++ b/tensorflow/java/src/gen/cc/op_gen_main.cc
@@ -67,10 +67,10 @@ int main(int argc, char* argv[]) {
   QCHECK(parsed_flags_ok && !output_dir.empty()) << usage;
   std::vector<tensorflow::string> api_dirs = tensorflow::str_util::Split(
       api_dirs_str, ",", tensorflow::str_util::SkipEmpty());
-  tensorflow::java::OpGenerator generator(base_package, output_dir, api_dirs);
+  tensorflow::java::OpGenerator generator(api_dirs);
   tensorflow::OpList ops;
   tensorflow::OpRegistry::Global()->Export(false, &ops);
-  TF_CHECK_OK(generator.Run(ops));
+  TF_CHECK_OK(generator.Run(ops, base_package, output_dir));
 
   return 0;
 }
diff --git a/tensorflow/java/src/gen/cc/op_generator.cc b/tensorflow/java/src/gen/cc/op_generator.cc
index 00f84bc9cd..2327a4daf1 100644
--- a/tensorflow/java/src/gen/cc/op_generator.cc
+++ b/tensorflow/java/src/gen/cc/op_generator.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include <list>
 #include <memory>
 #include <set>
+#include <ctime>
 
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/str_util.h"
@@ -38,23 +39,18 @@ namespace {
 const char* kLicenseSnippet =
     "tensorflow/java/src/gen/resources/license.java.snippet";
 
-const std::map<string, Type> kPrimitiveAttrTypes = {
-  { "Boolean", Type::Boolean() },
-  { "Byte", Type::Byte() },
-  { "Character", Type::Byte() },
-  { "Float", Type::Float() },
-  { "Integer", Type::Long() },
-  { "Long", Type::Long() },
-  { "Short", Type::Long() },
-  { "Double", Type::Float() },
-};
-
 enum RenderMode {
   DEFAULT,
   SINGLE_OUTPUT,
   SINGLE_LIST_OUTPUT
 };
 
+inline void AddArgument(const Variable& var, const string& description,
+    Method* method_out, Javadoc* javadoc_out) {
+  method_out->add_argument(var);
+  javadoc_out->add_param_tag(var.name(), description);
+}
+
 void CollectOpDependencies(const OpSpec& op, RenderMode mode,
     std::list<Type>* out) {
   out->push_back(Type::Class("Operation", "org.tensorflow"));
@@ -81,9 +77,7 @@ void CollectOpDependencies(const OpSpec& op, RenderMode mode,
   }
   for (const AttributeSpec& attribute : op.attributes()) {
     out->push_back(attribute.var().type());
-    if (attribute.var().type().name() == "Class") {
-      out->push_back(Type::Enum("DataType", "org.tensorflow"));
-    }
+    out->push_back(attribute.jni_type());
   }
   for (const AttributeSpec& optional_attribute : op.optional_attributes()) {
     out->push_back(optional_attribute.var().type());
@@ -92,45 +86,38 @@ void CollectOpDependencies(const OpSpec& op, RenderMode mode,
 
 void WriteSetAttrDirective(const AttributeSpec& attr, bool optional,
     SourceWriter* writer) {
-  string var = optional ? "opts." + attr.var().name() : attr.var().name();
+  string var_name = optional ? "opts." + attr.var().name() : attr.var().name();
   if (attr.iterable()) {
-    const Type& type = attr.type();
-    std::map<string, Type>::const_iterator it =
-      kPrimitiveAttrTypes.find(type.name());
-    if (it != kPrimitiveAttrTypes.end()) {
-      string array = attr.var().name() + "Array";
-      writer->AppendType(it->second)
-          .Append("[] " + array + " = new ")
-          .AppendType(it->second)
-          .Append("[" + var + ".size()];")
-          .EndLine();
-      writer->BeginBlock("for (int i = 0; i < " + array + ".length; ++i)")
-          .Append(array + "[i] = " + var + ".get(i);")
-          .EndLine()
-          .EndBlock()
-          .Append("opBuilder.setAttr(\"" + attr.op_def_name() + "\", " + array)
-          .Append(");")
-          .EndLine();
+    string array_name = attr.var().name() + "Array";
+    writer->AppendType(attr.jni_type())
+        .Append("[] " + array_name + " = new ")
+        .AppendType(attr.jni_type())
+        .Append("[" + var_name + ".size()];")
+        .EndLine()
+        .BeginBlock("for (int i = 0; i < " + array_name + ".length; ++i)")
+        .Append(array_name + "[i] = ");
+    if (attr.type().kind() == Type::GENERIC) {
+      writer->Append("DataType.fromClass(" + var_name + ".get(i));");
     } else {
-      writer->Append("opBuilder.setAttr(\"" + attr.op_def_name() + "\", " + var)
-          .Append(".toArray(new ")
-          .AppendType(type)
-          .Append("[" + var + ".size()]));")
-          .EndLine();
+      writer->Append(var_name + ".get(i);");
     }
+    writer->EndLine()
+        .EndBlock()
+        .Append("opBuilder.setAttr(\"" + attr.op_def_name() + "\", ")
+        .Append(array_name + ");")
+        .EndLine();
   } else {
-    Type type = attr.var().type();
     writer->Append("opBuilder.setAttr(\"" + attr.op_def_name() + "\", ");
-    if (type.name() == "Class") {
-      writer->Append("DataType.fromClass(" + attr.var().name() + "));");
+    if (attr.var().type().name() == "Class") {
+      writer->Append("DataType.fromClass(" + var_name + "));");
     } else {
-      writer->Append(var + ");");
+      writer->Append(var_name + ");");
     }
     writer->EndLine();
   }
 }
 
-void RenderFactoryMethod(const OpSpec& op, const Type& op_class,
+void RenderFactoryMethods(const OpSpec& op, const Type& op_class,
     SourceWriter* writer) {
   Method factory = Method::Create("create", op_class);
   Javadoc factory_doc = Javadoc::Create(
@@ -138,27 +125,24 @@ void RenderFactoryMethod(const OpSpec& op, const Type& op_class,
       + " operation to the graph.");
   Variable scope =
       Variable::Create("scope", Type::Class("Scope", "org.tensorflow.op"));
-  factory.add_argument(scope);
-  factory_doc.add_param_tag(scope.name(), "Current graph scope");
+  AddArgument(scope, "current graph scope", &factory, &factory_doc);
   for (const ArgumentSpec& input : op.inputs()) {
-    factory.add_argument(input.var());
-    factory_doc.add_param_tag(input.var().name(), input.description());
+    AddArgument(input.var(), input.description(), &factory, &factory_doc);
   }
-  for (const AttributeSpec& attribute : op.attributes()) {
-    factory.add_argument(attribute.var());
-    factory_doc.add_param_tag(attribute.var().name(), attribute.description());
+  for (const AttributeSpec& attr : op.attributes()) {
+    AddArgument(attr.var(), attr.description(), &factory, &factory_doc);
   }
   if (!op.optional_attributes().empty()) {
-    factory.add_argument(Variable::Varargs("options", Type::Class("Options")));
-    factory_doc.add_param_tag("options", "carries optional attributes values");
+    AddArgument(Variable::Varargs("options", Type::Class("Options")),
+        "carries optional attributes values", &factory, &factory_doc);
   }
   factory_doc.add_tag("return", "a new instance of " + op_class.name());
+
   writer->BeginMethod(factory, PUBLIC|STATIC, &factory_doc);
   writer->Append("OperationBuilder opBuilder = scope.graph().opBuilder(\""
       + op.graph_op_name() + "\", scope.makeOpName(\""
       + op_class.name() + "\"));");
   writer->EndLine();
-
   for (const ArgumentSpec& input : op.inputs()) {
     if (input.iterable()) {
       writer->Append("opBuilder.addInputList(Operands.asOutputs("
@@ -192,10 +176,9 @@ void RenderFactoryMethod(const OpSpec& op, const Type& op_class,
 
 void RenderConstructor(const OpSpec& op, const Type& op_class,
     SourceWriter* writer) {
-  Method constructor = Method::ConstructorFor(op_class)
-    .add_argument(
-        Variable::Create("operation",
-            Type::Class("Operation", "org.tensorflow")));
+  Variable operation =
+      Variable::Create("operation", Type::Class("Operation", "org.tensorflow"));
+  Method constructor = Method::ConstructorFor(op_class).add_argument(operation);
   for (const ArgumentSpec& output : op.outputs()) {
     if (output.iterable() && !output.type().unknown()) {
       constructor.add_annotation(
@@ -237,15 +220,14 @@ void RenderConstructor(const OpSpec& op, const Type& op_class,
 }
 
 void RenderGettersAndSetters(const OpSpec& op, SourceWriter* writer) {
-  for (const AttributeSpec& attribute : op.optional_attributes()) {
+  for (const AttributeSpec& attr : op.optional_attributes()) {
     Method setter =
-        Method::Create(attribute.var().name(), Type::Class("Options"))
-            .add_argument(attribute.var());
-    Javadoc setter_doc = Javadoc::Create()
-        .add_param_tag(attribute.var().name(), attribute.description());
+        Method::Create(attr.var().name(), Type::Class("Options"));
+    Javadoc setter_doc = Javadoc::Create();
+    AddArgument(attr.var(), attr.description(), &setter, &setter_doc);
     writer->BeginMethod(setter, PUBLIC|STATIC, &setter_doc)
-        .Append("return new Options()." + attribute.var().name() + "("
-            + attribute.var().name() + ");")
+        .Append("return new Options()." + attr.var().name() + "("
+            + attr.var().name() + ");")
         .EndLine()
         .EndMethod();
   }
@@ -311,14 +293,12 @@ void RenderOptionsClass(const OpSpec& op, const Type& op_class,
   Javadoc options_doc = Javadoc::Create(
       "Optional attributes for {@link " + op_class.full_name() + "}");
   writer->BeginInnerType(options_class, PUBLIC | STATIC, &options_doc);
-  for (const AttributeSpec& attribute : op.optional_attributes()) {
-    Method setter = Method::Create(attribute.var().name(), options_class)
-        .add_argument(attribute.var());
-    Javadoc setter_doc = Javadoc::Create()
-        .add_param_tag(attribute.var().name(), attribute.description());
+  for (const AttributeSpec& attr : op.optional_attributes()) {
+    Method setter = Method::Create(attr.var().name(), options_class);
+    Javadoc setter_doc = Javadoc::Create();
+    AddArgument(attr.var(), attr.description(), &setter, &setter_doc);
     writer->BeginMethod(setter, PUBLIC, &setter_doc)
-        .Append("this." + attribute.var().name() + " = "
-            + attribute.var().name() + ";")
+        .Append("this." + attr.var().name() + " = " + attr.var().name() + ";")
         .EndLine()
         .Append("return this;")
         .EndLine()
@@ -339,12 +319,13 @@ inline Type ClassOf(const EndpointSpec& endpoint, const string& base_package) {
 }
 
 void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
-    const string& base_package, const string& output_dir, Env* env) {
+    const string& base_package, const string& output_dir, Env* env,
+    const std::tm* timestamp) {
   Type op_class(ClassOf(endpoint, base_package)
       .add_supertype(Type::Class("PrimitiveOp", "org.tensorflow.op")));
   Javadoc op_javadoc(endpoint.javadoc());
 
-  // implement Operand (or Iterable<Operand>) if the op has only one output
+  // op interfaces
   RenderMode mode = DEFAULT;
   if (op.outputs().size() == 1) {
     const ArgumentSpec& output = op.outputs().front();
@@ -360,18 +341,22 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
       op_class.add_supertype(operand_inf);
     }
   }
-  // declare all outputs generics at the op class level
+  // op generic parameters
   std::set<string> generics;
   for (const ArgumentSpec& output : op.outputs()) {
     if (output.type().kind() == Type::GENERIC && !output.type().unknown()
         && generics.find(output.type().name()) == generics.end()) {
       op_class.add_parameter(output.type());
       op_javadoc.add_param_tag("<" + output.type().name() + ">",
-          "data type of output {@code " + output.var().name() + "}");
+          "data type for {@code " + output.var().name() + "()} output");
       generics.insert(output.type().name());
     }
   }
-  // handle endpoint deprecation
+  // op annotations
+  char date[20];
+  strftime(date, sizeof date, "%FT%TZ", timestamp);
+  op_class.add_annotation(Annotation::Create("Generated", "javax.annotation")
+      .attributes(string("value = \"op_generator\", date = \"") + date + "\""));
   if (endpoint.deprecated()) {
     op_class.add_annotation(Annotation::Create("Deprecated"));
     string explanation;
@@ -384,8 +369,8 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
     }
     op_javadoc.add_tag("deprecated", explanation);
   }
-  // expose the op in the Ops Graph API only if it is visible
   if (!op.hidden()) {
+    // expose the op in the Ops Graph API only if it is visible
     op_class.add_annotation(
         Annotation::Create("Operator", "org.tensorflow.op.annotation")
           .attributes("group = \"" + endpoint.package() + "\""));
@@ -405,15 +390,12 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
   std::list<Type> dependencies;
   CollectOpDependencies(op, mode, &dependencies);
   writer.WriteFromFile(kLicenseSnippet)
-      .EndLine()
-      .Append("// This file is machine generated, DO NOT EDIT!")
-      .EndLine()
       .EndLine()
       .BeginType(op_class, PUBLIC|FINAL, &dependencies, &op_javadoc);
   if (!op.optional_attributes().empty()) {
     RenderOptionsClass(op, op_class, &writer);
   }
-  RenderFactoryMethod(op, op_class, &writer);
+  RenderFactoryMethods(op, op_class, &writer);
   RenderGettersAndSetters(op, &writer);
   if (mode != DEFAULT) {
     RenderInterfaceImpl(op, mode, &writer);
@@ -428,13 +410,8 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
 
 }  // namespace
 
-OpGenerator::OpGenerator(const string& base_package, const string& output_dir,
-    const std::vector<string>& api_dirs, Env* env)
-  : base_package_(base_package), output_dir_(output_dir), api_dirs_(api_dirs),
-    env_(env) {
-}
-
-Status OpGenerator::Run(const OpList& op_list) {
+Status OpGenerator::Run(const OpList& op_list, const string& base_package,
+    const string& output_dir) {
   ApiDefMap api_map(op_list);
   if (!api_dirs_.empty()) {
     // Only load api files that correspond to the requested "op_list"
@@ -449,12 +426,14 @@ Status OpGenerator::Run(const OpList& op_list) {
     }
   }
   api_map.UpdateDocs();
+  time_t now;
+  time(&now);
   for (const auto& op_def : op_list.op()) {
     const ApiDef* api_def = api_map.GetApiDef(op_def.name());
     if (api_def->visibility() != ApiDef::SKIP) {
       OpSpec op(OpSpec::Create(op_def, *api_def));
       for (const EndpointSpec& endpoint : op.endpoints()) {
-        GenerateOp(op, endpoint, base_package_, output_dir_, env_);
+        GenerateOp(op, endpoint, base_package, output_dir, env_, gmtime(&now));
       }
     }
   }
diff --git a/tensorflow/java/src/gen/cc/op_generator.h b/tensorflow/java/src/gen/cc/op_generator.h
index 06b08e852a..b789e11fa9 100644
--- a/tensorflow/java/src/gen/cc/op_generator.h
+++ b/tensorflow/java/src/gen/cc/op_generator.h
@@ -36,18 +36,17 @@ namespace java {
 // ops definitions.
 class OpGenerator {
  public:
-  OpGenerator(const string& base_package, const string& output_dir,
-      const std::vector<string>& api_dirs, Env* env = Env::Default());
+  explicit OpGenerator(const std::vector<string>& api_dirs,
+      Env* env = Env::Default()) : api_dirs_(api_dirs), env_(env) {}
 
   // Generates wrappers for the given list of 'ops'.
   //
-  // Output files are generated in <output_dir>/<base_package>/<lib_package>,
-  // where 'lib_package' is derived from ops endpoints.
-  Status Run(const OpList& op_list);
+  // Output files are generated in <output_dir>/<base_package>/<op_package>,
+  // where 'op_package' is derived from ops endpoints.
+  Status Run(const OpList& op_list, const string& base_package,
+      const string& output_dir);
 
  private:
-  const string base_package_;
-  const string output_dir_;
   const std::vector<string> api_dirs_;
   Env* env_;
 };
diff --git a/tensorflow/java/src/gen/cc/op_specs.cc b/tensorflow/java/src/gen/cc/op_specs.cc
index a0e7a180f2..dcc6388614 100644
--- a/tensorflow/java/src/gen/cc/op_specs.cc
+++ b/tensorflow/java/src/gen/cc/op_specs.cc
@@ -46,14 +46,30 @@ class TypeResolver {
   explicit TypeResolver(const OpDef& op_def) : op_def_(op_def) {}
 
   Type TypeOf(const OpDef_ArgDef& arg_def, bool *iterable_out);
-  Type TypeOf(const OpDef_AttrDef& attr_def, bool *iterable_out);
+  std::pair<Type, Type> TypeOf(const OpDef_AttrDef& attr_def,
+      bool *iterable_out);
   bool IsAttributeVisited(const string& attr_name) {
     return visited_attrs_.find(attr_name) != visited_attrs_.cend();
   }
+
  private:
   const OpDef op_def_;
   std::map<std::string, Type> visited_attrs_;
-  char next_generic_ = 'T';
+  char next_generic_letter_ = 'T';
+
+  std::pair<Type, Type> MakeTypePair(const Type& type, const Type& jni_type) {
+    return std::make_pair(type, jni_type);
+  }
+  std::pair<Type, Type> MakeTypePair(const Type& type) {
+    return std::make_pair(type, type);
+  }
+  Type NextGeneric() {
+    char generic_letter = next_generic_letter_++;
+    if (next_generic_letter_ > 'Z') {
+      next_generic_letter_ = 'A';
+    }
+    return Type::Generic(string(1, generic_letter));
+  }
 };
 
 Type TypeResolver::TypeOf(const OpDef_ArgDef& arg_def,
@@ -107,7 +123,7 @@ Type TypeResolver::TypeOf(const OpDef_ArgDef& arg_def,
     } else {
       for (const auto& attr_def : op_def_.attr()) {
         if (attr_def.name() == arg_def.type_attr()) {
-          type = TypeOf(attr_def, iterable_out);
+          type = TypeOf(attr_def, iterable_out).first;
           break;
         }
       }
@@ -125,51 +141,47 @@ Type TypeResolver::TypeOf(const OpDef_ArgDef& arg_def,
   return type;
 }
 
-Type TypeResolver::TypeOf(const OpDef_AttrDef& attr_def,
+std::pair<Type, Type> TypeResolver::TypeOf(const OpDef_AttrDef& attr_def,
     bool* iterable_out) {
+  std::pair<Type, Type> types = MakeTypePair(Type::Wildcard());
   *iterable_out = false;
   StringPiece attr_type = attr_def.type();
   if (str_util::ConsumePrefix(&attr_type, "list(")) {
     attr_type.remove_suffix(1);  // remove closing brace
     *iterable_out = true;
   }
-  Type type = *iterable_out ? Type::Wildcard() : Type::Class("Object");
-  if (attr_type == "type") {
-    if (*iterable_out) {
-      type = Type::Enum("DataType", "org.tensorflow");
-    } else {
-      type = Type::Generic(string(1, next_generic_));
-      next_generic_ = (next_generic_ == 'Z') ? 'A' : next_generic_ + 1;
-      if (IsRealNumbers(attr_def.allowed_values())) {
-        // enforce real numbers datasets by extending java.lang.Number
-        type.add_supertype(Type::Class("Number"));
-      }
-    }
-  } else if (attr_type == "string") {
-    type = Type::Class("String");
+  if (attr_type == "string") {
+    types = MakeTypePair(Type::Class("String"));
 
   } else if (attr_type == "int") {
-    type = Type::Class("Integer");
+    types = MakeTypePair(Type::Class("Long"), Type::Long());
 
   } else if (attr_type == "float") {
-    type = Type::Class("Float");
+    types = MakeTypePair(Type::Class("Float"), Type::Float());
 
   } else if (attr_type == "bool") {
-    type = Type::Class("Boolean");
+    types = MakeTypePair(Type::Class("Boolean"), Type::Boolean());
 
   } else if (attr_type == "shape") {
-    type = Type::Class("Shape", "org.tensorflow");
+    types = MakeTypePair(Type::Class("Shape", "org.tensorflow"));
 
   } else if (attr_type == "tensor") {
-    type = Type::Class("Tensor", "org.tensorflow")
-        .add_parameter(Type::Wildcard());
+    types = MakeTypePair(Type::Class("Tensor", "org.tensorflow")
+        .add_parameter(Type::Wildcard()));
+
+  } else if (attr_type == "type") {
+    Type type = *iterable_out ? Type::Wildcard() : NextGeneric();
+    if (IsRealNumbers(attr_def.allowed_values())) {
+      type.add_supertype(Type::Class("Number"));
+    }
+    types = MakeTypePair(type, Type::Enum("DataType", "org.tensorflow"));
 
   } else {
     LOG(FATAL) << "Cannot resolve data type for attribute \"" << attr_type
         << "\" in operation \"" << op_def_.name() << "\"";
   }
-  visited_attrs_.insert(std::make_pair(attr_def.name(), type));
-  return type;
+  visited_attrs_.insert(std::make_pair(attr_def.name(), types.first));
+  return types;
 }
 
 string SnakeToCamelCase(const string& str, bool upper = false) {
@@ -307,19 +319,19 @@ ArgumentSpec CreateInput(const OpDef_ArgDef& input_def,
 AttributeSpec CreateAttribute(const OpDef_AttrDef& attr_def,
     const ApiDef::Attr& attr_api_def, TypeResolver* type_resolver) {
   bool iterable = false;
-  Type type = type_resolver->TypeOf(attr_def, &iterable);
-  // type attributes must be passed explicitly in methods as a Class<> parameter
-  bool is_explicit = type.kind() == Type::GENERIC && !iterable;
-  Type var_type = is_explicit ? Type::Class("Class").add_parameter(type) : type;
+  std::pair<Type, Type> types = type_resolver->TypeOf(attr_def, &iterable);
+  Type var_type = types.first.kind() == Type::GENERIC ?
+      Type::Class("Class").add_parameter(types.first) : types.first;
   if (iterable) {
-    var_type = Type::ListOf(type);
+    var_type = Type::ListOf(var_type);
   }
   return AttributeSpec(attr_api_def.name(),
       Variable::Create(SnakeToCamelCase(attr_api_def.rename_to()), var_type),
-      type,
+      types.first,
+      types.second,
       ParseDocumentation(attr_api_def.description()),
       iterable,
-      attr_api_def.has_default_value() && !is_explicit);
+      attr_api_def.has_default_value());
 }
 
 ArgumentSpec CreateOutput(const OpDef_ArgDef& output_def,
@@ -340,7 +352,6 @@ ArgumentSpec CreateOutput(const OpDef_ArgDef& output_def,
 
 EndpointSpec CreateEndpoint(const OpDef& op_def, const ApiDef& api_def,
     const ApiDef_Endpoint& endpoint_def) {
-
   std::vector<string> name_tokens = str_util::Split(endpoint_def.name(), ".");
   string package;
   string name;
@@ -381,7 +392,7 @@ OpSpec OpSpec::Create(const OpDef& op_def, const ApiDef& api_def) {
       AttributeSpec attr = CreateAttribute(op_def.attr(i), api_def.attr(i),
           &type_resolver);
       // attributes with a default value are optional
-      if (attr.optional()) {
+      if (attr.has_default_value() && attr.type().kind() != Type::GENERIC) {
         op.optional_attributes_.push_back(attr);
       } else {
         op.attributes_.push_back(attr);
diff --git a/tensorflow/java/src/gen/cc/op_specs.h b/tensorflow/java/src/gen/cc/op_specs.h
index 55c2c3f307..7d64391446 100644
--- a/tensorflow/java/src/gen/cc/op_specs.h
+++ b/tensorflow/java/src/gen/cc/op_specs.h
@@ -21,6 +21,7 @@ limitations under the License.
 
 #include "tensorflow/core/framework/op_def.pb.h"
 #include "tensorflow/core/framework/api_def.pb.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/java/src/gen/cc/java_defs.h"
 
 namespace tensorflow {
@@ -87,20 +88,23 @@ class AttributeSpec : public ArgumentSpec {
   // op_def_name: attribute name, as known by TensorFlow core
   // var: a variable to represent this attribute in Java
   // type: the type of this attribute
+  // jni_type: the type of this attribute in JNI layer (see OperationBuilder)
   // description: a description of this attribute, in javadoc
   // iterable: true if this attribute is a list
-  // optional: true if this attribute does not require to be set explicitly
+  // has_default_value: true if this attribute has a default value if not set
   AttributeSpec(const string& op_def_name, const Variable& var,
-      const Type& type, const string& description, bool iterable,
-      bool optional)
+      const Type& type, const Type& jni_type, const string& description,
+      bool iterable, bool has_default_value)
     : ArgumentSpec(op_def_name, var, type, description, iterable),
-      optional_(optional) {}
+      jni_type_(jni_type), has_default_value_(has_default_value) {}
   virtual ~AttributeSpec() = default;
 
-  bool optional() const { return optional_; }
+  const Type& jni_type() const { return jni_type_; }
+  bool has_default_value() const { return has_default_value_; }
 
  private:
-  const bool optional_;
+  const Type jni_type_;
+  const bool has_default_value_;
 };
 
 class OpSpec {
-- 
GitLab


From dd1ef8fa8f6861e53e8a7953c171b3e9253043ed Mon Sep 17 00:00:00 2001
From: "karl@kubx.ca" <karl@kubx.ca>
Date: Thu, 3 May 2018 22:39:35 -0400
Subject: [PATCH 224/395] Second code review

---
 tensorflow/core/api_def/BUILD                 |  7 ++
 .../java_api/api_def_FilterDataset.pbtxt      |  4 +
 .../java_api/api_def_FlatMapDataset.pbtxt     |  4 +
 .../core/api_def/java_api/api_def_For.pbtxt   |  4 +
 .../java_api/api_def_GeneratorDataset.pbtxt   |  4 +
 .../api_def_GroupByWindowDataset.pbtxt        |  4 +
 .../core/api_def/java_api/api_def_If.pbtxt    |  4 +
 .../java_api/api_def_InterleaveDataset.pbtxt  |  4 +
 .../java_api/api_def_MapAndBatchDataset.pbtxt |  4 +
 .../api_def/java_api/api_def_MapDataset.pbtxt |  4 +
 .../java_api/api_def_OneShotIterator.pbtxt    |  4 +
 .../api_def_ParallelInterleaveDataset.pbtxt   |  4 +
 .../java_api/api_def_ParallelMapDataset.pbtxt |  4 +
 .../api_def/java_api/api_def_RemoteCall.pbtxt |  4 +
 .../java_api/api_def_ScanDataset.pbtxt        |  4 +
 .../java_api/api_def_SymbolicGradient.pbtxt   |  4 +
 .../core/api_def/java_api/api_def_While.pbtxt |  4 +
 tensorflow/java/BUILD                         | 39 ++++------
 tensorflow/java/src/gen/cc/java_defs.h        |  6 +-
 tensorflow/java/src/gen/cc/op_gen_main.cc     |  2 +-
 tensorflow/java/src/gen/cc/op_generator.cc    | 77 +++++++++++--------
 tensorflow/java/src/gen/cc/op_generator.h     |  2 +-
 tensorflow/java/src/gen/cc/op_specs.cc        | 25 +++++-
 tensorflow/java/src/gen/cc/op_specs.h         | 17 +++-
 tensorflow/java/src/gen/cc/source_writer.cc   | 20 +++--
 tensorflow/java/src/gen/cc/source_writer.h    |  2 +-
 .../java/src/gen/cc/source_writer_test.cc     |  2 +-
 tensorflow/java/src/gen/gen_ops.bzl           | 41 +++-------
 28 files changed, 195 insertions(+), 109 deletions(-)
 create mode 100644 tensorflow/core/api_def/java_api/api_def_FilterDataset.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_FlatMapDataset.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_For.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_GeneratorDataset.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_GroupByWindowDataset.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_If.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_InterleaveDataset.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_MapAndBatchDataset.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_MapDataset.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_OneShotIterator.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_ParallelInterleaveDataset.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_ParallelMapDataset.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_RemoteCall.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_ScanDataset.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_SymbolicGradient.pbtxt
 create mode 100644 tensorflow/core/api_def/java_api/api_def_While.pbtxt

diff --git a/tensorflow/core/api_def/BUILD b/tensorflow/core/api_def/BUILD
index 19d6438809..06b797e32e 100644
--- a/tensorflow/core/api_def/BUILD
+++ b/tensorflow/core/api_def/BUILD
@@ -4,6 +4,7 @@
 # The following targets can be used to access ApiDefs:
 #   :base_api_def
 #   :python_api_def
+#   :java_api_def
 
 package(
     default_visibility = ["//visibility:private"],
@@ -29,6 +30,12 @@ filegroup(
     visibility = ["//tensorflow:internal"],
 )
 
+filegroup(
+    name = "java_api_def",
+    srcs = glob(["java_api/*"]),
+    visibility = ["//tensorflow:internal"],
+)
+
 cc_library(
     name = "excluded_ops_lib",
     srcs = ["excluded_ops.cc"],
diff --git a/tensorflow/core/api_def/java_api/api_def_FilterDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_FilterDataset.pbtxt
new file mode 100644
index 0000000000..debd7e5709
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_FilterDataset.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "FilterDataset"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_FlatMapDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_FlatMapDataset.pbtxt
new file mode 100644
index 0000000000..329ab15ef5
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_FlatMapDataset.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "FlatMapDataset"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_For.pbtxt b/tensorflow/core/api_def/java_api/api_def_For.pbtxt
new file mode 100644
index 0000000000..caabc947bb
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_For.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "For"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_GeneratorDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_GeneratorDataset.pbtxt
new file mode 100644
index 0000000000..a6e5167c30
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_GeneratorDataset.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "GeneratorDataset"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_GroupByWindowDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_GroupByWindowDataset.pbtxt
new file mode 100644
index 0000000000..4c0b2084a8
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_GroupByWindowDataset.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "GroupByWindowDataset"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_If.pbtxt b/tensorflow/core/api_def/java_api/api_def_If.pbtxt
new file mode 100644
index 0000000000..13b8635ca7
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_If.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "If"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_InterleaveDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_InterleaveDataset.pbtxt
new file mode 100644
index 0000000000..ed748d4d2a
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_InterleaveDataset.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "InterleaveDataset"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_MapAndBatchDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_MapAndBatchDataset.pbtxt
new file mode 100644
index 0000000000..cb96bf63d8
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_MapAndBatchDataset.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "MapAndBatchDataset"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_MapDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_MapDataset.pbtxt
new file mode 100644
index 0000000000..e0ab8dd9db
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_MapDataset.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "MapDataset"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_OneShotIterator.pbtxt b/tensorflow/core/api_def/java_api/api_def_OneShotIterator.pbtxt
new file mode 100644
index 0000000000..13130e6882
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_OneShotIterator.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "OneShotIterator"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_ParallelInterleaveDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_ParallelInterleaveDataset.pbtxt
new file mode 100644
index 0000000000..6a985d24fa
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_ParallelInterleaveDataset.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "ParallelInterleaveDataset"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_ParallelMapDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_ParallelMapDataset.pbtxt
new file mode 100644
index 0000000000..64f25b9e5e
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_ParallelMapDataset.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "ParallelMapDataset"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_RemoteCall.pbtxt b/tensorflow/core/api_def/java_api/api_def_RemoteCall.pbtxt
new file mode 100644
index 0000000000..2ccb5c8cf3
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_RemoteCall.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "RemoteCall"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_ScanDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_ScanDataset.pbtxt
new file mode 100644
index 0000000000..3463e60049
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_ScanDataset.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "ScanDataset"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_SymbolicGradient.pbtxt b/tensorflow/core/api_def/java_api/api_def_SymbolicGradient.pbtxt
new file mode 100644
index 0000000000..88c3acea74
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_SymbolicGradient.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "SymbolicGradient"
+  visibility: SKIP
+}
diff --git a/tensorflow/core/api_def/java_api/api_def_While.pbtxt b/tensorflow/core/api_def/java_api/api_def_While.pbtxt
new file mode 100644
index 0000000000..33756682c3
--- /dev/null
+++ b/tensorflow/core/api_def/java_api/api_def_While.pbtxt
@@ -0,0 +1,4 @@
+op {
+  graph_op_name: "While"
+  visibility: SKIP
+}
\ No newline at end of file
diff --git a/tensorflow/java/BUILD b/tensorflow/java/BUILD
index 17566e1a9c..7cd0208dbf 100644
--- a/tensorflow/java/BUILD
+++ b/tensorflow/java/BUILD
@@ -68,34 +68,27 @@ filegroup(
     ],
 )
 
-# Build the gen tool as a library, as it will be linked to a core/ops binary
-# files before making it an executable.
 tf_java_op_gen_srcjar(
     name = "java_op_gen_sources",
     api_def_srcs = [
         "//tensorflow/core/api_def:base_api_def",
+        "//tensorflow/core/api_def:java_api_def",
     ],
-    gen_base_package = "org.tensorflow.op",
-    gen_tool = "java_op_gen_tool",
-    ops_libs = [
-        "array_ops",
-        "candidate_sampling_ops",
-        "control_flow_ops",
-        "data_flow_ops",
-        "image_ops",
-        "io_ops",
-        "linalg_ops",
-        "logging_ops",
-        "math_ops",
-        "nn_ops",
-        "no_op",
-        "parsing_ops",
-        "random_ops",
-        "sparse_ops",
-        "state_ops",
-        "string_ops",
-        "training_ops",
-        "user_ops",
+    base_package = "org.tensorflow.op",
+    gen_tool = ":java_op_gen_tool",
+)
+
+tf_cc_binary(
+    name = "java_op_gen_tool",
+    srcs = [
+        "src/gen/cc/op_gen_main.cc",
+    ],
+    copts = tf_copts(),
+    linkopts = ["-lm"],
+    linkstatic = 1,
+    deps = [
+        ":java_op_gen_lib",
+        "//tensorflow/core:ops",
     ],
 )
 
diff --git a/tensorflow/java/src/gen/cc/java_defs.h b/tensorflow/java/src/gen/cc/java_defs.h
index 81ac67eb2f..62575f6683 100644
--- a/tensorflow/java/src/gen/cc/java_defs.h
+++ b/tensorflow/java/src/gen/cc/java_defs.h
@@ -1,4 +1,4 @@
-/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+/* 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.
@@ -102,10 +102,10 @@ class Type {
   const Kind& kind() const { return kind_; }
   const string& name() const { return name_; }
   const string& package() const { return package_; }
-  const string full_name() const {
+  const string canonical_name() const {
     return package_.empty() ? name_ : package_ + "." + name_;
   }
-  bool unknown() const { return name_.empty(); }  // only wildcards has no name
+  bool wildcard() const { return name_.empty(); }  // only wildcards has no name
   const std::list<Type>& parameters() const { return parameters_; }
   Type& add_parameter(const Type& parameter) {
     parameters_.push_back(parameter);
diff --git a/tensorflow/java/src/gen/cc/op_gen_main.cc b/tensorflow/java/src/gen/cc/op_gen_main.cc
index a508c96516..6c35cd9595 100644
--- a/tensorflow/java/src/gen/cc/op_gen_main.cc
+++ b/tensorflow/java/src/gen/cc/op_gen_main.cc
@@ -1,4 +1,4 @@
-/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+/* 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.
diff --git a/tensorflow/java/src/gen/cc/op_generator.cc b/tensorflow/java/src/gen/cc/op_generator.cc
index 2327a4daf1..7355b3a395 100644
--- a/tensorflow/java/src/gen/cc/op_generator.cc
+++ b/tensorflow/java/src/gen/cc/op_generator.cc
@@ -1,4 +1,4 @@
-/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+/* 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.
@@ -19,7 +19,6 @@ limitations under the License.
 #include <list>
 #include <memory>
 #include <set>
-#include <ctime>
 
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/str_util.h"
@@ -39,13 +38,26 @@ namespace {
 const char* kLicenseSnippet =
     "tensorflow/java/src/gen/resources/license.java.snippet";
 
+// There is three different modes to render an op class, depending on the
+// number and type of outputs it has:
+//
+// DEFAULT: This mode does not provide any specialization for the op class, it
+//          is applied when the operation does not comply with any other mode
+//
+// OPERAND: The op class implements the Operand<T> interface, allowing an
+//          instance to be passed directly in input to another operation
+//
+// LIST_OPERAND: The op class implements the Iterable<Operand<T>> interface,
+//          allowing an instance to be passed directly as a list input to
+//          another operation
+//
 enum RenderMode {
   DEFAULT,
-  SINGLE_OUTPUT,
-  SINGLE_LIST_OUTPUT
+  OPERAND,
+  LIST_OPERAND
 };
 
-inline void AddArgument(const Variable& var, const string& description,
+void AddArgument(const Variable& var, const string& description,
     Method* method_out, Javadoc* javadoc_out) {
   method_out->add_argument(var);
   javadoc_out->add_param_tag(var.name(), description);
@@ -56,9 +68,9 @@ void CollectOpDependencies(const OpSpec& op, RenderMode mode,
   out->push_back(Type::Class("Operation", "org.tensorflow"));
   out->push_back(Type::Class("OperationBuilder", "org.tensorflow"));
   out->push_back(Type::Class("Scope", "org.tensorflow.op"));
-  if (mode == SINGLE_OUTPUT) {
+  if (mode == OPERAND) {
     out->push_back(Type::Class("Output", "org.tensorflow"));
-  } else if (mode == SINGLE_LIST_OUTPUT) {
+  } else if (mode == LIST_OPERAND) {
     out->push_back(Type::Interface("Iterator", "java.util"));
   }
   // Don't pay attention to duplicate types in the dependency list, they will
@@ -180,7 +192,7 @@ void RenderConstructor(const OpSpec& op, const Type& op_class,
       Variable::Create("operation", Type::Class("Operation", "org.tensorflow"));
   Method constructor = Method::ConstructorFor(op_class).add_argument(operation);
   for (const ArgumentSpec& output : op.outputs()) {
-    if (output.iterable() && !output.type().unknown()) {
+    if (output.iterable() && !output.type().wildcard()) {
       constructor.add_annotation(
           Annotation::Create("SuppressWarnings").attributes("\"unchecked\""));
       break;
@@ -200,7 +212,7 @@ void RenderConstructor(const OpSpec& op, const Type& op_class,
                 + "\");")
             .EndLine()
             .Append(output.var().name() + " = Arrays.asList(");
-        if (!output.type().unknown()) {
+        if (!output.type().wildcard()) {
           writer->Append("(")
               .AppendType(output.var().type().parameters().front())
               .Append("[])");
@@ -245,8 +257,8 @@ void RenderInterfaceImpl(const OpSpec& op, RenderMode mode,
     SourceWriter* writer) {
   ArgumentSpec output = op.outputs().front();
 
-  if (mode == SINGLE_OUTPUT) {
-    bool cast2obj = output.type().unknown();
+  if (mode == OPERAND) {
+    bool cast2obj = output.type().wildcard();
     Type return_type = Type::Class("Output", "org.tensorflow")
         .add_parameter(cast2obj ? Type::Class("Object") : output.type());
     Method as_output = Method::Create("asOutput", return_type)
@@ -265,9 +277,9 @@ void RenderInterfaceImpl(const OpSpec& op, RenderMode mode,
         .EndLine()
         .EndMethod();
 
-  } else if (mode == SINGLE_LIST_OUTPUT) {
+  } else if (mode == LIST_OPERAND) {
     Type operand = Type::Interface("Operand", "org.tensorflow");
-    if (output.type().unknown()) {
+    if (output.type().wildcard()) {
       operand.add_parameter(Type::Class("Object"));
     } else {
       operand.add_parameter(output.type());
@@ -291,7 +303,7 @@ void RenderOptionsClass(const OpSpec& op, const Type& op_class,
     SourceWriter* writer) {
   Type options_class = Type::Class("Options");
   Javadoc options_doc = Javadoc::Create(
-      "Optional attributes for {@link " + op_class.full_name() + "}");
+      "Optional attributes for {@link " + op_class.canonical_name() + "}");
   writer->BeginInnerType(options_class, PUBLIC | STATIC, &options_doc);
   for (const AttributeSpec& attr : op.optional_attributes()) {
     Method setter = Method::Create(attr.var().name(), options_class);
@@ -319,8 +331,7 @@ inline Type ClassOf(const EndpointSpec& endpoint, const string& base_package) {
 }
 
 void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
-    const string& base_package, const string& output_dir, Env* env,
-    const std::tm* timestamp) {
+    const string& base_package, const string& output_dir, Env* env) {
   Type op_class(ClassOf(endpoint, base_package)
       .add_supertype(Type::Class("PrimitiveOp", "org.tensorflow.op")));
   Javadoc op_javadoc(endpoint.javadoc());
@@ -329,22 +340,22 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
   RenderMode mode = DEFAULT;
   if (op.outputs().size() == 1) {
     const ArgumentSpec& output = op.outputs().front();
-    Type operand_type(output.type().unknown() ?
+    Type operand_type(output.type().wildcard() ?
         Type::Class("Object") : output.type());
     Type operand_inf(Type::Interface("Operand", "org.tensorflow")
         .add_parameter(operand_type));
     if (output.iterable()) {
-      mode = SINGLE_LIST_OUTPUT;
+      mode = LIST_OPERAND;
       op_class.add_supertype(Type::IterableOf(operand_inf));
     } else {
-      mode = SINGLE_OUTPUT;
+      mode = OPERAND;
       op_class.add_supertype(operand_inf);
     }
   }
   // op generic parameters
   std::set<string> generics;
   for (const ArgumentSpec& output : op.outputs()) {
-    if (output.type().kind() == Type::GENERIC && !output.type().unknown()
+    if (output.type().kind() == Type::GENERIC && !output.type().wildcard()
         && generics.find(output.type().name()) == generics.end()) {
       op_class.add_parameter(output.type());
       op_javadoc.add_param_tag("<" + output.type().name() + ">",
@@ -353,16 +364,15 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
     }
   }
   // op annotations
-  char date[20];
-  strftime(date, sizeof date, "%FT%TZ", timestamp);
-  op_class.add_annotation(Annotation::Create("Generated", "javax.annotation")
-      .attributes(string("value = \"op_generator\", date = \"") + date + "\""));
+  op_class.add_annotation(
+      Annotation::Create("Generated", "javax.annotation")
+          .attributes("value = \"TensorFlow Java Op Generator\""));
   if (endpoint.deprecated()) {
     op_class.add_annotation(Annotation::Create("Deprecated"));
     string explanation;
     if (!op.endpoints().front().deprecated()) {
       explanation = "use {@link " +
-          ClassOf(op.endpoints().front(), base_package).full_name()
+          ClassOf(op.endpoints().front(), base_package).canonical_name()
           + "} instead";
     } else {
       explanation = op.deprecation_explanation();
@@ -376,14 +386,16 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
           .attributes("group = \"" + endpoint.package() + "\""));
   }
   // create op class file
-  string op_dir = io::JoinPath(output_dir,
+  const string op_dir_name = io::JoinPath(output_dir,
       str_util::StringReplace(op_class.package(), ".", "/", true));
-  if (!env->FileExists(op_dir).ok()) {
-    TF_CHECK_OK(Env::Default()->RecursivelyCreateDir(op_dir));
+  if (!env->FileExists(op_dir_name).ok()) {
+    TF_CHECK_OK(Env::Default()->RecursivelyCreateDir(op_dir_name))
+        << op_dir_name;
   }
+  const string op_file_name = op_class.name() + ".java";
   std::unique_ptr<tensorflow::WritableFile> op_file;
   TF_CHECK_OK(env->NewWritableFile(
-      io::JoinPath(op_dir, op_class.name() + ".java"), &op_file));
+      io::JoinPath(op_dir_name, op_file_name), &op_file)) << op_file_name;
 
   // render endpoint source code
   SourceFileWriter writer(op_file.get());
@@ -420,20 +432,19 @@ Status OpGenerator::Run(const OpList& op_list, const string& base_package,
         const std::string api_def_file_pattern =
             io::JoinPath(api_def_dir, "api_def_" + op.name() + ".pbtxt");
         if (env_->FileExists(api_def_file_pattern).ok()) {
-          TF_CHECK_OK(api_map.LoadFile(env_, api_def_file_pattern));
+          TF_CHECK_OK(api_map.LoadFile(env_, api_def_file_pattern))
+              << api_def_file_pattern;
         }
       }
     }
   }
   api_map.UpdateDocs();
-  time_t now;
-  time(&now);
   for (const auto& op_def : op_list.op()) {
     const ApiDef* api_def = api_map.GetApiDef(op_def.name());
     if (api_def->visibility() != ApiDef::SKIP) {
       OpSpec op(OpSpec::Create(op_def, *api_def));
       for (const EndpointSpec& endpoint : op.endpoints()) {
-        GenerateOp(op, endpoint, base_package, output_dir, env_, gmtime(&now));
+        GenerateOp(op, endpoint, base_package, output_dir, env_);
       }
     }
   }
diff --git a/tensorflow/java/src/gen/cc/op_generator.h b/tensorflow/java/src/gen/cc/op_generator.h
index b789e11fa9..cfe842070a 100644
--- a/tensorflow/java/src/gen/cc/op_generator.h
+++ b/tensorflow/java/src/gen/cc/op_generator.h
@@ -1,4 +1,4 @@
-/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+/* 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.
diff --git a/tensorflow/java/src/gen/cc/op_specs.cc b/tensorflow/java/src/gen/cc/op_specs.cc
index dcc6388614..081062ceaf 100644
--- a/tensorflow/java/src/gen/cc/op_specs.cc
+++ b/tensorflow/java/src/gen/cc/op_specs.cc
@@ -45,9 +45,26 @@ class TypeResolver {
  public:
   explicit TypeResolver(const OpDef& op_def) : op_def_(op_def) {}
 
+  // Returns the class type of an input/output argument
+  //
+  // For example, if the argument's datatype is DT_STRING, this method will
+  // return "java.lang.String", so the argument can become "Operand<String>"
+  // in the Ops API
   Type TypeOf(const OpDef_ArgDef& arg_def, bool *iterable_out);
-  std::pair<Type, Type> TypeOf(const OpDef_AttrDef& attr_def,
+
+  // Returns types of an input attribute
+  //
+  // The first element of the pair is the class type of this attribute while
+  // the second is its JNI/primitive type equivalent, required for explicit
+  // unboxing.
+  //
+  // For example, if the attribute is of type "float", this method will return
+  // <java.lang.Float, float>, so the attribute can be used as a "Float" object
+  // in the Ops API and casted to a "float" when passing through the JNI layer.
+  std::pair<Type, Type> TypesOf(const OpDef_AttrDef& attr_def,
       bool *iterable_out);
+
+  // Returns true if the type of this attribute has already been resolved
   bool IsAttributeVisited(const string& attr_name) {
     return visited_attrs_.find(attr_name) != visited_attrs_.cend();
   }
@@ -123,7 +140,7 @@ Type TypeResolver::TypeOf(const OpDef_ArgDef& arg_def,
     } else {
       for (const auto& attr_def : op_def_.attr()) {
         if (attr_def.name() == arg_def.type_attr()) {
-          type = TypeOf(attr_def, iterable_out).first;
+          type = TypesOf(attr_def, iterable_out).first;
           break;
         }
       }
@@ -141,7 +158,7 @@ Type TypeResolver::TypeOf(const OpDef_ArgDef& arg_def,
   return type;
 }
 
-std::pair<Type, Type> TypeResolver::TypeOf(const OpDef_AttrDef& attr_def,
+std::pair<Type, Type> TypeResolver::TypesOf(const OpDef_AttrDef& attr_def,
     bool* iterable_out) {
   std::pair<Type, Type> types = MakeTypePair(Type::Wildcard());
   *iterable_out = false;
@@ -319,7 +336,7 @@ ArgumentSpec CreateInput(const OpDef_ArgDef& input_def,
 AttributeSpec CreateAttribute(const OpDef_AttrDef& attr_def,
     const ApiDef::Attr& attr_api_def, TypeResolver* type_resolver) {
   bool iterable = false;
-  std::pair<Type, Type> types = type_resolver->TypeOf(attr_def, &iterable);
+  std::pair<Type, Type> types = type_resolver->TypesOf(attr_def, &iterable);
   Type var_type = types.first.kind() == Type::GENERIC ?
       Type::Class("Class").add_parameter(types.first) : types.first;
   if (iterable) {
diff --git a/tensorflow/java/src/gen/cc/op_specs.h b/tensorflow/java/src/gen/cc/op_specs.h
index 7d64391446..81582ea207 100644
--- a/tensorflow/java/src/gen/cc/op_specs.h
+++ b/tensorflow/java/src/gen/cc/op_specs.h
@@ -65,7 +65,6 @@ class ArgumentSpec {
       const Type& type, const string& description, bool iterable)
     : op_def_name_(op_def_name), var_(var), type_(type),
       description_(description), iterable_(iterable) {}
-  virtual ~ArgumentSpec() = default;
 
   const string& op_def_name() const { return op_def_name_; }
   const Variable& var() const { return var_; }
@@ -81,7 +80,7 @@ class ArgumentSpec {
   const bool iterable_;
 };
 
-class AttributeSpec : public ArgumentSpec {
+class AttributeSpec {
  public:
   // A specification for an operation attribute
   //
@@ -95,14 +94,24 @@ class AttributeSpec : public ArgumentSpec {
   AttributeSpec(const string& op_def_name, const Variable& var,
       const Type& type, const Type& jni_type, const string& description,
       bool iterable, bool has_default_value)
-    : ArgumentSpec(op_def_name, var, type, description, iterable),
+    : op_def_name_(op_def_name), var_(var), type_(type),
+      description_(description), iterable_(iterable),
       jni_type_(jni_type), has_default_value_(has_default_value) {}
-  virtual ~AttributeSpec() = default;
 
+  const string& op_def_name() const { return op_def_name_; }
+  const Variable& var() const { return var_; }
+  const Type& type() const { return type_; }
+  const string& description() const { return description_; }
+  bool iterable() const { return iterable_; }
   const Type& jni_type() const { return jni_type_; }
   bool has_default_value() const { return has_default_value_; }
 
  private:
+  const string op_def_name_;
+  const Variable var_;
+  const Type type_;
+  const string description_;
+  const bool iterable_;
   const Type jni_type_;
   const bool has_default_value_;
 };
diff --git a/tensorflow/java/src/gen/cc/source_writer.cc b/tensorflow/java/src/gen/cc/source_writer.cc
index 7e427787f9..56806cbb6d 100644
--- a/tensorflow/java/src/gen/cc/source_writer.cc
+++ b/tensorflow/java/src/gen/cc/source_writer.cc
@@ -1,4 +1,4 @@
-/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+/* 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.
@@ -83,17 +83,19 @@ SourceWriter& SourceWriter::Append(const StringPiece& str) {
 }
 
 SourceWriter& SourceWriter::AppendType(const Type& type) {
-  if (type.unknown()) {
+  if (type.wildcard()) {
     Append("?");
   } else {
     Append(type.name());
     if (!type.parameters().empty()) {
       Append("<");
+      bool first = true;
       for (const Type& t : type.parameters()) {
-        if (&t != &type.parameters().front()) {
+        if (!first) {
           Append(", ");
         }
         AppendType(t);
+        first = false;
       }
       Append(">");
     }
@@ -145,11 +147,13 @@ SourceWriter& SourceWriter::BeginMethod(const Method& method, int modifiers,
     AppendType(method.return_type()).Append(" ");
   }
   Append(method.name()).Append("(");
+  bool first = true;
   for (const Variable& v : method.arguments()) {
-    if (&v != &method.arguments().front()) {
+    if (!first) {
       Append(", ");
     }
     AppendType(v.type()).Append(v.variadic() ? "... " : " ").Append(v.name());
+    first = false;
   }
   return Append(")").BeginBlock();
 }
@@ -294,14 +298,16 @@ SourceWriter& SourceWriter::WriteAnnotations(
 SourceWriter& SourceWriter::WriteGenerics(
     const std::list<const Type*>& generics) {
   Append("<");
+  bool first = true;
   for (const Type* pt : generics) {
-    if (pt != generics.front()) {
+    if (!first) {
       Append(", ");
     }
     Append(pt->name());
     if (!pt->supertypes().empty()) {
       Append(" extends ").AppendType(pt->supertypes().front());
     }
+    first = false;
   }
   return Append(">");
 }
@@ -339,7 +345,7 @@ void SourceWriter::TypeVisitor::Visit(const Type& type) {
 
 void SourceWriter::GenericNamespace::DoVisit(const Type& type) {
   // ignore non-generic parameters, wildcards and generics already declared
-  if (type.kind() == Type::GENERIC && !type.unknown()
+  if (type.kind() == Type::GENERIC && !type.wildcard()
       && generic_names_.find(type.name()) == generic_names_.end()) {
     declared_types_.push_back(&type);
     generic_names_.insert(type.name());
@@ -348,7 +354,7 @@ void SourceWriter::GenericNamespace::DoVisit(const Type& type) {
 
 void SourceWriter::TypeImporter::DoVisit(const Type& type) {
   if (!type.package().empty() && type.package() != current_package_) {
-    imports_.insert(type.full_name());
+    imports_.insert(type.canonical_name());
   }
 }
 
diff --git a/tensorflow/java/src/gen/cc/source_writer.h b/tensorflow/java/src/gen/cc/source_writer.h
index bcae33ccce..1f0febe9a3 100644
--- a/tensorflow/java/src/gen/cc/source_writer.h
+++ b/tensorflow/java/src/gen/cc/source_writer.h
@@ -1,4 +1,4 @@
-/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+/* 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.
diff --git a/tensorflow/java/src/gen/cc/source_writer_test.cc b/tensorflow/java/src/gen/cc/source_writer_test.cc
index 875ad99ae2..b9a5fee9be 100644
--- a/tensorflow/java/src/gen/cc/source_writer_test.cc
+++ b/tensorflow/java/src/gen/cc/source_writer_test.cc
@@ -1,4 +1,4 @@
-/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+/* 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.
diff --git a/tensorflow/java/src/gen/gen_ops.bzl b/tensorflow/java/src/gen/gen_ops.bzl
index 7017b52649..f4ff34ea03 100644
--- a/tensorflow/java/src/gen/gen_ops.bzl
+++ b/tensorflow/java/src/gen/gen_ops.bzl
@@ -3,33 +3,26 @@
 load(
     "//tensorflow:tensorflow.bzl",
     "tf_binary_additional_srcs",
-    "tf_cc_binary",
-    "tf_copts",
 )
 
-# Given a list of "ops_libs" (a list of files in the core/ops directory
-# without their .cc extensions), generate Java wrapper code for all operations
-# found in the ops files.
-# Then, combine all those source files into a single archive (.srcjar).
+# Generate Java wrapper classes for all registered core operations and package
+# them into a single source archive (.srcjar).
 #
 # For example:
-#  tf_java_op_gen_srcjar("gen_sources", "gen_tool", "my.package", [ "array_ops", "math_ops" ])
+#  tf_java_op_gen_srcjar("gen_sources", ":gen_tool", "my.package")
 #
-# will create a genrule named "gen_sources" that first generate source files:
-#     ops/src/main/java/my/package/array/*.java
-#     ops/src/main/java/my/package/math/*.java
+# will create a genrule named "gen_sources" that generates source files under
+#     ops/src/main/java/my/package/**/*.java
 #
-# and then archive those source files in:
+# and then archive those source files into
 #     ops/gen_sources.srcjar
 #
 def tf_java_op_gen_srcjar(name,
                           gen_tool,
-                          gen_base_package,
-                          ops_libs=[],
-                          ops_libs_pkg="//tensorflow/core",
+                          base_package,
+                          api_def_srcs=[],
                           out_dir="ops/",
                           out_src_dir="src/main/java/",
-                          api_def_srcs=[],
                           visibility=["//tensorflow/java:__pkg__"]):
 
   gen_cmds = ["rm -rf $(@D)"]  # Always start from fresh when generating source files
@@ -48,23 +41,9 @@ def tf_java_op_gen_srcjar(name,
           ") | cut -d\" \" -f1))")
     api_def_args_str = ",".join(api_def_args)
 
-  gen_tool_deps = [":java_op_gen_lib"]
-  for ops_lib in ops_libs:
-    gen_tool_deps.append(ops_libs_pkg + ":" + ops_lib + "_op_lib")
-
-  tf_cc_binary(
-      name=gen_tool,
-      srcs=[
-          "src/gen/cc/op_gen_main.cc",
-      ],
-      copts=tf_copts(),
-      linkopts=["-lm"],
-      linkstatic=1,  # Faster to link this one-time-use binary dynamically
-      deps = gen_tool_deps)
-
-  gen_cmds += ["$(location :" + gen_tool + ")" +
+  gen_cmds += ["$(location " + gen_tool + ")" +
                " --output_dir=$(@D)/" + out_src_dir +
-               " --base_package=" + gen_base_package +
+               " --base_package=" + base_package +
                " --api_dirs=" + api_def_args_str]
 
   # Generate a source archive containing generated code for these ops.
-- 
GitLab


From aaa345f5a662aab524bbee3912c605919239bef6 Mon Sep 17 00:00:00 2001
From: wangsiyu <siyu.wsy@gmail.com>
Date: Fri, 4 May 2018 10:52:26 +0800
Subject: [PATCH 225/395] refine by using iterator of partitioned variable

---
 tensorflow/python/layers/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/layers/base.py b/tensorflow/python/layers/base.py
index c050e6be04..f7b2e471b2 100644
--- a/tensorflow/python/layers/base.py
+++ b/tensorflow/python/layers/base.py
@@ -358,7 +358,7 @@ def _add_elements_to_collection(elements, collection_list):
 def _should_add_regularizer(variable, existing_variable_set):
   result = True
   if isinstance(variable, tf_variables.PartitionedVariable):
-    for var in variable._get_variable_list():
+    for var in variable:
       if var in existing_variable_set:
         result = False
         break
-- 
GitLab


From de9256f61a9d71a30b175e46116fc5d87063ceaa Mon Sep 17 00:00:00 2001
From: "William D. Irons" <wdirons@us.ibm.com>
Date: Fri, 4 May 2018 08:19:03 -0500
Subject: [PATCH 226/395] Add conditions:default to mkl build (#19008)

If building on a system that is not darwin, linux_x86_64, or
windows, the select statement in third_party/mkl/BUILD fails to
find a match and fails. Need to use no mkl libraries for non-x86
systems

Fixes #18084
---
 third_party/mkl/BUILD | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/third_party/mkl/BUILD b/third_party/mkl/BUILD
index c2adf578c7..017613abb0 100644
--- a/third_party/mkl/BUILD
+++ b/third_party/mkl/BUILD
@@ -34,6 +34,7 @@ filegroup(
         "@org_tensorflow//tensorflow:windows": [
             "@mkl_windows//:LICENSE",
         ],
+        "//conditions:default": []
     }),
     visibility = ["//visibility:public"],
 )
@@ -54,5 +55,6 @@ cc_library(
             "@mkl_windows//:mkl_headers",
             "@mkl_windows//:mkl_libs_windows",
         ],
+        "//conditions:default": []
     }),
 )
-- 
GitLab


From 7f0d43c1b7462645767712cd5942d754a5f7adb7 Mon Sep 17 00:00:00 2001
From: manhyuk <manhyuk@kw.ac.kr>
Date: Sat, 5 May 2018 01:53:35 +0900
Subject: [PATCH 227/395] fix typo

---
 tensorflow/contrib/tpu/python/tpu/tpu_context.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_context.py b/tensorflow/contrib/tpu/python/tpu/tpu_context.py
index fbc1173e49..4d7bc6a5a6 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu_context.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu_context.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ===================================================================
-"""TPU system metdata and associated tooling."""
+"""TPU system metadata and associated tooling."""
 
 from __future__ import absolute_import
 from __future__ import division
-- 
GitLab


From 9b43bd6459e410fc8d3dd1beba9f9a6a254096ba Mon Sep 17 00:00:00 2001
From: Akshay Agrawal <akshayka@google.com>
Date: Thu, 3 May 2018 13:40:20 -0700
Subject: [PATCH 228/395] Documentation for tf.contrib.eager.py_func

PiperOrigin-RevId: 195303454
---
 tensorflow/python/ops/script_ops.py | 66 ++++++++++++++++++++++++++---
 1 file changed, 60 insertions(+), 6 deletions(-)

diff --git a/tensorflow/python/ops/script_ops.py b/tensorflow/python/ops/script_ops.py
index 9f1dd2c4fd..f87c5dc5e3 100644
--- a/tensorflow/python/ops/script_ops.py
+++ b/tensorflow/python/ops/script_ops.py
@@ -243,14 +243,68 @@ def _internal_py_func(func, inp, Tout, stateful=None, eager=False, name=None):
 
 
 def eager_py_func(func, inp, Tout, name=None):
-  """Wraps a python function into a TensorFlow op.
+  """Wraps a python function into a TensorFlow op that executes it eagerly.
 
-  When the returned op is executed, `func` is invoked with eager execution
-  enabled. Inputs are Tensor objects and func must return None or objects
-  that may be converted to Tensor objects.
+  This function allows expressing computations in a TensorFlow graph as
+  Python functions. In particular, it wraps a Python function `func`
+  in a TensorFlow operation that executes it with eager exeuction enabled. As a
+  consequence, `tf.contrib.eager.py_func` makes it possible to express control
+  flow using Python constructs (`if`, `while`, `for`, etc.), instead of
+  TensorFlow control flow constructs (@{tf.cond}, @{tf.while_loop}). For
+  example, you might use `tf.contrib.eager.py_func` to implement the log huber
+  function:
+
+  ```python
+  def log_huber(x, m):
+    if tf.abs(x) <= m:
+      return x ** 2
+    else:
+      return m ** 2 * (1 - 2 * tf.log(m) + tf.log(x ** 2))
+
+  x = tf.placeholder(tf.float32)
+  m = tf.placeholder(tf.float32)
+
+  y = tf.contrib.eager.py_func(func=log_huber, inp=[x, m], Tout=tf.float32)
+
+  with tf.Session() as sess:
+    # The session executes `log_huber` eagerly. Given the feed values below,
+    # it will take the second branch, so `output` evaluates to 7.24372.
+    output = sess.run(y, feed_dict={x: 3.0, m: 2.0})
+  ```
+
+  You can also use `tf.contrib.eager.py_func` to debug your models at runtime
+  using Python tools, i.e., you can isolate portions of your code that
+  you want to debug, wrap them in Python functions and insert `pdb` tracepoints
+  or print statements as desired, and wrap those functions in
+  `tf.contrib.eager.py_func`.
+
+  For more information on eager execution, see @{$programmers_guide/eager}.
+
+  `tf.contrib.eager.py_func` is similar in spirit to @{tf.py_func}, but unlike
+  the latter, the former lets you use TensorFlow operations in the wrapped
+  Python function. In particular, while @{tf.py_func} only runs on CPUs and
+  wraps functions that take NumPy arrays as inputs and return NumPy arrays as
+  outputs, `tf.contrib.eager.py_func` can be placed on GPUs and wraps functions
+  that take Tensors as inputs, execute TensorFlow operations in their bodies,
+  and return Tensors as outputs.
+
+  `tf.contrib.eager.py_func` is not differentiable, though a gradient may be
+  implemented in the future; if you would like to differentiate through it,
+  please file an issue on Github.
+
+  Like @{tf.py_func}, `tf.contrib.eager.py_func` has the following limitations
+  with respect to serialization and distribution:
+
+  * The body of the function (i.e. `func`) will not be serialized in a
+    `GraphDef`. Therefore, you should not use this function if you need to
+    serialize your model and restore it in a different environment.
+
+  * The operation must run in the same address space as the Python program
+    that calls `tf.contrib.eager.py_func()`. If you are using distributed
+    TensorFlow, you must run a `tf.train.Server` in the same process as the
+    program that calls `tf.contrib.eager.py_func()` and you must pin the created
+    operation to a device in that server (e.g. using `with tf.device():`).
 
-  This function has the same limitations as `py_func` with respect to
-  serialization and distribution.
 
   Args:
     func: A Python function which accepts a list of `Tensor` objects
-- 
GitLab


From 86d3435503e20e44ab37c87613481f7a35d0c14e Mon Sep 17 00:00:00 2001
From: Jianwei Xie <xiejw@google.com>
Date: Thu, 3 May 2018 13:54:29 -0700
Subject: [PATCH 229/395] Fix a typo.

PiperOrigin-RevId: 195305770
---
 tensorflow/python/estimator/training.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/estimator/training.py b/tensorflow/python/estimator/training.py
index 41ffa371aa..2f14a6f560 100644
--- a/tensorflow/python/estimator/training.py
+++ b/tensorflow/python/estimator/training.py
@@ -657,7 +657,7 @@ class _TrainingExecutor(object):
           hooks=train_hooks)
 
       if not self._continuous_eval_listener.before_eval():
-        logging.info('Exiting training and evaluation lopp, as requested by '
+        logging.info('Exiting training and evaluation loop, as requested by '
                      '_ContinuousEvalListener.before_eval.')
         break
 
-- 
GitLab


From 518dfea0d6d45448a360a49635fe815a28730c46 Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Thu, 3 May 2018 14:00:56 -0700
Subject: [PATCH 230/395] [XLA:CPU] Remove dead function + DCHECK, NFC

There isn't a lot of benefit to fixing the function to do that it says it does
since I'm adding support for lowering batch matmul which will break this
precondition anyway.

PiperOrigin-RevId: 195306803
---
 tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc | 4 ----
 tensorflow/compiler/xla/service/cpu/dot_op_emitter.h  | 4 ----
 2 files changed, 8 deletions(-)

diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
index 495fecc4aa..801c523908 100644
--- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
+++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
@@ -557,8 +557,6 @@ DotOpEmitter::DotOpEmitter(
   return dot_emitter.Emit();
 }
 
-bool DotOpEmitter::ShapesAreLegalForRuntimeDot() const { return true; }
-
 bool DotOpEmitter::EmitLlvmIrDotIfProfitable() {
   if (dot_.shape().dimensions_size() != 2) {
     return false;
@@ -908,8 +906,6 @@ tensorflow::Status DotOpEmitter::EmitScalarDot() {
 }
 
 tensorflow::Status DotOpEmitter::EmitCallToRuntime() {
-  DCHECK(ShapesAreLegalForRuntimeDot());
-
   // The signature of the Eigen runtime matmul function is:
   //
   //   (void)(void* run_options, float* out, float* lhs, float* rhs,
diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h
index 9d748eb81f..47e0924334 100644
--- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h
+++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h
@@ -99,10 +99,6 @@ class DotOpEmitter {
       llvm_ir::ForLoopNest* loop_nest, const llvm_ir::IrArray& operand_array,
       int64 reduction_dimension, tensorflow::StringPiece name_suffix);
 
-  // Our runtime operation requires that all arrays have the same layout,
-  // no padding, and a rank of two.
-  bool ShapesAreLegalForRuntimeDot() const;
-
   // Represents the dimensions of a matrix-matrix multiply operation.
   struct MatMultDims {
     // The number of rows in the LHS.
-- 
GitLab


From a4a9e372f6af694e91ef7aaae9f23867d0ec0fc2 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 3 May 2018 14:11:13 -0700
Subject: [PATCH 231/395] Optimize idempotent ops, e.g., Snapshot(Snapshot(x))
 => Snapshot(x)

PiperOrigin-RevId: 195308675
---
 tensorflow/core/grappler/op_types.cc          | 15 ++--
 tensorflow/core/grappler/op_types.h           |  5 ++
 .../optimizers/arithmetic_optimizer.cc        | 30 ++++++++
 .../optimizers/arithmetic_optimizer.h         |  1 +
 .../optimizers/arithmetic_optimizer_test.cc   | 68 +++++++++++++++++++
 tensorflow/python/grappler/cluster_test.py    |  4 +-
 .../profiler/internal/run_metadata_test.py    |  6 +-
 7 files changed, 121 insertions(+), 8 deletions(-)

diff --git a/tensorflow/core/grappler/op_types.cc b/tensorflow/core/grappler/op_types.cc
index c48dc00941..e633ecf789 100644
--- a/tensorflow/core/grappler/op_types.cc
+++ b/tensorflow/core/grappler/op_types.cc
@@ -323,6 +323,8 @@ bool IsSize(const NodeDef& node) { return node.op() == "Size"; }
 
 bool IsSlice(const NodeDef& node) { return node.op() == "Slice"; }
 
+bool IsSnapshot(const NodeDef& node) { return node.op() == "Snapshot"; }
+
 bool IsSoftplusGrad(const NodeDef& node) { return node.op() == "SoftplusGrad"; }
 
 bool IsSoftsignGrad(const NodeDef& node) { return node.op() == "SoftsignGrad"; }
@@ -488,14 +490,13 @@ bool IsValueAndOrderAndShapePreserving(const NodeDef& node) {
               "DeepCopy"
               "Enter",
               "Exit",
-              "Identity",
-              "IdentityN",
               "PreventGradient",
               "Print",
               "Snapshot",
               "StopGradient",
           }));
-  return value_and_order_and_shape_preserving_ops->count(node.op()) > 0;
+  return value_and_order_and_shape_preserving_ops->count(node.op()) > 0 ||
+         IsIdentity(node);
 }
 
 bool IsValueAndOrderPreserving(const NodeDef& node) {
@@ -505,7 +506,7 @@ bool IsValueAndOrderPreserving(const NodeDef& node) {
   static const std::unordered_set<string>* value_and_order_preserving_ops =
       CHECK_NOTNULL((new const std::unordered_set<string>{
           "ExpandDims",
-          "Snapshot",
+          "Reshape",
           "Squeeze",
       }));
   return value_and_order_preserving_ops->count(node.op()) > 0 ||
@@ -576,7 +577,7 @@ bool IsUnaryElementWise(const NodeDef& node) {
           "Tanh",
       }));
   return element_wise_ops->count(node.op()) > 0 ||
-         (!IsIdentityN(node) && IsValueAndOrderAndShapePreserving(node));
+         IsValueAndOrderAndShapePreserving(node);
 }
 
 bool HasOpDef(const NodeDef& node) {
@@ -584,5 +585,9 @@ bool HasOpDef(const NodeDef& node) {
   return OpRegistry::Global()->LookUpOpDef(node.op(), &op_def).ok();
 }
 
+bool IsIdempotent(const NodeDef& node) {
+  return IsValueAndOrderAndShapePreserving(node) && IsFreeOfSideEffect(node);
+}
+
 }  // namespace grappler
 }  // end namespace tensorflow
diff --git a/tensorflow/core/grappler/op_types.h b/tensorflow/core/grappler/op_types.h
index e33dd21538..f6105d710e 100644
--- a/tensorflow/core/grappler/op_types.h
+++ b/tensorflow/core/grappler/op_types.h
@@ -123,6 +123,7 @@ bool IsShape(const NodeDef& node);
 bool IsShapeN(const NodeDef& node);
 bool IsShuffle(const NodeDef& node);
 bool IsSigmoidGrad(const NodeDef& node);
+bool IsSnapshot(const NodeDef& node);
 bool IsSoftplusGrad(const NodeDef& node);
 bool IsSoftsignGrad(const NodeDef& node);
 bool IsSplit(const NodeDef& node);
@@ -187,6 +188,10 @@ bool IsValueAndOrderPreserving(const NodeDef& node);
 // function returns true if the op commutes with all element-wise operations.
 bool IsValuePreserving(const NodeDef& node);
 
+// Returns true if node is idempotent w.r.t. its first input, i.e. if
+// Op(Op(x, y, z), y, z) = Op(x, y, z).
+bool IsIdempotent(const NodeDef& node);
+
 bool IsUnaryElementWise(const NodeDef& node);
 
 // Returns true if we can find an opdef corresponding to the op of the node.
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
index 2a5654f752..29f49079c4 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
@@ -295,6 +295,7 @@ class ArithmeticOptimizerStage : public GraphOptimizerStage<string> {
         }
       }
     }
+    DedupControlInputs(target_node);
   }
 
   bool IsInPreserveSet(const NodeDef& node) const {
@@ -1690,6 +1691,32 @@ class HoistCWiseUnaryChainsStage : public ArithmeticOptimizerStage {
   std::unordered_set<string> optimized_nodes_;
 };
 
+class RemoveIdempotentStage : public ArithmeticOptimizerStage {
+ public:
+  explicit RemoveIdempotentStage(const GraphOptimizerContext& ctx,
+                                 const ArithmeticOptimizerContext& ctx_ext)
+      : ArithmeticOptimizerStage("RemoveIdempotent", ctx, ctx_ext) {}
+  ~RemoveIdempotentStage() override = default;
+
+  bool IsSupported(const NodeDef* node) const override {
+    return IsIdempotent(*node) && !IsInPreserveSet(*node);
+  }
+
+  Status TrySimplify(NodeDef* node, string* simplified_node_name) override {
+    NodeDef* input;
+    TF_RETURN_IF_ERROR(GetInputNode(node->input(0), &input));
+    auto root_scope_and_name = ParseNodeScopeAndName(node->name());
+    const string new_name = OptimizedNodeName(root_scope_and_name);
+    if (input->op() == node->op() && input->device() == node->device() &&
+        IsIdempotent(*input) && !ctx().node_map->NodeExists(new_name)) {
+      NodeDef* new_input_node = AddCopyNode(new_name, input);
+      ForwardControlDependencies(new_input_node, {node});
+      *simplified_node_name = new_input_node->name();
+    }
+    return Status::OK();
+  }
+};
+
 // Performs the conversion:
 // Div(x, Sqrt(y)) => Mul(x, Rsqrt(y))
 // TODO(srjoglekar): Generalize to optimize cases like (x / pow(y, z)).
@@ -1975,6 +2002,7 @@ void ArithmeticOptimizer::ForwardControlDependencies(
       }
     }
   }
+  DedupControlInputs(target_node);
 }
 
 // TODO(ezhulenev): extract each individual simplify rewrite into separate
@@ -2381,6 +2409,8 @@ Status ArithmeticOptimizer::SimplifyArithmeticOps(bool can_use_shapes) {
     pipeline.AddStage<HoistCWiseUnaryChainsStage>(ctx, ctx_ext);
   if (options_.convert_sqrt_div_to_rsqrt_mul)
     pipeline.AddStage<SqrtDivToRsqrtMulStage>(ctx, ctx_ext);
+  if (options_.remove_idempotent)
+    pipeline.AddStage<RemoveIdempotentStage>(ctx, ctx_ext);
 
   VLOG(1) << "Run " << pipeline.NumStages() << " arithmetic optimizer stages: "
           << str_util::Join(pipeline.StageNames(), ", ");
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
index 6309dc1a33..3f9feac55f 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
@@ -67,6 +67,7 @@ class ArithmeticOptimizer : public GraphOptimizer {
     bool remove_negation = true;
     bool hoist_cwise_unary_chains = true;
     bool convert_sqrt_div_to_rsqrt_mul = false;
+    bool remove_idempotent = true;
 
     // Choose which arithmetic optimizer stages will be enabled for a given
     // optimization level by default.
diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
index d32743f3f2..e109e66633 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
@@ -83,6 +83,7 @@ class ArithmeticOptimizerTest : public GrapplerTest {
                         GraphDef* output) {
     TF_EXPECT_OK(optimizer->Optimize(nullptr, *item, output));
     item->graph.Swap(output);
+    output->Clear();
     TF_EXPECT_OK(ModelPruner().Optimize(nullptr, *item, output));
   }
 
@@ -91,6 +92,7 @@ class ArithmeticOptimizerTest : public GrapplerTest {
                      GraphDef* output) {
     TF_EXPECT_OK(optimizer->Optimize(nullptr, *item, output));
     item->graph.Swap(output);
+    output->Clear();
     TF_EXPECT_OK(optimizer->Optimize(nullptr, *item, output));
   }
 
@@ -99,8 +101,10 @@ class ArithmeticOptimizerTest : public GrapplerTest {
                              GraphDef* output) {
     TF_EXPECT_OK(optimizer->Optimize(nullptr, *item, output));
     item->graph.Swap(output);
+    output->Clear();
     TF_EXPECT_OK(optimizer->Optimize(nullptr, *item, output));
     item->graph.Swap(output);
+    output->Clear();
     TF_EXPECT_OK(ModelPruner().Optimize(nullptr, *item, output));
   }
 
@@ -168,6 +172,11 @@ class ArithmeticOptimizerTest : public GrapplerTest {
     DisableAllStages(optimizer);
     optimizer->options_.convert_sqrt_div_to_rsqrt_mul = true;
   }
+
+  void EnableOnlyRemoveIdempotent(ArithmeticOptimizer* optimizer) {
+    DisableAllStages(optimizer);
+    optimizer->options_.remove_idempotent = true;
+  }
 };
 
 TEST_F(ArithmeticOptimizerTest, NoOp) {
@@ -2390,5 +2399,64 @@ TEST_F(ArithmeticOptimizerTest, HoistCWiseUnaryIntoSplit) {
   }
 }
 
+TEST_F(ArithmeticOptimizerTest, RemoveIdempotent) {
+  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
+  Output a = ops::Const(s.WithOpName("a"), 3.14f, {32});
+  Output ctrl1 = ops::Const(s.WithOpName("ctrl1"), 1, {});
+  Output ctrl2 = ops::Const(s.WithOpName("ctrl2"), 2, {});
+  Output sn1 =
+      ops::Snapshot(s.WithOpName("sn1").WithControlDependencies(ctrl1), a);
+  Output sn2 =
+      ops::Snapshot(s.WithOpName("sn2").WithControlDependencies(ctrl2), sn1);
+  Output out1 = ops::Identity(s.WithOpName("out1"), sn2);
+  Output id1 = ops::Identity(s.WithOpName("id1"), a);
+  Output id2 = ops::Identity(s.WithOpName("id2"), id1);
+  Output out2 = ops::Identity(s.WithOpName("out2"), id2);
+  GrapplerItem item;
+  item.fetch = {"out1", "out2"};
+  TF_CHECK_OK(s.ToGraphDef(&item.graph));
+
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch);
+
+  GraphDef output;
+  ArithmeticOptimizer optimizer;
+  EnableOnlyRemoveIdempotent(&optimizer);
+  OptimizeTwice(&optimizer, &item, &output);
+
+  EXPECT_EQ(11, output.node_size());
+  int found = 0;
+  for (const NodeDef& node : output.node()) {
+    if (node.name() == "out1") {
+      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ("ArithmeticOptimizer/RemoveIdempotent_sn2", node.input(0));
+      found++;
+    } else if (node.name() == "ArithmeticOptimizer/RemoveIdempotent_sn2") {
+      EXPECT_EQ(3, node.input_size());
+      EXPECT_EQ("Snapshot", node.op());
+      EXPECT_EQ("a", node.input(0));
+      EXPECT_EQ("^ctrl1", node.input(1));
+      EXPECT_EQ("^ctrl2", node.input(2));
+      found++;
+    } else if (node.name() == "out2") {
+      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ("ArithmeticOptimizer/RemoveIdempotent_id2", node.input(0));
+      found++;
+    } else if (node.name() == "ArithmeticOptimizer/RemoveIdempotent_id2") {
+      EXPECT_EQ("Identity", node.op());
+      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ("a", node.input(0));
+      found++;
+    }
+  }
+  EXPECT_EQ(4, found);
+
+  auto tensors = EvaluateNodes(output, item.fetch);
+  EXPECT_EQ(tensors.size(), tensors_expected.size());
+  EXPECT_EQ(tensors.size(), item.fetch.size());
+  for (int i = 0; i < item.fetch.size(); ++i) {
+    test::ExpectTensorNear<float>(tensors_expected[i], tensors[i], 1e-6);
+  }
+}
+
 }  // namespace grappler
 }  // namespace tensorflow
diff --git a/tensorflow/python/grappler/cluster_test.py b/tensorflow/python/grappler/cluster_test.py
index 26c6f22d34..541747867f 100644
--- a/tensorflow/python/grappler/cluster_test.py
+++ b/tensorflow/python/grappler/cluster_test.py
@@ -45,7 +45,7 @@ class ClusterTest(test.TestCase):
       op_perfs, run_time, step_stats = grappler_cluster.MeasureCosts(
           grappler_item)
       self.assertTrue(run_time > 0)
-      self.assertEqual(len(op_perfs), 8)
+      self.assertEqual(len(op_perfs), 4)
       self.assertTrue(step_stats.dev_stats)
 
   def testNoDetailedStats(self):
@@ -129,7 +129,7 @@ class ClusterTest(test.TestCase):
         disable_detailed_stats=False, disable_timeline=False) as gcluster:
       op_perfs, run_time, step_stats = gcluster.MeasureCosts(grappler_item)
       self.assertTrue(run_time > 0)
-      self.assertEqual(len(op_perfs), 8)
+      self.assertEqual(len(op_perfs), 4)
       self.assertTrue(step_stats.dev_stats)
 
   def testAvailableOps(self):
diff --git a/tensorflow/python/profiler/internal/run_metadata_test.py b/tensorflow/python/profiler/internal/run_metadata_test.py
index fd893d6cde..216cc3dd54 100644
--- a/tensorflow/python/profiler/internal/run_metadata_test.py
+++ b/tensorflow/python/profiler/internal/run_metadata_test.py
@@ -23,6 +23,7 @@ from collections import defaultdict
 import six
 
 from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import rewriter_config_pb2
 from tensorflow.python.client import session
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import math_ops
@@ -65,7 +66,10 @@ def _run_model():
   w = random_ops.random_normal(shape=[SIZE, 2 * SIZE])
   y = math_ops.matmul(x, w)
 
-  with session.Session() as sess:
+  config = config_pb2.ConfigProto()
+  config.graph_options.rewrite_options.arithmetic_optimization = (
+      rewriter_config_pb2.RewriterConfig.OFF)
+  with session.Session(config=config) as sess:
     run_metadata = config_pb2.RunMetadata()
     opts = builder.time_and_memory()
     opts['min_micros'] = 0
-- 
GitLab


From 05425f25ee1f8b83624127cf0f403b6751e7d70a Mon Sep 17 00:00:00 2001
From: Nick Desaulniers <ndesaulniers@google.com>
Date: Thu, 3 May 2018 14:16:27 -0700
Subject: [PATCH 232/395] [TF:XLA] clean up interface to xla::VerifyHloModule

It seems that the first argument, platform, is unused.

PiperOrigin-RevId: 195309504
---
 tensorflow/compiler/xla/tests/test_utils.cc | 3 +--
 tensorflow/compiler/xla/tests/test_utils.h  | 2 +-
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/xla/tests/test_utils.cc b/tensorflow/compiler/xla/tests/test_utils.cc
index 997a1d8273..810cc25f1b 100644
--- a/tensorflow/compiler/xla/tests/test_utils.cc
+++ b/tensorflow/compiler/xla/tests/test_utils.cc
@@ -339,8 +339,7 @@ StatusOr<std::vector<std::unique_ptr<Literal>>> MakeFakeArguments(
   return std::move(arguments);
 }
 
-Status VerifyHloModule(const se::Platform& platform, HloModule* const module,
-                       bool allow_mixed_precision) {
+Status VerifyHloModule(HloModule* const module, bool allow_mixed_precision) {
   return HloVerifier(allow_mixed_precision).Run(module).status();
 }
 
diff --git a/tensorflow/compiler/xla/tests/test_utils.h b/tensorflow/compiler/xla/tests/test_utils.h
index 30c147910c..f483cdebea 100644
--- a/tensorflow/compiler/xla/tests/test_utils.h
+++ b/tensorflow/compiler/xla/tests/test_utils.h
@@ -68,7 +68,7 @@ StatusOr<std::vector<std::unique_ptr<Literal>>> MakeFakeArguments(
 
 // Check that a given module satisfies various constraints before trying to
 // execute it.
-Status VerifyHloModule(const se::Platform& platform, HloModule* const module,
+Status VerifyHloModule(HloModule* const module,
                        bool allow_mixed_precision = false);
 
 }  // namespace xla
-- 
GitLab


From 316e0bab900d2a513e4e9622940181414e0d0596 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 3 May 2018 14:18:07 -0700
Subject: [PATCH 233/395] Add separate get_read and get_updated helpers that
 work on code exceprts. Handle corner case for AugAssign. Fix bug in
 _node_sets_self_attribute.

PiperOrigin-RevId: 195309809
---
 .../pyct/static_analysis/activity.py          |  79 ++++++++--
 .../pyct/static_analysis/activity_test.py     | 136 +++++++++++++++---
 2 files changed, 187 insertions(+), 28 deletions(-)

diff --git a/tensorflow/contrib/autograph/pyct/static_analysis/activity.py b/tensorflow/contrib/autograph/pyct/static_analysis/activity.py
index 2c14c2c8c2..4d7b0cbb7b 100644
--- a/tensorflow/contrib/autograph/pyct/static_analysis/activity.py
+++ b/tensorflow/contrib/autograph/pyct/static_analysis/activity.py
@@ -23,11 +23,12 @@ import copy
 import gast
 
 from tensorflow.contrib.autograph.pyct import anno
+from tensorflow.contrib.autograph.pyct import qual_names
 from tensorflow.contrib.autograph.pyct import transformer
-from tensorflow.contrib.autograph.pyct.qual_names import QN
 from tensorflow.contrib.autograph.pyct.static_analysis.annos import NodeAnno
 
 # TODO(mdan): Add support for PY3 (e.g. Param vs arg).
+# TODO(alexbw): Ignore named literals (e.g. None)
 
 
 class Scope(object):
@@ -43,16 +44,20 @@ class Scope(object):
     used: identifiers referenced in this scope
   """
 
-  def __init__(self, parent, isolated=True):
+  def __init__(self, parent, isolated=True, add_unknown_symbols=False):
     """Create a new scope.
 
     Args:
       parent: A Scope or None.
       isolated: Whether the scope is isolated, that is, whether variables
           created in this scope should be visible to the parent scope.
+      add_unknown_symbols: Whether to handle attributed and subscripts
+          without having first seen the base name.
+          E.g., analyzing the statement 'x.y = z' without first having seen 'x'.
     """
     self.isolated = isolated
     self.parent = parent
+    self.add_unknown_symbols = add_unknown_symbols
     self.modified = set()
     self.created = set()
     self.used = set()
@@ -134,13 +139,17 @@ class Scope(object):
     self.params.add(name)
 
   def mark_creation(self, name, writes_create_symbol=False):
+    """Mark a qualified name as created."""
     if name.is_composite():
       parent = name.parent
-      if self.has(parent):
-        if not writes_create_symbol:
-          return
+      if not writes_create_symbol:
+        return
       else:
-        raise ValueError('Unknown symbol "%s".' % parent)
+        if not self.has(parent):
+          if self.add_unknown_symbols:
+            self.mark_read(parent)
+          else:
+            raise ValueError('Unknown symbol "%s".' % parent)
     self.created.add(name)
 
   def mark_write(self, name):
@@ -163,17 +172,25 @@ class Scope(object):
 
 
 class ActivityAnalyzer(transformer.Base):
-  """Annotates nodes with local scope information. See Scope."""
+  """Annotates nodes with local scope information.
 
-  def __init__(self, context, parent_scope):
+  See Scope.
+
+  The use of this class requires that qual_names.resolve() has been called on
+  the node. This class will ignore nodes have not been
+  annotated with their qualified names.
+  """
+
+  def __init__(self, context, parent_scope=None, add_unknown_symbols=False):
     super(ActivityAnalyzer, self).__init__(context)
-    self.scope = Scope(parent_scope)
+    self.scope = Scope(parent_scope, None, add_unknown_symbols)
     self._in_return_statement = False
+    self._in_aug_assign = False
 
   @property
   def _in_constructor(self):
-    innermost = self.enclosing_entities[-1]
     if len(self.enclosing_entities) > 1:
+      innermost = self.enclosing_entities[-1]
       parent = self.enclosing_entities[-2]
       return isinstance(parent, gast.ClassDef) and innermost.name == '__init__'
     return False
@@ -184,6 +201,7 @@ class ActivityAnalyzer(transformer.Base):
       # TODO(mdan): The 'self' argument is not guaranteed to be called 'self'.
       if qn.has_attr and qn.parent.qn == ('self',):
         return True
+    return False
 
   def _track_symbol(self,
                     node,
@@ -201,12 +219,14 @@ class ActivityAnalyzer(transformer.Base):
         self.scope.mark_write(qn.parent)
       if writes_create_symbol:
         self.scope.mark_creation(qn, writes_create_symbol=True)
+      if self._in_aug_assign:
+        self.scope.mark_read(qn)
     elif isinstance(node.ctx, gast.Load):
       self.scope.mark_read(qn)
     elif isinstance(node.ctx, gast.Param):
       # Param contexts appear in function defs, so they have the meaning of
       # defining a variable.
-      # TODO(mdan): This bay be incorrect with nested functions.
+      # TODO(mdan): This may be incorrect with nested functions.
       # For nested functions, we'll have to add the notion of hiding args from
       # the parent scope, not writing to them.
       self.scope.mark_creation(qn)
@@ -222,6 +242,14 @@ class ActivityAnalyzer(transformer.Base):
     if self._in_return_statement:
       self.scope.mark_returned(qn)
 
+  def visit_AugAssign(self, node):
+    # Special rules for AugAssign. In Assign, the target is only written,
+    # but in AugAssig (e.g. a += b), the target is both read and written.
+    self._in_aug_assign = True
+    self.generic_visit(node)
+    self._in_aug_assign = False
+    return node
+
   def visit_Name(self, node):
     self.generic_visit(node)
     self._track_symbol(node)
@@ -295,7 +323,7 @@ class ActivityAnalyzer(transformer.Base):
 
   def visit_FunctionDef(self, node):
     if self.scope:
-      qn = QN(node.name)
+      qn = qual_names.QN(node.name)
       self.scope.mark_write(qn)
     current_scope = self.scope
     body_scope = Scope(current_scope, isolated=True)
@@ -355,5 +383,32 @@ class ActivityAnalyzer(transformer.Base):
     return node
 
 
+def get_read(node, context):
+  """Return the variable names as QNs (qual_names.py) read by this statement."""
+  analyzer = ActivityAnalyzer(context, None, True)
+  analyzer.visit(node)
+  return analyzer.scope.used
+
+
+def get_updated(node, context):
+  """Return the variable names created or mutated by this statement.
+
+  This function considers assign statements, augmented assign statements, and
+  the targets of for loops, as well as function arguments.
+  For example, `x[0] = 2` will return `x`, `x, y = 3, 4` will return `x` and
+  `y`, `for i in range(x)` will return `i`, etc.
+  Args:
+    node: An AST node
+    context: An EntityContext instance
+
+  Returns:
+    A set of variable names (QNs, see qual_names.py) of all the variables
+    created or mutated.
+  """
+  analyzer = ActivityAnalyzer(context, None, True)
+  analyzer.visit(node)
+  return analyzer.scope.created | analyzer.scope.modified
+
+
 def resolve(node, context, parent_scope=None):
   return ActivityAnalyzer(context, parent_scope).visit(node)
diff --git a/tensorflow/contrib/autograph/pyct/static_analysis/activity_test.py b/tensorflow/contrib/autograph/pyct/static_analysis/activity_test.py
index ef79a295bf..fdbd349af9 100644
--- a/tensorflow/contrib/autograph/pyct/static_analysis/activity_test.py
+++ b/tensorflow/contrib/autograph/pyct/static_analysis/activity_test.py
@@ -123,7 +123,7 @@ class ActivityAnalyzerTest(test.TestCase):
         recursive=True)
     node = qual_names.resolve(node)
     node = activity.resolve(node, ctx)
-    return node
+    return node, ctx
 
   def test_local_markers(self):
 
@@ -133,7 +133,7 @@ class ActivityAnalyzerTest(test.TestCase):
         b -= 1
       return b
 
-    node = self._parse_and_analyze(test_fn)
+    node, _ = self._parse_and_analyze(test_fn)
     self.assertFalse(
         anno.getanno(node.body[0].body[0].value,
                      NodeAnno.IS_LOCAL))  # c in b = c
@@ -156,6 +156,7 @@ class ActivityAnalyzerTest(test.TestCase):
                               expected - actual, actual - expected))
 
   def assertScopeIsRmc(self, scope, used, modified, created):
+    """Assert the scope contains specific used, modified & created variables."""
     self.assertSymbolSetsAre(used, scope.used, 'read')
     self.assertSymbolSetsAre(modified, scope.modified, 'modified')
     self.assertSymbolSetsAre(created, scope.created, 'created')
@@ -168,7 +169,7 @@ class ActivityAnalyzerTest(test.TestCase):
       print(a, b)
       return c
 
-    node = self._parse_and_analyze(test_fn)
+    node, _ = self._parse_and_analyze(test_fn)
     print_node = node.body[0].body[2]
     if isinstance(print_node, gast.Print):
       # Python 2
@@ -191,7 +192,7 @@ class ActivityAnalyzerTest(test.TestCase):
       foo(a, b)  # pylint:disable=undefined-variable
       return c
 
-    node = self._parse_and_analyze(test_fn)
+    node, _ = self._parse_and_analyze(test_fn)
     call_node = node.body[0].body[2].value
     # We basically need to detect which variables are captured by the call
     # arguments.
@@ -208,7 +209,7 @@ class ActivityAnalyzerTest(test.TestCase):
       foo(a.b, a.c)
       return a.d
 
-    node = self._parse_and_analyze(test_fn)
+    node, _ = self._parse_and_analyze(test_fn)
     call_node = node.body[0].body[1].value
     self.assertScopeIsRmc(
         anno.getanno(call_node, NodeAnno.ARGS_SCOPE),
@@ -234,7 +235,7 @@ class ActivityAnalyzerTest(test.TestCase):
       foo(a[0], a[b])
       return a[c]
 
-    node = self._parse_and_analyze(test_fn)
+    node, _ = self._parse_and_analyze(test_fn)
     call_node = node.body[0].body[2].value
     self.assertScopeIsRmc(
         anno.getanno(call_node, NodeAnno.ARGS_SCOPE),
@@ -258,7 +259,7 @@ class ActivityAnalyzerTest(test.TestCase):
         b -= 1
       return b, c
 
-    node = self._parse_and_analyze(test_fn)
+    node, _ = self._parse_and_analyze(test_fn)
     while_node = node.body[0].body[1]
     self.assertScopeIsRmc(
         anno.getanno(while_node, NodeAnno.BODY_SCOPE), ('b',), ('b', 'c'),
@@ -278,7 +279,7 @@ class ActivityAnalyzerTest(test.TestCase):
         b -= 1
       return b, c
 
-    node = self._parse_and_analyze(test_fn)
+    node, _ = self._parse_and_analyze(test_fn)
     for_node = node.body[0].body[1]
     self.assertScopeIsRmc(
         anno.getanno(for_node, NodeAnno.BODY_SCOPE), ('b',), ('b', 'c'), ('c',))
@@ -299,7 +300,7 @@ class ActivityAnalyzerTest(test.TestCase):
         u = -y
       return z, u
 
-    node = self._parse_and_analyze(test_fn)
+    node, _ = self._parse_and_analyze(test_fn)
     if_node = node.body[0].body[0]
     self.assertScopeIsRmc(
         anno.getanno(if_node, NodeAnno.BODY_SCOPE), ('x', 'y'), ('x', 'y', 'z'),
@@ -326,7 +327,7 @@ class ActivityAnalyzerTest(test.TestCase):
         d = 1
       return d
 
-    node = self._parse_and_analyze(test_fn)
+    node, _ = self._parse_and_analyze(test_fn)
     if_node = node.body[0].body[0]
     self.assertScopeIsRmc(
         anno.getanno(if_node, NodeAnno.BODY_SCOPE),
@@ -358,7 +359,7 @@ class ActivityAnalyzerTest(test.TestCase):
         d = 1
       return d
 
-    node = self._parse_and_analyze(test_fn)
+    node, _ = self._parse_and_analyze(test_fn)
     if_node = node.body[0].body[0]
     self.assertScopeIsRmc(
         anno.getanno(if_node, NodeAnno.BODY_SCOPE),
@@ -390,7 +391,7 @@ class ActivityAnalyzerTest(test.TestCase):
           a = b * b
       return a
 
-    node = self._parse_and_analyze(test_fn)
+    node, _ = self._parse_and_analyze(test_fn)
     inner_if_node = node.body[0].body[0].body[0]
     self.assertScopeIsRmc(
         anno.getanno(inner_if_node, NodeAnno.BODY_SCOPE), ('b',), ('a',),
@@ -413,7 +414,7 @@ class ActivityAnalyzerTest(test.TestCase):
         b -= f(i)
       return b, c
 
-    node = self._parse_and_analyze(test_fn)
+    node, _ = self._parse_and_analyze(test_fn)
     fn_def_node = node.body[0].body[0]
 
     self.assertScopeIsRmc(
@@ -434,7 +435,7 @@ class ActivityAnalyzerTest(test.TestCase):
         self.b = a
         self.b.c = 1
 
-    node = self._parse_and_analyze(TestClass)
+    node, _ = self._parse_and_analyze(TestClass)
     init_node = node.body[0].body[0]
     self.assertScopeIsRmc(
         anno.getanno(init_node, NodeAnno.BODY_SCOPE),
@@ -448,15 +449,118 @@ class ActivityAnalyzerTest(test.TestCase):
     def test_fn(a):
       a[0] += 1
 
-    node = self._parse_and_analyze(test_fn)
+    node, _ = self._parse_and_analyze(test_fn)
     fn_node = node.body[0]
     self.assertScopeIsRmc(
         anno.getanno(fn_node, NodeAnno.BODY_SCOPE),
-        ('a',),
+        ('a', 'a[0]'),
         ('a', 'a[0]'),
         ('a',),
     )
 
+  def test_return_vars_are_read(self):
+
+    def test_fn(a, b, c):  # pylint: disable=unused-argument
+      return c
+
+    node, _ = self._parse_and_analyze(test_fn)
+    fn_node = node.body[0]
+    self.assertScopeIsRmc(
+        anno.getanno(fn_node, NodeAnno.BODY_SCOPE),
+        ('c',),
+        (),
+        (
+            'a',
+            'b',
+            'c',
+        ),
+    )
+
+  def test_aug_assign(self):
+
+    def test_fn(a, b):
+      a += b
+
+    node, _ = self._parse_and_analyze(test_fn)
+    fn_node = node.body[0]
+    self.assertScopeIsRmc(
+        anno.getanno(fn_node, NodeAnno.BODY_SCOPE),
+        ('a', 'b'),
+        ('a'),
+        ('a', 'b'),
+    )
+
+  def test_aug_assign_rvalues(self):
+
+    a = dict(bar=3)
+
+    def foo():
+      return a
+
+    def test_fn(x):
+      foo()['bar'] += x
+
+    node, _ = self._parse_and_analyze(test_fn)
+    fn_node = node.body[0]
+    self.assertScopeIsRmc(
+        anno.getanno(fn_node, NodeAnno.BODY_SCOPE),
+        ('foo', 'x'),
+        (),
+        ('x',),
+    )
+
+  def test_params_created(self):
+
+    def test_fn(a, b):  # pylint: disable=unused-argument
+      return b
+
+    node, _ = self._parse_and_analyze(test_fn)
+    fn_node = node.body[0]
+    self.assertScopeIsRmc(
+        anno.getanno(fn_node, NodeAnno.BODY_SCOPE), ('b',), (('')),
+        (('a', 'b')))
+
+  def test_get_read(self):
+
+    def test_fn(x, y):
+      z = test_fn(x, y)
+      return z
+
+    node, ctx = self._parse_and_analyze(test_fn)
+    node = node.body[0].body[0]
+    read_vars = activity.get_read(node, ctx)
+    self.assertEqual(read_vars, set(map(qual_names.QN, ('test_fn', 'x', 'y'))))
+
+    def test_fn2(x, y, z):
+      z += test_fn2(x, y, z)
+      return z
+
+    node, ctx = self._parse_and_analyze(test_fn2)
+    node = node.body[0].body[0]
+    read_vars = activity.get_read(node, ctx)
+    self.assertEqual(read_vars,
+                     set(map(qual_names.QN, ('test_fn2', 'x', 'y', 'z'))))
+
+  def test_get_updated(self):
+
+    def test_fn(x, y):
+      z = test_fn(x, y)
+      return z
+
+    node, ctx = self._parse_and_analyze(test_fn)
+    node = node.body[0].body[0]
+    updated_vars = activity.get_updated(node, ctx)
+    self.assertEqual(updated_vars, set(map(qual_names.QN, ('z'))))
+
+    def test_fn2(x, y, z):
+      z += test_fn2(x, y, z)
+      return z
+
+    node, ctx = self._parse_and_analyze(test_fn2)
+    node = node.body[0].body[0]
+    updated_vars = activity.get_updated(node, ctx)
+    self.assertEqual(updated_vars, set(map(qual_names.QN, ('z'))))
+
 
 if __name__ == '__main__':
   test.main()
-- 
GitLab


From 28d43e5ada3c1e16b81c64b08cbbc273407a0347 Mon Sep 17 00:00:00 2001
From: Zhixian Yan <zhixianyan@google.com>
Date: Thu, 3 May 2018 14:34:27 -0700
Subject: [PATCH 234/395] Add tflite listed models with accuracy and
 performance numbers.

PiperOrigin-RevId: 195312636
---
 tensorflow/contrib/lite/g3doc/models.md | 87 +++++++++++++++++--------
 1 file changed, 61 insertions(+), 26 deletions(-)

diff --git a/tensorflow/contrib/lite/g3doc/models.md b/tensorflow/contrib/lite/g3doc/models.md
index d8134d5a00..c1c8ef049f 100644
--- a/tensorflow/contrib/lite/g3doc/models.md
+++ b/tensorflow/contrib/lite/g3doc/models.md
@@ -1,28 +1,63 @@
 # List of Hosted Models
 
-*   [NASNet large](https://storage.googleapis.com/download.tensorflow.org/models/tflite/nasnet_large_2018_03_27.zip)
-*   [NASNet mobile](https://storage.googleapis.com/download.tensorflow.org/models/tflite/nasnet_mobile_2018_03_27.zip)
-*   [ResNet v2 101](https://storage.googleapis.com/download.tensorflow.org/models/tflite/resnet_v2_101_2018_03_27.zip)
-*   [ResNet v2 50](https://storage.googleapis.com/download.tensorflow.org/models/tflite/resnet_v2_50_2018_03_27.zip)
-*   [Inception ResNet v2](https://storage.googleapis.com/download.tensorflow.org/models/tflite/inception_resnet_v2_2018_03_27.zip)
-*   [Inception v4](https://storage.googleapis.com/download.tensorflow.org/models/tflite/inception_v4_2018_03_27.zip)
-*   [Inception v3 2015](https://storage.googleapis.com/download.tensorflow.org/models/tflite/inception_v3_2015_2017_11_10.zip)
-*   [Inception v3 Slim 2016](https://storage.googleapis.com/download.tensorflow.org/models/tflite/inception_v3_slim_2016_android_2017_11_10.zip)
-*   [Mobilenet 0.25 128 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_0.25_128_float_2017_11_08.zip)
-*   [Mobilenet 0.25 160 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_0.25_160_float_2017_11_08.zip)
-*   [Mobilenet 0.25 192 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_0.25_192_float_2017_11_08.zip)
-*   [Mobilenet 0.25 224 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_0.25_224_float_2017_11_08.zip)
-*   [Mobilenet 0.50 128 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_0.50_128_float_2017_11_08.zip)
-*   [Mobilenet 0.50 160 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_0.50_160_float_2017_11_08.zip)
-*   [Mobilenet 0.50 192 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_0.50_192_float_2017_11_08.zip)
-*   [Mobilenet 0.50 224 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_0.50_224_float_2017_11_08.zip)
-*   [Mobilenet 0.75 128 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_0.75_128_float_2017_11_08.zip)
-*   [Mobilenet 0.75 160 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_0.75_160_float_2017_11_08.zip)
-*   [Mobilenet 0.75 192 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_0.75_192_float_2017_11_08.zip)
-*   [Mobilenet 0.75 224 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_0.75_224_float_2017_11_08.zip)
-*   [Mobilenet 1.0 128 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_1.0_128_float_2017_11_08.zip)
-*   [Mobilenet 1.0 160 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_1.0_160_float_2017_11_08.zip)
-*   [Mobilenet 1.0 192 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_1.0_192_float_2017_11_08.zip)
-*   [Mobilenet 1.0 224 Float](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_1.0_224_float_2017_11_08.zip)
-*   [Mobilenet 1.0 224 Quant](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_224_android_quant_2017_11_08.zip)
-*   [Smart Reply 1.0 Android ](https://storage.googleapis.com/download.tensorflow.org/models/tflite/smartreply_1.0_2017_11_01.zip)
+## Image classification (Float Models)
+
+Model Name          | Paper_Model_Files^                                                                                                                                                                        | Model_Size | Top-1 Accuracy | Top-5 Accuracy | TF Lite Performance^^ | Tensorflow Performance
+------------------- | :---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------: | ---------: | -------------: | -------------: | --------------------: | ---------------------:
+DenseNet            | [paper](https://arxiv.org/abs/1608.06993), [tflite&pb](https://storage.googleapis.com/download.tensorflow.org/models/tflite/model_zoo/upload_20180427/densenet_2018_04_27.tgz)            | 43.6 Mb    | 64.2%          | 85.6%          | 894 ms                | 1262 ms
+SqueezeNet          | [paper](https://arxiv.org/abs/1602.07360), [tflite&pb](https://storage.googleapis.com/download.tensorflow.org/models/tflite/model_zoo/upload_20180427/squeezenet_2018_04_27.tgz)          | 5.0 Mb     | 49.0%          | 72.9%          | 224 ms                | 255 ms
+NASNet mobile       | [paper](https://arxiv.org/abs/1707.07012), [tflite&pb](https://storage.googleapis.com/download.tensorflow.org/models/tflite/model_zoo/upload_20180427/nasnet_mobile_2018_04_27.tgz)       | 21.4 Mb    | 72.2%          | 90.6%          | 261 ms                | 389 ms
+NASNet large        | [paper](https://arxiv.org/abs/1707.07012), [tflite&pb](https://storage.googleapis.com/download.tensorflow.org/models/tflite/model_zoo/upload_20180427/nasnet_large_2018_04_27.tgz)        | 355.3 Mb   | 82.1%          | 95.8%          | 6697 ms               | 7940 ms
+ResNet_V2_50        | [paper](https://arxiv.org/abs/1603.05027), [tflite&pb](https://storage.googleapis.com/download.tensorflow.org/models/tflite/model_zoo/upload_20180427/resnet_v2_50_2018_04_27.tgz)        | 102.3 Mb   | 68.1%          | 88.4%          | 942 ms                | 1008 ms
+ResNet_V2_101       | [paper](https://arxiv.org/abs/1603.05027), [tflite&pb](https://storage.googleapis.com/download.tensorflow.org/models/tflite/model_zoo/upload_20180427/resnet_v2_101_2018_04_27.tgz)       | 178.3 Mb   | 70.4%          | 89.6%          | 1880 ms               | 1970 ms
+Inception_V3        | [paper](http://arxiv.org/abs/1512.00567), [tflite&pb](https://storage.googleapis.com/download.tensorflow.org/models/tflite/model_zoo/upload_20180427/inception_v3_2018_04_27.tgz)         | 95.3 Mb    | 76.9%          | 93.5%          | 1433 ms               | 1522 ms
+Inception_V4        | [paper](http://arxiv.org/abs/1602.07261), [tflite&pb](https://storage.googleapis.com/download.tensorflow.org/models/tflite/model_zoo/upload_20180427/inception_v4_2018_04_27.tgz)         | 170.7 Mb   | 79.6%          | 94.6%          | 2986 ms               | 3139 ms
+Inception_ResNet_V2 | [paper](https://arxiv.org/abs/1602.07261), [tflite&pb](https://storage.googleapis.com/download.tensorflow.org/models/tflite/model_zoo/upload_20180427/inception_resnet_v2_2018_04_27.tgz) | 121.0 Mb   | 76.8%          | 93.5%          | 2731 ms               | 2926 ms
+Mobilenet_0.25_128  | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.25_128.tgz)                                       | 1.9 Mb     | 41.5%          | 66.3%          | 6.2 ms                | 13.0 ms
+Mobilenet_0.25_160  | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.25_160.tgz)                                       | 1.9 Mb     | 45.5%          | 70.3%          | 8.6 ms                | 19.5 ms
+Mobilenet_0.25_192  | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.25_192.tgz)                                       | 1.9 Mb     | 47.7%          | 72.3%          | 12.1 ms               | 27.8 ms
+Mobilenet_0.25_224  | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.25_224.tgz)                                       | 1.9 Mb     | 49.8%          | 74.2%          | 16.2 ms               | 37.3 ms
+Mobilenet_0.50_128  | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.5_128.tgz)                                        | 5.3 Mb     | 56.3%          | 79.4%          | 18.1 ms               | 29.9 ms
+Mobilenet_0.50_160  | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.5_160.tgz)                                        | 5.3 Mb     | 59.1%          | 81.9%          | 26.8 ms               | 45.9 ms
+Mobilenet_0.50_192  | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.5_192.tgz)                                        | 5.3 Mb     | 61.7%          | 83.6%          | 35.6 ms               | 65.3 ms
+Mobilenet_0.50_224  | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.5_224.tgz)                                        | 5.3 Mb     | 63.3%          | 84.9%          | 47.6 ms               | 164.2 ms
+Mobilenet_0.75_128  | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.75_128.tgz)                                       | 10.3 Mb    | 62.1%          | 83.9%          | 34.6 ms               | 48.7 ms
+Mobilenet_0.75_160  | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.75_160.tgz)                                       | 10.3 Mb    | 65.3%          | 86.0%          | 51.3 ms               | 75.2 ms
+Mobilenet_0.75_192  | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.75_192.tgz)                                       | 10.3 Mb    | 67.2%          | 87.3%          | 71.7 ms               | 107.0 ms
+Mobilenet_0.75_224  | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.75_224.tgz)                                       | 10.3 Mb    | 68.4%          | 88.2%          | 95.7 ms               | 143.4 ms
+Mobilenet_1.0_128   | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_1.0_128.tgz)                                        | 16.9 Mb    | 65.2%          | 85.8%          | 57.4 ms               | 76.8 ms
+Mobilenet_1.0_160   | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_1.0_160.tgz)                                        | 16.9 Mb    | 68.0%          | 87.7%          | 86.0 ms               | 117.7 ms
+Mobilenet_1.0_192   | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_1.0_192.tgz)                                        | 16.9 Mb    | 70.0%          | 89.2%          | 118.6 ms              | 167.3 ms
+Mobilenet_1.0_224   | [paper](https://arxiv.org/pdf/1704.04861.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_1.0_224.tgz)                                        | 16.9 Mb    | 70.9%          | 89.9%          | 160.1 ms              | 224.3 ms
+
+^ The model files include both TF Lite FlatBuffer and Tensorflow frozen Graph.
+
+^^ The performance numbers are generated in the benchmark on Pixel-2 using
+single thread large core.
+
+## Image classification (Quantized Models)
+
+Model Name               | Paper_Model_Files                                                                                                                                         | Model_Size | Top-1 Accuracy | Top-5 Accuracy | TF Lite Performance
+------------------------ | :-------------------------------------------------------------------------------------------------------------------------------------------------------: | ---------: | -------------: | -------------: | ------------------:
+Mobilenet_0.25_128_quant | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.25_128_quant.tgz) | 0.5 Mb     | 39.9%          | 65.8%          | 3.7 ms
+Mobilenet_0.25_160_quant | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.25_160_quant.tgz) | 0.5 Mb     | 43.5%          | 69.1%          | 5.5 ms
+Mobilenet_0.25_192_quant | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.25_192_quant.tgz) | 0.5 Mb     | 45.8%          | 71.9%          | 7.9 ms
+Mobilenet_0.25_224_quant | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.25_224_quant.tgz) | 0.5 Mb     | 48.2%          | 73.8%          | 10.4 ms
+Mobilenet_0.50_128_quant | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.5_128_quant.tgz)  | 1.4 Mb     | 54.9%          | 78.9%          | 8.8 ms
+Mobilenet_0.50_160_quant | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.5_160_quant.tgz)  | 1.4 Mb     | 57.7%          | 81.3%          | 13.0 ms
+Mobilenet_0.50_192_quant | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.5_192_quant.tgz)  | 1.4 Mb     | 60.4%          | 83.2%          | 18.3 ms
+Mobilenet_0.50_224_quant | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.5_224_quant.tgz)  | 1.4 Mb     | 62.2%          | 84.5%          | 24.7 ms
+Mobilenet_0.75_128_quant | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.75_128_quant.tgz) | 2.6 Mb     | 59.8%          | 82.8%          | 16.2 ms
+Mobilenet_0.75_160_quant | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.75_160_quant.tgz) | 2.6 Mb     | 63.9%          | 85.5%          | 24.3 ms
+Mobilenet_0.75_192_quant | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.75_192_quant.tgz) | 2.6 Mb     | 66.2%          | 87.1%          | 33.8 ms
+Mobilenet_0.75_224_quant | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_0.75_224_quant.tgz) | 2.6 Mb     | 67.9%          | 88.1%          | 45.4 ms
+Mobilenet_1.0_128_quant  | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_1.0_128_quant.tgz)  | 4.3 Mb     | 64.0%          | 85.5%          | 24.9 ms
+Mobilenet_1.0_160_quant  | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_1.0_160_quant.tgz)  | 4.3 Mb     | 67.3%          | 87.7%          | 37.4 ms
+Mobilenet_1.0_192_quant  | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_1.0_192_quant.tgz)  | 4.3 Mb     | 69.0%          | 88.9%          | 51.9 ms
+Mobilenet_1.0_224_quant  | [paper](https://arxiv.org/pdf/1712.05877.pdf), [tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_1.0_224_quant.tgz)  | 4.3 Mb     | 69.7%          | 89.5%          | 70.2 ms
+
+## Other models
+
+Model                   | TF Lite FlatBuffer
+----------------------- | :----------------:
+Smart Reply 1.0 Android | [reference](https://research.googleblog.com/2017/11/on-device-conversational-modeling-with.html), [tflite](https://storage.googleapis.com/download.tensorflow.org/models/smartreply_1.0_2017_11_01.zip)
-- 
GitLab


From 4f4b15cece96c6cfa749c3fcf3288f1f47986210 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 3 May 2018 15:20:05 -0700
Subject: [PATCH 235/395] Fix bug that disabled loop invariant node motion
 optimizer. Disable it options, since it is broken in the presence of gradient
 stacks.

Get rid of an unnecessary copy of the graph.

PiperOrigin-RevId: 195319766
---
 .../core/grappler/optimizers/loop_optimizer.cc     | 14 ++++++--------
 .../core/grappler/optimizers/loop_optimizer.h      |  2 +-
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/loop_optimizer.cc b/tensorflow/core/grappler/optimizers/loop_optimizer.cc
index f7994221bb..5adc5b9227 100644
--- a/tensorflow/core/grappler/optimizers/loop_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/loop_optimizer.cc
@@ -474,15 +474,13 @@ std::vector<int> GetStackPushNodesToConvert(
   return nodes_to_convert;
 }
 
-Status RemoveStackOps(const GrapplerItem& item, GraphDef* optimized_graph) {
-  const std::unordered_set<string> nodes_to_preserve = item.NodesToPreserve();
-  const GraphDef& graph = item.graph;
-  *optimized_graph = graph;
+Status RemoveStackOps(const std::unordered_set<string>& nodes_to_preserve,
+                      GraphDef* optimized_graph) {
   NodeMap node_map(optimized_graph);
   SimpleGraphView graph_view;
-  TF_RETURN_IF_ERROR(graph_view.Initialize(graph));
-  for (int node_idx = 0; node_idx < graph.node_size(); ++node_idx) {
-    if (IsStackOp(graph.node(node_idx))) {
+  TF_RETURN_IF_ERROR(graph_view.Initialize(*optimized_graph));
+  for (int node_idx = 0; node_idx < optimized_graph->node_size(); ++node_idx) {
+    if (IsStackOp(optimized_graph->node(node_idx))) {
       for (int push_node_idx : GetStackPushNodesToConvert(
                graph_view, nodes_to_preserve, node_idx)) {
         // We found push nodes without corresponding pops. Convert them to
@@ -517,7 +515,7 @@ Status LoopOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
     TF_RETURN_IF_ERROR(linm_optimizer.Optimize());
   }
   if (options_.enable_stack_push_removal) {
-    TF_RETURN_IF_ERROR(RemoveStackOps(item, optimized_graph));
+    TF_RETURN_IF_ERROR(RemoveStackOps(item.NodesToPreserve(), optimized_graph));
   }
 
   return Status::OK();
diff --git a/tensorflow/core/grappler/optimizers/loop_optimizer.h b/tensorflow/core/grappler/optimizers/loop_optimizer.h
index a422505d23..764506f7c1 100644
--- a/tensorflow/core/grappler/optimizers/loop_optimizer.h
+++ b/tensorflow/core/grappler/optimizers/loop_optimizer.h
@@ -52,7 +52,7 @@ class LoopOptimizer : public GraphOptimizer {
 
   // Granular control for loop optimizer stages.
   struct LoopOptimizerOptions {
-    bool enable_loop_invariant_node_motion = true;
+    bool enable_loop_invariant_node_motion = false;
     bool enable_stack_push_removal = true;
 
     static LoopOptimizerOptions Default(RewriterConfig::Toggle opt_level) {
-- 
GitLab


From f25dd60858bc9ebe7b618aa966c2ddc1eef1f775 Mon Sep 17 00:00:00 2001
From: Dan Moldovan <mdan@google.com>
Date: Thu, 3 May 2018 15:39:46 -0700
Subject: [PATCH 236/395] Use tuple instead of list to reduce the chance of it
 being picked by the list conversions.

PiperOrigin-RevId: 195322522
---
 tensorflow/contrib/autograph/converters/asserts.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/autograph/converters/asserts.py b/tensorflow/contrib/autograph/converters/asserts.py
index 2d9e2c58e3..3b0db677ce 100644
--- a/tensorflow/contrib/autograph/converters/asserts.py
+++ b/tensorflow/contrib/autograph/converters/asserts.py
@@ -33,7 +33,7 @@ class AssertsTransformer(transformer.Base):
     # Note: The lone tf.Assert call will be wrapped with control_dependencies
     # by side_effect_guards.
     template = """
-      tf.Assert(test, [msg])
+      tf.Assert(test, (msg,))
     """
 
     if node.msg is None:
-- 
GitLab


From 549d63acd35872061ae42c36c94df6dbef18ee2b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 3 May 2018 15:42:23 -0700
Subject: [PATCH 237/395] Do not hoist nodes that modify frame info.

PiperOrigin-RevId: 195322927
---
 tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
index 29f49079c4..adfae2e1a3 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc
@@ -1541,7 +1541,7 @@ class HoistCWiseUnaryChainsStage : public ArithmeticOptimizerStage {
                          const ChainLinkSet& ops) const {
     if (ops.empty()) return true;
     const NodeDef* op0 = ops.begin()->node;
-    if (!IsUnaryElementWise(*op0)) return false;
+    if (ModifiesFrameInfo(*op0) || !IsUnaryElementWise(*op0)) return false;
     for (const auto& link : ops) {
       const NodeDef* op = link.node;
       if (op->device() != root_node.device() || op->op() != op0->op() ||
-- 
GitLab


From 200f4a2089cd4bef7832679cd121a2dbe85d6180 Mon Sep 17 00:00:00 2001
From: Justin Lebar <jlebar@google.com>
Date: Thu, 3 May 2018 15:58:43 -0700
Subject: [PATCH 238/395] Fix oom_test so that it doesn't try to allocate a
 giant host buffer when run without --config=cuda.  Sadly the best way I could
 come up with is pretty hacky.

PiperOrigin-RevId: 195325149
---
 tensorflow/compiler/tests/oom_test.py | 29 ++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/tensorflow/compiler/tests/oom_test.py b/tensorflow/compiler/tests/oom_test.py
index 1434e965e3..d68d32057a 100644
--- a/tensorflow/compiler/tests/oom_test.py
+++ b/tensorflow/compiler/tests/oom_test.py
@@ -22,6 +22,8 @@ from tensorflow.compiler.tests import xla_test
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
 from tensorflow.python.platform import googletest
 
 
@@ -42,20 +44,33 @@ class OutOfMemoryTest(xla_test.XLATestCase):
     """
 
     def test_loop():
-      size = 2e8
+      size = int(2e8)
       while True:
         with self.test_session():
-          # Force the compiled code to not be constant by feeding in an addend.
-          p = array_ops.placeholder(dtypes.float32, shape=[])
+          # Force the compiled code to not be constant by feeding in a
+          # parameter.
+          p = array_ops.placeholder(dtypes.float32, shape=[2, 1, 1])
           with self.test_scope():
-            # Create a large R1 tensor.
-            c = array_ops.zeros([size, 1]) + p
+            # Create a computation that produces a large R1 tensor as an
+            # intermediate result.  Reduce it down so that if this file was
+            # compiled without --config=cuda, we don't force a D2H copy of a
+            # large tensor and potentially OOM the host.
+            #
+            # This is a bit tricky because XLA:GPU doesn't currently support RNG
+            # ops.  Here we rely on the fact that XLA doesn't do algebraic
+            # simplifications on conv(<ones>, <filter>).
+            c = math_ops.reduce_sum(
+                nn_ops.convolution(
+                    array_ops.ones([1, size, 1]),
+                    p,
+                    padding='SAME',
+                    data_format='NWC'))
 
-            c.eval(feed_dict={p: 1.0})
+            c.eval(feed_dict={p: [[[1.0]], [[2.0]]]})
             size *= 2
 
     self.assertRaises(errors.ResourceExhaustedError, test_loop)
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
   googletest.main()
-- 
GitLab


From 04d5adbf848eba94e5c352cd0843a094b9fa0a4a Mon Sep 17 00:00:00 2001
From: Jeremy Lau <lauj@google.com>
Date: Thu, 3 May 2018 16:08:48 -0700
Subject: [PATCH 239/395] Fix bugs in LogicalBuffer::ToString and
 BufferValue::ToProto: these functions may be called before set_color(), but
 color() check fails when no color is set.

PiperOrigin-RevId: 195327063
---
 tensorflow/compiler/xla/service/buffer_value.cc   | 4 +++-
 tensorflow/compiler/xla/service/logical_buffer.cc | 6 +++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/xla/service/buffer_value.cc b/tensorflow/compiler/xla/service/buffer_value.cc
index df1a5ca435..2bc556a9e2 100644
--- a/tensorflow/compiler/xla/service/buffer_value.cc
+++ b/tensorflow/compiler/xla/service/buffer_value.cc
@@ -59,7 +59,9 @@ LogicalBufferProto BufferValue::ToProto(const SizeFunction& size_fn) const {
   LogicalBufferProto::Location proto_location =
       ToLocationProto(*instruction(), index());
   proto.mutable_defined_at()->Swap(&proto_location);
-  proto.set_color(color().value());
+  if (has_color()) {
+    proto.set_color(color().value());
+  }
   return proto;
 }
 
diff --git a/tensorflow/compiler/xla/service/logical_buffer.cc b/tensorflow/compiler/xla/service/logical_buffer.cc
index 1b3de8ad17..c742d35a7b 100644
--- a/tensorflow/compiler/xla/service/logical_buffer.cc
+++ b/tensorflow/compiler/xla/service/logical_buffer.cc
@@ -32,9 +32,13 @@ LogicalBuffer::LogicalBuffer(HloInstruction* instruction,
 LogicalBuffer::~LogicalBuffer() {}
 
 string LogicalBuffer::ToString() const {
+  string color_string;
+  if (has_color()) {
+    color_string = tensorflow::strings::StrCat(" @", color().value());
+  }
   return tensorflow::strings::StrCat(instruction_->name(), "[",
                                      tensorflow::str_util::Join(index_, ","),
-                                     "](#", id(), " @", color().value(), ")");
+                                     "](#", id(), color_string, ")");
 }
 
 }  // namespace xla
-- 
GitLab


From c9a92808ab8c1e19fd0a6bba5b9814c8c2c42511 Mon Sep 17 00:00:00 2001
From: Russell Power <power@google.com>
Date: Thu, 3 May 2018 16:16:05 -0700
Subject: [PATCH 240/395] Adjust worker shutdown hooks for TPUs

PiperOrigin-RevId: 195328247
---
 .../contrib/tpu/python/tpu/session_support.py | 47 +++++++++++++++----
 .../contrib/tpu/python/tpu/tpu_estimator.py   | 23 ++++++++-
 2 files changed, 58 insertions(+), 12 deletions(-)

diff --git a/tensorflow/contrib/tpu/python/tpu/session_support.py b/tensorflow/contrib/tpu/python/tpu/session_support.py
index 7c25f6693c..3455e0b4a6 100644
--- a/tensorflow/contrib/tpu/python/tpu/session_support.py
+++ b/tensorflow/contrib/tpu/python/tpu/session_support.py
@@ -126,12 +126,21 @@ class WorkerHeartbeatManager(object):
     return WorkerHeartbeatManager(self._session, bad_devices, bad_ops,
                                   self._request_placeholder)
 
+  def __repr__(self):
+    return 'HeartbeatManager(%s)' % ','.join(self._devices)
+
   def shutdown(self, timeout_ms=10000):
     """Shutdown all workers after `shutdown_timeout_secs`."""
+    logging.info('Shutting down %s.', self)
     req = event_pb2.WorkerHeartbeatRequest(
         watchdog_config=event_pb2.WatchdogConfig(timeout_ms=timeout_ms))
     self.configure(req)
 
+    # Wait for workers to shutdown.  This isn't strictly required
+    # but it avoids triggering multiple checkpoints with the same lame worker.
+    logging.info('Waiting %dms for worker shutdown.', timeout_ms)
+    time.sleep(timeout_ms / 1000)
+
 
 def all_worker_devices(session):
   """Return a list of devices for each worker in the system."""
@@ -250,6 +259,7 @@ class GracefulShutdownHook(session_run_hook.SessionRunHook):
           ' in your model definition to allow checkpointing.')
 
     with self._graph.as_default():
+      logging.info('Installing graceful shutdown hook.')
       self._session = session_lib.Session(
           target=training_session.sess_str, graph=self._graph)
       self._workers = WorkerHeartbeatManager.from_devices(
@@ -296,16 +306,33 @@ class GracefulShutdownHook(session_run_hook.SessionRunHook):
         fn(run_context, self._workers, lame_workers)
 
 
-def restart_computation(run_context, all_workers, lame_workers):
-  del run_context, lame_workers
-  logging.info('Shutting down all workers.')
-  all_workers.shutdown()
+class RestartComputation(object):
+  """Restart the entire computation.
+
+  This hook shuts down all workers and returns control to the top-level by
+  throwing a CoordinatorShutdownException.
+  """
+
+  def __init__(self, timeout_ms=10000):
+    self.timeout_ms = timeout_ms
+
+  def __call__(self, run_context, all_workers, lame_workers):
+    del run_context, lame_workers
+    all_workers.shutdown(timeout_ms=self.timeout_ms)
+
+    logging.info('Terminating coordinator.')
+    raise CoordinatorShutdownException()
+
 
-  logging.info('Terminating coordinator.')
-  raise CoordinatorShutdownException()
+class ShutdownLameWorkers(object):
+  """Shutdown lamed workers.
+
+  Processing will continue normally (typically by waiting for the down
+  workers to be restarted).
+  """
 
+  def __init__(self, timeout_ms=10000):
+    self.timeout_in_ms = timeout_ms
 
-def shutdown_lame_workers(run_context, all_workers, lame_workers):
-  del run_context, all_workers
-  logging.info('Shutting down %s', lame_workers)
-  lame_workers.shutdown()
+  def __call__(self, run_context, all_workers, lame_workers):
+    lame_workers.shutdown(timeout_ms=self.timeout_in_ms)
diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
index 534042b42c..a69bfa9a20 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
@@ -2049,9 +2049,28 @@ class TPUEstimator(estimator_lib.Estimator):
           host_ops = host_call.create_tpu_hostcall()
           if host_ops is None:
             host_ops = []
+
           shutdown_hooks = []
-          if os.environ.get('TF_TPU_GRACEFUL_SHUTDOWN', '0') != '0':
-            shutdown_hooks.append(session_support.GracefulShutdownHook())
+          shutdown_mode = os.environ.get('TF_TPU_GRACEFUL_SHUTDOWN_MODE',
+                                         'shutdown_worker')
+          if shutdown_mode:
+            if shutdown_mode == 'shutdown_worker':
+              finalizer_hooks = [
+                  session_support.ShutdownLameWorkers(timeout_ms=1000),
+              ]
+            elif shutdown_mode == 'shutdown_computation':
+              finalizer_hooks = [
+                  session_support.RestartComputation(timeout_ms=1000),
+              ]
+            else:
+              raise ValueError('Unknown TF_TPU_GRACEFUL_SHUTDOWN_MODE "%s"' %
+                               shutdown_mode)
+
+            shutdown_hooks.append(session_support.GracefulShutdownHook(
+                checkpoint_prefix=self.model_dir + '/model.ckpt',
+                on_shutdown_hooks=finalizer_hooks
+            ))
+
           with ops.control_dependencies([loss]):
             global_step = array_ops.identity(training.get_global_step())
           hooks = input_hooks + shutdown_hooks
-- 
GitLab


From 4a74a5058f7cb3ac096fa582941d9ab801ba6d65 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 3 May 2018 16:34:11 -0700
Subject: [PATCH 241/395] Fix flaky test time-outs for dnn_test and rnn_test.

PiperOrigin-RevId: 195331183
---
 tensorflow/contrib/estimator/BUILD | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/contrib/estimator/BUILD b/tensorflow/contrib/estimator/BUILD
index 41a817673d..571e2e3a5d 100644
--- a/tensorflow/contrib/estimator/BUILD
+++ b/tensorflow/contrib/estimator/BUILD
@@ -77,6 +77,7 @@ py_test(
     tags = [
         "no_pip",
         "notsan",
+        "optonly",  # times out http://b/79220679
     ],
     deps = [
         ":dnn",
@@ -450,6 +451,7 @@ py_test(
         "no_pip",
         "noasan",  # times out
         "notsan",
+        "optonly",  # times out http://b/79220679
     ],
     deps = [
         ":head",
-- 
GitLab


From 213a98d893105945540e0169faa124ac7e1200ba Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 3 May 2018 17:03:03 -0700
Subject: [PATCH 242/395] [XLA] Redesign: deprecate ComputationBuilder.

PiperOrigin-RevId: 195335330
---
 tensorflow/compiler/xla/client/computation.h  |  2 +
 .../compiler/xla/client/computation_builder.h |  2 +
 tensorflow/compiler/xla/client/lib/BUILD      |  5 +-
 .../compiler/xla/client/lib/arithmetic.cc     | 90 +----------------
 .../compiler/xla/client/lib/arithmetic.h      | 55 +----------
 tensorflow/compiler/xla/client/lib/testing.cc | 16 ++-
 tensorflow/compiler/xla/client/lib/testing.h  |  1 -
 tensorflow/compiler/xla/service/BUILD         | 10 +-
 tensorflow/compiler/xla/service/cpu/BUILD     |  4 +-
 .../xla/service/cpu/sample_harness.cc         | 10 +-
 .../xla/service/hlo_cost_analysis_test.cc     | 73 ++++++--------
 .../xla/service/hlo_evaluator_test.cc         | 10 +-
 .../xla/service/hlo_tfgraph_builder_test.cc   |  1 -
 .../xla/service/transpose_folding_test.cc     | 10 +-
 .../zero_sized_hlo_elimination_test.cc        |  1 -
 tensorflow/compiler/xla/tests/BUILD           | 18 ----
 .../xla/tests/local_client_aot_test_helper.cc | 19 ++--
 .../xla/tests/set_return_value_test.cc        | 98 -------------------
 .../xla/tests/vector_ops_simple_test.cc       |  3 +-
 19 files changed, 85 insertions(+), 343 deletions(-)
 delete mode 100644 tensorflow/compiler/xla/tests/set_return_value_test.cc

diff --git a/tensorflow/compiler/xla/client/computation.h b/tensorflow/compiler/xla/client/computation.h
index a53fc9e9cf..9a1bcde763 100644
--- a/tensorflow/compiler/xla/client/computation.h
+++ b/tensorflow/compiler/xla/client/computation.h
@@ -30,6 +30,8 @@ namespace xla {
 // Wraps a ComputationHandle protobuf with a lifetime. Computation is
 // movable and not copyable to capture the same kind of unique
 // ownership that std::unique_ptr represents.
+//
+// TODO(b/74197823): Deprecated. Use XlaComputation instead.
 class Computation {
  public:
   // Creates a null Computation.
diff --git a/tensorflow/compiler/xla/client/computation_builder.h b/tensorflow/compiler/xla/client/computation_builder.h
index 9431c2c459..ac1eb915cc 100644
--- a/tensorflow/compiler/xla/client/computation_builder.h
+++ b/tensorflow/compiler/xla/client/computation_builder.h
@@ -48,6 +48,8 @@ namespace xla {
 // deferred from being handled until Build() is called.
 //
 // Thread-compatible.
+//
+// TODO(b/74197823): Deprecated. Use XlaBuilder instead.
 class ComputationBuilder {
  public:
   // client: client in which to build the computation.
diff --git a/tensorflow/compiler/xla/client/lib/BUILD b/tensorflow/compiler/xla/client/lib/BUILD
index 59c4a53c05..d49d959a6c 100644
--- a/tensorflow/compiler/xla/client/lib/BUILD
+++ b/tensorflow/compiler/xla/client/lib/BUILD
@@ -22,8 +22,6 @@ cc_library(
         "//tensorflow/compiler/xla:status_macros",
         "//tensorflow/compiler/xla:types",
         "//tensorflow/compiler/xla:xla_data_proto",
-        "//tensorflow/compiler/xla/client:computation",
-        "//tensorflow/compiler/xla/client:computation_builder",
         "//tensorflow/compiler/xla/client/xla_client:xla_builder",
         "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/core:lib",
@@ -43,9 +41,8 @@ cc_library(
         "//tensorflow/compiler/xla:util",
         "//tensorflow/compiler/xla:xla_data_proto",
         "//tensorflow/compiler/xla/client",
-        "//tensorflow/compiler/xla/client:computation",
-        "//tensorflow/compiler/xla/client:computation_builder",
         "//tensorflow/compiler/xla/client:global_data",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
         "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/compiler/xla/tests:test_utils",
         "//tensorflow/core:lib",
diff --git a/tensorflow/compiler/xla/client/lib/arithmetic.cc b/tensorflow/compiler/xla/client/lib/arithmetic.cc
index 63df449e0b..a1d34796cc 100644
--- a/tensorflow/compiler/xla/client/lib/arithmetic.cc
+++ b/tensorflow/compiler/xla/client/lib/arithmetic.cc
@@ -17,7 +17,8 @@ limitations under the License.
 
 #include <string>
 
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 #include "tensorflow/compiler/xla/shape_util.h"
 #include "tensorflow/compiler/xla/status_macros.h"
 #include "tensorflow/compiler/xla/types.h"
@@ -27,28 +28,6 @@ limitations under the License.
 namespace xla {
 namespace {
 
-using InstructionGenerator =
-    ComputationDataHandle (*)(ComputationBuilder*, const ComputationDataHandle&,
-                              const ComputationDataHandle&);
-
-Computation CreateScalarComputation(const string& name, PrimitiveType type,
-                                    ComputationBuilder* builder,
-                                    InstructionGenerator generator) {
-  std::unique_ptr<ComputationBuilder> b;
-  if (type == PRED) {
-    b = builder->CreateSubBuilder(name);
-  } else {
-    b = builder->CreateSubBuilder(
-        tensorflow::strings::StrCat(name, "_", PrimitiveType_Name(type)));
-  }
-
-  const Shape scalar = ShapeUtil::MakeShape(type, {});
-  auto lhs = b->Parameter(0, scalar, "lhs");
-  auto rhs = b->Parameter(1, scalar, "rhs");
-  generator(b.get(), lhs, rhs);
-  return b->BuildAndNoteError();
-}
-
 using XlaOpGenerator = XlaOp (*)(XlaBuilder*, const XlaOp&, const XlaOp&);
 
 XlaComputation CreateScalarComputation(const string& name, PrimitiveType type,
@@ -71,71 +50,6 @@ XlaComputation CreateScalarComputation(const string& name, PrimitiveType type,
 
 }  // namespace
 
-Computation CreateScalarAddComputation(PrimitiveType type,
-                                       ComputationBuilder* builder) {
-  return CreateScalarComputation(
-      "add", type, builder,
-      [](ComputationBuilder* b, const ComputationDataHandle& lhs,
-         const ComputationDataHandle& rhs) { return b->Add(lhs, rhs); });
-}
-
-Computation CreateScalarMultiplyComputation(PrimitiveType type,
-                                            ComputationBuilder* builder) {
-  return CreateScalarComputation(
-      "mul", type, builder,
-      [](ComputationBuilder* b, const ComputationDataHandle& lhs,
-         const ComputationDataHandle& rhs) { return b->Mul(lhs, rhs); });
-}
-
-Computation CreateScalarGeComputation(PrimitiveType type,
-                                      ComputationBuilder* builder) {
-  return CreateScalarComputation(
-      "ge", type, builder,
-      [](ComputationBuilder* b, const ComputationDataHandle& lhs,
-         const ComputationDataHandle& rhs) { return b->Ge(lhs, rhs); });
-}
-
-Computation CreateScalarMaxComputation(PrimitiveType type,
-                                       ComputationBuilder* builder) {
-  return CreateScalarComputation(
-      "max", type, builder,
-      [](ComputationBuilder* b, const ComputationDataHandle& lhs,
-         const ComputationDataHandle& rhs) { return b->Max(lhs, rhs); });
-}
-
-Computation CreateScalarMinComputation(PrimitiveType type,
-                                       ComputationBuilder* builder) {
-  return CreateScalarComputation(
-      "min", type, builder,
-      [](ComputationBuilder* b, const ComputationDataHandle& lhs,
-         const ComputationDataHandle& rhs) { return b->Min(lhs, rhs); });
-}
-
-Computation CreateScalarAndComputation(ComputationBuilder* builder) {
-  return CreateScalarComputation(
-      "and", PRED, builder,
-      [](ComputationBuilder* b, const ComputationDataHandle& lhs,
-         const ComputationDataHandle& rhs) { return b->And(lhs, rhs); });
-}
-
-Computation CreateScalarOrComputation(ComputationBuilder* builder) {
-  return CreateScalarComputation(
-      "or", PRED, builder,
-      [](ComputationBuilder* b, const ComputationDataHandle& lhs,
-         const ComputationDataHandle& rhs) { return b->Or(lhs, rhs); });
-}
-
-StatusOr<ComputationDataHandle> Any(const ComputationDataHandle& predicates,
-                                    ComputationBuilder* builder) {
-  auto f = builder->ConstantR0<bool>(false);
-  Computation logical_or = CreateScalarOrComputation(builder);
-  TF_ASSIGN_OR_RETURN(std::unique_ptr<Shape> predicates_shape,
-                      builder->GetShape(predicates));
-  std::vector<int64> all_dimensions(ShapeUtil::Rank(*predicates_shape));
-  std::iota(all_dimensions.begin(), all_dimensions.end(), 0);
-  return builder->Reduce(predicates, f, logical_or, all_dimensions);
-}
-
 XlaComputation CreateScalarAddComputation(PrimitiveType type,
                                           XlaBuilder* builder) {
   return CreateScalarComputation(
diff --git a/tensorflow/compiler/xla/client/lib/arithmetic.h b/tensorflow/compiler/xla/client/lib/arithmetic.h
index f4d3fc8015..64b6b7d633 100644
--- a/tensorflow/compiler/xla/client/lib/arithmetic.h
+++ b/tensorflow/compiler/xla/client/lib/arithmetic.h
@@ -18,83 +18,38 @@ limitations under the License.
 
 #include <memory>
 
-#include "tensorflow/compiler/xla/client/computation.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
 #include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
 
 namespace xla {
 
-// Creates a scalar add computation and returns it.
-Computation CreateScalarAddComputation(PrimitiveType type,
-                                       ComputationBuilder* builder);
-
-// Creates a scalar multiply computation and returns it.
-Computation CreateScalarMultiplyComputation(PrimitiveType type,
-                                            ComputationBuilder* builder);
-
-// Creates a scalar ge computation and returns it.
-Computation CreateScalarGeComputation(PrimitiveType type,
-                                      ComputationBuilder* builder);
-
-// Creates a scalar max computation and returns it.
-Computation CreateScalarMaxComputation(PrimitiveType type,
-                                       ComputationBuilder* builder);
-
-// Creates a scalar min computation and returns it.
-Computation CreateScalarMinComputation(PrimitiveType type,
-                                       ComputationBuilder* builder);
-
-// Creates a scalar logical AND computation and returns it.
-Computation CreateScalarAndComputation(ComputationBuilder* builder);
-
-// Creates a scalar logical OR computation and returns it.
-Computation CreateScalarOrComputation(ComputationBuilder* builder);
-
-// Returns whether any predicate in "predicates" is set.
-//
-// Note: if predicates is zero-sized, Any() vacuously returns false.
-StatusOr<ComputationDataHandle> Any(const ComputationDataHandle& predicates,
-                                    ComputationBuilder* builder);
-
-// TODO(b/74197823): This is a part of a NOT YET ready refactor.
-//
 // Creates a scalar add computation and returns it.
 XlaComputation CreateScalarAddComputation(PrimitiveType type,
                                           XlaBuilder* builder);
-// TODO(b/74197823): This is a part of a NOT YET ready refactor.
-//
+
 // Creates a scalar multiply computation and returns it.
 XlaComputation CreateScalarMultiplyComputation(PrimitiveType type,
                                                XlaBuilder* builder);
-// TODO(b/74197823): This is a part of a NOT YET ready refactor.
-//
+
 // Creates a scalar ge computation and returns it.
 XlaComputation CreateScalarGeComputation(PrimitiveType type,
                                          XlaBuilder* builder);
-// TODO(b/74197823): This is a part of a NOT YET ready refactor.
-//
+
 // Creates a scalar max computation and returns it.
 XlaComputation CreateScalarMaxComputation(PrimitiveType type,
                                           XlaBuilder* builder);
-// TODO(b/74197823): This is a part of a NOT YET ready refactor.
-//
+
 // Creates a scalar min computation and returns it.
 XlaComputation CreateScalarMinComputation(PrimitiveType type,
                                           XlaBuilder* builder);
-// TODO(b/74197823): This is a part of a NOT YET ready refactor.
-//
+
 // Creates a scalar logical AND computation and returns it.
 XlaComputation CreateScalarAndComputation(XlaBuilder* builder);
 
-// TODO(b/74197823): This is a part of a NOT YET ready refactor.
-//
 // Creates a scalar logical OR computation and returns it.
 XlaComputation CreateScalarOrComputation(XlaBuilder* builder);
 
-// TODO(b/74197823): This is a part of a NOT YET ready refactor.
-//
 // Returns whether any predicate in "predicates" is set.
 //
 // Note: if predicates is zero-sized, Any() vacuously returns false.
diff --git a/tensorflow/compiler/xla/client/lib/testing.cc b/tensorflow/compiler/xla/client/lib/testing.cc
index 311dc4bdd7..9cd87f7473 100644
--- a/tensorflow/compiler/xla/client/lib/testing.cc
+++ b/tensorflow/compiler/xla/client/lib/testing.cc
@@ -15,8 +15,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/client/lib/testing.h"
 
-#include "tensorflow/compiler/xla/client/computation.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/execution_options_util.h"
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/compiler/xla/shape_util.h"
@@ -46,16 +45,14 @@ int64 DataSizeOfShape(const Shape& shape) {
   return total_size;
 }
 
-// Create a ComputationDataHandle for an op what generates fake data with the
-// given shape.
-ComputationDataHandle BuildFakeDataOpOnDevice(const Shape& shape,
-                                              ComputationBuilder* builder) {
+// Creates a XlaOp for an op what generates fake data with the given shape.
+XlaOp BuildFakeDataOpOnDevice(const Shape& shape, XlaBuilder* builder) {
   if (ShapeUtil::IsArray(shape)) {
     return builder->Broadcast(
         builder->ConstantLiteral(Literal::One(shape.element_type())),
         AsInt64Slice(shape.dimensions()));
   }
-  std::vector<ComputationDataHandle> parts;
+  std::vector<XlaOp> parts;
   for (const Shape& s : shape.tuple_shapes()) {
     parts.push_back(BuildFakeDataOpOnDevice(s, builder));
   }
@@ -64,11 +61,10 @@ ComputationDataHandle BuildFakeDataOpOnDevice(const Shape& shape,
 
 std::unique_ptr<GlobalData> MakeFakeDataViaDeviceOrDie(const Shape& shape,
                                                        Client* client) {
-  ComputationBuilder b(
-      client,
+  XlaBuilder b(
       tensorflow::strings::StrCat("make_fake_", ShapeUtil::HumanString(shape)));
   BuildFakeDataOpOnDevice(shape, &b);
-  Computation computation = b.Build().ConsumeValueOrDie();
+  XlaComputation computation = b.Build().ConsumeValueOrDie();
 
   auto execution_options = CreateDefaultExecutionOptions();
   *execution_options.mutable_shape_with_output_layout() = shape;
diff --git a/tensorflow/compiler/xla/client/lib/testing.h b/tensorflow/compiler/xla/client/lib/testing.h
index 1dc2622972..9e06141b1f 100644
--- a/tensorflow/compiler/xla/client/lib/testing.h
+++ b/tensorflow/compiler/xla/client/lib/testing.h
@@ -20,7 +20,6 @@ limitations under the License.
 #include <vector>
 
 #include "tensorflow/compiler/xla/client/client.h"
-#include "tensorflow/compiler/xla/client/computation.h"
 #include "tensorflow/compiler/xla/client/global_data.h"
 #include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD
index 0b8b22b44c..9c362d8cad 100644
--- a/tensorflow/compiler/xla/service/BUILD
+++ b/tensorflow/compiler/xla/service/BUILD
@@ -233,7 +233,7 @@ tf_cc_test(
         "//tensorflow/compiler/xla:types",
         "//tensorflow/compiler/xla:util",
         "//tensorflow/compiler/xla:xla_data_proto",
-        "//tensorflow/compiler/xla/client:computation_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
         "//tensorflow/compiler/xla/service:hlo_element_type_converter",
         "//tensorflow/compiler/xla/tests:hlo_verified_test_base",
         "//tensorflow/compiler/xla/tests:literal_test_util",
@@ -1669,10 +1669,10 @@ tf_cc_test(
         "//tensorflow/compiler/xla:test_helpers",
         "//tensorflow/compiler/xla/client",
         "//tensorflow/compiler/xla/client:client_library",
-        "//tensorflow/compiler/xla/client:computation",
-        "//tensorflow/compiler/xla/client:computation_builder",
         "//tensorflow/compiler/xla/client:local_client",
         "//tensorflow/compiler/xla/client:padding",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/compiler/xla/tests:hlo_test_base",
         "//tensorflow/compiler/xla/tests:xla_internal_test_main",
         "//tensorflow/core:lib",
@@ -2406,7 +2406,6 @@ tf_cc_test(
     srcs = ["hlo_tfgraph_builder_test.cc"],
     deps = [
         ":hlo_tfgraph_builder",
-        "//tensorflow/compiler/xla/client:computation_builder",
         "//tensorflow/compiler/xla/tests:hlo_test_base",
         "//tensorflow/compiler/xla/tests:xla_internal_test_main",
         "//tensorflow/core:protos_all_cc",
@@ -2475,7 +2474,7 @@ tf_cc_test(
         "//tensorflow/compiler/xla:test",
         "//tensorflow/compiler/xla:test_helpers",
         "//tensorflow/compiler/xla:xla_data_proto",
-        "//tensorflow/compiler/xla/client:computation_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
         "//tensorflow/compiler/xla/service/gpu:ir_emission_utils",
         "//tensorflow/compiler/xla/tests:hlo_test_base",
         "//tensorflow/compiler/xla/tests:xla_internal_test_main",
@@ -2512,6 +2511,7 @@ tf_cc_test(
         "//tensorflow/compiler/xla:test_helpers",
         "//tensorflow/compiler/xla:xla_data_proto",
         "//tensorflow/compiler/xla/client:computation_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
         "//tensorflow/compiler/xla/tests:hlo_test_base",
         "//tensorflow/compiler/xla/tests:xla_internal_test_main",
         "//tensorflow/core:lib",
diff --git a/tensorflow/compiler/xla/service/cpu/BUILD b/tensorflow/compiler/xla/service/cpu/BUILD
index cb81e413a3..7e6d58c7fa 100644
--- a/tensorflow/compiler/xla/service/cpu/BUILD
+++ b/tensorflow/compiler/xla/service/cpu/BUILD
@@ -365,10 +365,10 @@ tf_cc_binary(
         "//tensorflow/compiler/xla:xla_data_proto",
         "//tensorflow/compiler/xla/client",
         "//tensorflow/compiler/xla/client:client_library",
-        "//tensorflow/compiler/xla/client:computation",
-        "//tensorflow/compiler/xla/client:computation_builder",
         "//tensorflow/compiler/xla/client:global_data",
         "//tensorflow/compiler/xla/client:local_client",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/core:lib",
     ],
 )
diff --git a/tensorflow/compiler/xla/service/cpu/sample_harness.cc b/tensorflow/compiler/xla/service/cpu/sample_harness.cc
index b3f4609d46..167aa4adda 100644
--- a/tensorflow/compiler/xla/service/cpu/sample_harness.cc
+++ b/tensorflow/compiler/xla/service/cpu/sample_harness.cc
@@ -19,10 +19,10 @@ limitations under the License.
 #include "tensorflow/compiler/xla/array4d.h"
 #include "tensorflow/compiler/xla/client/client.h"
 #include "tensorflow/compiler/xla/client/client_library.h"
-#include "tensorflow/compiler/xla/client/computation.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
 #include "tensorflow/compiler/xla/client/global_data.h"
 #include "tensorflow/compiler/xla/client/local_client.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/compiler/xla/statusor.h"
 #include "tensorflow/compiler/xla/types.h"
@@ -48,13 +48,13 @@ int main(int argc, char** argv) {
       client->TransferToServer(*param1_literal).ConsumeValueOrDie();
 
   // Build computation.
-  xla::ComputationBuilder builder(client, "");
+  xla::XlaBuilder builder("");
   auto p0 = builder.Parameter(0, param0_literal->shape(), "param0");
   auto p1 = builder.Parameter(1, param1_literal->shape(), "param1");
   auto add = builder.Add(p1, p0, {0});
 
-  xla::StatusOr<xla::Computation> computation_status = builder.Build();
-  xla::Computation computation = computation_status.ConsumeValueOrDie();
+  xla::StatusOr<xla::XlaComputation> computation_status = builder.Build();
+  xla::XlaComputation computation = computation_status.ConsumeValueOrDie();
 
   // Execute and transfer result of computation.
   xla::ExecutionProfile profile;
diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis_test.cc b/tensorflow/compiler/xla/service/hlo_cost_analysis_test.cc
index 81cc7c4bdc..16fdda8a8b 100644
--- a/tensorflow/compiler/xla/service/hlo_cost_analysis_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_cost_analysis_test.cc
@@ -20,16 +20,13 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/client/client.h"
 #include "tensorflow/compiler/xla/client/client_library.h"
-#include "tensorflow/compiler/xla/client/computation.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
 #include "tensorflow/compiler/xla/client/local_client.h"
 #include "tensorflow/compiler/xla/client/padding.h"
-#include "tensorflow/compiler/xla/service/computation_tracker.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 #include "tensorflow/compiler/xla/service/hlo_module.h"
 #include "tensorflow/compiler/xla/service/local_service.h"
 #include "tensorflow/compiler/xla/service/service.h"
-#include "tensorflow/compiler/xla/service/user_computation.h"
-#include "tensorflow/compiler/xla/service/versioned_computation_handle.h"
 #include "tensorflow/compiler/xla/shape_util.h"
 #include "tensorflow/compiler/xla/tests/hlo_test_base.h"
 #include "tensorflow/core/platform/logging.h"
@@ -58,11 +55,10 @@ class HloCostAnalysisTest : public ::testing::Test {
         // whitebox accesses to the user computation built from the client,
         // as shown in the BuildHloGraph functions below.
         service_(static_cast<Service*>(ClientLibrary::GetXlaService(
-            static_cast<LocalClient*>(client_)->platform()))),
-        computation_tracker_(service_->computation_tracker()) {
+            static_cast<LocalClient*>(client_)->platform()))) {
     // Create a computation for a unary user function: x => exp(x + 0.5)
     {
-      ComputationBuilder builder(client_, "add_and_exp");
+      XlaBuilder builder("add_and_exp");
       auto x = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x");
       auto half = builder.ConstantR0<float>(0.5);
       builder.Exp(builder.Add(x, half));
@@ -73,7 +69,7 @@ class HloCostAnalysisTest : public ::testing::Test {
 
     // Create a computation for a binary user function: (x, y) => x + y
     {
-      ComputationBuilder builder(client_, "add");
+      XlaBuilder builder("add");
       auto x = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x");
       auto y = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "y");
       builder.Add(x, y);
@@ -84,7 +80,7 @@ class HloCostAnalysisTest : public ::testing::Test {
 
     // Create a computation for a sigmoid function: x => 1 / (1 + exp(-x))
     {
-      ComputationBuilder builder(client_, "sigmoid");
+      XlaBuilder builder("sigmoid");
       auto x = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x");
       auto one = builder.ConstantR0<float>(1.0);
       builder.Div(one, builder.Add(one, builder.Exp(builder.Neg(x))));
@@ -95,7 +91,7 @@ class HloCostAnalysisTest : public ::testing::Test {
 
     // Create a computation for a binary max function: (x, y) => max (x, y)
     {
-      ComputationBuilder builder(client_, "max");
+      XlaBuilder builder("max");
       auto x = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x");
       auto y = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "y");
       builder.Max(x, y);
@@ -106,7 +102,7 @@ class HloCostAnalysisTest : public ::testing::Test {
 
     // Create a computation for a binary GT function: (x, y) => x > y
     {
-      ComputationBuilder builder(client_, "gt");
+      XlaBuilder builder("gt");
       auto x = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x");
       auto y = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "y");
       builder.Gt(x, y);
@@ -117,35 +113,30 @@ class HloCostAnalysisTest : public ::testing::Test {
   }
 
   // Build HLO graph from the given builder and return the HLO module.
-  std::unique_ptr<HloModule> BuildHloGraph(ComputationBuilder* builder) {
+  std::unique_ptr<HloModule> BuildHloGraph(XlaBuilder* builder) {
     auto computation_status = builder->Build();
     TF_CHECK_OK(computation_status.status());
     auto computation = computation_status.ConsumeValueOrDie();
-    auto user_computation_status =
-        computation_tracker_.Resolve(computation.handle());
-    TF_CHECK_OK(user_computation_status.status());
-    auto user_computation = user_computation_status.ConsumeValueOrDie();
-    VersionedComputationHandle versioned_handle =
-        user_computation->GetVersionedHandle();
-    return std::move(
-        computation_tracker_.BuildHloModule(versioned_handle, HloModuleConfig())
-            .ValueOrDie());
+    auto config = HloModule::CreateModuleConfigFromProto(computation.proto(),
+                                                         DebugOptions())
+                      .ConsumeValueOrDie();
+    return HloModule::CreateFromProto(computation.proto(), config)
+        .ConsumeValueOrDie();
   }
 
   Client* client_;
   Service* service_;
-  const ComputationTracker& computation_tracker_;
 
   // User computations used for higher order operations (e.g., Map, Reduce).
-  Computation add_;
-  Computation add_and_exp_;
-  Computation sigmoid_;
-  Computation max_;
-  Computation gt_;
+  XlaComputation add_;
+  XlaComputation add_and_exp_;
+  XlaComputation sigmoid_;
+  XlaComputation max_;
+  XlaComputation gt_;
 };
 
 TEST_F(HloCostAnalysisTest, MatrixMultiply) {
-  ComputationBuilder builder(client_, "matrix_multiply");
+  XlaBuilder builder("matrix_multiply");
   auto lhs = builder.Parameter(0, ShapeUtil::MakeShape(F32, {10, 5}), "lhs");
   auto rhs = builder.Parameter(1, ShapeUtil::MakeShape(F32, {5, 30}), "rhs");
   auto result = builder.Dot(lhs, rhs);
@@ -167,7 +158,7 @@ TEST_F(HloCostAnalysisTest, MatrixMultiply) {
 }
 
 TEST_F(HloCostAnalysisTest, Map) {
-  ComputationBuilder builder(client_, "map");
+  XlaBuilder builder("map");
   auto input = builder.Parameter(0, ShapeUtil::MakeShape(F32, {10}), "in");
   auto result = builder.Map({input}, add_and_exp_, {0});
 
@@ -184,7 +175,7 @@ TEST_F(HloCostAnalysisTest, Map) {
 }
 
 TEST_F(HloCostAnalysisTest, Convolution) {
-  ComputationBuilder builder(client_, "convolution");
+  XlaBuilder builder("convolution");
   auto input = builder.Parameter(
       0,
       ShapeUtil::MakeShape(F32, {/*p_dim=*/1, /*z_dim=*/1, /*y_dim=*/10,
@@ -213,7 +204,7 @@ TEST_F(HloCostAnalysisTest, Convolution) {
 }
 
 TEST_F(HloCostAnalysisTest, Reduce) {
-  ComputationBuilder builder(client_, "reduce");
+  XlaBuilder builder("reduce");
   auto input =
       builder.Parameter(0, ShapeUtil::MakeShape(F32, {10, 20}), "input");
   auto result =
@@ -231,7 +222,7 @@ TEST_F(HloCostAnalysisTest, Reduce) {
 }
 
 TEST_F(HloCostAnalysisTest, ReduceWindow) {
-  ComputationBuilder builder(client_, "reduce_window");
+  XlaBuilder builder("reduce_window");
   auto input =
       builder.Parameter(0, ShapeUtil::MakeShape(F32, {10, 20}), "input");
   auto result = builder.ReduceWindow(input, builder.ConstantR0<float>(0), add_,
@@ -248,7 +239,7 @@ TEST_F(HloCostAnalysisTest, ReduceWindow) {
 }
 
 TEST_F(HloCostAnalysisTest, SelectAndScatter) {
-  ComputationBuilder builder(client_, "select_and_scatter");
+  XlaBuilder builder("select_and_scatter");
   auto operand =
       builder.Parameter(0, ShapeUtil::MakeShape(F32, {10, 20}), "input");
   auto source =
@@ -269,7 +260,7 @@ TEST_F(HloCostAnalysisTest, SelectAndScatter) {
 }
 
 TEST_F(HloCostAnalysisTest, Broadcast) {
-  ComputationBuilder b(client_, "broadcast");
+  XlaBuilder b("broadcast");
   b.Broadcast(b.ConstantR0<float>(42), {10, 7});
   auto hlo_module = BuildHloGraph(&b);
   HloCostAnalysis analysis(ShapeSize);
@@ -280,7 +271,7 @@ TEST_F(HloCostAnalysisTest, Broadcast) {
 
 // Calculates the computation cost of a graph with more than one HLO node.
 TEST_F(HloCostAnalysisTest, FullyConnectedForward) {
-  ComputationBuilder builder(client_, "fully_connected_forward");
+  XlaBuilder builder("fully_connected_forward");
   auto input =
       builder.Parameter(0, ShapeUtil::MakeShape(F32, {10, 5}), "input");
   auto weight =
@@ -305,7 +296,7 @@ TEST_F(HloCostAnalysisTest, FullyConnectedForward) {
 TEST_F(HloCostAnalysisTest, MatmulAndConvolutionCanBeTheSameComputation) {
   HloCostAnalysis conv_analysis(ShapeSize);
   {
-    ComputationBuilder builder(client_, "conv_looking_matmul");
+    XlaBuilder builder("conv_looking_matmul");
     auto lhs = builder.Parameter(0, ShapeUtil::MakeShape(F32, {64, 64, 1, 1}),
                                  "input");
     auto rhs = builder.Parameter(1, ShapeUtil::MakeShape(F32, {64, 64, 1, 1}),
@@ -318,7 +309,7 @@ TEST_F(HloCostAnalysisTest, MatmulAndConvolutionCanBeTheSameComputation) {
 
   HloCostAnalysis matmul_analysis(ShapeSize);
   {
-    ComputationBuilder builder(client_, "matmul");
+    XlaBuilder builder("matmul");
     auto lhs =
         builder.Parameter(0, ShapeUtil::MakeShape(F32, {64, 64}), "input");
     auto rhs =
@@ -427,7 +418,7 @@ TEST_F(FusionCostAnalysis, NoLayout) {
 TEST_F(HloCostAnalysisTest, TupleCost) {
   HloCostAnalysis analysis(ShapeSize);
   {
-    ComputationBuilder builder(client_, "matmul");
+    XlaBuilder builder("matmul");
     auto x = builder.Parameter(0, ShapeUtil::MakeShape(F32, {123}), "x");
     auto y = builder.Parameter(1, ShapeUtil::MakeShape(F32, {42}), "y");
     auto tuple = builder.Tuple({x, y});
@@ -443,7 +434,7 @@ TEST_F(HloCostAnalysisTest, TupleCost) {
 }
 
 TEST_F(HloCostAnalysisTest, BaseDilatedConvolution) {
-  ComputationBuilder builder(client_, "BaseDilatedConvolution");
+  XlaBuilder builder("BaseDilatedConvolution");
   auto input = builder.Parameter(
       0,
       ShapeUtil::MakeShape(F32, {/*p_dim=*/1, /*z_dim=*/1, /*y_dim=*/10,
@@ -458,7 +449,7 @@ TEST_F(HloCostAnalysisTest, BaseDilatedConvolution) {
   auto result = builder.ConvGeneralDilated(
       input, kernel, /*window_strides=*/{1, 1}, /*padding=*/{{1, 1}, {1, 1}},
       /*lhs_dilation=*/{3, 5}, /*rhs_dilation=*/{7, 11},
-      ComputationBuilder::CreateDefaultConvDimensionNumbers(2));
+      XlaBuilder::CreateDefaultConvDimensionNumbers(2));
 
   // Run HLO cost analysis.
   auto hlo_module = BuildHloGraph(&builder);
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_test.cc b/tensorflow/compiler/xla/service/hlo_evaluator_test.cc
index 230147abfe..cc16446778 100644
--- a/tensorflow/compiler/xla/service/hlo_evaluator_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_test.cc
@@ -21,7 +21,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/compiler/xla/reference_util.h"
 #include "tensorflow/compiler/xla/service/hlo_computation.h"
@@ -827,7 +827,7 @@ TEST_P(HloEvaluatorTest, Simple4x4Conv2DWith2x2Kernel) {
   *window.add_dimensions() = dim;
 
   ConvolutionDimensionNumbers dnums =
-      ComputationBuilder::CreateDefaultConvDimensionNumbers(2);
+      XlaBuilder::CreateDefaultConvDimensionNumbers(2);
 
   const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 4, 4});
   b.AddInstruction(HloInstruction::CreateConvolve(
@@ -1046,7 +1046,7 @@ TEST_P(HloEvaluatorTest, DilatedBaseConv2DWithHighPadding) {
   *window.add_dimensions() = dim;
 
   ConvolutionDimensionNumbers dnums =
-      ComputationBuilder::CreateDefaultConvDimensionNumbers(2);
+      XlaBuilder::CreateDefaultConvDimensionNumbers(2);
 
   const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 7, 7});
   b.AddInstruction(HloInstruction::CreateConvolve(
@@ -1109,7 +1109,7 @@ TEST_P(HloEvaluatorTest, DilatedBaseConv2DWithLowAndHighPadding) {
   *window.add_dimensions() = dim;
 
   ConvolutionDimensionNumbers dnums =
-      ComputationBuilder::CreateDefaultConvDimensionNumbers(2);
+      XlaBuilder::CreateDefaultConvDimensionNumbers(2);
 
   const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 8, 8});
   b.AddInstruction(HloInstruction::CreateConvolve(
@@ -1180,7 +1180,7 @@ TEST_P(HloEvaluatorTest,
   *window.add_dimensions() = dim;
 
   ConvolutionDimensionNumbers dnums =
-      ComputationBuilder::CreateDefaultConvDimensionNumbers(2);
+      XlaBuilder::CreateDefaultConvDimensionNumbers(2);
 
   const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 9, 3});
   b.AddInstruction(HloInstruction::CreateConvolve(
diff --git a/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc b/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc
index f8d98f0678..be156d765d 100644
--- a/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc
@@ -14,7 +14,6 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/compiler/xla/service/hlo_tfgraph_builder.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
 #include "tensorflow/compiler/xla/tests/hlo_test_base.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/core/framework/tensor_shape.pb.h"
diff --git a/tensorflow/compiler/xla/service/transpose_folding_test.cc b/tensorflow/compiler/xla/service/transpose_folding_test.cc
index c7c4160345..0319109f7f 100644
--- a/tensorflow/compiler/xla/service/transpose_folding_test.cc
+++ b/tensorflow/compiler/xla/service/transpose_folding_test.cc
@@ -19,7 +19,7 @@ limitations under the License.
 #include <unordered_set>
 #include <vector>
 
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h"
 #include "tensorflow/compiler/xla/service/hlo_computation.h"
@@ -222,7 +222,7 @@ TEST_F(TransposeFoldingTest, FoldConvDimSwapTransposeRhs) {
   HloInstruction* transpose_y =
       builder.AddInstruction(HloInstruction::CreateTranspose(
           ShapeUtil::MakeShape(F32, {2, 3, 1, 1}), y, {1, 0, 2, 3}));
-  auto dnums = ComputationBuilder::CreateDefaultConvDimensionNumbers();
+  auto dnums = XlaBuilder::CreateDefaultConvDimensionNumbers();
   Window window;
   for (int i = 0; i < 2; ++i) {
     WindowDimension* dim = window.add_dimensions();
@@ -275,7 +275,7 @@ TEST_F(TransposeFoldingTest, FoldConvComplexTransposeRhs) {
   HloInstruction* transpose_y =
       builder.AddInstruction(HloInstruction::CreateTranspose(
           ShapeUtil::MakeShape(F32, {2, 3, 1, 1}), y, {1, 3, 0, 2}));
-  auto dnums = ComputationBuilder::CreateDefaultConvDimensionNumbers();
+  auto dnums = XlaBuilder::CreateDefaultConvDimensionNumbers();
   Window window;
   for (int i = 0; i < 2; ++i) {
     WindowDimension* dim = window.add_dimensions();
@@ -334,7 +334,7 @@ TEST_F(TransposeFoldingTest, FoldConvTransposeLhs) {
   HloInstruction* transpose_x =
       builder.AddInstruction(HloInstruction::CreateTranspose(
           ShapeUtil::MakeShape(F32, {2, 3, 1, 1}), x, {1, 0, 2, 3}));
-  auto dnums = ComputationBuilder::CreateDefaultConvDimensionNumbers();
+  auto dnums = XlaBuilder::CreateDefaultConvDimensionNumbers();
   Window window;
   for (int i = 0; i < 2; ++i) {
     WindowDimension* dim = window.add_dimensions();
@@ -398,7 +398,7 @@ TEST_F(TransposeFoldingTest, FoldConvComplexTransposeLhs) {
   HloInstruction* transpose_x =
       builder.AddInstruction(HloInstruction::CreateTranspose(
           ShapeUtil::MakeShape(F32, {2, 3, 1, 1}), x, {1, 0, 3, 2}));
-  auto dnums = ComputationBuilder::CreateDefaultConvDimensionNumbers();
+  auto dnums = XlaBuilder::CreateDefaultConvDimensionNumbers();
   Window window;
   for (int i = 0; i < 2; ++i) {
     WindowDimension* dim = window.add_dimensions();
diff --git a/tensorflow/compiler/xla/service/zero_sized_hlo_elimination_test.cc b/tensorflow/compiler/xla/service/zero_sized_hlo_elimination_test.cc
index a4e67cc9d9..f5331280ee 100644
--- a/tensorflow/compiler/xla/service/zero_sized_hlo_elimination_test.cc
+++ b/tensorflow/compiler/xla/service/zero_sized_hlo_elimination_test.cc
@@ -19,7 +19,6 @@ limitations under the License.
 #include <unordered_set>
 #include <vector>
 
-#include "tensorflow/compiler/xla/client/computation_builder.h"
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/compiler/xla/service/hlo_computation.h"
 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD
index 54cf0543b8..0571ff5055 100644
--- a/tensorflow/compiler/xla/tests/BUILD
+++ b/tensorflow/compiler/xla/tests/BUILD
@@ -1933,24 +1933,6 @@ xla_test(
     ],
 )
 
-xla_test(
-    name = "set_return_value_test",
-    srcs = ["set_return_value_test.cc"],
-    deps = [
-        "//tensorflow/compiler/xla:shape_util",
-        "//tensorflow/compiler/xla/client:computation_builder",
-        "//tensorflow/compiler/xla/client:local_client",
-        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
-        "//tensorflow/compiler/xla/client/xla_client:xla_computation",
-        "//tensorflow/compiler/xla/tests:client_library_test_base",
-        "//tensorflow/compiler/xla/tests:hlo_test_base",
-        "//tensorflow/compiler/xla/tests:literal_test_util",
-        "//tensorflow/compiler/xla/tests:xla_internal_test_main",
-        "//tensorflow/core:lib",
-        "//tensorflow/core:test",
-    ],
-)
-
 xla_test(
     name = "reshape_motion_test",
     srcs = ["reshape_motion_test.cc"],
diff --git a/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc b/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc
index 3704ddd801..a366afe826 100644
--- a/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc
+++ b/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc
@@ -21,7 +21,8 @@ limitations under the License.
 
 #include "llvm/ADT/Triple.h"
 #include "tensorflow/compiler/xla/client/client_library.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 #include "tensorflow/compiler/xla/service/cpu/cpu_compiler.h"
 #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h"
 #include "tensorflow/compiler/xla/types.h"
@@ -29,27 +30,31 @@ limitations under the License.
 #include "tensorflow/core/platform/init_main.h"
 #include "tensorflow/core/platform/logging.h"
 
+namespace {
+
 using xla::string;
 
-xla::Computation Doubler(xla::Client* client) {
-  xla::ComputationBuilder builder(client, "doubler");
+xla::XlaComputation Doubler() {
+  xla::XlaBuilder builder("doubler");
   auto r0f32 = xla::ShapeUtil::MakeShape(xla::F32, {});
   auto x = builder.Parameter(0, r0f32, "x");
   builder.Mul(x, builder.ConstantR0<float>(2.0));
   return std::move(builder.Build().ValueOrDie());
 }
 
+}  // namespace
+
 int main(int argc, char** argv) {
   tensorflow::port::InitMain(argv[0], &argc, &argv);
 
   auto client = xla::ClientLibrary::GetOrCreateCompileOnlyClient().ValueOrDie();
 
-  xla::ComputationBuilder builder(client, "aot_test_helper");
+  xla::XlaBuilder builder("aot_test_helper");
   auto opaque_shape = xla::ShapeUtil::MakeOpaqueShape();
   auto opaque_param = builder.Parameter(0, opaque_shape, "x");
   auto r0f32 = xla::ShapeUtil::MakeShape(xla::F32, {});
   auto sum = builder.CustomCall("SumStructElements", {opaque_param}, r0f32);
-  builder.Call(Doubler(client), {sum});
+  builder.Call(Doubler(), {sum});
 
   if (argc != 2) {
     LOG(FATAL) << "local_client_aot_test_helper TARGET_CPU";
@@ -71,8 +76,8 @@ int main(int argc, char** argv) {
 
   llvm::Triple triple(xla::llvm_ir::AsStringRef(triple_string));
 
-  xla::Computation computation = builder.Build().ConsumeValueOrDie();
-  xla::CompileOnlyClient::AotComputationInstance instance{
+  xla::XlaComputation computation = builder.Build().ConsumeValueOrDie();
+  xla::CompileOnlyClient::AotXlaComputationInstance instance{
       &computation, /*argument_layouts=*/{&opaque_shape}, &r0f32};
 
   xla::cpu::CpuAotCompilationOptions options(
diff --git a/tensorflow/compiler/xla/tests/set_return_value_test.cc b/tensorflow/compiler/xla/tests/set_return_value_test.cc
deleted file mode 100644
index 29f79ec28a..0000000000
--- a/tensorflow/compiler/xla/tests/set_return_value_test.cc
+++ /dev/null
@@ -1,98 +0,0 @@
-/* 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.
-==============================================================================*/
-
-#include <vector>
-
-#include "tensorflow/compiler/xla/client/computation_builder.h"
-#include "tensorflow/compiler/xla/client/local_client.h"
-#include "tensorflow/compiler/xla/tests/client_library_test_base.h"
-#include "tensorflow/compiler/xla/tests/literal_test_util.h"
-#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/platform/test.h"
-
-namespace xla {
-namespace {
-
-class SetReturnValueTest : public ClientLibraryTestBase {};
-
-TEST_F(SetReturnValueTest, NoSetValue) {
-  ComputationBuilder builder(client_, "no_set_value");
-  auto alpha = builder.ConstantR0<float>(1.0);
-  auto x = builder.ConstantR1<float>(
-      {-1.0, 1.0, 2.0, -2.0, -3.0, 3.0, 4.0, -4.0, -5.0, 5.0});
-  auto ax = builder.Add(alpha, x);
-  auto aax = builder.Add(alpha, ax);
-
-  std::vector<float> expected = {1.0, 3.0, 4.0,  0.0,  -1.0,
-                                 5.0, 6.0, -2.0, -3.0, 7.0};
-
-  ComputeAndCompareR1<float>(&builder, expected, {}, ErrorSpec(0.0001));
-}
-
-TEST_F(SetReturnValueTest, SetValue) {
-  ComputationBuilder builder(client_, "set_value");
-  auto alpha = builder.ConstantR0<float>(1.0);
-  auto x = builder.ConstantR1<float>(
-      {-1.0, 1.0, 2.0, -2.0, -3.0, 3.0, 4.0, -4.0, -5.0, 5.0});
-  auto ax = builder.Add(alpha, x);
-  auto aax = builder.Add(alpha, ax);
-  auto builder_status = builder.SetReturnValue(ax);
-  EXPECT_TRUE(builder_status.ok());
-
-  std::vector<float> expected = {0.0, 2.0, 3.0,  -1.0, -2.0,
-                                 4.0, 5.0, -3.0, -4.0, 6.0};
-
-  ComputeAndCompareR1<float>(&builder, expected, {}, ErrorSpec(0.0001));
-}
-
-TEST_F(SetReturnValueTest, SetValueAndModify) {
-  ComputationBuilder builder(client_, "set_value_and_modify");
-  auto alpha = builder.ConstantR0<float>(1.0);
-  auto x = builder.ConstantR1<float>(
-      {-1.0, 1.0, 2.0, -2.0, -3.0, 3.0, 4.0, -4.0, -5.0, 5.0});
-  auto ax = builder.Add(alpha, x);
-  auto aax = builder.Add(alpha, ax);
-  auto builder_status = builder.SetReturnValue(ax);
-  EXPECT_TRUE(builder_status.ok());
-  auto aaax = builder.Add(alpha, aax);
-
-  std::vector<float> expected = {0.0, 2.0, 3.0,  -1.0, -2.0,
-                                 4.0, 5.0, -3.0, -4.0, 6.0};
-
-  ComputeAndCompareR1<float>(&builder, expected, {}, ErrorSpec(0.0001));
-}
-
-TEST_F(SetReturnValueTest, SetValueMultipleTimesAndModify) {
-  ComputationBuilder builder(client_, "set_value_multiple_times_and_modify");
-  auto alpha = builder.ConstantR0<float>(1.0);
-  auto x = builder.ConstantR1<float>(
-      {-1.0, 1.0, 2.0, -2.0, -3.0, 3.0, 4.0, -4.0, -5.0, 5.0});
-  auto ax = builder.Add(alpha, x);
-  auto aax = builder.Add(alpha, ax);
-  auto builder_status = builder.SetReturnValue(aax);
-  EXPECT_TRUE(builder_status.ok());
-  auto aaax = builder.Add(alpha, aax);
-  builder_status = builder.SetReturnValue(ax);
-  EXPECT_TRUE(builder_status.ok());
-  auto aaaax = builder.Add(alpha, aaax);
-
-  std::vector<float> expected = {0.0, 2.0, 3.0,  -1.0, -2.0,
-                                 4.0, 5.0, -3.0, -4.0, 6.0};
-
-  ComputeAndCompareR1<float>(&builder, expected, {}, ErrorSpec(0.0001));
-}
-
-}  // namespace
-}  // namespace xla
diff --git a/tensorflow/compiler/xla/tests/vector_ops_simple_test.cc b/tensorflow/compiler/xla/tests/vector_ops_simple_test.cc
index 3dded3f715..5cce7a2bf8 100644
--- a/tensorflow/compiler/xla/tests/vector_ops_simple_test.cc
+++ b/tensorflow/compiler/xla/tests/vector_ops_simple_test.cc
@@ -18,7 +18,6 @@ limitations under the License.
 #include <vector>
 
 #include "tensorflow/compiler/xla/array4d.h"
-#include "tensorflow/compiler/xla/client/computation.h"
 #include "tensorflow/compiler/xla/client/global_data.h"
 #include "tensorflow/compiler/xla/client/lib/arithmetic.h"
 #include "tensorflow/compiler/xla/client/local_client.h"
@@ -350,7 +349,7 @@ XLA_TEST_F(VecOpsSimpleTest, ClampTenValuesConstantNonzeroLower) {
 }
 
 XLA_TEST_F(VecOpsSimpleTest, ClampValuesConstantS64) {
-  ComputationBuilder builder(client_, TestName());
+  XlaBuilder builder(TestName());
   auto zero = builder.ConstantR0<int64>(0);
   auto one = builder.ConstantR0<int64>(10);
   auto x = builder.ConstantR1<int64>({-3, 3, 9, 13});
-- 
GitLab


From fc7b593cda65f4a3a3de0cc733270f0864f820e2 Mon Sep 17 00:00:00 2001
From: Ruoxin Sang <rxsang@google.com>
Date: Thu, 3 May 2018 17:21:26 -0700
Subject: [PATCH 243/395] Clear the stat cache of the target when renaming the
 file.

PiperOrigin-RevId: 195337886
---
 .../core/platform/cloud/gcs_file_system.cc    |  4 +-
 .../platform/cloud/gcs_file_system_test.cc    | 72 +++++++++++++++++++
 2 files changed, 74 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/platform/cloud/gcs_file_system.cc b/tensorflow/core/platform/cloud/gcs_file_system.cc
index 488f9cc75d..e44e897434 100644
--- a/tensorflow/core/platform/cloud/gcs_file_system.cc
+++ b/tensorflow/core/platform/cloud/gcs_file_system.cc
@@ -1375,9 +1375,9 @@ Status GcsFileSystem::RenameObject(const string& src, const string& target) {
   request->SetResultBuffer(&output_buffer);
   TF_RETURN_WITH_CONTEXT_IF_ERROR(request->Send(), " when renaming ", src,
                                   " to ", target);
-  // Flush the target from the block cache.  The source will be flushed in the
+  // Flush the target from the caches.  The source will be flushed in the
   // DeleteFile call below.
-  file_block_cache_->RemoveFile(target);
+  ClearFileCaches(target);
   Json::Value root;
   TF_RETURN_IF_ERROR(ParseJson(output_buffer, &root));
   bool done;
diff --git a/tensorflow/core/platform/cloud/gcs_file_system_test.cc b/tensorflow/core/platform/cloud/gcs_file_system_test.cc
index c639299954..28be13869b 100644
--- a/tensorflow/core/platform/cloud/gcs_file_system_test.cc
+++ b/tensorflow/core/platform/cloud/gcs_file_system_test.cc
@@ -1902,6 +1902,78 @@ TEST(GcsFileSystemTest, RenameFile_Object) {
   EXPECT_EQ("fedcba98", result);
 }
 
+TEST(GcsFileSystemTest, RenameFile_Object_FlushTargetStatCache) {
+  std::vector<HttpRequest*> requests(
+      {// Stat the target file.
+       new FakeHttpRequest(
+           "Uri: https://www.googleapis.com/storage/v1/b/bucket/o/"
+           "path%2Fdst.txt?fields=size%2Cupdated\n"
+           "Auth Token: fake_token\n"
+           "Timeouts: 5 1 10\n",
+           strings::StrCat("{\"size\": \"1000\","
+                           "\"updated\": \"2016-04-29T23:15:24.896Z\"}")),
+       // IsDirectory is checking whether there are children objects.
+       new FakeHttpRequest(
+           "Uri: https://www.googleapis.com/storage/v1/b/bucket/o?"
+           "fields=items%2Fname%2CnextPageToken&prefix=path%2Fsrc.txt%2F"
+           "&maxResults=1\n"
+           "Auth Token: fake_token\n"
+           "Timeouts: 5 1 10\n",
+           "{}"),
+       // IsDirectory is checking if the path exists as an object.
+       new FakeHttpRequest(
+           "Uri: https://www.googleapis.com/storage/v1/b/bucket/o/"
+           "path%2Fsrc.txt?fields=size%2Cupdated\n"
+           "Auth Token: fake_token\n"
+           "Timeouts: 5 1 10\n",
+           strings::StrCat("{\"size\": \"1010\","
+                           "\"updated\": \"2016-04-29T23:15:24.896Z\"}")),
+       // Copying to the new location.
+       new FakeHttpRequest(
+           "Uri: https://www.googleapis.com/storage/v1/b/bucket/o/"
+           "path%2Fsrc.txt/rewriteTo/b/bucket/o/path%2Fdst.txt\n"
+           "Auth Token: fake_token\n"
+           "Post: yes\n"
+           "Timeouts: 5 1 10\n",
+           "{\"done\": true}"),
+       // Deleting the original file.
+       new FakeHttpRequest(
+           "Uri: https://www.googleapis.com/storage/v1/b/bucket/o/"
+           "path%2Fsrc.txt\n"
+           "Auth Token: fake_token\n"
+           "Timeouts: 5 1 10\n"
+           "Delete: yes\n",
+           ""),
+       new FakeHttpRequest(
+           "Uri: https://www.googleapis.com/storage/v1/b/bucket/o/"
+           "path%2Fdst.txt?fields=size%2Cupdated\n"
+           "Auth Token: fake_token\n"
+           "Timeouts: 5 1 10\n",
+           strings::StrCat("{\"size\": \"1010\","
+                           "\"updated\": \"2016-04-29T23:15:24.896Z\"}"))});
+  GcsFileSystem fs(
+      std::unique_ptr<AuthProvider>(new FakeAuthProvider),
+      std::unique_ptr<HttpRequest::Factory>(
+          new FakeHttpRequestFactory(&requests)),
+      0 /* block size */, 0 /* max bytes */, 0 /* max staleness */,
+      3600 /* stat cache max age */, 0 /* stat cache max entries */,
+      0 /* matching paths cache max age */,
+      0 /* matching paths cache max entries */, 0 /* initial retry delay*/,
+      kTestTimeoutConfig, nullptr /* gcs additional header */);
+  // Do an initial stat of the destination file to load their contents into the
+  // stat cache.
+  FileStatistics stat_before_renaming;
+  TF_EXPECT_OK(fs.Stat("gs://bucket/path/dst.txt", &stat_before_renaming));
+  EXPECT_EQ(1000, stat_before_renaming.length);
+
+  TF_EXPECT_OK(
+      fs.RenameFile("gs://bucket/path/src.txt", "gs://bucket/path/dst.txt"));
+
+  FileStatistics stat_after_renaming;
+  TF_EXPECT_OK(fs.Stat("gs://bucket/path/dst.txt", &stat_after_renaming));
+  EXPECT_EQ(1010, stat_after_renaming.length);
+}
+
 /// Tests the scenario when deletion returns a failure, but actually succeeds.
 TEST(GcsFileSystemTest, RenameFile_Object_DeletionRetried) {
   std::vector<HttpRequest*> requests(
-- 
GitLab


From fa7b5a9d1ab654bbd466487e39a8b3f83c17f3f0 Mon Sep 17 00:00:00 2001
From: Chris Leary <leary@google.com>
Date: Thu, 3 May 2018 18:08:34 -0700
Subject: [PATCH 244/395] [XLA] Make LocalShapedBuffer::FromLiteral fallible by
 passing StatusOr wrapper.

PiperOrigin-RevId: 195345724
---
 .../xla/python/local_computation_builder.cc      | 16 ++++++++--------
 .../xla/python/local_computation_builder.h       |  6 ++++--
 .../xla/python/local_computation_builder.i       | 13 +++++++++++++
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/tensorflow/compiler/xla/python/local_computation_builder.cc b/tensorflow/compiler/xla/python/local_computation_builder.cc
index 7102f46737..044458164f 100644
--- a/tensorflow/compiler/xla/python/local_computation_builder.cc
+++ b/tensorflow/compiler/xla/python/local_computation_builder.cc
@@ -104,25 +104,25 @@ static StatusOr<ScopedShapedBuffer> ToBuffer(LocalClient* client,
 }
 
 /* static */
-LocalShapedBuffer* LocalShapedBuffer::FromLiteral(
+StatusOr<LocalShapedBuffer*> LocalShapedBuffer::FromLiteral(
     const Literal& argument,
     const tensorflow::gtl::optional<Shape>& shape_with_layout) {
   LocalClient* client = GetOrCreateLocalClient();
-  ScopedShapedBuffer buf = [&] {
+  StatusOr<ScopedShapedBuffer> buf = [&] {
     if (shape_with_layout) {
       std::unique_ptr<Literal> relaid =
           argument.Relayout(shape_with_layout.value());
-      return ToBuffer(client, /*device_ordinal=*/0, *relaid)
-          .ConsumeValueOrDie();
+      return ToBuffer(client, /*device_ordinal=*/0, *relaid);
     }
-    return ToBuffer(client, /*device_ordinal=*/0, argument).ConsumeValueOrDie();
+    return ToBuffer(client, /*device_ordinal=*/0, argument);
   }();
-  return new LocalShapedBuffer(std::move(buf));
+  TF_RETURN_IF_ERROR(buf.status());
+  return new LocalShapedBuffer(std::move(buf).ValueOrDie());
 }
 
-std::unique_ptr<Literal> LocalShapedBuffer::ToLiteral() const {
+StatusOr<std::unique_ptr<Literal>> LocalShapedBuffer::ToLiteral() const {
   LocalClient* client = GetOrCreateLocalClient();
-  return client->ShapedBufferToLiteral(*shaped_buffer()).ConsumeValueOrDie();
+  return client->ShapedBufferToLiteral(*shaped_buffer());
 }
 
 CompiledLocalComputation::CompiledLocalComputation(
diff --git a/tensorflow/compiler/xla/python/local_computation_builder.h b/tensorflow/compiler/xla/python/local_computation_builder.h
index e1048909ab..5ec097846a 100644
--- a/tensorflow/compiler/xla/python/local_computation_builder.h
+++ b/tensorflow/compiler/xla/python/local_computation_builder.h
@@ -59,12 +59,14 @@ StatusOr<std::unique_ptr<Literal> > TransferFromOutfeedLocalReplica(
 // client.
 class LocalShapedBuffer {
  public:
-  static LocalShapedBuffer* FromLiteral(
+  static StatusOr<LocalShapedBuffer*> FromLiteral(
       const Literal& argument,
       const tensorflow::gtl::optional<Shape>& shape_with_layout);
+
   LocalShapedBuffer(ScopedShapedBuffer shaped_buffer);
   const ScopedShapedBuffer* shaped_buffer() const;
-  std::unique_ptr<Literal> ToLiteral() const;
+
+  StatusOr<std::unique_ptr<Literal> > ToLiteral() const;
 
  private:
   ScopedShapedBuffer shaped_buffer_;
diff --git a/tensorflow/compiler/xla/python/local_computation_builder.i b/tensorflow/compiler/xla/python/local_computation_builder.i
index ac792e8189..b8cce5a5f7 100644
--- a/tensorflow/compiler/xla/python/local_computation_builder.i
+++ b/tensorflow/compiler/xla/python/local_computation_builder.i
@@ -205,6 +205,19 @@ tensorflow::ImportNumpy();
   }
 }
 
+%typemap(out) StatusOr<xla::swig::LocalShapedBuffer*> {
+  if ($1.ok()) {
+    auto* value = $1.ValueOrDie();
+    {
+      auto* $1 = value;
+      $typemap(out, xla::swig::LocalShapedBuffer*)
+    }
+  } else {
+    PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str());
+    SWIG_fail;
+  }
+}
+
 %typemap(out) StatusOr< std::unique_ptr<Literal> > {
   if ($1.ok()) {
     std::unique_ptr<Literal> value = $1.ConsumeValueOrDie();
-- 
GitLab


From 0abbff6c0bdf0ee4690def786513298afc8b772a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 3 May 2018 19:45:59 -0700
Subject: [PATCH 245/395] [XLA] Redesign: cleanup client_library_test_base.

PiperOrigin-RevId: 195357555
---
 .../xla/tests/client_library_test_base.cc     | 221 +---------------
 .../xla/tests/client_library_test_base.h      | 243 ++++++------------
 2 files changed, 92 insertions(+), 372 deletions(-)

diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.cc b/tensorflow/compiler/xla/tests/client_library_test_base.cc
index 22660c35dc..c09e7eaf2b 100644
--- a/tensorflow/compiler/xla/tests/client_library_test_base.cc
+++ b/tensorflow/compiler/xla/tests/client_library_test_base.cc
@@ -94,27 +94,13 @@ string ClientLibraryTestBase::TestName() const {
   return ::testing::UnitTest::GetInstance()->current_test_info()->name();
 }
 
-template <typename BuilderT>
 StatusOr<std::unique_ptr<GlobalData>> ClientLibraryTestBase::Execute(
-    BuilderT* builder, tensorflow::gtl::ArraySlice<GlobalData*> arguments) {
+    XlaBuilder* builder, tensorflow::gtl::ArraySlice<GlobalData*> arguments) {
   // Build the computation, as a convenience.
   TF_ASSIGN_OR_RETURN(auto computation, builder->Build());
   return client_->Execute(computation, arguments, &execution_options_);
 }
 
-StatusOr<std::unique_ptr<Literal>> ClientLibraryTestBase::ExecuteAndTransfer(
-    const Computation& computation,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments,
-    const Shape* shape_with_output_layout) {
-  ExecutionOptions execution_options = execution_options_;
-  if (shape_with_output_layout != nullptr) {
-    *execution_options.mutable_shape_with_output_layout() =
-        *shape_with_output_layout;
-  }
-  return client_->ExecuteAndTransfer(computation, arguments,
-                                     &execution_options);
-}
-
 StatusOr<std::unique_ptr<Literal>> ClientLibraryTestBase::ExecuteAndTransfer(
     const XlaComputation& computation,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments,
@@ -128,17 +114,6 @@ StatusOr<std::unique_ptr<Literal>> ClientLibraryTestBase::ExecuteAndTransfer(
                                      &execution_options);
 }
 
-template <>
-StatusOr<std::unique_ptr<Literal>> ClientLibraryTestBase::ExecuteAndTransfer(
-    ComputationBuilder* builder,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments,
-    const Shape* shape_with_output_layout) {
-  // Build the computation, as a convenience.
-  TF_ASSIGN_OR_RETURN(auto computation, builder->Build());
-  return ExecuteAndTransfer(computation, arguments, shape_with_output_layout);
-}
-
-template <>
 StatusOr<std::unique_ptr<Literal>> ClientLibraryTestBase::ExecuteAndTransfer(
     XlaBuilder* builder, tensorflow::gtl::ArraySlice<GlobalData*> arguments,
     const Shape* shape_with_output_layout) {
@@ -162,18 +137,6 @@ ClientLibraryTestBase::ExecuteAndTransferReference(
                                          &execution_options);
 }
 
-std::unique_ptr<GlobalData> ClientLibraryTestBase::ExecuteOrDie(
-    ComputationBuilder* builder,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments) {
-  return Execute(builder, arguments).ConsumeValueOrDie();
-}
-
-std::unique_ptr<Literal> ClientLibraryTestBase::ExecuteAndTransferOrDie(
-    ComputationBuilder* builder,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments) {
-  return ExecuteAndTransfer(builder, arguments).ConsumeValueOrDie();
-}
-
 string ClientLibraryTestBase::ExecuteToString(
     XlaBuilder* builder, tensorflow::gtl::ArraySlice<GlobalData*> arguments) {
   auto computation_status = builder->Build();
@@ -191,32 +154,6 @@ string ClientLibraryTestBase::ExecuteToString(
   }
 }
 
-string ClientLibraryTestBase::ExecuteToString(
-    ComputationBuilder* builder,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments) {
-  auto computation_status = builder->Build();
-  if (!computation_status.ok()) {
-    return computation_status.status().ToString();
-  }
-  auto computation = computation_status.ConsumeValueOrDie();
-
-  auto result =
-      client_->ExecuteAndTransfer(computation, arguments, &execution_options_);
-  if (!result.ok()) {
-    return result.status().ToString();
-  } else {
-    return result.ValueOrDie()->ToString();
-  }
-}
-
-void ClientLibraryTestBase::ComputeAndCompareR1(
-    ComputationBuilder* builder, const tensorflow::core::Bitmap& expected,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments) {
-  std::unique_ptr<Literal> expected_literal = Literal::CreateR1(expected);
-  ClientLibraryTestBase::ComputeAndCompareLiteral(builder, *expected_literal,
-                                                  arguments);
-}
-
 void ClientLibraryTestBase::ComputeAndCompareR1(
     XlaBuilder* builder, const tensorflow::core::Bitmap& expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments) {
@@ -225,18 +162,16 @@ void ClientLibraryTestBase::ComputeAndCompareR1(
                                                   arguments);
 }
 
-template <typename BuilderT>
 void ClientLibraryTestBase::ComputeAndCompareLiteral(
-    BuilderT* builder, const Literal& expected,
+    XlaBuilder* builder, const Literal& expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments,
     const Shape* shape_with_layout) {
   EXPECT_IS_OK(ComputeAndCompareLiteralWithStatus(builder, expected, arguments,
                                                   shape_with_layout));
 }
 
-template <typename BuilderT>
 void ClientLibraryTestBase::ComputeAndCompareLiteral(
-    BuilderT* builder, const Literal& expected,
+    XlaBuilder* builder, const Literal& expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments, ErrorSpec error,
     const Shape* shape_with_layout) {
   EXPECT_IS_OK(ComputeAndCompareLiteralWithStatus(builder, expected, arguments,
@@ -245,7 +180,7 @@ void ClientLibraryTestBase::ComputeAndCompareLiteral(
 
 tensorflow::Status
 ClientLibraryTestBase::ComputeAndCompareLiteralWithAllOutputLayouts(
-    const xla::Computation& computation, const Literal& expected,
+    const xla::XlaComputation& computation, const Literal& expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments,
     const std::function<void(const Literal& actual,
                              const string& error_message)>& verify_output) {
@@ -271,7 +206,7 @@ ClientLibraryTestBase::ComputeAndCompareLiteralWithAllOutputLayouts(
 
 tensorflow::Status
 ClientLibraryTestBase::ComputeAndCompareLiteralWithAllInputLayouts(
-    const xla::Computation& computation, const Literal& expected,
+    const xla::XlaComputation& computation, const Literal& /*expected*/,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments,
     const std::function<void(const Literal& actual,
                              const string& error_message)>& verify_output,
@@ -334,28 +269,8 @@ ClientLibraryTestBase::ComputeAndCompareLiteralWithAllInputLayouts(
   return choose(0);
 }
 
-tensorflow::Status
-ClientLibraryTestBase::ComputeAndCompareLiteralWithAllOutputLayouts(
-    const xla::XlaComputation& /*computation*/, const Literal& /*expected*/,
-    tensorflow::gtl::ArraySlice<GlobalData*> /*arguments*/,
-    const std::function<void(const Literal& actual,
-                             const string& error_message)>& /*verify_output*/) {
-  return Unimplemented("not yet implemented for XlaComputation");
-}
-
-tensorflow::Status
-ClientLibraryTestBase::ComputeAndCompareLiteralWithAllInputLayouts(
-    const xla::XlaComputation& /*computation*/, const Literal& /*expected*/,
-    tensorflow::gtl::ArraySlice<GlobalData*> /*arguments*/,
-    const std::function<void(const Literal& actual,
-                             const string& error_message)>& /*verify_output*/,
-    const Shape* /*output_with_layout*/) {
-  return Unimplemented("not yet implemented for XlaComputation");
-}
-
-template <typename BuilderT>
 tensorflow::Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus(
-    BuilderT* builder, const Literal& expected,
+    XlaBuilder* builder, const Literal& expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments_passed_in,
     const Shape* shape_with_layout) {
   std::vector<GlobalData*> arguments(arguments_passed_in.begin(),
@@ -412,9 +327,8 @@ tensorflow::Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus(
   return tensorflow::Status::OK();
 }
 
-template <typename BuilderT>
 tensorflow::Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus(
-    BuilderT* builder, const Literal& expected,
+    XlaBuilder* builder, const Literal& expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments_passed_in,
     ErrorSpec error, const Shape* shape_with_layout) {
   std::vector<GlobalData*> arguments(arguments_passed_in.begin(),
@@ -484,9 +398,8 @@ void ClientLibraryTestBase::ComputeAndCompareR1U8(
   EXPECT_EQ(expected, actual->GetR1U8AsString());
 }
 
-template <typename BuilderT>
 void ClientLibraryTestBase::ComputeAndCompareTuple(
-    BuilderT* builder, const Literal& expected,
+    XlaBuilder* builder, const Literal& expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments) {
   auto actual_status = ExecuteAndTransfer(builder, arguments);
   EXPECT_IS_OK(actual_status.status());
@@ -497,9 +410,8 @@ void ClientLibraryTestBase::ComputeAndCompareTuple(
   LiteralTestUtil::ExpectEqual(expected, *actual);
 }
 
-template <typename BuilderT>
 void ClientLibraryTestBase::ComputeAndCompareTuple(
-    BuilderT* builder, const Literal& expected,
+    XlaBuilder* builder, const Literal& expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments, ErrorSpec error) {
   auto actual_status = ExecuteAndTransfer(builder, arguments);
   EXPECT_IS_OK(actual_status.status());
@@ -510,60 +422,6 @@ void ClientLibraryTestBase::ComputeAndCompareTuple(
   LiteralTestUtil::ExpectNear(expected, *actual, error);
 }
 
-void ClientLibraryTestBase::ComputeAndCompare(
-    ComputationBuilder* builder, const ComputationDataHandle& operand,
-    tensorflow::gtl::ArraySlice<Literal> arguments) {
-  auto status_or_data = ComputeValueAndReference(builder, operand, arguments);
-  EXPECT_IS_OK(status_or_data);
-  if (!status_or_data.ok()) {
-    return;
-  }
-  std::unique_ptr<Literal> reference, result;
-  std::tie(reference, result) = status_or_data.ConsumeValueOrDie();
-  LiteralTestUtil::ExpectEqual(*reference, *result);
-}
-
-void ClientLibraryTestBase::ComputeAndCompare(
-    ComputationBuilder* builder, const ComputationDataHandle& operand,
-    tensorflow::gtl::ArraySlice<Literal> arguments, ErrorSpec error) {
-  auto status_or_data = ComputeValueAndReference(builder, operand, arguments);
-  EXPECT_IS_OK(status_or_data);
-  if (!status_or_data.ok()) {
-    return;
-  }
-  std::unique_ptr<Literal> reference, result;
-  std::tie(reference, result) = status_or_data.ConsumeValueOrDie();
-  LiteralTestUtil::ExpectNear(*reference, *result, error);
-}
-
-StatusOr<std::pair<std::unique_ptr<Literal>, std::unique_ptr<Literal>>>
-ClientLibraryTestBase::ComputeValueAndReference(
-    ComputationBuilder* builder, const ComputationDataHandle& operand,
-    tensorflow::gtl::ArraySlice<Literal> arguments) {
-  // Transfer the arguments to the executor service. We put the unique_ptr's
-  // into a vector to keep the data alive on the service until the end of this
-  // function.
-  std::vector<std::unique_ptr<GlobalData>> argument_data;
-  for (const auto& arg : arguments) {
-    TF_ASSIGN_OR_RETURN(auto data, client_->TransferToServer(arg));
-    argument_data.push_back(std::move(data));
-  }
-
-  // Create raw pointers to the GlobalData for the rest of the call stack.
-  std::vector<GlobalData*> argument_data_ptr;
-  std::transform(
-      argument_data.begin(), argument_data.end(),
-      std::back_inserter(argument_data_ptr),
-      [](const std::unique_ptr<GlobalData>& data) { return data.get(); });
-
-  TF_ASSIGN_OR_RETURN(
-      auto reference,
-      builder->ComputeConstant(operand, /*output_layout=*/nullptr, arguments));
-  TF_ASSIGN_OR_RETURN(auto result,
-                      ExecuteAndTransfer(builder, argument_data_ptr));
-  return std::make_pair(std::move(reference), std::move(result));
-}
-
 void ClientLibraryTestBase::ComputeAndCompare(
     XlaBuilder* builder, tensorflow::gtl::ArraySlice<Literal> arguments) {
   auto status_or_data = ComputeValueAndReference(builder, arguments);
@@ -651,8 +509,8 @@ XlaComputation ClientLibraryTestBase::CreateScalarMax() {
   return computation_status.ConsumeValueOrDie();
 }
 
-Computation ClientLibraryTestBase::CreateScalarReluSensitivity() {
-  ComputationBuilder builder(client_, "relu_sensitivity");
+XlaComputation ClientLibraryTestBase::CreateScalarReluSensitivity() {
+  XlaBuilder builder("relu_sensitivity");
   auto shape = ShapeUtil::MakeShape(use_bfloat16_ ? BF16 : F32, {});
   auto activation = builder.Parameter(0, shape, "activation");
   auto backprop = builder.Parameter(1, shape, "backprop");
@@ -693,14 +551,6 @@ ClientLibraryTestBase::CreatePatternedMatrixWithZeroPadding(int rows, int cols,
   return array;
 }
 
-ComputationDataHandle ClientLibraryTestBase::AddParam(
-    const Literal& argument, ComputationBuilder* builder) {
-  ComputationDataHandle data_handle;
-  arguments_.push_back(CreateParameterAndTransferLiteral(
-      arguments_.size(), argument, "", builder, &data_handle));
-  return data_handle;
-}
-
 XlaOp ClientLibraryTestBase::AddParam(const Literal& argument,
                                       XlaBuilder* builder) {
   XlaOp data_handle;
@@ -709,59 +559,10 @@ XlaOp ClientLibraryTestBase::AddParam(const Literal& argument,
   return data_handle;
 }
 
-ComputationDataHandle ClientLibraryTestBase::CreateConstantFromLiteral(
-    const Literal& literal, ComputationBuilder* builder) {
-  return builder->ConstantLiteral(
-      use_bfloat16_ ? *LiteralTestUtil::ConvertF32ToBF16(literal) : literal);
-}
-
 XlaOp ClientLibraryTestBase::CreateConstantFromLiteral(const Literal& literal,
                                                        XlaBuilder* builder) {
   return builder->ConstantLiteral(
       use_bfloat16_ ? *LiteralTestUtil::ConvertF32ToBF16(literal) : literal);
 }
 
-template void ClientLibraryTestBase::ComputeAndCompareLiteral(
-    ComputationBuilder* builder, const Literal& expected,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments,
-    const Shape* shape_with_layout);
-
-template void ClientLibraryTestBase::ComputeAndCompareLiteral(
-    XlaBuilder* builder, const Literal& expected,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments,
-    const Shape* shape_with_layout);
-
-template void ClientLibraryTestBase::ComputeAndCompareLiteral(
-    ComputationBuilder* builder, const Literal& expected,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments, ErrorSpec error,
-    const Shape* shape_with_layout);
-
-template void ClientLibraryTestBase::ComputeAndCompareLiteral(
-    XlaBuilder* builder, const Literal& expected,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments, ErrorSpec error,
-    const Shape* shape_with_layout);
-
-template void ClientLibraryTestBase::ComputeAndCompareTuple(
-    ComputationBuilder* builder, const Literal& expected,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments);
-
-template void ClientLibraryTestBase::ComputeAndCompareTuple(
-    XlaBuilder* builder, const Literal& expected,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments);
-
-template void ClientLibraryTestBase::ComputeAndCompareTuple(
-    ComputationBuilder* builder, const Literal& expected,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments, ErrorSpec error);
-
-template void ClientLibraryTestBase::ComputeAndCompareTuple(
-    XlaBuilder* builder, const Literal& expected,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments, ErrorSpec error);
-
-template StatusOr<std::unique_ptr<GlobalData>> ClientLibraryTestBase::Execute(
-    ComputationBuilder* builder,
-    tensorflow::gtl::ArraySlice<GlobalData*> arguments);
-
-template StatusOr<std::unique_ptr<GlobalData>> ClientLibraryTestBase::Execute(
-    XlaBuilder* builder, tensorflow::gtl::ArraySlice<GlobalData*> arguments);
-
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.h b/tensorflow/compiler/xla/tests/client_library_test_base.h
index 32eea7c2f3..e58979a303 100644
--- a/tensorflow/compiler/xla/tests/client_library_test_base.h
+++ b/tensorflow/compiler/xla/tests/client_library_test_base.h
@@ -25,10 +25,9 @@ limitations under the License.
 #include "tensorflow/compiler/xla/array3d.h"
 #include "tensorflow/compiler/xla/array4d.h"
 #include "tensorflow/compiler/xla/client/client_library.h"
-#include "tensorflow/compiler/xla/client/computation.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
 #include "tensorflow/compiler/xla/client/global_data.h"
 #include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 #include "tensorflow/compiler/xla/literal_util.h"
 #include "tensorflow/compiler/xla/ptr_util.h"
 #include "tensorflow/compiler/xla/statusor.h"
@@ -91,21 +90,11 @@ class ClientLibraryTestBase : public ::testing::Test {
   // Convenience methods for building and running a computation with the member
   // execution options. Modify execution_options_ in your test if you want to
   // customize the options.
-  template <typename BuilderT>
   StatusOr<std::unique_ptr<GlobalData>> Execute(
-      BuilderT* builder, tensorflow::gtl::ArraySlice<GlobalData*> arguments);
+      XlaBuilder* builder, tensorflow::gtl::ArraySlice<GlobalData*> arguments);
 
-  // TODO(b/74197823): Remove the template type 'BuilderT' in all methods once
-  // the migration to XlaBuilder is complete.
-
-  template <typename BuilderT>
   StatusOr<std::unique_ptr<Literal>> ExecuteAndTransfer(
-      BuilderT* builder, tensorflow::gtl::ArraySlice<GlobalData*> arguments,
-      const Shape* shape_with_output_layout = nullptr);
-
-  StatusOr<std::unique_ptr<Literal>> ExecuteAndTransfer(
-      const Computation& computation,
-      tensorflow::gtl::ArraySlice<GlobalData*> arguments,
+      XlaBuilder* builder, tensorflow::gtl::ArraySlice<GlobalData*> arguments,
       const Shape* shape_with_output_layout = nullptr);
 
   StatusOr<std::unique_ptr<Literal>> ExecuteAndTransfer(
@@ -121,101 +110,90 @@ class ClientLibraryTestBase : public ::testing::Test {
       tensorflow::gtl::ArraySlice<GlobalData*> arguments,
       const Shape* shape_with_output_layout = nullptr);
 
-  // Convenience OrDie variants of above methods.
-  std::unique_ptr<GlobalData> ExecuteOrDie(
-      ComputationBuilder* builder,
-      tensorflow::gtl::ArraySlice<GlobalData*> arguments);
-  std::unique_ptr<Literal> ExecuteAndTransferOrDie(
-      ComputationBuilder* builder,
-      tensorflow::gtl::ArraySlice<GlobalData*> arguments);
-
   // Run a computation and return its value as a string. If an error
   // occurs, then instead return the error as a string.
   string ExecuteToString(XlaBuilder* builder,
                          tensorflow::gtl::ArraySlice<GlobalData*> arguments);
-  string ExecuteToString(ComputationBuilder* builder,
-                         tensorflow::gtl::ArraySlice<GlobalData*> arguments);
 
   // Convenience methods for building and running a computation, transferring
   // the result, and comparing it to the expected value(s). Methods are
   // templated on the native host type which maps to specific XLA types (See
-  // ComputationBuilder/XlaBuilder for details). For each rank, two forms are
+  // XlaBuilder for details). For each rank, two forms are
   // provided: one for floating point types with an ErrorSpec parameter, and one
   // for integral types without the ErrorSpec parameter.
-  template <typename NativeT, typename BuilderT>
-  void ComputeAndCompareR0(BuilderT* builder, NativeT expected,
+  template <typename NativeT>
+  void ComputeAndCompareR0(XlaBuilder* builder, NativeT expected,
                            tensorflow::gtl::ArraySlice<GlobalData*> arguments);
-  template <typename NativeT, typename BuilderT>
-  void ComputeAndCompareR0(BuilderT* builder, NativeT expected,
+  template <typename NativeT>
+  void ComputeAndCompareR0(XlaBuilder* builder, NativeT expected,
                            tensorflow::gtl::ArraySlice<GlobalData*> arguments,
                            ErrorSpec error);
 
-  template <typename NativeT, typename BuilderT>
-  void ComputeAndCompareR1(BuilderT* builder,
+  template <typename NativeT>
+  void ComputeAndCompareR1(XlaBuilder* builder,
                            tensorflow::gtl::ArraySlice<NativeT> expected,
                            tensorflow::gtl::ArraySlice<GlobalData*> arguments);
-  template <typename NativeT, typename BuilderT>
-  void ComputeAndCompareR1(BuilderT* builder,
+  template <typename NativeT>
+  void ComputeAndCompareR1(XlaBuilder* builder,
                            tensorflow::gtl::ArraySlice<NativeT> expected,
                            tensorflow::gtl::ArraySlice<GlobalData*> arguments,
                            ErrorSpec error);
 
   // As above, but uses a bitmap to hold the predicate vector to avoid
   // deficiencies of vector<bool>.
-  void ComputeAndCompareR1(ComputationBuilder* builder,
-                           const tensorflow::core::Bitmap& expected,
-                           tensorflow::gtl::ArraySlice<GlobalData*> arguments);
   void ComputeAndCompareR1(XlaBuilder* builder,
                            const tensorflow::core::Bitmap& expected,
                            tensorflow::gtl::ArraySlice<GlobalData*> arguments);
 
-  template <typename NativeT, typename BuilderT>
-  void ComputeAndCompareR2(BuilderT* builder, const Array2D<NativeT>& expected,
+  template <typename NativeT>
+  void ComputeAndCompareR2(XlaBuilder* builder,
+                           const Array2D<NativeT>& expected,
                            tensorflow::gtl::ArraySlice<GlobalData*> arguments);
-  template <typename NativeT, typename BuilderT>
-  void ComputeAndCompareR2(BuilderT* builder, const Array2D<NativeT>& expected,
+  template <typename NativeT>
+  void ComputeAndCompareR2(XlaBuilder* builder,
+                           const Array2D<NativeT>& expected,
                            tensorflow::gtl::ArraySlice<GlobalData*> arguments,
                            ErrorSpec error);
 
-  template <typename NativeT, typename BuilderT>
-  void ComputeAndCompareR3(BuilderT* builder, const Array3D<NativeT>& expected,
+  template <typename NativeT>
+  void ComputeAndCompareR3(XlaBuilder* builder,
+                           const Array3D<NativeT>& expected,
                            tensorflow::gtl::ArraySlice<GlobalData*> arguments);
-  template <typename NativeT, typename BuilderT>
-  void ComputeAndCompareR3(BuilderT* builder, const Array3D<NativeT>& expected,
+  template <typename NativeT>
+  void ComputeAndCompareR3(XlaBuilder* builder,
+                           const Array3D<NativeT>& expected,
                            tensorflow::gtl::ArraySlice<GlobalData*> arguments,
                            ErrorSpec error);
 
-  template <typename NativeT, typename BuilderT>
-  void ComputeAndCompareR4(BuilderT* builder, const Array4D<NativeT>& expected,
+  template <typename NativeT>
+  void ComputeAndCompareR4(XlaBuilder* builder,
+                           const Array4D<NativeT>& expected,
                            tensorflow::gtl::ArraySlice<GlobalData*> arguments);
-  template <typename NativeT, typename BuilderT>
-  void ComputeAndCompareR4(BuilderT* builder, const Array4D<NativeT>& expected,
+  template <typename NativeT>
+  void ComputeAndCompareR4(XlaBuilder* builder,
+                           const Array4D<NativeT>& expected,
                            tensorflow::gtl::ArraySlice<GlobalData*> arguments,
                            ErrorSpec error);
 
   // Build and run the computation and compare the result with the given
   // literal. shape_with_layout indicates the result layout to request when
   // calling Execute.
-  template <typename BuilderT>
   void ComputeAndCompareLiteral(
-      BuilderT* builder, const Literal& expected,
+      XlaBuilder* builder, const Literal& expected,
       tensorflow::gtl::ArraySlice<GlobalData*> arguments,
       const Shape* shape_with_layout = nullptr);
-  template <typename BuilderT>
   void ComputeAndCompareLiteral(
-      BuilderT* builder, const Literal& expected,
+      XlaBuilder* builder, const Literal& expected,
       tensorflow::gtl::ArraySlice<GlobalData*> arguments, ErrorSpec error,
       const Shape* shape_with_layout = nullptr);
 
   // ComputeAndCompare variant which returns an error status.
-  template <typename BuilderT>
   tensorflow::Status ComputeAndCompareLiteralWithStatus(
-      BuilderT* builder, const Literal& expected,
+      XlaBuilder* builder, const Literal& expected,
       tensorflow::gtl::ArraySlice<GlobalData*> arguments,
       const Shape* shape_with_layout = nullptr);
-  template <typename BuilderT>
   tensorflow::Status ComputeAndCompareLiteralWithStatus(
-      BuilderT* builder, const Literal& expected,
+      XlaBuilder* builder, const Literal& expected,
       tensorflow::gtl::ArraySlice<GlobalData*> arguments, ErrorSpec error,
       const Shape* shape_with_layout = nullptr);
 
@@ -227,25 +205,13 @@ class ClientLibraryTestBase : public ::testing::Test {
 
   // Convenience method for running a built computation, transferring the
   // result, and comparing it to the expected tuple literal.
-  template <typename BuilderT>
   void ComputeAndCompareTuple(
-      BuilderT* builder, const Literal& expected,
+      XlaBuilder* builder, const Literal& expected,
       tensorflow::gtl::ArraySlice<GlobalData*> arguments);
-  template <typename BuilderT>
   void ComputeAndCompareTuple(
-      BuilderT* builder, const Literal& expected,
+      XlaBuilder* builder, const Literal& expected,
       tensorflow::gtl::ArraySlice<GlobalData*> arguments, ErrorSpec error);
 
-  // Convenience method for running a built computation and comparing the result
-  // with the HloEvaluator.
-  void ComputeAndCompare(ComputationBuilder* builder,
-                         const ComputationDataHandle& operand,
-                         tensorflow::gtl::ArraySlice<Literal> arguments);
-  void ComputeAndCompare(ComputationBuilder* builder,
-                         const ComputationDataHandle& operand,
-                         tensorflow::gtl::ArraySlice<Literal> arguments,
-                         ErrorSpec error);
-
   // Convenience method for running a built computation and comparing the result
   // with the reference result.
   void ComputeAndCompare(XlaBuilder* builder,
@@ -257,7 +223,7 @@ class ClientLibraryTestBase : public ::testing::Test {
   // Create scalar operations for use in reductions.
   XlaComputation CreateScalarRelu();
   XlaComputation CreateScalarMax();
-  Computation CreateScalarReluSensitivity();
+  XlaComputation CreateScalarReluSensitivity();
 
   // Special case convenience functions for creating filled arrays.
 
@@ -297,34 +263,25 @@ class ClientLibraryTestBase : public ::testing::Test {
   // server, then stores into "data_handle" the global handle for that
   // parameter. When the use_bfloat16 flag is set but the literal has F32
   // elements, the literal will be converted to BF16 before being transferred.
-  template <typename BuilderT, typename HandleT>
   std::unique_ptr<GlobalData> CreateParameterAndTransferLiteral(
       int64 parameter_number, const Literal& literal, const string& name,
-      BuilderT* builder, HandleT* data_handle);
+      XlaBuilder* builder, XlaOp* data_handle);
 
   // As above, but the caller can specify the device that the literal is
   // transferred to. If device_handle is nullptr, the literal will be
   // transferred to the default device.
-  template <typename BuilderT, typename HandleT>
   std::unique_ptr<GlobalData> CreateParameterAndTransferLiteral(
       int64 parameter_number, const Literal& literal, const string& name,
-      const DeviceHandle* device_handle, BuilderT* builder,
-      HandleT* data_handle);
+      const DeviceHandle* device_handle, XlaBuilder* builder,
+      XlaOp* data_handle);
 
   // Creates a parameter instruction and sets the value that will be passed to
   // the computation as specified. This function must be used for all parameters
   // or none and no parameters must be passed when invoking the computation if
   // using this mechanism. If using this mechanism, then each parameter must be
   // set exactly once. The first added parameter gets index 0, then 1 and so on.
-  ComputationDataHandle AddParam(const Literal& argument,
-                                 ComputationBuilder* builder);
   XlaOp AddParam(const Literal& argument, XlaBuilder* builder);
 
-  template <class T>
-  ComputationDataHandle AddParam(const Array<T>& argument,
-                                 ComputationBuilder* builder) {
-    return AddParam(*Literal::CreateFromArray(argument), builder);
-  }
   template <class T>
   XlaOp AddParam(const Array<T>& argument, XlaBuilder* builder) {
     return AddParam(*Literal::CreateFromArray(argument), builder);
@@ -333,18 +290,11 @@ class ClientLibraryTestBase : public ::testing::Test {
   // Creates a constant instruction with the given literal. When the
   // use_bfloat16 flag is set but the literal has F32 elements, the elements
   // will be converted to BF16s.
-  ComputationDataHandle CreateConstantFromLiteral(const Literal& literal,
-                                                  ComputationBuilder* builder);
   XlaOp CreateConstantFromLiteral(const Literal& literal, XlaBuilder* builder);
 
   // Creates a constant instruction with the given array. When the use_bfloat16
   // flag is set but the array has float elements, the elements will be
   // converted to bfloat16s.
-  template <typename NativeT>
-  ComputationDataHandle CreateConstantFromArray(const Array<NativeT>& array,
-                                                ComputationBuilder* builder) {
-    return CreateConstantFromLiteral(*Literal::CreateFromArray(array), builder);
-  }
 
   template <typename NativeT>
   XlaOp CreateConstantFromArray(const Array<NativeT>& array,
@@ -353,13 +303,6 @@ class ClientLibraryTestBase : public ::testing::Test {
   }
 
   // Same as CreateConstantFromArray, but for scalars.
-  template <typename NativeT>
-  ComputationDataHandle CreateConstantFromScalar(NativeT value,
-                                                 ComputationBuilder* builder) {
-    return CreateConstantFromLiteral(*Literal::CreateR0<NativeT>(value),
-                                     builder);
-  }
-
   template <typename NativeT>
   XlaOp CreateConstantFromScalar(NativeT value, XlaBuilder* builder) {
     return CreateConstantFromLiteral(*Literal::CreateR0<NativeT>(value),
@@ -374,12 +317,12 @@ class ClientLibraryTestBase : public ::testing::Test {
   //
   // When the use_bfloat16 flag is set but NativeT is float, the data will be
   // converted to bfloat16.
-  template <typename NativeT, typename BuilderT, typename HandleT>
+  template <typename NativeT>
   std::unique_ptr<GlobalData> CreateR0Parameter(NativeT value,
                                                 int64 parameter_number,
                                                 const string& name,
-                                                BuilderT* builder,
-                                                HandleT* data_handle);
+                                                XlaBuilder* builder,
+                                                XlaOp* data_handle);
 
   // Creates a parameter instruction that wraps the given values and then stores
   // into "data_handle" the global handle for that parameter.
@@ -389,10 +332,10 @@ class ClientLibraryTestBase : public ::testing::Test {
   //
   // When the use_bfloat16 flag is set but NativeT is float, the data will be
   // converted to bfloat16.
-  template <typename NativeT, typename BuilderT, typename HandleT>
+  template <typename NativeT>
   std::unique_ptr<GlobalData> CreateR1Parameter(
       tensorflow::gtl::ArraySlice<NativeT> values, int64 parameter_number,
-      const string& name, BuilderT* builder, HandleT* data_handle);
+      const string& name, XlaBuilder* builder, XlaOp* data_handle);
 
   // Creates a parameter instruction that wraps the given constant array
   // "array_2d" and then stores to "data_handle" the global handle for that
@@ -403,10 +346,10 @@ class ClientLibraryTestBase : public ::testing::Test {
   //
   // When the use_bfloat16 flag is set but NativeT is float, the data will be
   // converted to bfloat16.
-  template <typename NativeT, typename BuilderT, typename HandleT>
+  template <typename NativeT>
   std::unique_ptr<GlobalData> CreateR2Parameter(
       const Array2D<NativeT>& array_2d, int64 parameter_number,
-      const string& name, BuilderT* builder, HandleT* data_handle);
+      const string& name, XlaBuilder* builder, XlaOp* data_handle);
 
   // Creates a parameter instruction that wraps the given constant array
   // "array_3d" and then stores to "data_handle" the global handle for that
@@ -417,10 +360,10 @@ class ClientLibraryTestBase : public ::testing::Test {
   //
   // When the use_bfloat16 flag is set but NativeT is float, the data will be
   // converted to bfloat16.
-  template <typename NativeT, typename BuilderT, typename HandleT>
+  template <typename NativeT>
   std::unique_ptr<GlobalData> CreateR3Parameter(
       const Array3D<NativeT>& array_3d, int64 parameter_number,
-      const string& name, BuilderT* builder, HandleT* data_handle);
+      const string& name, XlaBuilder* builder, XlaOp* data_handle);
 
   // Getter and setter for the use_bfloat16 flag, which indicates whether to run
   // tests with all float-type input/output converted to bfloat16.
@@ -435,21 +378,6 @@ class ClientLibraryTestBase : public ::testing::Test {
   ExecutionOptions execution_options_;
 
  private:
-  // Build and run the computation with all permutations of output layouts.
-  tensorflow::Status ComputeAndCompareLiteralWithAllOutputLayouts(
-      const xla::Computation& computation, const Literal& expected,
-      tensorflow::gtl::ArraySlice<GlobalData*> arguments,
-      const std::function<void(const Literal& actual,
-                               const string& error_message)>& verify_output);
-  // Build and run the computation with all permutations of layouts of all input
-  // arguments.
-  tensorflow::Status ComputeAndCompareLiteralWithAllInputLayouts(
-      const xla::Computation& computation, const Literal& expected,
-      tensorflow::gtl::ArraySlice<GlobalData*> arguments,
-      const std::function<void(const Literal& actual,
-                               const string& error_message)>& verify_output,
-      const Shape* output_with_layout = nullptr);
-
   tensorflow::Status ComputeAndCompareLiteralWithAllOutputLayouts(
       const xla::XlaComputation& computation, const Literal& expected,
       tensorflow::gtl::ArraySlice<GlobalData*> arguments,
@@ -462,13 +390,6 @@ class ClientLibraryTestBase : public ::testing::Test {
                                const string& error_message)>& verify_output,
       const Shape* output_with_layout = nullptr);
 
-  // Executes the computation and calculates the expected reference value using
-  // the HloEvaluator. Returns two literals in the order of (expected, actual).
-  StatusOr<std::pair<std::unique_ptr<Literal>, std::unique_ptr<Literal>>>
-  ComputeValueAndReference(ComputationBuilder* builder,
-                           const ComputationDataHandle& operand,
-                           tensorflow::gtl::ArraySlice<Literal> arguments);
-
   // Executes the computation and calculates the expected reference value using
   // the reference client. Returns two literals in the order of (expected,
   // actual).
@@ -484,9 +405,9 @@ class ClientLibraryTestBase : public ::testing::Test {
   std::vector<std::unique_ptr<GlobalData>> arguments_;
 };
 
-template <typename NativeT, typename BuilderT>
+template <typename NativeT>
 void ClientLibraryTestBase::ComputeAndCompareR0(
-    BuilderT* builder, NativeT expected,
+    XlaBuilder* builder, NativeT expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments) {
   std::unique_ptr<Literal> expected_literal =
       Literal::CreateR0<NativeT>(expected);
@@ -494,9 +415,9 @@ void ClientLibraryTestBase::ComputeAndCompareR0(
                                                   arguments);
 }
 
-template <typename NativeT, typename BuilderT>
+template <typename NativeT>
 void ClientLibraryTestBase::ComputeAndCompareR0(
-    BuilderT* builder, NativeT expected,
+    XlaBuilder* builder, NativeT expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments, ErrorSpec error) {
   static_assert(std::is_same<NativeT, float>::value ||
                     std::is_same<NativeT, double>::value ||
@@ -510,9 +431,9 @@ void ClientLibraryTestBase::ComputeAndCompareR0(
                                                   arguments, error);
 }
 
-template <typename NativeT, typename BuilderT>
+template <typename NativeT>
 void ClientLibraryTestBase::ComputeAndCompareR1(
-    BuilderT* builder, tensorflow::gtl::ArraySlice<NativeT> expected,
+    XlaBuilder* builder, tensorflow::gtl::ArraySlice<NativeT> expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments) {
   std::unique_ptr<Literal> expected_literal =
       Literal::CreateR1<NativeT>(expected);
@@ -520,9 +441,9 @@ void ClientLibraryTestBase::ComputeAndCompareR1(
                                                   arguments);
 }
 
-template <typename NativeT, typename BuilderT>
+template <typename NativeT>
 void ClientLibraryTestBase::ComputeAndCompareR1(
-    BuilderT* builder, tensorflow::gtl::ArraySlice<NativeT> expected,
+    XlaBuilder* builder, tensorflow::gtl::ArraySlice<NativeT> expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments, ErrorSpec error) {
   static_assert(std::is_same<NativeT, float>::value ||
                     std::is_same<NativeT, double>::value ||
@@ -536,9 +457,9 @@ void ClientLibraryTestBase::ComputeAndCompareR1(
                                                   arguments, error);
 }
 
-template <typename NativeT, typename BuilderT>
+template <typename NativeT>
 void ClientLibraryTestBase::ComputeAndCompareR2(
-    BuilderT* builder, const Array2D<NativeT>& expected,
+    XlaBuilder* builder, const Array2D<NativeT>& expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments) {
   std::unique_ptr<Literal> expected_literal =
       Literal::CreateR2FromArray2D<NativeT>(expected);
@@ -546,9 +467,9 @@ void ClientLibraryTestBase::ComputeAndCompareR2(
                                                   arguments);
 }
 
-template <typename NativeT, typename BuilderT>
+template <typename NativeT>
 void ClientLibraryTestBase::ComputeAndCompareR2(
-    BuilderT* builder, const Array2D<NativeT>& expected,
+    XlaBuilder* builder, const Array2D<NativeT>& expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments, ErrorSpec error) {
   static_assert(std::is_same<NativeT, float>::value ||
                     std::is_same<NativeT, double>::value ||
@@ -562,9 +483,9 @@ void ClientLibraryTestBase::ComputeAndCompareR2(
                                                   arguments, error);
 }
 
-template <typename NativeT, typename BuilderT>
+template <typename NativeT>
 void ClientLibraryTestBase::ComputeAndCompareR3(
-    BuilderT* builder, const Array3D<NativeT>& expected,
+    XlaBuilder* builder, const Array3D<NativeT>& expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments) {
   std::unique_ptr<Literal> expected_literal =
       Literal::CreateR3FromArray3D<NativeT>(expected);
@@ -572,9 +493,9 @@ void ClientLibraryTestBase::ComputeAndCompareR3(
                                                   arguments);
 }
 
-template <typename NativeT, typename BuilderT>
+template <typename NativeT>
 void ClientLibraryTestBase::ComputeAndCompareR3(
-    BuilderT* builder, const Array3D<NativeT>& expected,
+    XlaBuilder* builder, const Array3D<NativeT>& expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments, ErrorSpec error) {
   static_assert(std::is_same<NativeT, float>::value ||
                     std::is_same<NativeT, double>::value ||
@@ -588,9 +509,9 @@ void ClientLibraryTestBase::ComputeAndCompareR3(
                                                   arguments, error);
 }
 
-template <typename NativeT, typename BuilderT>
+template <typename NativeT>
 void ClientLibraryTestBase::ComputeAndCompareR4(
-    BuilderT* builder, const Array4D<NativeT>& expected,
+    XlaBuilder* builder, const Array4D<NativeT>& expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments) {
   std::unique_ptr<Literal> expected_literal =
       Literal::CreateR4FromArray4D<NativeT>(expected);
@@ -598,9 +519,9 @@ void ClientLibraryTestBase::ComputeAndCompareR4(
                                                   arguments);
 }
 
-template <typename NativeT, typename BuilderT>
+template <typename NativeT>
 void ClientLibraryTestBase::ComputeAndCompareR4(
-    BuilderT* builder, const Array4D<NativeT>& expected,
+    XlaBuilder* builder, const Array4D<NativeT>& expected,
     tensorflow::gtl::ArraySlice<GlobalData*> arguments, ErrorSpec error) {
   static_assert(std::is_same<NativeT, float>::value ||
                     std::is_same<NativeT, double>::value ||
@@ -614,10 +535,10 @@ void ClientLibraryTestBase::ComputeAndCompareR4(
                                                   arguments, error);
 }
 
-template <typename NativeT, typename BuilderT, typename HandleT>
+template <typename NativeT>
 std::unique_ptr<GlobalData> ClientLibraryTestBase::CreateR0Parameter(
     NativeT value, int64 parameter_number, const string& name,
-    BuilderT* builder, HandleT* data_handle) {
+    XlaBuilder* builder, XlaOp* data_handle) {
   std::unique_ptr<Literal> literal = Literal::CreateR0(value);
   if (use_bfloat16_ && literal->shape().element_type() == F32) {
     literal = LiteralTestUtil::ConvertF32ToBF16(*literal);
@@ -628,10 +549,10 @@ std::unique_ptr<GlobalData> ClientLibraryTestBase::CreateR0Parameter(
   return data;
 }
 
-template <typename NativeT, typename BuilderT, typename HandleT>
+template <typename NativeT>
 std::unique_ptr<GlobalData> ClientLibraryTestBase::CreateR1Parameter(
     tensorflow::gtl::ArraySlice<NativeT> values, int64 parameter_number,
-    const string& name, BuilderT* builder, HandleT* data_handle) {
+    const string& name, XlaBuilder* builder, XlaOp* data_handle) {
   std::unique_ptr<Literal> literal = Literal::CreateR1(values);
   if (use_bfloat16_ && literal->shape().element_type() == F32) {
     literal = LiteralTestUtil::ConvertF32ToBF16(*literal);
@@ -642,10 +563,10 @@ std::unique_ptr<GlobalData> ClientLibraryTestBase::CreateR1Parameter(
   return data;
 }
 
-template <typename NativeT, typename BuilderT, typename HandleT>
+template <typename NativeT>
 std::unique_ptr<GlobalData> ClientLibraryTestBase::CreateR2Parameter(
     const Array2D<NativeT>& array_2d, int64 parameter_number,
-    const string& name, BuilderT* builder, HandleT* data_handle) {
+    const string& name, XlaBuilder* builder, XlaOp* data_handle) {
   std::unique_ptr<Literal> literal = Literal::CreateR2FromArray2D(array_2d);
   if (use_bfloat16_ && literal->shape().element_type() == F32) {
     literal = LiteralTestUtil::ConvertF32ToBF16(*literal);
@@ -656,10 +577,10 @@ std::unique_ptr<GlobalData> ClientLibraryTestBase::CreateR2Parameter(
   return data;
 }
 
-template <typename NativeT, typename BuilderT, typename HandleT>
+template <typename NativeT>
 std::unique_ptr<GlobalData> ClientLibraryTestBase::CreateR3Parameter(
     const Array3D<NativeT>& array_3d, int64 parameter_number,
-    const string& name, BuilderT* builder, HandleT* data_handle) {
+    const string& name, XlaBuilder* builder, XlaOp* data_handle) {
   std::unique_ptr<Literal> literal = Literal::CreateR3FromArray3D(array_3d);
   if (use_bfloat16_ && literal->shape().element_type() == F32) {
     literal = LiteralTestUtil::ConvertF32ToBF16(*literal);
@@ -695,23 +616,21 @@ std::unique_ptr<Array2D<NativeT>> ClientLibraryTestBase::CreatePseudorandomR2(
   return result;
 }
 
-template <typename BuilderT, typename HandleT>
 std::unique_ptr<GlobalData>
 ClientLibraryTestBase::CreateParameterAndTransferLiteral(int64 parameter_number,
                                                          const Literal& literal,
                                                          const string& name,
-                                                         BuilderT* builder,
-                                                         HandleT* data_handle) {
+                                                         XlaBuilder* builder,
+                                                         XlaOp* data_handle) {
   return CreateParameterAndTransferLiteral(parameter_number, literal, name,
                                            nullptr, builder, data_handle);
 }
 
-template <typename BuilderT, typename HandleT>
 std::unique_ptr<GlobalData>
 ClientLibraryTestBase::CreateParameterAndTransferLiteral(
     int64 parameter_number, const Literal& literal, const string& name,
-    const DeviceHandle* device_handle, BuilderT* builder,
-    HandleT* data_handle) {
+    const DeviceHandle* device_handle, XlaBuilder* builder,
+    XlaOp* data_handle) {
   const Literal* param_literal = &literal;
   std::unique_ptr<Literal> converted_literal;
   if (use_bfloat16_) {
-- 
GitLab


From 8ec11ae8eb7b97caced73ed3971209236e2aef5c Mon Sep 17 00:00:00 2001
From: Yuefeng Zhou <yuefengz@google.com>
Date: Thu, 3 May 2018 22:01:39 -0700
Subject: [PATCH 246/395] Add the MultiWorkerMirroredStrategy

PiperOrigin-RevId: 195368876
---
 tensorflow/contrib/distribute/python/BUILD    |  13 ++
 .../distribute/python/mirrored_strategy.py    |   1 +
 .../python/multi_worker_strategy.py           | 141 ++++++++++++++++++
 .../python/multi_worker_strategy_test.py      |  64 ++++++++
 .../distribute/python/one_device_strategy.py  |   1 +
 tensorflow/python/training/distribute.py      |  20 ++-
 6 files changed, 238 insertions(+), 2 deletions(-)
 create mode 100644 tensorflow/contrib/distribute/python/multi_worker_strategy.py
 create mode 100644 tensorflow/contrib/distribute/python/multi_worker_strategy_test.py

diff --git a/tensorflow/contrib/distribute/python/BUILD b/tensorflow/contrib/distribute/python/BUILD
index aaafc184bf..8dfcaf6032 100644
--- a/tensorflow/contrib/distribute/python/BUILD
+++ b/tensorflow/contrib/distribute/python/BUILD
@@ -86,6 +86,19 @@ py_library(
     ],
 )
 
+py_library(
+    name = "multi_worker_strategy",
+    srcs = ["multi_worker_strategy.py"],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        ":mirrored_strategy",
+        ":values",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python:training",
+        "//tensorflow/python:util",
+    ],
+)
+
 py_library(
     name = "one_device_strategy",
     srcs = ["one_device_strategy.py"],
diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy.py b/tensorflow/contrib/distribute/python/mirrored_strategy.py
index 2e57b02583..8237b23dbb 100644
--- a/tensorflow/contrib/distribute/python/mirrored_strategy.py
+++ b/tensorflow/contrib/distribute/python/mirrored_strategy.py
@@ -80,6 +80,7 @@ class MirroredStrategy(distribute_lib.DistributionStrategy):
         dict((d, i) for i, d in enumerate(devices)))
     self._cross_tower_ops = cross_tower_ops
     self._prefetch_on_device = prefetch_on_device
+    # TODO(yuefengz): consider setting the default device.
 
   def _create_variable(self, next_creator, *args, **kwargs):
     """Create a mirrored variable. See `DistributionStrategy.scope`."""
diff --git a/tensorflow/contrib/distribute/python/multi_worker_strategy.py b/tensorflow/contrib/distribute/python/multi_worker_strategy.py
new file mode 100644
index 0000000000..a552b370eb
--- /dev/null
+++ b/tensorflow/contrib/distribute/python/multi_worker_strategy.py
@@ -0,0 +1,141 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Classes implementing a mirrored DistributionStrategy for multiple workers."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from functools import partial
+
+from tensorflow.contrib.distribute.python import values
+from tensorflow.contrib.distribute.python.mirrored_strategy import MirroredStrategy
+from tensorflow.core.protobuf import cluster_pb2
+from tensorflow.python.training import device_util
+from tensorflow.python.training import server_lib
+from tensorflow.python.util import nest
+
+
+# TODO(yuefengz): support between-graph replication.
+# TODO(yuefengz): merge this class into its base class.
+# TODO(yuefengz): in some cases, we probably want to use configure method to
+# configure this class.
+# TODO(yuefengz): MirroredStrategy.worker_devices may be confusing after the
+# class is introduced.
+class MultiWorkerMirroredStrategy(MirroredStrategy):
+  """Mirrored strategy that works on multiple workers with in-graph replication.
+
+  There are several important concepts for distributed TensorFlow, e.g.
+  `client`, `job`, 'task', `cluster`, `in-graph replication` and
+  'synchronous training' and they have already been defined in the
+  [TensorFlow's documentation](https://www.tensorflow.org/deploy/distributed).
+  The distribution strategy inherits these concepts as well and in addition to
+  that we also clarify several more concepts:
+    * **In-graph replication**: the `client` creates a single `tf.Graph` that
+    specifies tasks for devices on all workers. The `client` then creates a
+    client session which will talk to the `master` service of a `worker`. Then
+    the `master` will parition the graph and distribute the work to all
+    participating workers.
+    * **Worker**: A `worker` is a TensorFlow `task` that usually maps to one
+    physical machine. We will have multiple `worker`s with different `task`
+    index. They all do similar things except for one worker checkpointing model
+    variables, writing summaries, etc. in addition to its ordinary work.
+
+  This class maps one tower to one device on a worker. It mirrors all model
+  variables on all towers. For example, if you have two `worker`s and each
+  `worker` has 4 GPUs, it will create 8 copies of the model variables on these 8
+  GPUs. Then like in MirroredStrategy, each tower performs their computation
+  with their own copy of variables unless in cross-tower model where variable or
+  tensor reduction happens.
+  """
+
+  def __init__(self,
+               num_gpus_per_worker=1,
+               worker_job_name=None,
+               num_workers=None,
+               cluster=None,
+               cross_tower_ops=None,
+               prefetch_on_device=None):
+    """Initialize the strategy object.
+
+    Args:
+      num_gpus_per_worker: number of GPUs per work. If it is zero, the local
+        CPU will be used.
+      worker_job_name: the job name for `worker`, typically just 'worker'.
+      num_workers: the number of workers. If it is 0, it regenerates to
+        single-worker MirroredStrategy.
+      cluster: a `tf.train.ClusterSpec` object or a dict that can be used to
+        construct a `tf.train.ClusterSpec` object or a `tf.train.ClusterDef`
+        proto buffer. It is an alternative way to initialize this object.
+      cross_tower_ops: the cross tower ops to use. If None, a default one will
+        be used. If configure method is called, a best one for the configuration
+        will be chosen.
+      prefetch_on_device: a boolean to specify whether to prefetech input to
+        each worker's devices.
+
+    Raises:
+      ValueError: if got an unexpected `cluster`.
+    """
+    if cluster is None:
+      self._workers = [
+          '/job:%s/task:%d' % (worker_job_name, task_index)
+          for task_index in range(num_workers)
+      ]
+    else:
+      if isinstance(cluster, (dict, cluster_pb2.ClusterDef)):
+        cluster_spec = server_lib.ClusterSpec(cluster)
+      elif isinstance(cluster, server_lib.ClusterSpec):
+        cluster_spec = cluster
+      else:
+        raise ValueError(
+            "`cluster_spec' should be dict or a `tf.train.ClusterSpec` or a "
+            '`tf.train.ClusterDef` object')
+
+      self._workers = []
+      for job in sorted(cluster_spec.jobs):
+        for task in range(cluster_spec.num_tasks(job)):
+          self._workers.append('/job:%s/task:%d' % (job, task))
+
+    self._num_gpus_per_worker = num_gpus_per_worker
+    if num_gpus_per_worker > 0:
+      self._worker_device_map = {
+          worker: [
+              device_util.canonicalize(worker + '/device:GPU:%d' % gpu)
+              for gpu in range(num_gpus_per_worker)
+          ] for worker in self._workers
+      }
+    else:
+      self._worker_device_map = {
+          worker: [device_util.canonicalize(worker, '/device:CPU:0')]
+          for worker in self._workers
+      }
+    self._devices = nest.flatten(self._worker_device_map.values())
+
+    super(MultiWorkerMirroredStrategy, self).__init__(
+        devices=self._devices, prefetch_on_device=prefetch_on_device)
+
+    # Setting `_default_device` will add a device scope in the
+    # distribution.scope. We set the default device to the first worker. When
+    # users specify device under distribution.scope by
+    #   with tf.device("/cpu:0"):
+    #     ...
+    # their ops will end up on the cpu device of its first worker, e.g.
+    # "/job:worker/task:0/device:CPU:0". Note this is not used in tower mode.
+    self._default_device = self._workers[0]
+
+  def distribute_dataset(self, dataset_fn):
+    return values.MultiWorkerDataset(
+        partial(self._call_dataset_fn, dataset_fn), self._worker_device_map,
+        self._prefetch_on_device)
diff --git a/tensorflow/contrib/distribute/python/multi_worker_strategy_test.py b/tensorflow/contrib/distribute/python/multi_worker_strategy_test.py
new file mode 100644
index 0000000000..ee7588163e
--- /dev/null
+++ b/tensorflow/contrib/distribute/python/multi_worker_strategy_test.py
@@ -0,0 +1,64 @@
+# Copyright 2018 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 MultiWorkerMirroredStrategy."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.contrib.distribute.python import multi_worker_strategy
+from tensorflow.contrib.distribute.python import multi_worker_test_base
+from tensorflow.contrib.distribute.python import strategy_test_lib
+from tensorflow.python.eager import context
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.training import server_lib
+
+
+@test_util.with_c_api
+class MultiWorkerStrategyTest(multi_worker_test_base.MultiWorkerTestBase,
+                              strategy_test_lib.DistributionTestBase):
+
+  def _get_distribution_strategy(self):
+    return multi_worker_strategy.MultiWorkerMirroredStrategy(
+        cluster=server_lib.ClusterSpec({
+            'worker': ['/job:worker/task:0', '/job:worker/task:1']
+        }),
+        num_gpus_per_worker=context.num_gpus())
+
+  def testMinimizeLossGraph(self):
+    self._test_minimize_loss_graph(self._get_distribution_strategy())
+
+
+class DeviceScopeTest(test.TestCase):
+  """Test the device scope of MultiWorkerMirroredStrategy."""
+
+  def testDeviceScope(self):
+    with context.graph_mode():
+      strategy = multi_worker_strategy.MultiWorkerMirroredStrategy(
+          cluster={'worker': ['/job:worker/task:0', '/job:worker/task:1']},
+          num_gpus_per_worker=context.num_gpus())
+      with strategy.scope():
+        a = constant_op.constant(1.)
+        with ops.device('/cpu:0'):
+          b = constant_op.constant(1.)
+        self.assertEqual(a.device, '/job:worker/task:0')
+        self.assertEqual(b.device, '/job:worker/task:0/device:CPU:0')
+
+
+if __name__ == '__main__':
+  test.main()
diff --git a/tensorflow/contrib/distribute/python/one_device_strategy.py b/tensorflow/contrib/distribute/python/one_device_strategy.py
index 64aa369201..09b6d4a515 100644
--- a/tensorflow/contrib/distribute/python/one_device_strategy.py
+++ b/tensorflow/contrib/distribute/python/one_device_strategy.py
@@ -40,6 +40,7 @@ class OneDeviceStrategy(distribute_lib.DistributionStrategy):
     super(OneDeviceStrategy, self).__init__()
     self._device = device
     self._prefetch_on_device = prefetch_on_device
+    self._default_device = device
 
   def _create_variable(self, next_creator, *args, **kwargs):
     # No need to distinguish tower-local variables when not mirroring,
diff --git a/tensorflow/python/training/distribute.py b/tensorflow/python/training/distribute.py
index c16b05102e..21f81ee187 100644
--- a/tensorflow/python/training/distribute.py
+++ b/tensorflow/python/training/distribute.py
@@ -290,19 +290,31 @@ def _require_distribution_strategy_scope(distribution_strategy):
 class _CurrentDistributionContext(object):
   """Context manager for setting the `DistributionStrategy` and var creator."""
 
-  def __init__(self, distribution_strategy, var_creator_scope, var_scope=None):
+  def __init__(self,
+               distribution_strategy,
+               var_creator_scope,
+               var_scope=None,
+               default_device=None):
     self._context = _CrossTowerThreadMode(distribution_strategy)
     self._var_creator_scope = var_creator_scope
     self._var_scope = var_scope
+    if default_device:
+      self._device_scope = ops.device(default_device)
+    else:
+      self._device_scope = None
 
   def __enter__(self):
     _push_per_thread_mode(self._context)
     if self._var_scope:
       self._var_scope.__enter__()
     self._var_creator_scope.__enter__()
+    if self._device_scope:
+      self._device_scope.__enter__()
     return self._context.distribution_strategy
 
   def __exit__(self, exception_type, exception_value, traceback):
+    if self._device_scope:
+      self._device_scope.__exit__(exception_type, exception_value, traceback)
     self._var_creator_scope.__exit__(exception_type, exception_value, traceback)
     if self._var_scope:
       self._var_scope.__exit__(exception_type, exception_value, traceback)
@@ -557,6 +569,9 @@ class DistributionStrategy(object):
   # TODO(josh11b): List of towers with their worker and parameter devices
   #   (where the parameter devices may overlap in the ps case).
 
+  def __init__(self):
+    self._default_device = None
+
   def scope(self):
     """Returns a context manager selecting this DistributionStrategy as current.
 
@@ -587,7 +602,8 @@ class DistributionStrategy(object):
         self, variable_scope.variable_creator_scope(creator_with_resource_vars),
         variable_scope.variable_scope(
             variable_scope.get_variable_scope(),
-            custom_getter=disable_partitioned_variables))
+            custom_getter=disable_partitioned_variables),
+        self._default_device)
 
   def _create_variable(self, next_creator, *args, **kwargs):
     # Note: should support "colocate_with" argument.
-- 
GitLab


From da0dcb21501b765932e392ae710ebbecefeb309c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 3 May 2018 23:36:02 -0700
Subject: [PATCH 247/395] Internal change.

PiperOrigin-RevId: 195374319
---
 .../kernels/bidirectional_sequence_lstm.cc    |  6 ++---
 tensorflow/contrib/lite/kernels/kernel_util.h |  4 +++
 tensorflow/contrib/lite/kernels/lstm.cc       |  4 +--
 tensorflow/contrib/lite/kernels/mean.cc       | 16 ++++++------
 tensorflow/contrib/lite/kernels/svdf.cc       | 26 +++++++++----------
 .../kernels/unidirectional_sequence_lstm.cc   |  4 +--
 6 files changed, 30 insertions(+), 30 deletions(-)

diff --git a/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm.cc b/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm.cc
index 3ac0210f36..a35ba23ced 100644
--- a/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm.cc
+++ b/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm.cc
@@ -365,8 +365,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TfLiteIntArrayFree(node->temporaries);
   node->temporaries = TfLiteIntArrayCreate(2);
   node->temporaries->data[0] = *scratch_tensor_index;
-  TfLiteTensor* fw_scratch_buffer =
-      &context->tensors[node->temporaries->data[0]];
+  TfLiteTensor* fw_scratch_buffer = GetTemporary(context, node, /*index=*/0);
   fw_scratch_buffer->type = input->type;
   fw_scratch_buffer->allocation_type = kTfLiteArenaRw;
 
@@ -434,8 +433,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   // Create a scratch buffer tensor.
   node->temporaries->data[1] = *(scratch_tensor_index) + 1;
-  TfLiteTensor* bw_scratch_buffer =
-      &context->tensors[node->temporaries->data[1]];
+  TfLiteTensor* bw_scratch_buffer = GetTemporary(context, node, /*index=*/1);
   bw_scratch_buffer->type = input->type;
   bw_scratch_buffer->allocation_type = kTfLiteArenaRw;
 
diff --git a/tensorflow/contrib/lite/kernels/kernel_util.h b/tensorflow/contrib/lite/kernels/kernel_util.h
index 2f407b5da3..e225443a67 100644
--- a/tensorflow/contrib/lite/kernels/kernel_util.h
+++ b/tensorflow/contrib/lite/kernels/kernel_util.h
@@ -32,6 +32,10 @@ inline TfLiteTensor* GetOutput(TfLiteContext* context, TfLiteNode* node,
                                int index) {
   return &context->tensors[node->outputs->data[index]];
 }
+inline TfLiteTensor* GetTemporary(TfLiteContext* context, TfLiteNode* node,
+                                  int index) {
+  return &context->tensors[node->temporaries->data[index]];
+}
 inline int NumInputs(const TfLiteNode* node) { return node->inputs->size; }
 inline int NumOutputs(const TfLiteNode* node) { return node->outputs->size; }
 
diff --git a/tensorflow/contrib/lite/kernels/lstm.cc b/tensorflow/contrib/lite/kernels/lstm.cc
index 668226e674..a1521efbb4 100644
--- a/tensorflow/contrib/lite/kernels/lstm.cc
+++ b/tensorflow/contrib/lite/kernels/lstm.cc
@@ -290,7 +290,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TfLiteIntArrayFree(node->temporaries);
   node->temporaries = TfLiteIntArrayCreate(1);
   node->temporaries->data[0] = *scratch_tensor_index;
-  TfLiteTensor* scratch_buffer = &context->tensors[node->temporaries->data[0]];
+  TfLiteTensor* scratch_buffer = GetTemporary(context, node, /*index=*/0);
   scratch_buffer->type = input->type;
   scratch_buffer->allocation_type = kTfLiteArenaRw;
 
@@ -378,7 +378,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   const bool use_peephole = (cell_to_output_weights != nullptr);
 
   // Index the scratch buffers pointers to the global scratch buffer.
-  TfLiteTensor* scratch_buffer = &context->tensors[node->temporaries->data[0]];
+  TfLiteTensor* scratch_buffer = GetTemporary(context, node, /*index=*/0);
 
   float* input_gate_scratch = nullptr;
   float* cell_scratch = nullptr;
diff --git a/tensorflow/contrib/lite/kernels/mean.cc b/tensorflow/contrib/lite/kernels/mean.cc
index 047bdd1039..98f80e32d9 100644
--- a/tensorflow/contrib/lite/kernels/mean.cc
+++ b/tensorflow/contrib/lite/kernels/mean.cc
@@ -146,7 +146,7 @@ TfLiteStatus InitializeTemporaries(TfLiteContext* context, TfLiteNode* node,
   TfLiteIntArrayFree(node->temporaries);
   node->temporaries = TfLiteIntArrayCreate(3);
   node->temporaries->data[0] = *scratch_tensor_index;
-  TfLiteTensor* scratch_tensor = &context->tensors[node->temporaries->data[0]];
+  TfLiteTensor* scratch_tensor = GetTemporary(context, node, /*index=*/0);
   scratch_tensor->type = kTfLiteInt32;
   scratch_tensor->allocation_type = kTfLiteArenaRw;
   TfLiteIntArray* index_size = TfLiteIntArrayCreate(1);
@@ -156,11 +156,11 @@ TfLiteStatus InitializeTemporaries(TfLiteContext* context, TfLiteNode* node,
 
   // Creates a temp tensor to store resolved axis given input data.
   node->temporaries->data[1] = *scratch_tensor_index + 1;
-  TfLiteTensor* resolved_axis = &context->tensors[node->temporaries->data[1]];
+  TfLiteTensor* resolved_axis = GetTemporary(context, node, /*index=*/1);
   resolved_axis->type = kTfLiteInt32;
   // Creates a temp tensor to store temp sums when calculating mean.
   node->temporaries->data[2] = *scratch_tensor_index + 2;
-  TfLiteTensor* temp_sum = &context->tensors[node->temporaries->data[2]];
+  TfLiteTensor* temp_sum = GetTemporary(context, node, /*index=*/2);
   switch (op_context->input->type) {
     case kTfLiteFloat32:
       temp_sum->type = kTfLiteFloat32;
@@ -187,8 +187,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   MeanContext op_context(context, node);
   TF_LITE_ENSURE_OK(context, InitializeTemporaries(context, node, &op_context));
 
-  TfLiteTensor* resolved_axis = &context->tensors[node->temporaries->data[1]];
-  TfLiteTensor* temp_sum = &context->tensors[node->temporaries->data[2]];
+  TfLiteTensor* resolved_axis = GetTemporary(context, node, /*index=*/1);
+  TfLiteTensor* temp_sum = GetTemporary(context, node, /*index=*/2);
   // Leaves work to Eval if axis is not constant; else resizes output.
   if (!IsConstantTensor(op_context.axis)) {
     SetTensorToDynamic(op_context.output);
@@ -208,9 +208,9 @@ template <KernelType kernel_type>
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   MeanContext op_context(context, node);
   int num_axis = static_cast<int>(NumElements(op_context.axis));
-  TfLiteTensor* temp_index = &context->tensors[node->temporaries->data[0]];
-  TfLiteTensor* resolved_axis = &context->tensors[node->temporaries->data[1]];
-  TfLiteTensor* temp_sum = &context->tensors[node->temporaries->data[2]];
+  TfLiteTensor* temp_index = GetTemporary(context, node, /*index=*/0);
+  TfLiteTensor* resolved_axis = GetTemporary(context, node, /*index=*/1);
+  TfLiteTensor* temp_sum = GetTemporary(context, node, /*index=*/2);
   // Resize the output tensor if the output tensor is dynamic.
   if (IsDynamicTensor(op_context.output)) {
     TF_LITE_ENSURE_OK(context,
diff --git a/tensorflow/contrib/lite/kernels/svdf.cc b/tensorflow/contrib/lite/kernels/svdf.cc
index c69755447d..13da51c7a7 100644
--- a/tensorflow/contrib/lite/kernels/svdf.cc
+++ b/tensorflow/contrib/lite/kernels/svdf.cc
@@ -37,7 +37,7 @@ constexpr int kWeightsFeatureTensor = 1;
 constexpr int kWeightsTimeTensor = 2;
 constexpr int kBiasTensor = 3;
 constexpr int kStateTensor = 0;
-constexpr int KOutputTensor = 1;
+constexpr int kOutputTensor = 1;
 
 void* Init(TfLiteContext* context, const char* buffer, size_t length) {
   auto* scratch_tensor_index = new int;
@@ -59,9 +59,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   TfLiteTensor* input = &context->tensors[node->inputs->data[kInputTensor]];
   TfLiteTensor* weights_feature =
-      &context->tensors[node->inputs->data[kWeightsFeatureTensor]];
-  TfLiteTensor* weights_time =
-      &context->tensors[node->inputs->data[kWeightsTimeTensor]];
+      GetInput(context, node, kWeightsFeatureTensor);
+  TfLiteTensor* weights_time = GetInput(context, node, kWeightsTimeTensor);
 
   // Check all the parameters of tensor match within themselves and match the
   // input configuration.
@@ -79,8 +78,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     TF_LITE_ASSERT_EQ(bias->dims->data[0], num_units);
   }
 
-  TfLiteTensor* state = &context->tensors[node->outputs->data[kStateTensor]];
-  TfLiteTensor* output = &context->tensors[node->outputs->data[KOutputTensor]];
+  TfLiteTensor* state = GetOutput(context, node, kStateTensor);
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
 
   // Resize state.
   // For each batch, the state is a 2-D tensor: memory_size * num_filters
@@ -112,7 +111,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   scratch_size_array->data[0] = batch_size;
   scratch_size_array->data[1] = num_filters;
 
-  TfLiteTensor* scratch_tensor = &context->tensors[node->temporaries->data[0]];
+  TfLiteTensor* scratch_tensor = GetTemporary(context, node, /*index=*/0);
   scratch_tensor->type = input->type;
   scratch_tensor->allocation_type = kTfLiteArenaRw;
   TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_tensor,
@@ -124,15 +123,14 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   auto* params = reinterpret_cast<TfLiteSVDFParams*>(node->builtin_data);
 
-  TfLiteTensor* input = &context->tensors[node->inputs->data[kInputTensor]];
+  TfLiteTensor* input = GetInput(context, node, kInputTensor);
   TfLiteTensor* weights_feature =
-      &context->tensors[node->inputs->data[kWeightsFeatureTensor]];
-  TfLiteTensor* weights_time =
-      &context->tensors[node->inputs->data[kWeightsTimeTensor]];
+      GetInput(context, node, kWeightsFeatureTensor);
+  TfLiteTensor* weights_time = GetInput(context, node, kWeightsTimeTensor);
 
-  TfLiteTensor* state = &context->tensors[node->outputs->data[kStateTensor]];
-  TfLiteTensor* output = &context->tensors[node->outputs->data[KOutputTensor]];
-  TfLiteTensor* scratch = &context->tensors[node->temporaries->data[0]];
+  TfLiteTensor* state = GetOutput(context, node, kStateTensor);
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+  TfLiteTensor* scratch = GetTemporary(context, node, /*index=*/0);
 
   TfLiteTensor* bias = GetOptionalInputTensor(context, node, kBiasTensor);
 
diff --git a/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.cc b/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.cc
index 3c1256d3a6..5987bf68b5 100644
--- a/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.cc
+++ b/tensorflow/contrib/lite/kernels/unidirectional_sequence_lstm.cc
@@ -292,7 +292,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TfLiteIntArrayFree(node->temporaries);
   node->temporaries = TfLiteIntArrayCreate(1);
   node->temporaries->data[0] = *scratch_tensor_index;
-  TfLiteTensor* scratch_buffer = &context->tensors[node->temporaries->data[0]];
+  TfLiteTensor* scratch_buffer = GetTemporary(context, node, /*index=*/0);
   scratch_buffer->type = input->type;
   scratch_buffer->allocation_type = kTfLiteArenaRw;
 
@@ -381,7 +381,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   const bool use_peephole = (cell_to_output_weights != nullptr);
 
   // Index the scratch buffers pointers to the global scratch buffer.
-  TfLiteTensor* scratch_buffer = &context->tensors[node->temporaries->data[0]];
+  TfLiteTensor* scratch_buffer = GetTemporary(context, node, /*index=*/0);
   float* input_gate_scratch = nullptr;
   float* cell_scratch = nullptr;
   float* forget_gate_scratch = nullptr;
-- 
GitLab


From 0bb55f02022e88affefc111cf9a8cf70a046d1da Mon Sep 17 00:00:00 2001
From: HyoukJoong Lee <hyouklee@google.com>
Date: Fri, 4 May 2018 00:51:58 -0700
Subject: [PATCH 248/395] Automated g4 rollback of changelist 194829761

PiperOrigin-RevId: 195379693
---
 .../xla/service/hlo_module_group_metadata.cc        |  7 -------
 .../xla/service/hlo_module_group_metadata.h         |  3 ---
 tensorflow/compiler/xla/service/service.cc          | 13 +++----------
 3 files changed, 3 insertions(+), 20 deletions(-)

diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc
index 3367d76ded..54c34ce116 100644
--- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc
+++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc
@@ -194,13 +194,6 @@ int64 HloModuleGroupMetadata::GetModuleId(const HloModule* module) const {
   LOG(FATAL) << "unknown module";
 }
 
-int64 HloModuleGroupMetadata::GetDeviceModulesCount() const {
-  return std::count_if(modules_.begin(), modules_.end(),
-                       [](const HloModule* module) {
-                         return !module->config().is_host_module();
-                       });
-}
-
 Status HloModuleGroupMetadata::RecordInstructions() {
   const auto visitor = [this](HloInstruction* hlo) -> Status {
     if (hlo->opcode() == HloOpcode::kWhile) {
diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.h b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h
index d619082616..c48a7ab0b5 100644
--- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.h
+++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h
@@ -147,9 +147,6 @@ class HloModuleGroupMetadata {
   // the module in the module vector.
   int64 GetModuleId(const HloModule* module) const;
 
-  // Returns the number of modules for devices (excluding the host module).
-  int64 GetDeviceModulesCount() const;
-
   // Returns the companion instructions for the given instruction.
   //
   // Precondition: IsCompanionWhile(instruction) is true.
diff --git a/tensorflow/compiler/xla/service/service.cc b/tensorflow/compiler/xla/service/service.cc
index 6ce03ab39d..495f8801ba 100644
--- a/tensorflow/compiler/xla/service/service.cc
+++ b/tensorflow/compiler/xla/service/service.cc
@@ -626,16 +626,9 @@ Service::ExecuteParallelAndRegisterResult(
   // profiled.
   std::map<int64, se::Stream*> index_to_profiled_streams;
 
-  // Build DeviceAssignment for all cores based on the provided device handles.
-  DeviceAssignment device_assignment(options_.number_of_replicas(),
-                                     executables.size());
-  for (int64 i = 0; i < executables.size(); i++) {
-    TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*backend, device_handles[i]));
-    CHECK_EQ(replicas.size(), arguments[i].size());
-    for (int64 replica = 0; replica < replicas.size(); ++replica) {
-      device_assignment(replica, i) = replicas[replica]->device_ordinal();
-    }
-  }
+  TF_ASSIGN_OR_RETURN(DeviceAssignment device_assignment,
+                      backend->computation_placer()->AssignDevices(
+                          options_.number_of_replicas(), executables.size()));
 
   for (int64 i = 0; i < executables.size(); i++) {
     // Stream executors for the replicas of the current computation.
-- 
GitLab


From 1284047dca0dd58745a31cd2fd68da3173c7e120 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 01:47:12 -0700
Subject: [PATCH 249/395]   * Don't copy on-host and on-device shapes locally. 
  * Use ForEachMutableElement rather than the iterators, as it is much
 quicker.

There is still room for improvement; ForEachMutableElement is linear in the number of nodes in the shape tree but we want to be linear in the number of nodes in the sub shape tree. But I feel this is a good enough improvement.

PiperOrigin-RevId: 195384423
---
 tensorflow/compiler/jit/BUILD                 | 25 ++++++++
 tensorflow/compiler/jit/xla_launch_util.cc    | 22 ++++---
 tensorflow/compiler/jit/xla_launch_util.h     | 11 ++++
 .../compiler/jit/xla_launch_util_test.cc      | 64 +++++++++++++++++++
 4 files changed, 113 insertions(+), 9 deletions(-)
 create mode 100644 tensorflow/compiler/jit/xla_launch_util_test.cc

diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD
index af2965bba5..07136d6a74 100644
--- a/tensorflow/compiler/jit/BUILD
+++ b/tensorflow/compiler/jit/BUILD
@@ -360,6 +360,31 @@ tf_cc_test(
     ],
 )
 
+tf_cc_test(
+    name = "xla_launch_util_test",
+    size = "small",
+    srcs = ["xla_launch_util_test.cc"],
+    deps = [
+        ":common",
+        ":xla_compilation_cache",
+        ":xla_launch_util",
+        ":xla_tensor",
+        "//tensorflow/compiler/tf2xla:common",
+        "//tensorflow/compiler/tf2xla:xla_compiler",
+        "//tensorflow/compiler/xla:statusor",
+        "//tensorflow/compiler/xla/client:client_library",
+        "//tensorflow/compiler/xla/client:local_client",
+        "//tensorflow/core:core_cpu_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:gpu_runtime",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core/kernels:variable_ops",
+    ],
+)
+
 # This target can be used by XLA device plugins to prevent circular dependencies, and provides access to all of the required headers for building a device library.
 cc_header_only_library(
     name = "xla_jit_headers_lib",
diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc
index 2a7f04271d..33e53612b9 100644
--- a/tensorflow/compiler/jit/xla_launch_util.cc
+++ b/tensorflow/compiler/jit/xla_launch_util.cc
@@ -77,16 +77,16 @@ Status XlaAllocator::Deallocate(int device_ordinal, se::DeviceMemoryBase* mem) {
   return Status::OK();
 }
 
-namespace {
+namespace internal {
 // Return the 'index''th subtree of the given ShapedBuffer as a
 // ScopedShapedBuffer. The returned ScopedShapedBuffer takes ownership of the
 // subtree, and sets the input's buffer pointers to nullptr for the subtree.
 ScopedShapedBuffer ExtractSubShapedBuffer(
     ShapedBuffer* shaped_buffer, int index,
     xla::DeviceMemoryAllocator* allocator) {
-  xla::Shape on_host_shape = xla::ShapeUtil::GetTupleElementShape(
+  const xla::Shape& on_host_shape = xla::ShapeUtil::GetTupleElementShape(
       shaped_buffer->on_host_shape(), index);
-  xla::Shape on_device_shape = xla::ShapeUtil::GetTupleElementShape(
+  const xla::Shape& on_device_shape = xla::ShapeUtil::GetTupleElementShape(
       shaped_buffer->on_device_shape(), index);
 
   ShapedBuffer sub_shaped_buffer(on_host_shape, on_device_shape,
@@ -98,14 +98,18 @@ ScopedShapedBuffer ExtractSubShapedBuffer(
   sub_shape_tree.CopySubtreeFrom(shape_tree,
                                  /*source_base_index=*/{index},
                                  /*target_base_index=*/{});
-  for (auto& index_to_buffer : shape_tree) {
-    if (!index_to_buffer.first.empty() && index_to_buffer.first[0] == index) {
-      index_to_buffer.second = se::DeviceMemoryBase(nullptr, 0);
-    }
-  }
+  shape_tree.ForEachMutableElement(
+      [index](const xla::ShapeIndex& shape_index,
+              tensorflow::se::DeviceMemoryBase* data) {
+        // shape_index is empty for the root node. Ignore that.
+        if (!shape_index.empty() && shape_index[0] == index) {
+          *data = tensorflow::se::DeviceMemoryBase(nullptr, 0);
+        }
+      });
   return ScopedShapedBuffer(std::move(sub_shaped_buffer), allocator);
 }
-}  // namespace
+}  // namespace internal
+using internal::ExtractSubShapedBuffer;
 
 XlaComputationLaunchContext::XlaComputationLaunchContext(
     int64 num_resource_args, xla::LocalClient* client,
diff --git a/tensorflow/compiler/jit/xla_launch_util.h b/tensorflow/compiler/jit/xla_launch_util.h
index 8a6ff3b0c7..38291b0bd4 100644
--- a/tensorflow/compiler/jit/xla_launch_util.h
+++ b/tensorflow/compiler/jit/xla_launch_util.h
@@ -140,6 +140,17 @@ class XlaTensorBuffer : public TensorBuffer {
   Allocator* allocator_;
 };
 
+// Exposed in this header file for microbenchmarking purposes, but this is an
+// internal implementation detail.
+namespace internal {
+// Return the 'index''th subtree of the given ShapedBuffer as a
+// ScopedShapedBuffer. The returned ScopedShapedBuffer takes ownership of the
+// subtree, and sets the input's buffer pointers to nullptr for the subtree.
+xla::ScopedShapedBuffer ExtractSubShapedBuffer(
+    xla::ShapedBuffer* shaped_buffer, int index,
+    xla::DeviceMemoryAllocator* allocator);
+}  // namespace internal
+
 }  // namespace tensorflow
 
 #endif
diff --git a/tensorflow/compiler/jit/xla_launch_util_test.cc b/tensorflow/compiler/jit/xla_launch_util_test.cc
new file mode 100644
index 0000000000..27813efc0b
--- /dev/null
+++ b/tensorflow/compiler/jit/xla_launch_util_test.cc
@@ -0,0 +1,64 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+// Contains microbenchmarks for performance critical functions in
+// xla_launch_util.cc.
+
+#include "tensorflow/compiler/jit/xla_launch_util.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/test_benchmark.h"
+
+// Test ExtractSubBuffer with different depths (depth of ShapeTree) and fan-outs
+// (cardinality of each non-leaf node's children).
+void BM_ExtractSubBuffer(int iters, int depth, int fan_out) {
+  tensorflow::testing::StopTiming();
+  xla::Shape shape = xla::ShapeUtil::MakeShape(xla::F32, {32, 64, 128});
+  for (int i = 0; i < depth; ++i) {
+    std::vector<xla::Shape> shapes(fan_out, shape);
+    shape = xla::ShapeUtil::MakeTupleShape(shapes);
+  }
+  xla::ShapedBuffer shaped_buffer(shape, shape, /*platform=*/nullptr,
+                                  /*device_ordinal=*/0);
+  tensorflow::testing::StartTiming();
+  for (int i = 0; i < iters; ++i) {
+    // Extract a buffer from approximately the middle of the first level of the
+    // tree.
+    tensorflow::internal::ExtractSubShapedBuffer(&shaped_buffer,
+                                                 /*index=*/fan_out / 2,
+                                                 /*allocator=*/nullptr)
+        .release();
+  }
+}
+
+BENCHMARK(BM_ExtractSubBuffer)
+    ->ArgPair(1, 4)
+    ->ArgPair(1, 8)
+    ->ArgPair(1, 32)
+    ->ArgPair(1, 64)
+    ->ArgPair(1, 128)
+    ->ArgPair(1, 256)
+    ->ArgPair(1, 512)
+    ->ArgPair(2, 4)
+    ->ArgPair(2, 8)
+    ->ArgPair(2, 32)
+    ->ArgPair(2, 64)
+    ->ArgPair(2, 128);
+
+int main(int argc, char** argv) {
+  testing::InitGoogleTest(&argc, argv);
+  tensorflow::testing::RunBenchmarks();
+  return RUN_ALL_TESTS();
+}
-- 
GitLab


From 73a1908b3c50d2f665a3a9af491e217d814edb40 Mon Sep 17 00:00:00 2001
From: Tom Hennigan <tomhennigan@google.com>
Date: Fri, 4 May 2018 01:57:02 -0700
Subject: [PATCH 250/395] Prefer non-nested GradientTape.gradient call when
 only one source is passed.

PiperOrigin-RevId: 195385406
---
 tensorflow/docs_src/programmers_guide/eager.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/docs_src/programmers_guide/eager.md b/tensorflow/docs_src/programmers_guide/eager.md
index 595e6be4af..5926e9f7f4 100644
--- a/tensorflow/docs_src/programmers_guide/eager.md
+++ b/tensorflow/docs_src/programmers_guide/eager.md
@@ -227,8 +227,8 @@ w = tfe.Variable([[1.0]])
 with tf.GradientTape() as tape:
   loss = w * w
 
-grad = tape.gradient(loss, [w])
-print(grad)  # => [tf.Tensor([[ 2.]], shape=(1, 1), dtype=float32)]
+grad = tape.gradient(loss, w)
+print(grad)  # => tf.Tensor([[ 2.]], shape=(1, 1), dtype=float32)
 ```
 
 Here's an example of `tf.GradientTape` that records forward-pass operations
@@ -596,7 +596,7 @@ def line_search_step(fn, init_x, rate=1.0):
     # Variables are automatically recorded, but manually watch a tensor
     tape.watch(init_x)
     value = fn(init_x)
-  grad, = tape.gradient(value, [init_x])
+  grad = tape.gradient(value, init_x)
   grad_norm = tf.reduce_sum(grad * grad)
   init_value = value
   while value > init_value - rate * grad_norm:
-- 
GitLab


From c183c5600b1393767c8c85aad34a436feb3bbe75 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 02:04:33 -0700
Subject: [PATCH 251/395] Fixing some linter errors in TF documentation (Github
 > GitHub, the the > the).

PiperOrigin-RevId: 195386172
---
 tensorflow/docs_src/deploy/index.md                          | 2 +-
 tensorflow/docs_src/get_started/get_started_for_beginners.md | 2 +-
 tensorflow/docs_src/mobile/android_build.md                  | 4 ++--
 tensorflow/docs_src/mobile/linking_libs.md                   | 2 +-
 tensorflow/docs_src/mobile/mobile_intro.md                   | 4 ++--
 tensorflow/docs_src/mobile/tflite/index.md                   | 2 +-
 6 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensorflow/docs_src/deploy/index.md b/tensorflow/docs_src/deploy/index.md
index 07b1bc9257..61edba04b4 100644
--- a/tensorflow/docs_src/deploy/index.md
+++ b/tensorflow/docs_src/deploy/index.md
@@ -14,4 +14,4 @@ the following documents:
     designed for production environments. TensorFlow Serving provides
     out-of-the-box integration with TensorFlow models.
     [Source code for TensorFlow Serving](https://github.com/tensorflow/serving)
-    is available on Github.
+    is available on GitHub.
diff --git a/tensorflow/docs_src/get_started/get_started_for_beginners.md b/tensorflow/docs_src/get_started/get_started_for_beginners.md
index fbe0ed74f8..d5a80e22c5 100644
--- a/tensorflow/docs_src/get_started/get_started_for_beginners.md
+++ b/tensorflow/docs_src/get_started/get_started_for_beginners.md
@@ -233,7 +233,7 @@ The Iris program requires the data from the following two .csv files:
 *   `http://download.tensorflow.org/data/iris_training.csv`, which contains
     the training set.
 *   `http://download.tensorflow.org/data/iris_test.csv`, which contains the
-    the test set.
+    test set.
 
 The **training set** contains the examples that we'll use to train the model;
 the **test set** contains the examples that we'll use to evaluate the trained
diff --git a/tensorflow/docs_src/mobile/android_build.md b/tensorflow/docs_src/mobile/android_build.md
index c35530061d..f4b07db459 100644
--- a/tensorflow/docs_src/mobile/android_build.md
+++ b/tensorflow/docs_src/mobile/android_build.md
@@ -26,7 +26,7 @@ If you haven't already, do the following two things:
 - Install [Android Studio](https://developer.android.com/studio/index.html),
   following the instructions on their website.
 
-- Clone the TensorFlow repository from Github:
+- Clone the TensorFlow repository from GitHub:
 
         git clone https://github.com/tensorflow/tensorflow
 
@@ -37,7 +37,7 @@ If you haven't already, do the following two things:
 
 2. From the **Open File or Project** window that appears, navigate to and select
     the `tensorflow/examples/android` directory from wherever you cloned the
-    TensorFlow Github repo.  Click OK.
+    TensorFlow GitHub repo.  Click OK.
 
     If it asks you to do a Gradle Sync, click OK.
 
diff --git a/tensorflow/docs_src/mobile/linking_libs.md b/tensorflow/docs_src/mobile/linking_libs.md
index 2a0a77c92d..cf0db59021 100644
--- a/tensorflow/docs_src/mobile/linking_libs.md
+++ b/tensorflow/docs_src/mobile/linking_libs.md
@@ -32,7 +32,7 @@ include this functionality in your program:
 2. Download the nightly precompiled version from
 [ci.tensorflow.org](http://ci.tensorflow.org/view/Nightly/job/nightly-android/lastSuccessfulBuild/artifact/out/).
 
-3. Build the JAR file yourself using the instructions [in our Android Github repo](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/android)
+3. Build the JAR file yourself using the instructions [in our Android GitHub repo](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/android)
 
 ### iOS
 
diff --git a/tensorflow/docs_src/mobile/mobile_intro.md b/tensorflow/docs_src/mobile/mobile_intro.md
index 69b63ae7d2..1b0b9b44b4 100644
--- a/tensorflow/docs_src/mobile/mobile_intro.md
+++ b/tensorflow/docs_src/mobile/mobile_intro.md
@@ -80,7 +80,7 @@ tracking is especially important for applications where you’re trying to count
 how many objects are present over time, since it gives you a good idea when a
 new object enters or leaves the scene. We have some sample code for this
 available for Android [on
-Github](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android),
+GitHub](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android),
 and also a [more general object detection
 model](https://github.com/tensorflow/models/tree/master/research/object_detection/README.md)
 available as well.
@@ -231,7 +231,7 @@ process.
 The next step is to pick an effective model to use. You might be able to avoid
 training a model from scratch if someone else has already implemented a model
 similar to what you need; we have a repository of models implemented in
-TensorFlow [on Github](https://github.com/tensorflow/models) that you can look
+TensorFlow [on GitHub](https://github.com/tensorflow/models) that you can look
 through. Lean towards the simplest model you can find, and try to get started as
 soon as you have even a small amount of labelled data, since you’ll get the best
 results when you’re able to iterate quickly. The shorter the time it takes to
diff --git a/tensorflow/docs_src/mobile/tflite/index.md b/tensorflow/docs_src/mobile/tflite/index.md
index 11f11ea4dc..01881ccf3b 100644
--- a/tensorflow/docs_src/mobile/tflite/index.md
+++ b/tensorflow/docs_src/mobile/tflite/index.md
@@ -11,7 +11,7 @@ optimizing the kernels for mobile apps, pre-fused activations, and quantized
 kernels that allow smaller and faster (fixed-point math) models.
 
 Most of our TensorFlow Lite documentation is [on
-Github](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite)
+GitHub](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite)
 for the time being.
 
 ## What does TensorFlow Lite contain?
-- 
GitLab


From 7a7bbc303c451fea5b3dd93109028531a89a18ab Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 02:22:14 -0700
Subject: [PATCH 252/395] Do not crash on ROOT outfeed operations.

PiperOrigin-RevId: 195388075
---
 .../compiler/xla/service/cpu/ir_emitter.cc    |  8 ++-
 .../compiler/xla/service/cpu/tests/BUILD      | 14 +++++
 .../cpu/tests/cpu_literal_caching_test.cc     | 16 +-----
 .../xla/service/cpu/tests/cpu_outfeed_test.cc | 57 +++++++++++++++++++
 4 files changed, 79 insertions(+), 16 deletions(-)
 create mode 100644 tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc

diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
index e473389a29..6347ee2a2a 100644
--- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
+++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
@@ -2563,8 +2563,12 @@ Status IrEmitter::FinishVisit(HloInstruction* root) {
   // nothing to do since the result was already written directly into the output
   // buffer.
   VLOG(2) << "FinishVisit root: " << root->ToString();
-  llvm::Value* root_value = GetEmittedValueFor(root);
-  VLOG(2) << "  value: " << llvm_ir::DumpToString(*root_value);
+  if (root->opcode() == HloOpcode::kOutfeed) {
+    VLOG(2) << "  outfeed with value: "
+            << llvm_ir::DumpToString(*GetEmittedValueFor(root->operand(0)));
+  } else {
+    VLOG(2) << "  value: " << llvm_ir::DumpToString(*GetEmittedValueFor(root));
+  }
 
   auto record_complete_computation = [&](llvm::Value* prof_counter) {
     if (prof_counter) {
diff --git a/tensorflow/compiler/xla/service/cpu/tests/BUILD b/tensorflow/compiler/xla/service/cpu/tests/BUILD
index 4ddb7a85bc..18a915e533 100644
--- a/tensorflow/compiler/xla/service/cpu/tests/BUILD
+++ b/tensorflow/compiler/xla/service/cpu/tests/BUILD
@@ -161,3 +161,17 @@ tf_cc_test(
         "//tensorflow/core:test_main",
     ],
 )
+
+tf_cc_test(
+    name = "cpu_outfeed_test",
+    srcs = ["cpu_outfeed_test.cc"],
+    deps = [
+        "//tensorflow/compiler/xla/service:hlo",
+        "//tensorflow/compiler/xla/service/cpu:cpu_compiler",
+        "//tensorflow/compiler/xla/service/cpu/tests:cpu_codegen_test",
+        "//tensorflow/compiler/xla/tools/parser:hlo_parser",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc
index b10eb74635..d6e0425c55 100644
--- a/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc
+++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc
@@ -50,16 +50,10 @@ ENTRY main {
   const_b = f32[2,3,2] while(f32[2,3,2] const_a), condition=while_cond, body=while_body
 
   out0 = () outfeed(f32[2,3,2] const_a)
-  out1 = () outfeed(f32[2,3,2] const_b)
-
-  ROOT root = f32[] constant(1)
+  ROOT out1 = () outfeed(f32[2,3,2] const_b)
 }
 )";
 
-  // TODO(b/78879738): The fake "f32[] constant(1)" root is only needed to work
-  // around b/78879738.  Once b/78879738 is fixed, we can set one of the
-  // outfeeds as the root.
-
   string filecheck_pattern = R"(
 CHECK: private constant [2 x [3 x [2 x float]]]
 CHECK-NOT: private constant [2 x [3 x [2 x float]]]
@@ -99,16 +93,10 @@ ENTRY main {
   const_b = (f32[2,1]{1,0}, f32[2]{0}) while((f32[2,1]{1,0}, f32[2]{0}) const_a), condition=while_cond, body=while_body
 
   out0 = () outfeed((f32[2,1]{1,0}, f32[2]{0}) const_a)
-  out1 = () outfeed((f32[2,1]{1,0}, f32[2]{0}) const_b)
-
-  ROOT root = f32[] constant(1)
+  ROOT out1 = () outfeed((f32[2,1]{1,0}, f32[2]{0}) const_b)
 }
 )";
 
-  // TODO(b/78879738): The fake "f32[] constant(1)" root is only needed to work
-  // around b/78879738.  Once b/78879738 is fixed, we can set one of the
-  // outfeeds as the root.
-
   string filecheck_pattern = R"(
 CHECK: private constant [2 x float]
 CHECK: private constant [2 x [1 x float]]
diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc
new file mode 100644
index 0000000000..879372eb13
--- /dev/null
+++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc
@@ -0,0 +1,57 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/cpu/cpu_compiler.h"
+#include "tensorflow/compiler/xla/service/cpu/tests/cpu_codegen_test.h"
+#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h"
+
+namespace xla {
+namespace cpu {
+namespace {
+class CpuOutfeedTest : public CpuCodegenTest {};
+
+TEST_F(CpuOutfeedTest, OutfeedRoot) {
+  const string hlo_text = R"(
+HloModule Outfeed
+
+ENTRY main {
+  const_a = f32[2,3,2] constant(
+  f32[2,3,2]
+    {{{1, 2}, {1001, 1002}, {2001, 2002}},
+     {{2, 1}, {2001, 3002}, {2001, 2002}}})
+
+  ROOT out = () outfeed(f32[2,3,2] const_a)
+}
+)";
+
+  string filecheck_pattern = R"(
+CHECK: private constant [2 x [3 x [2 x float]]]
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_text));
+
+  CpuAotCompilationOptions options{
+      /*triple=*/"x86_64-pc-linux", /*cpu_name=*/"", /*features=*/"",
+      /*entry_point_name=*/"entry",
+      /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static};
+
+  CompileAheadOfTimeAndVerifyIr(std::move(module), options, filecheck_pattern,
+                                /*match_optimized_ir=*/false);
+}
+
+}  // namespace
+}  // namespace cpu
+}  // namespace xla
-- 
GitLab


From 34bb6643654b9a207b93d046d5fde807eb7ee499 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 03:27:18 -0700
Subject: [PATCH 253/395] Fix HloSharding::GetSubSharding to return correct
 array shardings

Previously it always returned a tuple sharding even if the specified
index was referenceing a non-tuple element.

PiperOrigin-RevId: 195393313
---
 tensorflow/compiler/xla/service/hlo_sharding.cc | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/xla/service/hlo_sharding.cc b/tensorflow/compiler/xla/service/hlo_sharding.cc
index 994de44123..7f7e3f7dab 100644
--- a/tensorflow/compiler/xla/service/hlo_sharding.cc
+++ b/tensorflow/compiler/xla/service/hlo_sharding.cc
@@ -367,10 +367,14 @@ HloSharding HloSharding::GetSubSharding(const Shape& shape,
                                         const ShapeIndex& index) const {
   CHECK(IsTuple());
 
-  ShapeTree<HloSharding> sub_shape_tree(ShapeUtil::GetSubshape(shape, index),
-                                        Replicate());
+  Shape sub_shape = ShapeUtil::GetSubshape(shape, index);
+  ShapeTree<HloSharding> sub_shape_tree(sub_shape, Replicate());
   sub_shape_tree.CopySubtreeFrom(GetAsShapeTree(shape), index, {});
-  return Tuple(sub_shape_tree);
+  if (ShapeUtil::IsTuple(sub_shape)) {
+    return Tuple(sub_shape_tree);
+  } else {
+    return sub_shape_tree.element({});
+  }
 }
 
 std::ostream& operator<<(std::ostream& out, const HloSharding& sharding) {
-- 
GitLab


From 2d6170fc0afee7269cab7f84647f2a65b86e7020 Mon Sep 17 00:00:00 2001
From: Benjamin Kramer <kramerb@google.com>
Date: Fri, 4 May 2018 03:43:00 -0700
Subject: [PATCH 254/395] [XLA] Remove template keyword on non-template
 methods.

This is an error with clang trunk.

PiperOrigin-RevId: 195394277
---
 tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc | 9 +++------
 third_party/libxsmm.BUILD                               | 2 +-
 2 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc b/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc
index 6cb470caf8..464cc01214 100644
--- a/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc
+++ b/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc
@@ -67,8 +67,7 @@ XLA_TYPED_TEST(MatOpsSimpleTest_F16F32, ExpTwoByTwoValues) {
       Literal::CreateR2FromArray2D<T>({{2.71828f, 1.00000f},    // row 0
                                        {0.36788f, 1.64872f}});  // row 1
 
-  this->template ComputeAndCompareLiteral(&builder, *expected, {},
-                                          ErrorSpec(1e-5));
+  this->ComputeAndCompareLiteral(&builder, *expected, {}, ErrorSpec(1e-5));
 }
 
 XLA_TYPED_TEST(MatOpsSimpleTest_F16F32, MapTwoByTwo) {
@@ -96,8 +95,7 @@ XLA_TYPED_TEST(MatOpsSimpleTest_F16F32, MapTwoByTwo) {
   std::unique_ptr<Literal> expected =
       Literal::CreateR2FromArray2D<T>({{1.5f, 0.5f},     // row 0
                                        {-0.5f, 1.0f}});  // row 1
-  this->template ComputeAndCompareLiteral(&builder, *expected, {},
-                                          ErrorSpec(1e-5));
+  this->ComputeAndCompareLiteral(&builder, *expected, {}, ErrorSpec(1e-5));
 }
 
 XLA_TYPED_TEST(MatOpsSimpleTest_F16F32, MaxTwoByTwoValues) {
@@ -116,8 +114,7 @@ XLA_TYPED_TEST(MatOpsSimpleTest_F16F32, MaxTwoByTwoValues) {
   std::unique_ptr<Literal> expected =
       Literal::CreateR2FromArray2D<T>({{7.0f, 6.0f},     // row 0
                                        {3.0f, -4.0f}});  // row 1
-  this->template ComputeAndCompareLiteral(&builder, *expected, {},
-                                          ErrorSpec(1e-6));
+  this->ComputeAndCompareLiteral(&builder, *expected, {}, ErrorSpec(1e-6));
 }
 
 struct TestLinspaceMaxParam {
diff --git a/third_party/libxsmm.BUILD b/third_party/libxsmm.BUILD
index 78ed1f4e16..4124f2db63 100644
--- a/third_party/libxsmm.BUILD
+++ b/third_party/libxsmm.BUILD
@@ -38,8 +38,8 @@ genrule(
         ":libxsmm_interface",
     ],
     visibility = [
-        "//third_party/eigen3:__pkg__",
         "//tensorflow/core/kernels:__pkg__",
+        "//third_party/eigen3:__pkg__",
     ],
 )
 
-- 
GitLab


From 3db0e545d2460be0392dfcaa304231cd2105648e Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Fri, 4 May 2018 10:18:46 -0700
Subject: [PATCH 255/395] Change RecvTensor RPC implementation to use
 DeviceContext::CopyDeviceTensorToCPU rather than calling
 GPUUtil::CopyGPUTensorToCPU. The direct call into the GPU code is problematic
 for non-GPU devices.

PiperOrigin-RevId: 195433287
---
 tensorflow/core/distributed_runtime/rpc/BUILD |  1 -
 .../rpc/grpc_worker_service.cc                | 24 +++++++------------
 2 files changed, 9 insertions(+), 16 deletions(-)

diff --git a/tensorflow/core/distributed_runtime/rpc/BUILD b/tensorflow/core/distributed_runtime/rpc/BUILD
index e973a22f45..c2719f5462 100644
--- a/tensorflow/core/distributed_runtime/rpc/BUILD
+++ b/tensorflow/core/distributed_runtime/rpc/BUILD
@@ -169,7 +169,6 @@ tf_cuda_library(
         ":grpc_worker_service_impl",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
-        "//tensorflow/core:gpu_runtime",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:worker_proto_cc",
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_worker_service.cc b/tensorflow/core/distributed_runtime/rpc/grpc_worker_service.cc
index bbf7391377..26fad1fc3c 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_worker_service.cc
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_worker_service.cc
@@ -23,9 +23,6 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/device.h"
 #include "tensorflow/core/common_runtime/device_mgr.h"
 #include "tensorflow/core/common_runtime/dma_helper.h"
-#if GOOGLE_CUDA
-#include "tensorflow/core/common_runtime/gpu/gpu_util.h"
-#endif  // GOOGLE_CUDA
 #include "tensorflow/core/common_runtime/local_device.h"
 #include "tensorflow/core/common_runtime/process_util.h"
 #include "tensorflow/core/common_runtime/step_stats_collector.h"
@@ -439,10 +436,10 @@ void GrpcWorker::GrpcRecvTensorAsync(CallOptions* opts,
   opts->SetCancelCallback([this, step_id]() { AbortStep(step_id); });
   env_->rendezvous_mgr->RecvLocalAsync(
       step_id, parsed,
-      [opts, response, done, src_dev](const Status& status,
-                                      const Rendezvous::Args& send_args,
-                                      const Rendezvous::Args& recv_args,
-                                      const Tensor& val, const bool is_dead) {
+      [opts, response, done, src_dev, request](
+          const Status& status, const Rendezvous::Args& send_args,
+          const Rendezvous::Args& recv_args, const Tensor& val,
+          const bool is_dead) {
         opts->ClearCancelCallback();
         if (status.ok()) {
           // DMA can only be used for Tensors that do not fall into
@@ -455,8 +452,7 @@ void GrpcWorker::GrpcRecvTensorAsync(CallOptions* opts,
           {
             // Non-DMA cases.
             if (src_dev->tensorflow_gpu_device_info() && (!on_host)) {
-#if GOOGLE_CUDA
-              const DeviceContext* send_dev_context = send_args.device_context;
+              DeviceContext* send_dev_context = send_args.device_context;
               AllocatorAttributes alloc_attrs;
               alloc_attrs.set_gpu_compatible(true);
               alloc_attrs.set_on_host(true);
@@ -465,7 +461,8 @@ void GrpcWorker::GrpcRecvTensorAsync(CallOptions* opts,
               CHECK(send_dev_context)
                   << "send dev name: " << src_dev->name()
                   << " gpu_info: " << src_dev->tensorflow_gpu_device_info();
-              // "val" is on a GPU. Uses GPUUtil to fill the copy on host.
+              // "val" is on an accelerator device. Uses the device_context to
+              // fill the copy on host.
               StatusCallback copy_ready = [response, done, copy,
                                            is_dead](const Status& s) {
                 // The value is now ready to be returned on the wire.
@@ -474,11 +471,8 @@ void GrpcWorker::GrpcRecvTensorAsync(CallOptions* opts,
                 delete copy;
               };
 
-              GPUUtil::CopyGPUTensorToCPU(src_dev, send_dev_context, &val, copy,
-                                          copy_ready);
-#else
-              done(errors::Internal("No GPU device in process"));
-#endif  // GOOGLE_CUDA
+              send_dev_context->CopyDeviceTensorToCPU(
+                  &val, request->rendezvous_key(), src_dev, copy, copy_ready);
             } else {
               grpc::EncodeTensorToByteBuffer(is_dead, val, response);
               done(Status::OK());
-- 
GitLab


From a5f44b3519627859fb476a9cad1acc354bfa649f Mon Sep 17 00:00:00 2001
From: Alan Chiao <alanchiao@google.com>
Date: Fri, 4 May 2018 10:31:01 -0700
Subject: [PATCH 256/395] Implement neg op

PiperOrigin-RevId: 195435079
---
 tensorflow/contrib/lite/builtin_ops.h         |   1 +
 .../lite/g3doc/tf_ops_compatibility.md        |  11 ++
 tensorflow/contrib/lite/kernels/BUILD         |  14 ++
 tensorflow/contrib/lite/kernels/neg.cc        |  79 +++++++++++
 tensorflow/contrib/lite/kernels/neg_test.cc   |  80 +++++++++++
 tensorflow/contrib/lite/kernels/register.cc   |   2 +
 tensorflow/contrib/lite/model.cc              |   1 +
 tensorflow/contrib/lite/nnapi_delegate.cc     |   1 +
 tensorflow/contrib/lite/schema/schema.fbs     |   5 +
 .../contrib/lite/schema/schema_generated.h    | 124 +++++++++++++++++-
 tensorflow/contrib/lite/testing/BUILD         |   1 +
 .../contrib/lite/testing/generate_examples.py |  25 ++++
 .../testing/generated_examples_zip_test.cc    |   1 +
 .../contrib/lite/toco/tflite/operator.cc      |   3 +-
 .../contrib/lite/toco/tflite/operator_test.cc |   1 +
 15 files changed, 341 insertions(+), 8 deletions(-)
 create mode 100644 tensorflow/contrib/lite/kernels/neg.cc
 create mode 100644 tensorflow/contrib/lite/kernels/neg_test.cc

diff --git a/tensorflow/contrib/lite/builtin_ops.h b/tensorflow/contrib/lite/builtin_ops.h
index 21e0e04ef6..962a7a8970 100644
--- a/tensorflow/contrib/lite/builtin_ops.h
+++ b/tensorflow/contrib/lite/builtin_ops.h
@@ -84,6 +84,7 @@ typedef enum {
   kTfLiteBuiltinArgMax = 56,
   kTfLiteBuiltinMinimum = 57,
   kTfLiteBuiltinLess = 58,
+  kTfLiteBuiltinNeg = 59,
 } TfLiteBuiltinOperator;
 
 #ifdef __cplusplus
diff --git a/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md b/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md
index aa28f8d050..0051ee84ec 100644
--- a/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md
+++ b/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md
@@ -397,6 +397,17 @@ Options {
 }
 ```
 
+**NEG**
+
+```
+Inputs {
+  0: a tensor
+}
+Outputs {
+  0: elementwise negation of the input tensor
+}
+```
+
 **PAD**
 
 ```
diff --git a/tensorflow/contrib/lite/kernels/BUILD b/tensorflow/contrib/lite/kernels/BUILD
index 57b3136cce..feab18b5c2 100644
--- a/tensorflow/contrib/lite/kernels/BUILD
+++ b/tensorflow/contrib/lite/kernels/BUILD
@@ -158,6 +158,7 @@ cc_library(
         "mean.cc",
         "mfcc.cc",
         "mul.cc",
+        "neg.cc",
         "pad.cc",
         "pooling.cc",
         "register.cc",
@@ -856,6 +857,19 @@ tf_cc_test(
     ],
 )
 
+tf_cc_test(
+    name = "neg_test",
+    size = "small",
+    srcs = ["neg_test.cc"],
+    tags = ["tflite_not_portable_ios"],
+    deps = [
+        ":builtin_ops",
+        "//tensorflow/contrib/lite:framework",
+        "//tensorflow/contrib/lite/kernels:test_util",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
 filegroup(
     name = "all_files",
     srcs = glob(
diff --git a/tensorflow/contrib/lite/kernels/neg.cc b/tensorflow/contrib/lite/kernels/neg.cc
new file mode 100644
index 0000000000..692da81727
--- /dev/null
+++ b/tensorflow/contrib/lite/kernels/neg.cc
@@ -0,0 +1,79 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/contrib/lite/context.h"
+#include "tensorflow/contrib/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace neg {
+
+constexpr int kInputTensor = 0;
+constexpr int kOutputTensor = 0;
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+  TfLiteTensor* input = GetInput(context, node, kInputTensor);
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+
+  output->type = input->type;
+  return context->ResizeTensor(context, output,
+                               TfLiteIntArrayCopy(input->dims));
+}
+
+template <typename T>
+void Negate(const T* in_data, int num_elements, T* out_data) {
+  // TODO(alanchiao): add vectorized version.
+  for (int i = 0; i < num_elements; ++i) {
+    out_data[i] = -in_data[i];
+  }
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  TfLiteTensor* input = GetInput(context, node, kInputTensor);
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+  const int num_elements = NumElements(input);
+  switch (input->type) {
+    case kTfLiteInt64:
+      Negate(input->data.i64, num_elements, output->data.i64);
+      break;
+    case kTfLiteInt32:
+      Negate(input->data.i32, num_elements, output->data.i32);
+      break;
+    case kTfLiteFloat32:
+      Negate(input->data.f, num_elements, output->data.f);
+      break;
+    default:
+      context->ReportError(
+          context, "Neg only currently supports int64, int32, and float32.",
+          input->type);
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+}  // namespace neg
+
+TfLiteRegistration* Register_NEG() {
+  static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr,
+                                 neg::Prepare, neg::Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite
diff --git a/tensorflow/contrib/lite/kernels/neg_test.cc b/tensorflow/contrib/lite/kernels/neg_test.cc
new file mode 100644
index 0000000000..3c95ac8cc2
--- /dev/null
+++ b/tensorflow/contrib/lite/kernels/neg_test.cc
@@ -0,0 +1,80 @@
+/* Copyright 2018 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.
+==============================================================================*/
+#include <gtest/gtest.h>
+#include "tensorflow/contrib/lite/interpreter.h"
+#include "tensorflow/contrib/lite/kernels/register.h"
+#include "tensorflow/contrib/lite/kernels/test_util.h"
+#include "tensorflow/contrib/lite/model.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+class NegOpModel : public SingleOpModel {
+ public:
+  NegOpModel(const TensorData& input, const TensorData& output) {
+    input_ = AddInput(input);
+    output_ = AddOutput(output);
+    SetBuiltinOp(BuiltinOperator_NEG, BuiltinOptions_NegOptions,
+                 CreateNegOptions(builder_).Union());
+    BuildInterpreter({GetShape(input_)});
+  }
+
+  template <class T>
+  void SetInput(std::initializer_list<T> data) {
+    PopulateTensor<T>(input_, data);
+  }
+
+  template <class T>
+  std::vector<T> GetOutput() {
+    return ExtractVector<T>(output_);
+  }
+
+ protected:
+  int input_;
+  int output_;
+};
+
+TEST(NegOpModel, NegFloat) {
+  NegOpModel m({TensorType_FLOAT32, {2, 3}}, {TensorType_FLOAT32, {2, 3}});
+  m.SetInput<float>({-2.0f, -1.0f, 0.f, 1.0f, 2.0f, 3.0f});
+  m.Invoke();
+  EXPECT_THAT(m.GetOutput<float>(),
+              ElementsAreArray({2.0f, 1.0f, 0.f, -1.0f, -2.0f, -3.0f}));
+}
+
+TEST(NegOpModel, NegInt32) {
+  NegOpModel m({TensorType_INT32, {2, 3}}, {TensorType_INT32, {2, 3}});
+  m.SetInput<int32>({-2, -1, 0, 1, 2, 3});
+  m.Invoke();
+  EXPECT_THAT(m.GetOutput<int32>(), ElementsAreArray({2, 1, 0, -1, -2, -3}));
+}
+
+TEST(NegOpModel, NegInt64) {
+  NegOpModel m({TensorType_INT64, {2, 3}}, {TensorType_INT64, {2, 3}});
+  m.SetInput<int64_t>({-2, -1, 0, 1, 2, 3});
+  m.Invoke();
+  EXPECT_THAT(m.GetOutput<int64_t>(), ElementsAreArray({2, 1, 0, -1, -2, -3}));
+}
+
+}  // namespace
+}  // namespace tflite
+
+int main(int argc, char** argv) {
+  ::tflite::LogToStderr();
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}
diff --git a/tensorflow/contrib/lite/kernels/register.cc b/tensorflow/contrib/lite/kernels/register.cc
index f91d188ffa..29ea718a96 100644
--- a/tensorflow/contrib/lite/kernels/register.cc
+++ b/tensorflow/contrib/lite/kernels/register.cc
@@ -81,6 +81,7 @@ TfLiteRegistration* Register_MINIMUM();
 TfLiteRegistration* Register_ARG_MAX();
 TfLiteRegistration* Register_LESS();
 TfLiteRegistration* Register_FLOOR();
+TfLiteRegistration* Register_NEG();
 
 BuiltinOpResolver::BuiltinOpResolver() {
   AddBuiltin(BuiltinOperator_RELU, Register_RELU());
@@ -143,6 +144,7 @@ BuiltinOpResolver::BuiltinOpResolver() {
   AddBuiltin(BuiltinOperator_ARG_MAX, Register_ARG_MAX());
   AddBuiltin(BuiltinOperator_LESS, Register_LESS());
   AddBuiltin(BuiltinOperator_FLOOR, Register_FLOOR());
+  AddBuiltin(BuiltinOperator_NEG, Register_NEG());
 
   // TODO(andrewharp, ahentz): Move these somewhere more appropriate so that
   // custom ops aren't always included by default.
diff --git a/tensorflow/contrib/lite/model.cc b/tensorflow/contrib/lite/model.cc
index e15f1be7d3..590f042e21 100644
--- a/tensorflow/contrib/lite/model.cc
+++ b/tensorflow/contrib/lite/model.cc
@@ -351,6 +351,7 @@ TfLiteStatus ParseOpData(const Operator* op, BuiltinOperator op_type,
     case BuiltinOperator_DEQUANTIZE:
     case BuiltinOperator_PRELU:
     case BuiltinOperator_FLOOR:
+    case BuiltinOperator_NEG:
       break;
     case BuiltinOperator_CAST: {
       TfLiteCastParams* params = MallocPOD<TfLiteCastParams>();
diff --git a/tensorflow/contrib/lite/nnapi_delegate.cc b/tensorflow/contrib/lite/nnapi_delegate.cc
index e1895dd38e..6eac18c4f5 100644
--- a/tensorflow/contrib/lite/nnapi_delegate.cc
+++ b/tensorflow/contrib/lite/nnapi_delegate.cc
@@ -372,6 +372,7 @@ void AddOpsAndParams(tflite::Interpreter* interpreter,
       case tflite::BuiltinOperator_MINIMUM:
       case tflite::BuiltinOperator_ARG_MAX:
       case tflite::BuiltinOperator_LESS:
+      case tflite::BuiltinOperator_NEG:
         FATAL("Op code %d is currently not delegated to NNAPI", builtin);
         nn_op_type = -1;  // set to invalid
         break;
diff --git a/tensorflow/contrib/lite/schema/schema.fbs b/tensorflow/contrib/lite/schema/schema.fbs
index b16baf02dc..265b1dd3fe 100644
--- a/tensorflow/contrib/lite/schema/schema.fbs
+++ b/tensorflow/contrib/lite/schema/schema.fbs
@@ -136,6 +136,7 @@ enum BuiltinOperator : byte {
   ARG_MAX = 56,
   MINIMUM = 57,
   LESS = 58,
+  NEG = 59,
 }
 
 // Options for the builtin operators.
@@ -181,6 +182,7 @@ union BuiltinOptions {
   MaximumMinimumOptions,
   ArgMaxOptions,
   LessOptions,
+  NegOptions,
 }
 
 enum Padding : byte { SAME, VALID }
@@ -406,6 +408,9 @@ table ArgMaxOptions {
 table LessOptions {
 }
 
+table NegOptions {
+}
+
 // An OperatorCode can be an enum value (BuiltinOperator) if the operator is a
 // builtin, or a string if the operator is custom.
 table OperatorCode {
diff --git a/tensorflow/contrib/lite/schema/schema_generated.h b/tensorflow/contrib/lite/schema/schema_generated.h
index 57af973460..c172f77aa9 100755
--- a/tensorflow/contrib/lite/schema/schema_generated.h
+++ b/tensorflow/contrib/lite/schema/schema_generated.h
@@ -154,6 +154,9 @@ struct ArgMaxOptionsT;
 struct LessOptions;
 struct LessOptionsT;
 
+struct NegOptions;
+struct NegOptionsT;
+
 struct OperatorCode;
 struct OperatorCodeT;
 
@@ -272,11 +275,12 @@ enum BuiltinOperator {
   BuiltinOperator_ARG_MAX = 56,
   BuiltinOperator_MINIMUM = 57,
   BuiltinOperator_LESS = 58,
+  BuiltinOperator_NEG = 59,
   BuiltinOperator_MIN = BuiltinOperator_ADD,
-  BuiltinOperator_MAX = BuiltinOperator_LESS
+  BuiltinOperator_MAX = BuiltinOperator_NEG
 };
 
-inline BuiltinOperator (&EnumValuesBuiltinOperator())[58] {
+inline BuiltinOperator (&EnumValuesBuiltinOperator())[59] {
   static BuiltinOperator values[] = {
     BuiltinOperator_ADD,
     BuiltinOperator_AVERAGE_POOL_2D,
@@ -335,7 +339,8 @@ inline BuiltinOperator (&EnumValuesBuiltinOperator())[58] {
     BuiltinOperator_MAXIMUM,
     BuiltinOperator_ARG_MAX,
     BuiltinOperator_MINIMUM,
-    BuiltinOperator_LESS
+    BuiltinOperator_LESS,
+    BuiltinOperator_NEG
   };
   return values;
 }
@@ -401,6 +406,7 @@ inline const char **EnumNamesBuiltinOperator() {
     "ARG_MAX",
     "MINIMUM",
     "LESS",
+    "NEG",
     nullptr
   };
   return names;
@@ -454,11 +460,12 @@ enum BuiltinOptions {
   BuiltinOptions_MaximumMinimumOptions = 39,
   BuiltinOptions_ArgMaxOptions = 40,
   BuiltinOptions_LessOptions = 41,
+  BuiltinOptions_NegOptions = 42,
   BuiltinOptions_MIN = BuiltinOptions_NONE,
-  BuiltinOptions_MAX = BuiltinOptions_LessOptions
+  BuiltinOptions_MAX = BuiltinOptions_NegOptions
 };
 
-inline BuiltinOptions (&EnumValuesBuiltinOptions())[42] {
+inline BuiltinOptions (&EnumValuesBuiltinOptions())[43] {
   static BuiltinOptions values[] = {
     BuiltinOptions_NONE,
     BuiltinOptions_Conv2DOptions,
@@ -501,7 +508,8 @@ inline BuiltinOptions (&EnumValuesBuiltinOptions())[42] {
     BuiltinOptions_DequantizeOptions,
     BuiltinOptions_MaximumMinimumOptions,
     BuiltinOptions_ArgMaxOptions,
-    BuiltinOptions_LessOptions
+    BuiltinOptions_LessOptions,
+    BuiltinOptions_NegOptions
   };
   return values;
 }
@@ -550,6 +558,7 @@ inline const char **EnumNamesBuiltinOptions() {
     "MaximumMinimumOptions",
     "ArgMaxOptions",
     "LessOptions",
+    "NegOptions",
     nullptr
   };
   return names;
@@ -728,6 +737,10 @@ template<> struct BuiltinOptionsTraits<LessOptions> {
   static const BuiltinOptions enum_value = BuiltinOptions_LessOptions;
 };
 
+template<> struct BuiltinOptionsTraits<NegOptions> {
+  static const BuiltinOptions enum_value = BuiltinOptions_NegOptions;
+};
+
 struct BuiltinOptionsUnion {
   BuiltinOptions type;
   void *value;
@@ -1087,6 +1100,14 @@ struct BuiltinOptionsUnion {
     return type == BuiltinOptions_LessOptions ?
       reinterpret_cast<const LessOptionsT *>(value) : nullptr;
   }
+  NegOptionsT *AsNegOptions() {
+    return type == BuiltinOptions_NegOptions ?
+      reinterpret_cast<NegOptionsT *>(value) : nullptr;
+  }
+  const NegOptionsT *AsNegOptions() const {
+    return type == BuiltinOptions_NegOptions ?
+      reinterpret_cast<const NegOptionsT *>(value) : nullptr;
+  }
 };
 
 bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *obj, BuiltinOptions type);
@@ -4014,6 +4035,46 @@ inline flatbuffers::Offset<LessOptions> CreateLessOptions(
 
 flatbuffers::Offset<LessOptions> CreateLessOptions(flatbuffers::FlatBufferBuilder &_fbb, const LessOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
 
+struct NegOptionsT : public flatbuffers::NativeTable {
+  typedef NegOptions TableType;
+  NegOptionsT() {
+  }
+};
+
+struct NegOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
+  typedef NegOptionsT NativeTableType;
+  bool Verify(flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           verifier.EndTable();
+  }
+  NegOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  void UnPackTo(NegOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  static flatbuffers::Offset<NegOptions> Pack(flatbuffers::FlatBufferBuilder &_fbb, const NegOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
+};
+
+struct NegOptionsBuilder {
+  flatbuffers::FlatBufferBuilder &fbb_;
+  flatbuffers::uoffset_t start_;
+  explicit NegOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  NegOptionsBuilder &operator=(const NegOptionsBuilder &);
+  flatbuffers::Offset<NegOptions> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = flatbuffers::Offset<NegOptions>(end);
+    return o;
+  }
+};
+
+inline flatbuffers::Offset<NegOptions> CreateNegOptions(
+    flatbuffers::FlatBufferBuilder &_fbb) {
+  NegOptionsBuilder builder_(_fbb);
+  return builder_.Finish();
+}
+
+flatbuffers::Offset<NegOptions> CreateNegOptions(flatbuffers::FlatBufferBuilder &_fbb, const NegOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
+
 struct OperatorCodeT : public flatbuffers::NativeTable {
   typedef OperatorCode TableType;
   BuiltinOperator builtin_code;
@@ -4254,6 +4315,9 @@ struct Operator FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
   const LessOptions *builtin_options_as_LessOptions() const {
     return builtin_options_type() == BuiltinOptions_LessOptions ? static_cast<const LessOptions *>(builtin_options()) : nullptr;
   }
+  const NegOptions *builtin_options_as_NegOptions() const {
+    return builtin_options_type() == BuiltinOptions_NegOptions ? static_cast<const NegOptions *>(builtin_options()) : nullptr;
+  }
   const flatbuffers::Vector<uint8_t> *custom_options() const {
     return GetPointer<const flatbuffers::Vector<uint8_t> *>(VT_CUSTOM_OPTIONS);
   }
@@ -4444,6 +4508,10 @@ template<> inline const LessOptions *Operator::builtin_options_as<LessOptions>()
   return builtin_options_as_LessOptions();
 }
 
+template<> inline const NegOptions *Operator::builtin_options_as<NegOptions>() const {
+  return builtin_options_as_NegOptions();
+}
+
 struct OperatorBuilder {
   flatbuffers::FlatBufferBuilder &fbb_;
   flatbuffers::uoffset_t start_;
@@ -6070,6 +6138,29 @@ inline flatbuffers::Offset<LessOptions> CreateLessOptions(flatbuffers::FlatBuffe
       _fbb);
 }
 
+inline NegOptionsT *NegOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const {
+  auto _o = new NegOptionsT();
+  UnPackTo(_o, _resolver);
+  return _o;
+}
+
+inline void NegOptions::UnPackTo(NegOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const {
+  (void)_o;
+  (void)_resolver;
+}
+
+inline flatbuffers::Offset<NegOptions> NegOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const NegOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) {
+  return CreateNegOptions(_fbb, _o, _rehasher);
+}
+
+inline flatbuffers::Offset<NegOptions> CreateNegOptions(flatbuffers::FlatBufferBuilder &_fbb, const NegOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) {
+  (void)_rehasher;
+  (void)_o;
+  struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const NegOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va;
+  return tflite::CreateNegOptions(
+      _fbb);
+}
+
 inline OperatorCodeT *OperatorCode::UnPack(const flatbuffers::resolver_function_t *_resolver) const {
   auto _o = new OperatorCodeT();
   UnPackTo(_o, _resolver);
@@ -6417,6 +6508,10 @@ inline bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *ob
       auto ptr = reinterpret_cast<const LessOptions *>(obj);
       return verifier.VerifyTable(ptr);
     }
+    case BuiltinOptions_NegOptions: {
+      auto ptr = reinterpret_cast<const NegOptions *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
     default: return false;
   }
 }
@@ -6599,6 +6694,10 @@ inline void *BuiltinOptionsUnion::UnPack(const void *obj, BuiltinOptions type, c
       auto ptr = reinterpret_cast<const LessOptions *>(obj);
       return ptr->UnPack(resolver);
     }
+    case BuiltinOptions_NegOptions: {
+      auto ptr = reinterpret_cast<const NegOptions *>(obj);
+      return ptr->UnPack(resolver);
+    }
     default: return nullptr;
   }
 }
@@ -6769,6 +6868,10 @@ inline flatbuffers::Offset<void> BuiltinOptionsUnion::Pack(flatbuffers::FlatBuff
       auto ptr = reinterpret_cast<const LessOptionsT *>(value);
       return CreateLessOptions(_fbb, ptr, _rehasher).Union();
     }
+    case BuiltinOptions_NegOptions: {
+      auto ptr = reinterpret_cast<const NegOptionsT *>(value);
+      return CreateNegOptions(_fbb, ptr, _rehasher).Union();
+    }
     default: return 0;
   }
 }
@@ -6939,6 +7042,10 @@ inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) FL
       value = new LessOptionsT(*reinterpret_cast<LessOptionsT *>(u.value));
       break;
     }
+    case BuiltinOptions_NegOptions: {
+      value = new NegOptionsT(*reinterpret_cast<NegOptionsT *>(u.value));
+      break;
+    }
     default:
       break;
   }
@@ -7151,6 +7258,11 @@ inline void BuiltinOptionsUnion::Reset() {
       delete ptr;
       break;
     }
+    case BuiltinOptions_NegOptions: {
+      auto ptr = reinterpret_cast<NegOptionsT *>(value);
+      delete ptr;
+      break;
+    }
     default: break;
   }
   value = nullptr;
diff --git a/tensorflow/contrib/lite/testing/BUILD b/tensorflow/contrib/lite/testing/BUILD
index a1162cef38..211de63d58 100644
--- a/tensorflow/contrib/lite/testing/BUILD
+++ b/tensorflow/contrib/lite/testing/BUILD
@@ -43,6 +43,7 @@ gen_zipped_test_files(
         "mean.zip",
         "minimum.zip",
         "mul.zip",
+        "neg.zip",
         "pad.zip",
         "relu.zip",
         "relu1.zip",
diff --git a/tensorflow/contrib/lite/testing/generate_examples.py b/tensorflow/contrib/lite/testing/generate_examples.py
index 2f8f7a1a79..7e892769bf 100644
--- a/tensorflow/contrib/lite/testing/generate_examples.py
+++ b/tensorflow/contrib/lite/testing/generate_examples.py
@@ -2061,6 +2061,31 @@ def make_floor_tests(zip_path):
   make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs)
 
 
+def make_neg_tests(zip_path):
+  """Make a set of tests to do neg."""
+
+  test_parameters = [{
+      "input_dtype": [tf.float32, tf.int32],
+      "input_shape": [[1, 3, 4, 3], [5]],
+  }]
+
+  def build_graph(parameters):
+    """Build the neg op testing graph."""
+    input_tensor = tf.placeholder(
+        dtype=parameters["input_dtype"],
+        name="input",
+        shape=parameters["input_shape"])
+    out = tf.negative(input_tensor)
+    return [input_tensor], [out]
+
+  def build_inputs(parameters, sess, inputs, outputs):
+    values = create_tensor_data(parameters["input_dtype"],
+                                parameters["input_shape"])
+    return [values], sess.run(outputs, feed_dict=dict(zip(inputs, [values])))
+
+  make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs)
+
+
 # Toco binary path provided by the generate rule.
 bin_path = None
 
diff --git a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
index 34abb213c9..0673a3bb46 100644
--- a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
+++ b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
@@ -266,6 +266,7 @@ INSTANTIATE_TESTS(maximum)
 INSTANTIATE_TESTS(mean)
 INSTANTIATE_TESTS(minimum)
 INSTANTIATE_TESTS(mul)
+INSTANTIATE_TESTS(neg)
 INSTANTIATE_TESTS(pad)
 // INSTANTIATE_TESTS(prelu)
 INSTANTIATE_TESTS(relu)
diff --git a/tensorflow/contrib/lite/toco/tflite/operator.cc b/tensorflow/contrib/lite/toco/tflite/operator.cc
index d2e14ac5e0..e18ae805c0 100644
--- a/tensorflow/contrib/lite/toco/tflite/operator.cc
+++ b/tensorflow/contrib/lite/toco/tflite/operator.cc
@@ -872,7 +872,6 @@ std::vector<std::unique_ptr<BaseOperator>> BuildOperatorList() {
   // attributes.
   ops.emplace_back(
       new SimpleOperator<AddNOperator>("ADDN", OperatorType::kAddN));
-  ops.emplace_back(new SimpleOperator<NegOperator>("NEG", OperatorType::kNeg));
   ops.emplace_back(new SimpleOperator<TensorFlowRsqrtOperator>(
       "RSQRT", OperatorType::kTensorFlowRsqrt));
   // Simple Operators.
@@ -901,7 +900,7 @@ std::vector<std::unique_ptr<BaseOperator>> BuildOperatorList() {
       "MINIMUM", OperatorType::kTensorFlowMinimum));
   ops.emplace_back(new SimpleOperator<TensorFlowLessOperator>(
       "LESS", OperatorType::kTensorFlowLess));
-
+  ops.emplace_back(new SimpleOperator<NegOperator>("NEG", OperatorType::kNeg));
   return ops;
 }
 }  // namespace
diff --git a/tensorflow/contrib/lite/toco/tflite/operator_test.cc b/tensorflow/contrib/lite/toco/tflite/operator_test.cc
index 36ed741541..2b6c32b07c 100644
--- a/tensorflow/contrib/lite/toco/tflite/operator_test.cc
+++ b/tensorflow/contrib/lite/toco/tflite/operator_test.cc
@@ -115,6 +115,7 @@ TEST_F(OperatorTest, SimpleOperators) {
       "MINIMUM", OperatorType::kTensorFlowMinimum);
   CheckSimpleOperator<TensorFlowLessOperator>("LESS",
                                               OperatorType::kTensorFlowLess);
+  CheckSimpleOperator<NegOperator>("NEG", OperatorType::kNeg);
 }
 
 TEST_F(OperatorTest, BuiltinAdd) {
-- 
GitLab


From 47f1bd90658dd6858fb4bbefd4ef8acbef4ca931 Mon Sep 17 00:00:00 2001
From: Allen Lavoie <allenl@google.com>
Date: Fri, 4 May 2018 10:37:42 -0700
Subject: [PATCH 257/395] TFTS: Make it easier to swap in different
 autoregressive models.

Adds a very simple LSTM encoder/decoder option as an example.

ARModel's new constructor argument is a bit awkward, since Estimator's new graphs mean we need a Model factory rather than a Model (or to un-build the model?). It's still a much more pleasant way to write autoregressive models than fiddling with ARModel directly, since ARModel handles collecting all the features (and the prediction loop, etc.). Happy to hear other ideas for an API.

PiperOrigin-RevId: 195436186
---
 .../timeseries/python/timeseries/ar_model.py  | 284 +++++++++++++++---
 .../python/timeseries/ar_model_test.py        |  86 ++++--
 .../python/timeseries/estimators.py           |  17 +-
 .../python/timeseries/estimators_test.py      |  17 +-
 4 files changed, 319 insertions(+), 85 deletions(-)

diff --git a/tensorflow/contrib/timeseries/python/timeseries/ar_model.py b/tensorflow/contrib/timeseries/python/timeseries/ar_model.py
index 558d9480b4..ce96180c92 100644
--- a/tensorflow/contrib/timeseries/python/timeseries/ar_model.py
+++ b/tensorflow/contrib/timeseries/python/timeseries/ar_model.py
@@ -20,6 +20,7 @@ from __future__ import print_function
 
 from tensorflow.contrib import distributions
 
+from tensorflow.contrib.rnn.python.ops import lstm_ops
 from tensorflow.contrib.timeseries.python.timeseries import model
 from tensorflow.contrib.timeseries.python.timeseries import model_utils
 from tensorflow.contrib.timeseries.python.timeseries.feature_keys import PredictionFeatures
@@ -29,6 +30,9 @@ from tensorflow.python.estimator import estimator_lib
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
+from tensorflow.python.keras._impl.keras.engine import sequential
+from tensorflow.python.keras._impl.keras.engine import training
+from tensorflow.python.keras._impl.keras.layers import core
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import check_ops
 from tensorflow.python.ops import control_flow_ops
@@ -40,12 +44,150 @@ from tensorflow.python.ops import tensor_array_ops
 from tensorflow.python.ops import variable_scope
 
 
+class FlatPredictionModel(training.Model):
+  """Flattens input and output windows and puts them through dense layers.
+
+  This model does not operate on its own, but rather is a plugin to
+  `ARModel`. See `ARModel`'s constructor documentation
+  (`prediction_model_factory`) for a usage example.
+  """
+
+  def __init__(self,
+               num_features,
+               input_window_size,
+               output_window_size,
+               hidden_layer_sizes=None):
+    """Construct the flat prediction model.
+
+    Args:
+      num_features: number of input features per time step.
+      input_window_size: Number of past time steps of data to look at when doing
+        the regression.
+      output_window_size: Number of future time steps to predict. Note that
+        setting it to > 1 empirically seems to give a better fit.
+      hidden_layer_sizes: list of sizes of hidden layers.
+    """
+    super(FlatPredictionModel, self).__init__()
+    self._input_flatten = core.Flatten()
+    self._output_flatten = core.Flatten()
+    if hidden_layer_sizes:
+      self._hidden_layers = sequential.Sequential([
+          core.Dense(layer_size, activation=nn_ops.relu)
+          for layer_size in hidden_layer_sizes])
+    else:
+      self._hidden_layers = None
+    self._mean_transform = core.Dense(num_features * output_window_size,
+                                      name="predicted_mean")
+    self._covariance_transform = core.Dense(num_features * output_window_size,
+                                            name="log_sigma_square")
+    self._prediction_shape = [-1, output_window_size, num_features]
+
+  def call(self, input_window_features, output_window_features):
+    """Compute predictions from input and output windows.
+
+    Args:
+      input_window_features: A floating point Tensor with shape [batch size,
+        input window size, input features]. The batch dimension may not have
+        static shape information, but the window size and number of input
+        features are known at graph construction time and recorded in the static
+        shape information for the `input_window_features` `Tensor`. Note that
+        `input_window_size` may be zero.
+      output_window_features: A floating point Tensor with shape [batch size,
+        output window size, output features]. As with `input_window_features`,
+        the last two dimensions have static shape information. If there are no
+        output features, the size of the last dimension will be zero.
+    Returns:
+      A dictionary of predictions with keys "mean" and "covariance" (only
+      diagonal covariances are currently supported). Each has shape
+      [batch size, output window size, num_features], where num_features is the
+      same as the constructor argument.
+    """
+    if input_window_features.shape[1].value == 0:
+      # TODO(allenl): Make reshape()'s static shape information work on
+      # zero-size Tensors? Currently this special case is required because
+      # otherwise the Dense layers get unknown last dimensions.
+      activation = self._output_flatten(output_window_features)
+    elif output_window_features.shape[2].value == 0:
+      activation = self._input_flatten(input_window_features)
+    else:
+      activation = array_ops.concat(
+          [self._input_flatten(input_window_features),
+           self._output_flatten(output_window_features)],
+          axis=1)
+    if self._hidden_layers:
+      activation = self._hidden_layers(activation)
+    predicted_mean = array_ops.reshape(
+        self._mean_transform(activation),
+        self._prediction_shape)
+    predicted_covariance = array_ops.reshape(
+        gen_math_ops.exp(self._covariance_transform(activation)),
+        self._prediction_shape)
+    return {"mean": predicted_mean,
+            "covariance": predicted_covariance}
+
+
+class LSTMPredictionModel(training.Model):
+  """A simple encoder/decoder model using an LSTM.
+
+  This model does not operate on its own, but rather is a plugin to
+  `ARModel`. See `ARModel`'s constructor documentation
+  (`prediction_model_factory`) for a usage example.
+  """
+
+  def __init__(self,
+               num_features,
+               input_window_size,
+               output_window_size,
+               num_units=128):
+    """Construct the LSTM prediction model.
+
+    Args:
+      num_features: number of input features per time step.
+      input_window_size: Number of past time steps of data to look at when doing
+        the regression.
+      output_window_size: Number of future time steps to predict. Note that
+        setting it to > 1 empirically seems to give a better fit.
+      num_units: The number of units in the encoder and decoder LSTM cells.
+    """
+    super(LSTMPredictionModel, self).__init__()
+    self._encoder = lstm_ops.LSTMBlockFusedCell(
+        num_units=num_units, name="encoder")
+    self._decoder = lstm_ops.LSTMBlockFusedCell(
+        num_units=num_units, name="decoder")
+    self._mean_transform = core.Dense(num_features,
+                                      name="mean_transform")
+    self._covariance_transform = core.Dense(num_features,
+                                            name="covariance_transform")
+
+  def call(self, input_window_features, output_window_features):
+    """Compute predictions from input and output windows."""
+    # Convert to time major
+    input_window_features = array_ops.transpose(input_window_features,
+                                                [1, 0, 2])
+    output_window_features = array_ops.transpose(output_window_features,
+                                                 [1, 0, 2])
+    _, encoder_state = self._encoder(
+        input_window_features, dtype=self.dtype)
+    decoder_output, _ = self._decoder(
+        output_window_features, dtype=self.dtype,
+        initial_state=encoder_state)
+
+    # Switch back to batch major
+    decoder_output = array_ops.transpose(decoder_output, [1, 0, 2])
+    predicted_mean = self._mean_transform(decoder_output)
+    predicted_covariance = gen_math_ops.exp(
+        self._covariance_transform(decoder_output))
+    return {"mean": predicted_mean,
+            "covariance": predicted_covariance}
+
+
 class ARModel(model.TimeSeriesModel):
   """Auto-regressive model, both linear and non-linear.
 
   Features to the model include time and values of input_window_size timesteps,
-  and times for output_window_size timesteps. These are passed through zero or
-  more hidden layers, and then fed to a loss function (e.g. squared loss).
+  and times for output_window_size timesteps. These are passed through a
+  configurable prediction model, and then fed to a loss function (e.g. squared
+  loss).
 
   Note that this class can also be used to regress against time only by setting
   the input_window_size to zero.
@@ -58,9 +200,9 @@ class ARModel(model.TimeSeriesModel):
                input_window_size,
                output_window_size,
                num_features,
+               prediction_model_factory=FlatPredictionModel,
                num_time_buckets=10,
                loss=NORMAL_LIKELIHOOD_LOSS,
-               hidden_layer_sizes=None,
                exogenous_feature_columns=None):
     """Constructs an auto-regressive model.
 
@@ -73,6 +215,22 @@ class ARModel(model.TimeSeriesModel):
       output_window_size: Number of future time steps to predict. Note that
         setting it to > 1 empirically seems to give a better fit.
       num_features: number of input features per time step.
+      prediction_model_factory: A callable taking arguments `num_features`,
+        `input_window_size`, and `output_window_size` and returning a
+        `tf.keras.Model`. The `Model`'s `call()` takes two arguments: an input
+        window and an output window, and returns a dictionary of
+        predictions. See `FlatPredictionModel` for an example. Example usage:
+
+        ```python
+        model = ar_model.ARModel(
+          periodicities=2, num_features=3,
+          prediction_model_factory=functools.partial(
+              FlatPredictionModel,
+              hidden_layer_sizes=[10, 10]))
+        ```
+
+        The default model computes predictions as a linear function of flattened
+        input and output windows.
       num_time_buckets: Number of buckets into which to divide (time %
         periodicity) for generating time based features.
       loss: Loss function to use for training. Currently supported values are
@@ -81,18 +239,15 @@ class ARModel(model.TimeSeriesModel):
         SQUARED_LOSS, the evaluation loss is reported based on un-scaled
         observations and predictions, while the training loss is computed on
         normalized data (if input statistics are available).
-      hidden_layer_sizes: list of sizes of hidden layers.
       exogenous_feature_columns: A list of `tf.feature_column`s (for example
           `tf.feature_column.embedding_column`) corresponding to exogenous
           features which provide extra information to the model but are not part
           of the series to be predicted. Passed to
           `tf.feature_column.input_layer`.
     """
+    self._model_factory = prediction_model_factory
     self.input_window_size = input_window_size
     self.output_window_size = output_window_size
-    if hidden_layer_sizes is None:
-      hidden_layer_sizes = []
-    self.hidden_layer_sizes = hidden_layer_sizes
     self.window_size = self.input_window_size + self.output_window_size
     self.loss = loss
     super(ARModel, self).__init__(
@@ -115,6 +270,19 @@ class ARModel(model.TimeSeriesModel):
     assert len(self._periods) or self.input_window_size
     assert output_window_size > 0
 
+  def initialize_graph(self, input_statistics=None):
+    super(ARModel, self).initialize_graph(input_statistics=input_statistics)
+    self._model_scope = variable_scope.variable_scope(
+        # The trailing slash means we strip all enclosing variable_scopes, which
+        # unfortunately is necessary because the model gets called inside and
+        # outside a "while" scope (for prediction and training respectively),
+        # and the variables names need to match.
+        "model/", use_resource=True)
+    self._model_instance = self._model_factory(
+        num_features=self.num_features,
+        input_window_size=self.input_window_size,
+        output_window_size=self.output_window_size)
+
   def get_start_state(self):
     # State which matches the format we'll return later. Typically this will not
     # be used by the model directly, but the shapes and dtypes should match so
@@ -166,17 +334,6 @@ class ARModel(model.TimeSeriesModel):
     return array_ops.reshape(predicted_mean,
                              [-1, self.output_window_size, self.num_features])
 
-  def _create_hidden_stack(self, activation, activation_size):
-    activations = []
-    for layer_number, layer_size in enumerate(self.hidden_layer_sizes):
-      # TODO(agarwal): Migrate to fully_connected in tf slim
-      activation = model_utils.fully_connected(
-          activation, activation_size, layer_size,
-          name="layer_{}".format(layer_number))
-      activation_size = layer_size
-      activations.append((activation, activation_size))
-    return activations
-
   def prediction_ops(self, times, values, exogenous_regressors):
     """Compute model predictions given input data.
 
@@ -195,7 +352,7 @@ class ARModel(model.TimeSeriesModel):
       self.num_features].
     """
     times.get_shape().assert_is_compatible_with([None, self.window_size])
-    activations = []
+    batch_size = array_ops.shape(times)[0]
     if self.input_window_size:
       values.get_shape().assert_is_compatible_with(
           [None, self.input_window_size, self.num_features])
@@ -203,39 +360,66 @@ class ARModel(model.TimeSeriesModel):
       exogenous_regressors.get_shape().assert_is_compatible_with(
           [None, self.window_size, self.exogenous_size])
     # Create input features.
-    activation_components = []
+    input_window_features = []
+    input_feature_size = 0
+    output_window_features = []
+    output_feature_size = 0
     if self._periods:
       _, time_features = self._compute_time_features(times)
-      activation_size = self.window_size * self._buckets * len(self._periods)
-      activation_components.append(
-          array_ops.reshape(time_features, [-1, activation_size]))
-    else:
-      activation_size = 0
+      num_time_features = self._buckets * len(self._periods)
+      time_features = array_ops.reshape(
+          time_features,
+          [batch_size,
+           self.window_size,
+           num_time_features])
+      input_time_features, output_time_features = array_ops.split(
+          time_features, (self.input_window_size, self.output_window_size),
+          axis=1)
+      input_feature_size += num_time_features
+      output_feature_size += num_time_features
+      input_window_features.append(input_time_features)
+      output_window_features.append(output_time_features)
     if self.input_window_size:
       inp = array_ops.slice(values, [0, 0, 0], [-1, self.input_window_size, -1])
-      inp_size = self.input_window_size * self.num_features
-      inp = array_ops.reshape(inp, [-1, inp_size])
-      activation_components.append(inp)
-      activation_size += inp_size
+      input_window_features.append(
+          array_ops.reshape(
+              inp,
+              [batch_size, self.input_window_size, self.num_features]))
+      input_feature_size += self.num_features
     if self.exogenous_size:
-      exogenous_size = self.window_size * self.exogenous_size
-      activation_size += exogenous_size
-      exogenous_flattened = array_ops.reshape(
-          exogenous_regressors, [-1, exogenous_size])
-      activation_components.append(exogenous_flattened)
-    assert activation_size
-    assert activation_components
-    activation = array_ops.concat(activation_components, axis=1)
-    activations.append((activation, activation_size))
-    # Create hidden layers.
-    activations += self._create_hidden_stack(activation, activation_size)
-    # Create mean and convariance ops.
-    predicted_mean = self._predicted_mean_op(activations)
-    predicted_covariance = self._predicted_covariance_op(activations,
-                                                         self.num_features)
-    return {"activations": activations,
-            "mean": predicted_mean,
-            "covariance": predicted_covariance}
+      input_exogenous_features, output_exogenous_features = array_ops.split(
+          exogenous_regressors,
+          (self.input_window_size, self.output_window_size),
+          axis=1)
+      input_feature_size += self.exogenous_size
+      output_feature_size += self.exogenous_size
+      input_window_features.append(input_exogenous_features)
+      output_window_features.append(output_exogenous_features)
+    assert input_window_features
+    input_window_features = array_ops.concat(input_window_features, axis=2)
+    if output_window_features:
+      output_window_features = array_ops.concat(output_window_features, axis=2)
+    else:
+      output_window_features = array_ops.zeros(
+          [batch_size, self.output_window_size, 0],
+          dtype=self.dtype)
+    static_batch_size = times.get_shape()[0].value
+    input_window_features.set_shape(
+        [static_batch_size, self.input_window_size, input_feature_size])
+    output_window_features.set_shape(
+        [static_batch_size, self.output_window_size, output_feature_size])
+    return self._output_window_predictions(input_window_features,
+                                           output_window_features)
+
+  def _output_window_predictions(
+      self, input_window_features, output_window_features):
+    with self._model_scope:
+      predictions = self._model_instance(
+          input_window_features, output_window_features)
+      result_shape = [None, self.output_window_size, self.num_features]
+      for v in predictions.values():
+        v.set_shape(result_shape)
+      return predictions
 
   def loss_op(self, targets, prediction_ops):
     """Create loss_op."""
@@ -286,6 +470,8 @@ class ARModel(model.TimeSeriesModel):
       values are Tensors of shape [batch_size, predict window size,
       num_features] and correspond to the values passed in `TIMES`.
     """
+    if not self._graph_initialized:
+      self.initialize_graph()
     predict_times = math_ops.cast(
         ops.convert_to_tensor(features[PredictionFeatures.TIMES]), dtypes.int32)
     exogenous_regressors = self._process_exogenous_features(
@@ -701,9 +887,9 @@ class AnomalyMixtureARModel(ARModel):
                input_window_size,
                output_window_size,
                num_features,
+               prediction_model_factory=FlatPredictionModel,
                anomaly_distribution=GAUSSIAN_ANOMALY,
                num_time_buckets=10,
-               hidden_layer_sizes=None,
                exogenous_feature_columns=None):
     assert (anomaly_prior_probability < 1.0 and
             anomaly_prior_probability > 0.0)
@@ -719,7 +905,7 @@ class AnomalyMixtureARModel(ARModel):
         input_window_size=input_window_size,
         output_window_size=output_window_size,
         loss=ARModel.NORMAL_LIKELIHOOD_LOSS,
-        hidden_layer_sizes=hidden_layer_sizes,
+        prediction_model_factory=prediction_model_factory,
         exogenous_feature_columns=exogenous_feature_columns)
 
   def _create_anomaly_ops(self, times, values, prediction_ops_dict):
diff --git a/tensorflow/contrib/timeseries/python/timeseries/ar_model_test.py b/tensorflow/contrib/timeseries/python/timeseries/ar_model_test.py
index d078ac8d46..63f5d3568b 100644
--- a/tensorflow/contrib/timeseries/python/timeseries/ar_model_test.py
+++ b/tensorflow/contrib/timeseries/python/timeseries/ar_model_test.py
@@ -18,12 +18,13 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import functools
+
 import numpy as np
 
+from tensorflow.contrib.timeseries.python.timeseries import ar_model
 from tensorflow.contrib.timeseries.python.timeseries import input_pipeline
 from tensorflow.contrib.timeseries.python.timeseries import test_utils
-from tensorflow.contrib.timeseries.python.timeseries.ar_model import AnomalyMixtureARModel
-from tensorflow.contrib.timeseries.python.timeseries.ar_model import ARModel
 from tensorflow.contrib.timeseries.python.timeseries.estimators import ARRegressor
 from tensorflow.contrib.timeseries.python.timeseries.feature_keys import PredictionFeatures
 from tensorflow.contrib.timeseries.python.timeseries.feature_keys import TrainEvalFeatures
@@ -91,7 +92,7 @@ class ARModelTest(test.TestCase):
     np.random.seed(3)
     data_noise_stddev = 0.2
     if max_loss is None:
-      if loss == ARModel.NORMAL_LIKELIHOOD_LOSS:
+      if loss == ar_model.ARModel.NORMAL_LIKELIHOOD_LOSS:
         max_loss = 1.0
       else:
         max_loss = 0.05 / (data_noise_stddev ** 2)
@@ -137,7 +138,7 @@ class ARModelTest(test.TestCase):
     test_loss = test_evaluation["loss"]
     logging.info("Final test loss: %f", test_loss)
     self.assertLess(test_loss, max_loss)
-    if loss == ARModel.SQUARED_LOSS:
+    if loss == ar_model.ARModel.SQUARED_LOSS:
       # Test that the evaluation loss is reported without input scaling.
       self.assertAllClose(
           test_loss,
@@ -169,7 +170,7 @@ class ARModelTest(test.TestCase):
     predicted_mean = predictions["mean"][:, 0]
     true_values = predict_true_values[0, :, 0]
 
-    if loss == ARModel.NORMAL_LIKELIHOOD_LOSS:
+    if loss == ar_model.ARModel.NORMAL_LIKELIHOOD_LOSS:
       variances = predictions["covariance"][:, 0]
       standard_deviations = np.sqrt(variances)
       # Note that we may get tighter bounds with more training steps.
@@ -180,26 +181,26 @@ class ARModelTest(test.TestCase):
   def test_time_regression_squared(self):
     self.train_helper(input_window_size=0,
                       train_steps=350,
-                      loss=ARModel.SQUARED_LOSS)
+                      loss=ar_model.ARModel.SQUARED_LOSS)
 
   def test_autoregression_squared(self):
     self.train_helper(input_window_size=15,
-                      loss=ARModel.SQUARED_LOSS)
+                      loss=ar_model.ARModel.SQUARED_LOSS)
 
   def test_autoregression_short_input_window(self):
     self.train_helper(input_window_size=8,
-                      loss=ARModel.SQUARED_LOSS)
+                      loss=ar_model.ARModel.SQUARED_LOSS)
 
   def test_autoregression_normal(self):
     self.train_helper(input_window_size=10,
-                      loss=ARModel.NORMAL_LIKELIHOOD_LOSS,
+                      loss=ar_model.ARModel.NORMAL_LIKELIHOOD_LOSS,
                       train_steps=300,
                       max_loss=1.5,
                       anomaly_distribution=None)
 
   def test_autoregression_normal_multiple_periods(self):
     self.train_helper(input_window_size=10,
-                      loss=ARModel.NORMAL_LIKELIHOOD_LOSS,
+                      loss=ar_model.ARModel.NORMAL_LIKELIHOOD_LOSS,
                       max_loss=2.0,
                       multiple_periods=True,
                       anomaly_distribution=None)
@@ -207,15 +208,15 @@ class ARModelTest(test.TestCase):
   def test_autoregression_normal_anomalies_normal(self):
     self.train_helper(
         input_window_size=10,
-        loss=ARModel.NORMAL_LIKELIHOOD_LOSS,
-        anomaly_distribution=AnomalyMixtureARModel.GAUSSIAN_ANOMALY)
+        loss=ar_model.ARModel.NORMAL_LIKELIHOOD_LOSS,
+        anomaly_distribution=ar_model.AnomalyMixtureARModel.GAUSSIAN_ANOMALY)
 
   def test_autoregression_normal_anomalies_cauchy(self):
     self.train_helper(
         input_window_size=10,
         max_loss=1.5,
-        loss=ARModel.NORMAL_LIKELIHOOD_LOSS,
-        anomaly_distribution=AnomalyMixtureARModel.CAUCHY_ANOMALY)
+        loss=ar_model.ARModel.NORMAL_LIKELIHOOD_LOSS,
+        anomaly_distribution=ar_model.AnomalyMixtureARModel.CAUCHY_ANOMALY)
 
   def test_wrong_window_size(self):
     estimator = ARRegressor(
@@ -237,15 +238,38 @@ class ARModelTest(test.TestCase):
     with self.assertRaisesRegexp(ValueError, "requires a window of at least"):
       estimator.evaluate(input_fn=_bad_window_size_input_fn, steps=1)
 
-  def test_predictions_direct(self):
+  def test_predictions_direct_flat(self):
+    g = ops.Graph()
+    with g.as_default():
+      model = ar_model.ARModel(periodicities=2,
+                               num_features=1,
+                               num_time_buckets=10,
+                               input_window_size=2,
+                               output_window_size=2,
+                               prediction_model_factory=functools.partial(
+                                   ar_model.FlatPredictionModel,
+                                   hidden_layer_sizes=[40, 10]))
+      with session.Session():
+        predicted_values = model.predict({
+            PredictionFeatures.TIMES: [[4, 6, 10]],
+            PredictionFeatures.STATE_TUPLE: (
+                [[1, 2]], [[[1.], [2.]]], [[[], []]])
+        })
+        variables.global_variables_initializer().run()
+        self.assertAllEqual(predicted_values["mean"].eval().shape,
+                            [1, 3, 1])
+
+  def test_predictions_direct_lstm(self):
     g = ops.Graph()
     with g.as_default():
-      model = ARModel(periodicities=2,
-                      num_features=1,
-                      num_time_buckets=10,
-                      input_window_size=2,
-                      output_window_size=2,
-                      hidden_layer_sizes=[40, 10])
+      model = ar_model.ARModel(periodicities=2,
+                               num_features=1,
+                               num_time_buckets=10,
+                               input_window_size=2,
+                               output_window_size=2,
+                               prediction_model_factory=functools.partial(
+                                   ar_model.LSTMPredictionModel,
+                                   num_units=16))
       with session.Session():
         predicted_values = model.predict({
             PredictionFeatures.TIMES: [[4, 6, 10]],
@@ -259,11 +283,11 @@ class ARModelTest(test.TestCase):
   def test_long_eval(self):
     g = ops.Graph()
     with g.as_default():
-      model = ARModel(periodicities=2,
-                      num_features=1,
-                      num_time_buckets=10,
-                      input_window_size=2,
-                      output_window_size=1)
+      model = ar_model.ARModel(periodicities=2,
+                               num_features=1,
+                               num_time_buckets=10,
+                               input_window_size=2,
+                               output_window_size=1)
       raw_features = {
           TrainEvalFeatures.TIMES: [[1, 3, 5, 7, 11]],
           TrainEvalFeatures.VALUES: [[[1.], [2.], [3.], [4.], [5.]]]}
@@ -309,11 +333,11 @@ class ARModelTest(test.TestCase):
   def test_long_eval_discard_indivisible(self):
     g = ops.Graph()
     with g.as_default():
-      model = ARModel(periodicities=2,
-                      num_features=1,
-                      num_time_buckets=10,
-                      input_window_size=2,
-                      output_window_size=2)
+      model = ar_model.ARModel(periodicities=2,
+                               num_features=1,
+                               num_time_buckets=10,
+                               input_window_size=2,
+                               output_window_size=2)
       raw_features = {
           TrainEvalFeatures.TIMES: [[1, 3, 5, 7, 11]],
           TrainEvalFeatures.VALUES: [[[1.], [2.], [3.], [4.], [5.]]]}
diff --git a/tensorflow/contrib/timeseries/python/timeseries/estimators.py b/tensorflow/contrib/timeseries/python/timeseries/estimators.py
index f4608ca2d1..4ec8d26116 100644
--- a/tensorflow/contrib/timeseries/python/timeseries/estimators.py
+++ b/tensorflow/contrib/timeseries/python/timeseries/estimators.py
@@ -18,6 +18,8 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import functools
+
 from tensorflow.contrib.timeseries.python.timeseries import ar_model
 from tensorflow.contrib.timeseries.python.timeseries import feature_keys
 from tensorflow.contrib.timeseries.python.timeseries import head as ts_head_lib
@@ -61,7 +63,10 @@ class TimeSeriesRegressor(estimator_lib.Estimator):
     input_statistics_generator = math_utils.InputStatisticsFromMiniBatch(
         dtype=model.dtype, num_features=model.num_features)
     if state_manager is None:
-      state_manager = state_management.PassthroughStateManager()
+      if isinstance(model, ar_model.ARModel):
+        state_manager = state_management.FilteringOnlyStateManager()
+      else:
+        state_manager = state_management.PassthroughStateManager()
     if optimizer is None:
       optimizer = train.AdamOptimizer(0.02)
     self._model = model
@@ -246,11 +251,13 @@ class ARRegressor(TimeSeriesRegressor):
         anomaly_distribution = ar_model.AnomalyMixtureARModel.GAUSSIAN_ANOMALY
       model = ar_model.ARModel(
           periodicities=periodicities, num_features=num_features,
+          prediction_model_factory=functools.partial(
+              ar_model.FlatPredictionModel,
+              hidden_layer_sizes=hidden_layer_sizes),
           exogenous_feature_columns=exogenous_feature_columns,
           num_time_buckets=num_time_buckets,
           input_window_size=input_window_size,
-          output_window_size=output_window_size, loss=loss,
-          hidden_layer_sizes=hidden_layer_sizes)
+          output_window_size=output_window_size, loss=loss)
     else:
       if loss != ar_model.ARModel.NORMAL_LIKELIHOOD_LOSS:
         raise ValueError(
@@ -261,9 +268,11 @@ class ARRegressor(TimeSeriesRegressor):
           input_window_size=input_window_size,
           output_window_size=output_window_size,
           num_features=num_features,
+          prediction_model_factory=functools.partial(
+              ar_model.FlatPredictionModel,
+              hidden_layer_sizes=hidden_layer_sizes),
           exogenous_feature_columns=exogenous_feature_columns,
           num_time_buckets=num_time_buckets,
-          hidden_layer_sizes=hidden_layer_sizes,
           anomaly_prior_probability=anomaly_prior_probability,
           anomaly_distribution=anomaly_distribution)
     state_manager = state_management.FilteringOnlyStateManager()
diff --git a/tensorflow/contrib/timeseries/python/timeseries/estimators_test.py b/tensorflow/contrib/timeseries/python/timeseries/estimators_test.py
index eebee053f8..706742ca28 100644
--- a/tensorflow/contrib/timeseries/python/timeseries/estimators_test.py
+++ b/tensorflow/contrib/timeseries/python/timeseries/estimators_test.py
@@ -16,6 +16,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import functools
 import tempfile
 
 import numpy
@@ -178,7 +179,7 @@ class TimeSeriesRegressorTest(test.TestCase):
             session=sess)
         self.assertAllEqual([10, 15, 1], predictions["mean"].shape)
 
-  def test_fit_restore_fit_ar_regressor(self):
+  def test_fit_restore_fit_ar_flat(self):
     def _estimator_fn(model_dir, exogenous_feature_columns):
       return estimators.ARRegressor(
           periodicities=10, input_window_size=10, output_window_size=6,
@@ -189,6 +190,20 @@ class TimeSeriesRegressorTest(test.TestCase):
           exogenous_feature_columns=exogenous_feature_columns)
     self._fit_restore_fit_test_template(_estimator_fn, dtype=dtypes.float32)
 
+  def test_fit_restore_fit_ar_lstm(self):
+    def _estimator_fn(model_dir, exogenous_feature_columns):
+      return estimators.TimeSeriesRegressor(
+          model=ar_model.ARModel(
+              periodicities=10, input_window_size=10, output_window_size=6,
+              num_features=1,
+              exogenous_feature_columns=exogenous_feature_columns,
+              prediction_model_factory=functools.partial(
+                  ar_model.LSTMPredictionModel,
+                  num_units=10)),
+          config=_SeedRunConfig(),
+          model_dir=model_dir)
+    self._fit_restore_fit_test_template(_estimator_fn, dtype=dtypes.float32)
+
   def test_fit_restore_fit_structural_ensemble_regressor(self):
     dtype = dtypes.float32
     def _estimator_fn(model_dir, exogenous_feature_columns):
-- 
GitLab


From e32c42a6deed1f8ed1dcdeaaba0acf74685c18e3 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 10:47:38 -0700
Subject: [PATCH 258/395] Improve broadcast add implementation.

PiperOrigin-RevId: 195437679
---
 .../internal/optimized/optimized_ops.h        | 87 ++++++++++++++++++-
 .../internal/reference/reference_ops.h        | 81 +++++++++++++++++
 2 files changed, 165 insertions(+), 3 deletions(-)

diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
index 3d6042c31f..4776726972 100644
--- a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
@@ -2593,7 +2593,7 @@ inline void Add(int left_shift, const uint8* input1_data,
   }
 #endif  // NEON
 
-  for (; i < size; i++) {
+  for (; i < size; ++i) {
     const int32 input1_val = input1_offset + input1_data[i];
     const int32 input2_val = input2_offset + input2_data[i];
     const int32 shifted_input1_val = input1_val * (1 << left_shift);
@@ -2750,7 +2750,7 @@ inline void BroadcastAdd(int left_shift, const uint8* input1_data,
                          int32 output_activation_min,
                          int32 output_activation_max, uint8* output_data,
                          const Dims<4>& output_dims) {
-  gemmlowp::ScopedProfilingLabel label("BroadcastAdd/8bit");
+  gemmlowp::ScopedProfilingLabel label("BroadcastAddGeneric/8bit");
 
   NdArrayDesc<4> desc1;
   NdArrayDesc<4> desc2;
@@ -2799,6 +2799,60 @@ inline void BroadcastAdd(int left_shift, const uint8* input1_data,
   }
 }
 
+inline void BroadcastAddFivefold(
+    int y0, int y1, int y2, int y3, int y4, int left_shift,
+    const uint8* input1_data, const Dims<4>& input1_dims, int32 input1_offset,
+    int32 input1_multiplier, int input1_shift, const uint8* input2_data,
+    const Dims<4>& input2_dims, int32 input2_offset, int32 input2_multiplier,
+    int input2_shift, int32 output_offset, int32 output_multiplier,
+    int output_shift, int32 output_activation_min, int32 output_activation_max,
+    uint8* output_data, const Dims<4>& output_dims) {
+  gemmlowp::ScopedProfilingLabel label("BroadcastAddFivefold/8bit");
+
+  // Fivefold nested loops. The second input resets its position for each
+  // iteration of the second loop. The first input resets its position at the
+  // beginning of the fourth loop. The innermost loop is an elementwise add of
+  // sections of the arrays.
+  uint8* output_data_ptr = output_data;
+  const uint8* input1_data_ptr = input1_data;
+  const uint8* input2_data_reset = input2_data;
+  for (int i4 = 0; i4 < y4; ++i4) {
+    const uint8* input2_data_ptr;
+    for (int i3 = 0; i3 < y3; ++i3) {
+      input2_data_ptr = input2_data_reset;
+      for (int i2 = 0; i2 < y2; ++i2) {
+        for (int i1 = 0; i1 < y1; ++i1) {
+          for (int i0 = 0; i0 < y0; ++i0) {
+            const int32 input1_val = input1_offset + input1_data_ptr[i0];
+            const int32 input2_val = input2_offset + input2_data_ptr[i0];
+            const int32 shifted_input1_val = input1_val * (1 << left_shift);
+            const int32 shifted_input2_val = input2_val * (1 << left_shift);
+            const int32 scaled_input1_val =
+                MultiplyByQuantizedMultiplierSmallerThanOne(
+                    shifted_input1_val, input1_multiplier, input1_shift);
+            const int32 scaled_input2_val =
+                MultiplyByQuantizedMultiplierSmallerThanOne(
+                    shifted_input2_val, input2_multiplier, input2_shift);
+            const int32 raw_sum = scaled_input1_val + scaled_input2_val;
+            const int32 raw_output =
+                MultiplyByQuantizedMultiplierSmallerThanOne(
+                    raw_sum, output_multiplier, output_shift) +
+                output_offset;
+            const int32 clamped_output =
+                std::min(output_activation_max,
+                         std::max(output_activation_min, raw_output));
+            output_data_ptr[i0] = static_cast<uint8>(clamped_output);
+          }
+          input2_data_ptr += y0;
+          output_data_ptr += y0;
+        }
+        input1_data_ptr += y0;
+      }
+    }
+    input2_data_reset = input2_data_ptr;
+  }
+}
+
 template <FusedActivationFunctionType Ac>
 inline void BroadcastAdd(int left_shift, const uint8* input1_data,
                          const Dims<4>& input1_dims, int32 input1_offset,
@@ -2827,6 +2881,33 @@ inline void BroadcastAdd(int left_shift, const uint8* input1_data,
                output_activation_max, output_data, output_dims);
 }
 
+template <FusedActivationFunctionType Ac>
+inline void BroadcastAddFivefold(
+    int y0, int y1, int y2, int y3, int y4, int left_shift,
+    const uint8* input1_data, const Dims<4>& input1_dims, int32 input1_offset,
+    int32 input1_multiplier, int input1_shift, const uint8* input2_data,
+    const Dims<4>& input2_dims, int32 input2_offset, int32 input2_multiplier,
+    int input2_shift, int32 output_offset, int32 output_multiplier,
+    int output_shift, int32 output_activation_min, int32 output_activation_max,
+    uint8* output_data, const Dims<4>& output_dims) {
+  static_assert(Ac == FusedActivationFunctionType::kNone ||
+                    Ac == FusedActivationFunctionType::kRelu ||
+                    Ac == FusedActivationFunctionType::kRelu6 ||
+                    Ac == FusedActivationFunctionType::kRelu1,
+                "");
+  TFLITE_DCHECK_LE(output_activation_min, output_activation_max);
+  if (Ac == FusedActivationFunctionType::kNone) {
+    TFLITE_DCHECK_EQ(output_activation_min, 0);
+    TFLITE_DCHECK_EQ(output_activation_max, 255);
+  }
+  BroadcastAddFivefold(y0, y1, y2, y3, y4, left_shift, input1_data, input1_dims,
+                       input1_offset, input1_multiplier, input1_shift,
+                       input2_data, input2_dims, input2_offset,
+                       input2_multiplier, input2_shift, output_offset,
+                       output_multiplier, output_shift, output_activation_min,
+                       output_activation_max, output_data, output_dims);
+}
+
 inline void Mul(const float* input1_data, const Dims<4>& input1_dims,
                 const float* input2_data, const Dims<4>& input2_dims,
                 float output_activation_min, float output_activation_max,
@@ -4375,7 +4456,7 @@ inline void Softmax(const uint8* input_data, const Dims<4>& input_dims,
   using FixedPointAccum = gemmlowp::FixedPoint<int32, kAccumulationIntegerBits>;
   using FixedPoint0 = gemmlowp::FixedPoint<int32, 0>;
 
-gemmlowp::ScopedProfilingLabel label("Softmax/8bit");
+  gemmlowp::ScopedProfilingLabel label("Softmax/8bit");
   const int batches = MatchingArraySize(input_dims, 3, output_dims, 3);
   const int height = MatchingArraySize(input_dims, 2, output_dims, 2);
   const int width = MatchingArraySize(input_dims, 1, output_dims, 1);
diff --git a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
index d41ade4c9d..c6ed614593 100644
--- a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
@@ -1189,6 +1189,60 @@ inline void BroadcastAdd(int left_shift, const uint8* input1_data,
   }
 }
 
+inline void BroadcastAddFivefold(
+    int y0, int y1, int y2, int y3, int y4, int left_shift,
+    const uint8* input1_data, const Dims<4>& input1_dims, int32 input1_offset,
+    int32 input1_multiplier, int input1_shift, const uint8* input2_data,
+    const Dims<4>& input2_dims, int32 input2_offset, int32 input2_multiplier,
+    int input2_shift, int32 output_offset, int32 output_multiplier,
+    int output_shift, int32 output_activation_min, int32 output_activation_max,
+    uint8* output_data, const Dims<4>& output_dims) {
+  gemmlowp::ScopedProfilingLabel label("BroadcastAddFivefold/8bit");
+
+  int sb1 = y0;
+  int sa2 = y0;
+  int sb2 = y0 * y1;
+  int sa3 = y0 * y2;
+  int sa4 = y0 * y2 * y3;
+  int sb4 = y0 * y1 * y2;
+
+  uint8* output_data_ptr = output_data;
+  for (int i4 = 0; i4 < y4; ++i4) {
+    for (int i3 = 0; i3 < y3; ++i3) {
+      for (int i2 = 0; i2 < y2; ++i2) {
+        for (int i1 = 0; i1 < y1; ++i1) {
+          for (int i0 = 0; i0 < y0; ++i0) {
+            const int32 input1_val =
+                input1_offset +
+                input1_data[i4 * sa4 + i3 * sa3 + i2 * sa2 + i0];
+            const int32 input2_val =
+                input2_offset +
+                input2_data[i4 * sb4 + i2 * sb2 + i1 * sb1 + i0];
+            const int32 shifted_input1_val = input1_val * (1 << left_shift);
+            const int32 shifted_input2_val = input2_val * (1 << left_shift);
+            const int32 scaled_input1_val =
+                MultiplyByQuantizedMultiplierSmallerThanOne(
+                    shifted_input1_val, input1_multiplier, input1_shift);
+            const int32 scaled_input2_val =
+                MultiplyByQuantizedMultiplierSmallerThanOne(
+                    shifted_input2_val, input2_multiplier, input2_shift);
+            const int32 raw_sum = scaled_input1_val + scaled_input2_val;
+            const int32 raw_output =
+                MultiplyByQuantizedMultiplierSmallerThanOne(
+                    raw_sum, output_multiplier, output_shift) +
+                output_offset;
+            const int32 clamped_output =
+                std::min(output_activation_max,
+                         std::max(output_activation_min, raw_output));
+            *output_data_ptr = static_cast<uint8>(clamped_output);
+            ++output_data_ptr;
+          }
+        }
+      }
+    }
+  }
+}
+
 template <FusedActivationFunctionType Ac>
 inline void BroadcastAdd(int left_shift, const uint8* input1_data,
                          const Dims<4>& input1_dims, int32 input1_offset,
@@ -1217,6 +1271,33 @@ inline void BroadcastAdd(int left_shift, const uint8* input1_data,
                output_activation_max, output_data, output_dims);
 }
 
+template <FusedActivationFunctionType Ac>
+inline void BroadcastAddFivefold(
+    int y0, int y1, int y2, int y3, int y4, int left_shift,
+    const uint8* input1_data, const Dims<4>& input1_dims, int32 input1_offset,
+    int32 input1_multiplier, int input1_shift, const uint8* input2_data,
+    const Dims<4>& input2_dims, int32 input2_offset, int32 input2_multiplier,
+    int input2_shift, int32 output_offset, int32 output_multiplier,
+    int output_shift, int32 output_activation_min, int32 output_activation_max,
+    uint8* output_data, const Dims<4>& output_dims) {
+  static_assert(Ac == FusedActivationFunctionType::kNone ||
+                    Ac == FusedActivationFunctionType::kRelu ||
+                    Ac == FusedActivationFunctionType::kRelu6 ||
+                    Ac == FusedActivationFunctionType::kRelu1,
+                "");
+  TFLITE_DCHECK_LE(output_activation_min, output_activation_max);
+  if (Ac == FusedActivationFunctionType::kNone) {
+    TFLITE_DCHECK_EQ(output_activation_min, 0);
+    TFLITE_DCHECK_EQ(output_activation_max, 255);
+  }
+  BroadcastAddFivefold(y0, y1, y2, y3, y4, left_shift, input1_data, input1_dims,
+                       input1_offset, input1_multiplier, input1_shift,
+                       input2_data, input2_dims, input2_offset,
+                       input2_multiplier, input2_shift, output_offset,
+                       output_multiplier, output_shift, output_activation_min,
+                       output_activation_max, output_data, output_dims);
+}
+
 inline void Mul(const float* input1_data, const Dims<4>& input1_dims,
                 const float* input2_data, const Dims<4>& input2_dims,
                 float output_activation_min, float output_activation_max,
-- 
GitLab


From 09d0e300f9a136838102c94e4b5cf0d4d0876ace Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 10:50:29 -0700
Subject: [PATCH 259/395] Internal clean up: change scanf to use int64_t
 instead of int64

PiperOrigin-RevId: 195438212
---
 tensorflow/core/lib/strings/numbers.cc     | 4 ++--
 tensorflow/core/util/command_line_flags.cc | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/lib/strings/numbers.cc b/tensorflow/core/lib/strings/numbers.cc
index e4b909296e..987e4fe733 100644
--- a/tensorflow/core/lib/strings/numbers.cc
+++ b/tensorflow/core/lib/strings/numbers.cc
@@ -392,8 +392,8 @@ string FpToString(Fprint fp) {
 
 bool StringToFp(const string& s, Fprint* fp) {
   char junk;
-  uint64 result;
-  if (sscanf(s.c_str(), "%llx%c", &result, &junk) == 1) {
+  uint64_t result;
+  if (sscanf(s.c_str(), "%lx%c", &result, &junk) == 1) {
     *fp = result;
     return true;
   } else {
diff --git a/tensorflow/core/util/command_line_flags.cc b/tensorflow/core/util/command_line_flags.cc
index 8c27d01917..b281acb2b0 100644
--- a/tensorflow/core/util/command_line_flags.cc
+++ b/tensorflow/core/util/command_line_flags.cc
@@ -69,8 +69,8 @@ bool ParseInt64Flag(tensorflow::StringPiece arg, tensorflow::StringPiece flag,
       str_util::ConsumePrefix(&arg, flag) &&
       str_util::ConsumePrefix(&arg, "=")) {
     char extra;
-    int64 parsed_int64;
-    if (sscanf(arg.data(), "%lld%c", &parsed_int64, &extra) != 1) {
+    int64_t parsed_int64;
+    if (sscanf(arg.data(), "%ld%c", &parsed_int64, &extra) != 1) {
       LOG(ERROR) << "Couldn't interpret value " << arg << " for flag " << flag
                  << ".";
       *value_parsing_ok = false;
-- 
GitLab


From 01a70dc43d32eb5add5f1cb5de2d6c98ed88dd83 Mon Sep 17 00:00:00 2001
From: Suharsh Sivakumar <suharshs@google.com>
Date: Fri, 4 May 2018 11:17:17 -0700
Subject: [PATCH 260/395] Add operations before Identity operations should be
 quantized.

Fixes #19014

PiperOrigin-RevId: 195443326
---
 .../contrib/quantize/python/quantize.py       |  6 ++++--
 .../contrib/quantize/python/quantize_test.py  | 20 +++++++++----------
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/tensorflow/contrib/quantize/python/quantize.py b/tensorflow/contrib/quantize/python/quantize.py
index efc1a94b3c..60616ea749 100644
--- a/tensorflow/contrib/quantize/python/quantize.py
+++ b/tensorflow/contrib/quantize/python/quantize.py
@@ -33,7 +33,7 @@ from tensorflow.python.platform import tf_logging as logging
 _QUANTIZABLE_TYPES = {'Conv2D', 'MatMul', 'DepthwiseConv2dNative'}
 
 # Activations that are supported by the quantization rewrite.
-_ACTIVATION_TYPES = {'Relu', 'Relu6', 'Identity'}
+_ACTIVATION_TYPES = {'Relu', 'Relu6'}
 
 
 def Quantize(graph,
@@ -267,8 +267,10 @@ def _FindLayersToQuantize(graph):
 
   # The input to the activation can come from bias add, fold bias add, the
   # bypasses.
+  # TODO(suharshs): We should ideally skip Identity operations instead of
+  # treating them as an activation.
   activation_pattern = graph_matcher.OpTypePattern(
-      '|'.join(_ACTIVATION_TYPES),
+      '|'.join(_ACTIVATION_TYPES) + '|Identity',
       inputs=[
           graph_matcher.OneofPattern([
               bias_add_pattern, folded_bias_add_pattern, bypass_pattern_a,
diff --git a/tensorflow/contrib/quantize/python/quantize_test.py b/tensorflow/contrib/quantize/python/quantize_test.py
index 5e479f3946..e7360ae03c 100644
--- a/tensorflow/contrib/quantize/python/quantize_test.py
+++ b/tensorflow/contrib/quantize/python/quantize_test.py
@@ -74,7 +74,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
                     weights_initializer=self._WeightInit(0.09),
                     activation_fn=None, scope='test/test')
       node = math_ops.add(conv, input2, name='test/add')
-      node = array_ops.identity(node, name='test/identity')
+      node = nn_ops.relu6(node, name='test/relu6')
       update_barrier = control_flow_ops.no_op(name='update_barrier')
       with ops.control_dependencies([update_barrier]):
         array_ops.identity(node, name='control_dependency')
@@ -97,7 +97,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
         for output in quant_op.outputs:
           consumers.extend(output.consumers())
 
-        self.assertNotIn('test/identity', [c.name for c in consumers])
+        self.assertNotIn('test/relu6', [c.name for c in consumers])
 
   def testInsertQuantOpForAddAfterSeparableConv2d(self):
     self._RunTestOverParameters(
@@ -114,7 +114,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
                               weights_initializer=self._WeightInit(0.09),
                               activation_fn=None, scope='test/test')
       node = math_ops.add(conv, input2, name='test/add')
-      node = array_ops.identity(node, name='test/identity')
+      node = nn_ops.relu6(node, name='test/relu6')
       update_barrier = control_flow_ops.no_op(name='update_barrier')
       with ops.control_dependencies([update_barrier]):
         array_ops.identity(node, name='control_dependency')
@@ -135,7 +135,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
         for output in quant_op.outputs:
           consumers.extend(output.consumers())
 
-        self.assertNotIn('test/identity', [c.name for c in consumers])
+        self.assertNotIn('test/relu6', [c.name for c in consumers])
 
   def testFinalLayerQuantized(self):
     self._RunTestOverParameters(self._TestFinalLayerQuantized)
@@ -174,7 +174,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
           stride=2,
           padding='SAME',
           weights_initializer=self._WeightInit(0.09),
-          activation_fn=array_ops.identity,
+          activation_fn=nn_ops.relu6,
           scope='test/test')
       bypass_tensor = math_ops.add(conv, input2, name='test/add')
       # The output of the post_activation bypass will be another layer.
@@ -184,7 +184,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
           stride=2,
           padding='SAME',
           weights_initializer=self._WeightInit(0.09),
-          activation_fn=array_ops.identity,
+          activation_fn=nn_ops.relu6,
           scope='test/unused')
 
       quantize.Quantize(graph, is_training, weight_bits=8, activation_bits=8)
@@ -212,7 +212,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
           stride=2,
           padding='SAME',
           weights_initializer=self._WeightInit(0.09),
-          activation_fn=array_ops.identity,
+          activation_fn=nn_ops.relu6,
           scope='test/test1')
 
       # The bypass of this conv is the post activation bypass of the previous
@@ -227,7 +227,7 @@ class QuantizeTest(test_util.TensorFlowTestCase):
           scope='test/test2')
 
       bypass_tensor = math_ops.add(conv1, conv2, name='test/add')
-      _ = array_ops.identity(bypass_tensor, name='test/output')
+      _ = nn_ops.relu6(bypass_tensor, name='test/output')
 
       quantize.Quantize(graph, is_training, weight_bits=8, activation_bits=8)
 
@@ -248,11 +248,11 @@ class QuantizeTest(test_util.TensorFlowTestCase):
           'test/test1/act_quant/FakeQuantWithMinMaxVars' in op_names)
       self.assertTrue('test/act_quant/FakeQuantWithMinMaxVars' in op_names)
       self.assertEqual(
-          'Identity',
+          'Relu6',
           graph.get_operation_by_name(
               'test/test1/act_quant/FakeQuantWithMinMaxVars').inputs[0].op.type)
       self.assertEqual(
-          'Identity',
+          'Relu6',
           graph.get_operation_by_name(
               'test/act_quant/FakeQuantWithMinMaxVars').inputs[0].op.type)
 
-- 
GitLab


From 2314acf98fb874317dd17ef3daf438d7af87f900 Mon Sep 17 00:00:00 2001
From: Anya Petrova <a.petrova.ds@gmail.com>
Date: Fri, 4 May 2018 13:22:03 -0700
Subject: [PATCH 261/395] Fix a small typo.

---
 SECURITY.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/SECURITY.md b/SECURITY.md
index a5ce3a62ee..01886b613e 100644
--- a/SECURITY.md
+++ b/SECURITY.md
@@ -173,7 +173,7 @@ the progress being made towards a fix and announcement.
 In addition, please include the following information along with your report:
 
 * Your name and affiliation (if any).
-* A description the technical details of the vulnerabilities. It is very
+* A description of the technical details of the vulnerabilities. It is very
   important to let us know how we can reproduce your findings.
 * An explanation who can exploit this vulnerability, and what they gain when
   doing so -- write an attack scenario. This will help us evaluate your report
-- 
GitLab


From f368558429f5ebdbc0a187c3801dccf1ca6963c7 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 11:40:01 -0700
Subject: [PATCH 262/395] [XLA] Cleanup client_library_test_base: move
 definition of CreateParameterAndTransferLiteral to .cc file

PiperOrigin-RevId: 195446864
---
 .../xla/tests/client_library_test_base.cc     | 29 +++++++++++++++++++
 .../xla/tests/client_library_test_base.h      | 29 -------------------
 2 files changed, 29 insertions(+), 29 deletions(-)

diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.cc b/tensorflow/compiler/xla/tests/client_library_test_base.cc
index c09e7eaf2b..41f9a5f666 100644
--- a/tensorflow/compiler/xla/tests/client_library_test_base.cc
+++ b/tensorflow/compiler/xla/tests/client_library_test_base.cc
@@ -565,4 +565,33 @@ XlaOp ClientLibraryTestBase::CreateConstantFromLiteral(const Literal& literal,
       use_bfloat16_ ? *LiteralTestUtil::ConvertF32ToBF16(literal) : literal);
 }
 
+std::unique_ptr<GlobalData>
+ClientLibraryTestBase::CreateParameterAndTransferLiteral(int64 parameter_number,
+                                                         const Literal& literal,
+                                                         const string& name,
+                                                         XlaBuilder* builder,
+                                                         XlaOp* data_handle) {
+  return CreateParameterAndTransferLiteral(parameter_number, literal, name,
+                                           nullptr, builder, data_handle);
+}
+
+std::unique_ptr<GlobalData>
+ClientLibraryTestBase::CreateParameterAndTransferLiteral(
+    int64 parameter_number, const Literal& literal, const string& name,
+    const DeviceHandle* device_handle, XlaBuilder* builder,
+    XlaOp* data_handle) {
+  const Literal* param_literal = &literal;
+  std::unique_ptr<Literal> converted_literal;
+  if (use_bfloat16_) {
+    converted_literal = LiteralTestUtil::ConvertF32ToBF16(literal);
+    param_literal = converted_literal.get();
+  }
+  std::unique_ptr<GlobalData> data =
+      client_->TransferToServer(*param_literal, device_handle)
+          .ConsumeValueOrDie();
+  *data_handle =
+      builder->Parameter(parameter_number, param_literal->shape(), name);
+  return data;
+}
+
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.h b/tensorflow/compiler/xla/tests/client_library_test_base.h
index e58979a303..16e838e60f 100644
--- a/tensorflow/compiler/xla/tests/client_library_test_base.h
+++ b/tensorflow/compiler/xla/tests/client_library_test_base.h
@@ -616,35 +616,6 @@ std::unique_ptr<Array2D<NativeT>> ClientLibraryTestBase::CreatePseudorandomR2(
   return result;
 }
 
-std::unique_ptr<GlobalData>
-ClientLibraryTestBase::CreateParameterAndTransferLiteral(int64 parameter_number,
-                                                         const Literal& literal,
-                                                         const string& name,
-                                                         XlaBuilder* builder,
-                                                         XlaOp* data_handle) {
-  return CreateParameterAndTransferLiteral(parameter_number, literal, name,
-                                           nullptr, builder, data_handle);
-}
-
-std::unique_ptr<GlobalData>
-ClientLibraryTestBase::CreateParameterAndTransferLiteral(
-    int64 parameter_number, const Literal& literal, const string& name,
-    const DeviceHandle* device_handle, XlaBuilder* builder,
-    XlaOp* data_handle) {
-  const Literal* param_literal = &literal;
-  std::unique_ptr<Literal> converted_literal;
-  if (use_bfloat16_) {
-    converted_literal = LiteralTestUtil::ConvertF32ToBF16(literal);
-    param_literal = converted_literal.get();
-  }
-  std::unique_ptr<GlobalData> data =
-      client_->TransferToServer(*param_literal, device_handle)
-          .ConsumeValueOrDie();
-  *data_handle =
-      builder->Parameter(parameter_number, param_literal->shape(), name);
-  return data;
-}
-
 }  // namespace xla
 
 #endif  // TENSORFLOW_COMPILER_XLA_TESTS_CLIENT_LIBRARY_TEST_BASE_H_
-- 
GitLab


From be9b87375adecad9bd8bb12c81b2566c77a68ad7 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 11:40:20 -0700
Subject: [PATCH 263/395] [XLA] Redesign: migrate the SWIG wrapped xla client.
 Added LocalOp that wraps XlaOp, so that it's fully visible to swig.

PiperOrigin-RevId: 195446939
---
 tensorflow/compiler/xla/python/BUILD          |   3 +-
 .../xla/python/local_computation_builder.cc   | 315 ++++++++-------
 .../xla/python/local_computation_builder.h    | 206 +++++-----
 .../xla/python/local_computation_builder.i    |  53 +--
 tensorflow/compiler/xla/python/xla_client.py  | 362 +++++++-----------
 5 files changed, 415 insertions(+), 524 deletions(-)

diff --git a/tensorflow/compiler/xla/python/BUILD b/tensorflow/compiler/xla/python/BUILD
index ecb87bd889..932cce943f 100644
--- a/tensorflow/compiler/xla/python/BUILD
+++ b/tensorflow/compiler/xla/python/BUILD
@@ -49,9 +49,10 @@ cc_library(
         "//tensorflow/compiler/xla:util",
         "//tensorflow/compiler/xla:xla_data_proto",
         "//tensorflow/compiler/xla/client:client_library",
-        "//tensorflow/compiler/xla/client:computation_builder",
         "//tensorflow/compiler/xla/client:executable_build_options",
         "//tensorflow/compiler/xla/client:local_client",
+        "//tensorflow/compiler/xla/client/xla_client:xla_builder",
+        "//tensorflow/compiler/xla/client/xla_client:xla_computation",
         "//tensorflow/compiler/xla/service:shaped_buffer",
         "//tensorflow/core:framework_lite",
         "//tensorflow/core:lib",
diff --git a/tensorflow/compiler/xla/python/local_computation_builder.cc b/tensorflow/compiler/xla/python/local_computation_builder.cc
index 044458164f..df262c97bf 100644
--- a/tensorflow/compiler/xla/python/local_computation_builder.cc
+++ b/tensorflow/compiler/xla/python/local_computation_builder.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/python/local_computation_builder.h"
 #include "tensorflow/compiler/xla/executable_run_options.h"
+#include "tensorflow/compiler/xla/ptr_util.h"
 #include "tensorflow/compiler/xla/util.h"
 #include "tensorflow/core/platform/default/thread_annotations.h"
 
@@ -248,7 +249,7 @@ LocalShapedBuffer* CompiledLocalComputation::ExecuteWithShapedBuffers(
   return new LocalShapedBuffer(std::move(result_buffer));
 }
 
-LocalComputation::LocalComputation(Computation computation)
+LocalComputation::LocalComputation(XlaComputation computation)
     : computation_(std::move(computation)) {}
 
 StatusOr<CompiledLocalComputation*> LocalComputation::Compile(
@@ -271,7 +272,7 @@ StatusOr<CompiledLocalComputation*> LocalComputation::Compile(
   return new CompiledLocalComputation(std::move(local_executable));
 }
 
-const Computation& LocalComputation::computation() const {
+const XlaComputation& LocalComputation::computation() const {
   return computation_;
 }
 
@@ -281,8 +282,12 @@ StatusOr<Shape> LocalComputation::GetReturnValueShape() const {
   return std::move(*program_shape.mutable_result());
 }
 
+LocalOp::LocalOp(const XlaOp& op) : op_(op) {}
+
+const XlaOp& LocalOp::op() const { return op_; }
+
 LocalComputationBuilder::LocalComputationBuilder(const string& computation_name)
-    : builder_(GetOrCreateLocalClient(), computation_name) {}
+    : builder_(computation_name) {}
 
 void LocalComputationBuilder::SetOpMetadata(const OpMetadata& metadata) {
   builder_.SetOpMetadata(metadata);
@@ -291,19 +296,21 @@ void LocalComputationBuilder::SetOpMetadata(const OpMetadata& metadata) {
 void LocalComputationBuilder::ClearOpMetadata() { builder_.ClearOpMetadata(); }
 
 StatusOr<LocalComputation*> LocalComputationBuilder::Build() {
-  TF_ASSIGN_OR_RETURN(Computation computation, builder_.Build());
+  TF_ASSIGN_OR_RETURN(XlaComputation computation, builder_.Build());
   return new LocalComputation(std::move(computation));
 }
 
-ComputationDataHandle LocalComputationBuilder::Parameter(int64 parameter_number,
-                                                         const Shape& shape,
-                                                         const string& name) {
+LocalOp LocalComputationBuilder::Parameter(int64 parameter_number,
+                                           const Shape& shape,
+                                           const string& name) {
   return builder_.Parameter(parameter_number, shape, name);
 }
 
 std::unique_ptr<Shape> LocalComputationBuilder::GetShape(
-    const ComputationDataHandle& operand) {
-  return builder_.GetShape(operand).ConsumeValueOrDie();
+    const LocalOp& operand) {
+  auto result = MakeUnique<Shape>();
+  *result = builder_.GetShape(operand.op()).ValueOrDie();
+  return result;
 }
 
 StatusOr<Shape> LocalComputationBuilder::GetReturnValueShape() {
@@ -311,222 +318,236 @@ StatusOr<Shape> LocalComputationBuilder::GetReturnValueShape() {
   return program_shape.result();
 }
 
-ComputationDataHandle LocalComputationBuilder::Infeed(const Shape& shape) {
+LocalOp LocalComputationBuilder::Infeed(const Shape& shape) {
   return builder_.Infeed(shape);
 }
 
-void LocalComputationBuilder::Outfeed(const ComputationDataHandle& operand,
+void LocalComputationBuilder::Outfeed(const LocalOp& operand,
                                       const Shape& shape,
                                       const string& outfeed_config) {
-  builder_.Outfeed(operand, shape, outfeed_config);
+  builder_.Outfeed(operand.op(), shape, outfeed_config);
 }
 
-ComputationDataHandle LocalComputationBuilder::ConstantLiteral(
-    const Literal& literal) {
+LocalOp LocalComputationBuilder::ConstantLiteral(const Literal& literal) {
   return builder_.ConstantLiteral(literal);
 }
 
-ComputationDataHandle LocalComputationBuilder::Broadcast(
-    const ComputationDataHandle& operand,
+LocalOp LocalComputationBuilder::Broadcast(
+    const LocalOp& operand,
     tensorflow::gtl::ArraySlice<int64> broadcast_sizes) {
-  return builder_.Broadcast(operand, broadcast_sizes);
+  return builder_.Broadcast(operand.op(), broadcast_sizes);
 }
 
-ComputationDataHandle LocalComputationBuilder::Pad(
-    const ComputationDataHandle& operand,
-    const ComputationDataHandle& padding_value,
-    const PaddingConfig& padding_config) {
-  return builder_.Pad(operand, padding_value, padding_config);
+LocalOp LocalComputationBuilder::Pad(const LocalOp& operand,
+                                     const LocalOp& padding_value,
+                                     const PaddingConfig& padding_config) {
+  return builder_.Pad(operand.op(), padding_value.op(), padding_config);
 }
 
-ComputationDataHandle LocalComputationBuilder::Reshape(
-    const ComputationDataHandle& operand,
-    tensorflow::gtl::ArraySlice<int64> dimensions,
+LocalOp LocalComputationBuilder::Reshape(
+    const LocalOp& operand, tensorflow::gtl::ArraySlice<int64> dimensions,
     tensorflow::gtl::ArraySlice<int64> new_sizes) {
-  return builder_.Reshape(operand, dimensions, new_sizes);
+  return builder_.Reshape(operand.op(), dimensions, new_sizes);
 }
 
-ComputationDataHandle LocalComputationBuilder::Collapse(
-    const ComputationDataHandle& operand,
-    tensorflow::gtl::ArraySlice<int64> dimensions) {
-  return builder_.Collapse(operand, dimensions);
+LocalOp LocalComputationBuilder::Collapse(
+    const LocalOp& operand, tensorflow::gtl::ArraySlice<int64> dimensions) {
+  return builder_.Collapse(operand.op(), dimensions);
 }
 
-ComputationDataHandle LocalComputationBuilder::CrossReplicaSum(
-    const ComputationDataHandle& operand) {
-  return builder_.CrossReplicaSum(operand);
+LocalOp LocalComputationBuilder::CrossReplicaSum(const LocalOp& operand) {
+  return builder_.CrossReplicaSum(operand.op());
 }
 
-ComputationDataHandle LocalComputationBuilder::Slice(
-    const ComputationDataHandle& operand,
-    tensorflow::gtl::ArraySlice<int64> start_indices,
+LocalOp LocalComputationBuilder::Slice(
+    const LocalOp& operand, tensorflow::gtl::ArraySlice<int64> start_indices,
     tensorflow::gtl::ArraySlice<int64> limit_indices,
     tensorflow::gtl::ArraySlice<int64> strides) {
-  return builder_.Slice(operand, start_indices, limit_indices, strides);
+  return builder_.Slice(operand.op(), start_indices, limit_indices, strides);
 }
 
-ComputationDataHandle LocalComputationBuilder::SliceInDim(
-    const ComputationDataHandle& operand, int64 start_index, int64 limit_index,
-    int64 stride, int64 dimno) {
-  return builder_.SliceInDim(operand, start_index, limit_index, stride, dimno);
+LocalOp LocalComputationBuilder::SliceInDim(const LocalOp& operand,
+                                            int64 start_index,
+                                            int64 limit_index, int64 stride,
+                                            int64 dimno) {
+  return builder_.SliceInDim(operand.op(), start_index, limit_index, stride,
+                             dimno);
 }
 
-ComputationDataHandle LocalComputationBuilder::DynamicSlice(
-    const ComputationDataHandle& operand,
-    const ComputationDataHandle& start_indices,
+LocalOp LocalComputationBuilder::DynamicSlice(
+    const LocalOp& operand, const LocalOp& start_indices,
     tensorflow::gtl::ArraySlice<int64> slice_sizes) {
-  return builder_.DynamicSlice(operand, start_indices, slice_sizes);
+  return builder_.DynamicSlice(operand.op(), start_indices.op(), slice_sizes);
 }
 
-ComputationDataHandle LocalComputationBuilder::DynamicUpdateSlice(
-    const ComputationDataHandle& operand, const ComputationDataHandle& update,
-    const ComputationDataHandle& start_indices) {
-  return builder_.DynamicUpdateSlice(operand, update, start_indices);
+LocalOp LocalComputationBuilder::DynamicUpdateSlice(
+    const LocalOp& operand, const LocalOp& update,
+    const LocalOp& start_indices) {
+  return builder_.DynamicUpdateSlice(operand.op(), update.op(),
+                                     start_indices.op());
 }
 
-ComputationDataHandle LocalComputationBuilder::ConcatInDim(
-    tensorflow::gtl::ArraySlice<ComputationDataHandle> operands,
-    int64 dimension) {
-  return builder_.ConcatInDim(operands, dimension);
+LocalOp LocalComputationBuilder::ConcatInDim(
+    tensorflow::gtl::ArraySlice<LocalOp> operands, int64 dimension) {
+  std::vector<XlaOp> xla_ops;
+  xla_ops.reserve(operands.size());
+  for (const auto& op : operands) {
+    xla_ops.push_back(op.op());
+  }
+  return builder_.ConcatInDim(xla_ops, dimension);
 }
 
-ComputationDataHandle
-LocalComputationBuilder::SelectAndScatterWithGeneralPadding(
-    const ComputationDataHandle& operand, const LocalComputation& select,
+LocalOp LocalComputationBuilder::SelectAndScatterWithGeneralPadding(
+    const LocalOp& operand, const LocalComputation& select,
     tensorflow::gtl::ArraySlice<int64> window_dimensions,
     tensorflow::gtl::ArraySlice<int64> window_strides,
     tensorflow::gtl::ArraySlice<std::pair<int64, int64>> padding,
-    const ComputationDataHandle& source,
-    const ComputationDataHandle& init_value, const LocalComputation& scatter) {
+    const LocalOp& source, const LocalOp& init_value,
+    const LocalComputation& scatter) {
   return builder_.SelectAndScatterWithGeneralPadding(
-      operand, select.computation(), window_dimensions, window_strides, padding,
-      source, init_value, scatter.computation());
+      operand.op(), select.computation(), window_dimensions, window_strides,
+      padding, source.op(), init_value.op(), scatter.computation());
 }
 
-ComputationDataHandle LocalComputationBuilder::Tuple(
-    tensorflow::gtl::ArraySlice<ComputationDataHandle> elements) {
-  return builder_.Tuple(elements);
+LocalOp LocalComputationBuilder::Tuple(
+    tensorflow::gtl::ArraySlice<LocalOp> elements) {
+  std::vector<XlaOp> xla_ops;
+  xla_ops.reserve(elements.size());
+  for (const auto& op : elements) {
+    xla_ops.push_back(op.op());
+  }
+
+  return builder_.Tuple(xla_ops);
 }
 
-ComputationDataHandle LocalComputationBuilder::GetTupleElement(
-    const ComputationDataHandle& tuple_data, int64 index) {
-  return builder_.GetTupleElement(tuple_data, index);
+LocalOp LocalComputationBuilder::GetTupleElement(const LocalOp& tuple_data,
+                                                 int64 index) {
+  return builder_.GetTupleElement(tuple_data.op(), index);
 }
 
-ComputationDataHandle LocalComputationBuilder::Dot(
-    const ComputationDataHandle& lhs, const ComputationDataHandle& rhs) {
-  return builder_.Dot(lhs, rhs);
+LocalOp LocalComputationBuilder::Dot(const LocalOp& lhs, const LocalOp& rhs) {
+  return builder_.Dot(lhs.op(), rhs.op());
 }
 
-ComputationDataHandle LocalComputationBuilder::DotGeneral(
-    const ComputationDataHandle& lhs, const ComputationDataHandle& rhs,
+LocalOp LocalComputationBuilder::DotGeneral(
+    const LocalOp& lhs, const LocalOp& rhs,
     const DotDimensionNumbers& dimension_numbers) {
-  return builder_.DotGeneral(lhs, rhs, dimension_numbers);
+  return builder_.DotGeneral(lhs.op(), rhs.op(), dimension_numbers);
 }
 
-ComputationDataHandle LocalComputationBuilder::ConvGeneralDilated(
-    const ComputationDataHandle& lhs, const ComputationDataHandle& rhs,
+LocalOp LocalComputationBuilder::ConvGeneralDilated(
+    const LocalOp& lhs, const LocalOp& rhs,
     tensorflow::gtl::ArraySlice<int64> window_strides,
     tensorflow::gtl::ArraySlice<std::pair<int64, int64>> padding,
     tensorflow::gtl::ArraySlice<int64> lhs_dilation,
     tensorflow::gtl::ArraySlice<int64> rhs_dilation,
     const ConvolutionDimensionNumbers& dimension_numbers) {
-  return builder_.ConvGeneralDilated(lhs, rhs, window_strides, padding,
-                                     lhs_dilation, rhs_dilation,
+  return builder_.ConvGeneralDilated(lhs.op(), rhs.op(), window_strides,
+                                     padding, lhs_dilation, rhs_dilation,
                                      dimension_numbers);
 }
 
-ComputationDataHandle LocalComputationBuilder::ConvertElementType(
-    const ComputationDataHandle& operand, PrimitiveType new_element_type) {
-  return builder_.ConvertElementType(operand, new_element_type);
+LocalOp LocalComputationBuilder::ConvertElementType(
+    const LocalOp& operand, PrimitiveType new_element_type) {
+  return builder_.ConvertElementType(operand.op(), new_element_type);
 }
 
-ComputationDataHandle LocalComputationBuilder::Call(
+LocalOp LocalComputationBuilder::Call(
     const LocalComputation& local_computation,
-    tensorflow::gtl::ArraySlice<ComputationDataHandle> operands) {
-  return builder_.Call(local_computation.computation(), operands);
+    tensorflow::gtl::ArraySlice<LocalOp> operands) {
+  std::vector<XlaOp> xla_ops;
+  xla_ops.reserve(operands.size());
+  for (const auto& op : operands) {
+    xla_ops.push_back(op.op());
+  }
+  return builder_.Call(local_computation.computation(), xla_ops);
 }
 
-ComputationDataHandle LocalComputationBuilder::Transpose(
-    const ComputationDataHandle& operand,
-    tensorflow::gtl::ArraySlice<int64> permutation) {
-  return builder_.Transpose(operand, permutation);
+LocalOp LocalComputationBuilder::Transpose(
+    const LocalOp& operand, tensorflow::gtl::ArraySlice<int64> permutation) {
+  return builder_.Transpose(operand.op(), permutation);
 }
 
-ComputationDataHandle LocalComputationBuilder::Rev(
-    const ComputationDataHandle& operand,
-    tensorflow::gtl::ArraySlice<int64> dimensions) {
-  return builder_.Rev(operand, dimensions);
+LocalOp LocalComputationBuilder::Rev(
+    const LocalOp& operand, tensorflow::gtl::ArraySlice<int64> dimensions) {
+  return builder_.Rev(operand.op(), dimensions);
 }
 
-ComputationDataHandle LocalComputationBuilder::Map(
-    tensorflow::gtl::ArraySlice<ComputationDataHandle> operands,
+LocalOp LocalComputationBuilder::Map(
+    tensorflow::gtl::ArraySlice<LocalOp> operands,
     const LocalComputation& local_computation,
     tensorflow::gtl::ArraySlice<int64> dimensions,
-    tensorflow::gtl::ArraySlice<ComputationDataHandle> static_operands) {
-  return builder_.Map(operands, local_computation.computation(), dimensions,
-                      static_operands);
+    tensorflow::gtl::ArraySlice<LocalOp> static_operands) {
+  std::vector<XlaOp> xla_ops;
+  xla_ops.reserve(operands.size());
+  for (const auto& op : operands) {
+    xla_ops.push_back(op.op());
+  }
+
+  std::vector<XlaOp> static_xla_ops;
+  static_xla_ops.reserve(static_operands.size());
+  for (const auto& op : static_operands) {
+    static_xla_ops.push_back(op.op());
+  }
+
+  return builder_.Map(xla_ops, local_computation.computation(), dimensions,
+                      static_xla_ops);
 }
 
-ComputationDataHandle LocalComputationBuilder::Reduce(
-    const ComputationDataHandle& operand,
-    const ComputationDataHandle& init_value,
+LocalOp LocalComputationBuilder::Reduce(
+    const LocalOp& operand, const LocalOp& init_value,
     const LocalComputation& local_computation,
     tensorflow::gtl::ArraySlice<int64> dimensions_to_reduce) {
-  return builder_.Reduce(operand, init_value, local_computation.computation(),
-                         dimensions_to_reduce);
+  return builder_.Reduce(operand.op(), init_value.op(),
+                         local_computation.computation(), dimensions_to_reduce);
 }
 
-ComputationDataHandle LocalComputationBuilder::ReduceWindowWithGeneralPadding(
-    const ComputationDataHandle& operand,
-    const ComputationDataHandle& init_value,
+LocalOp LocalComputationBuilder::ReduceWindowWithGeneralPadding(
+    const LocalOp& operand, const LocalOp& init_value,
     const LocalComputation& local_computation,
     tensorflow::gtl::ArraySlice<int64> window_dimensions,
     tensorflow::gtl::ArraySlice<int64> window_strides,
     tensorflow::gtl::ArraySlice<std::pair<int64, int64>> padding) {
   return builder_.ReduceWindowWithGeneralPadding(
-      operand, init_value, local_computation.computation(), window_dimensions,
-      window_strides, padding);
+      operand.op(), init_value.op(), local_computation.computation(),
+      window_dimensions, window_strides, padding);
 }
 
-ComputationDataHandle LocalComputationBuilder::RngNormal(
-    const ComputationDataHandle& mu, const ComputationDataHandle& sigma,
-    const Shape& shape) {
-  return builder_.RngNormal(mu, sigma, shape);
+LocalOp LocalComputationBuilder::RngNormal(const LocalOp& mu,
+                                           const LocalOp& sigma,
+                                           const Shape& shape) {
+  return builder_.RngNormal(mu.op(), sigma.op(), shape);
 }
 
-ComputationDataHandle LocalComputationBuilder::RngUniform(
-    const ComputationDataHandle& a, const ComputationDataHandle& b,
-    const Shape& shape) {
-  return builder_.RngUniform(a, b, shape);
+LocalOp LocalComputationBuilder::RngUniform(const LocalOp& a, const LocalOp& b,
+                                            const Shape& shape) {
+  return builder_.RngUniform(a.op(), b.op(), shape);
 }
 
-ComputationDataHandle LocalComputationBuilder::While(
-    const LocalComputation& condition, const LocalComputation& body,
-    const ComputationDataHandle& init) {
-  return builder_.While(condition.computation(), body.computation(), init);
+LocalOp LocalComputationBuilder::While(const LocalComputation& condition,
+                                       const LocalComputation& body,
+                                       const LocalOp& init) {
+  return builder_.While(condition.computation(), body.computation(), init.op());
 }
 
-ComputationDataHandle LocalComputationBuilder::Conditional(
-    const ComputationDataHandle& predicate,
-    const ComputationDataHandle& true_operand,
-    const LocalComputation& true_computation,
-    const ComputationDataHandle& false_operand,
+LocalOp LocalComputationBuilder::Conditional(
+    const LocalOp& predicate, const LocalOp& true_operand,
+    const LocalComputation& true_computation, const LocalOp& false_operand,
     const LocalComputation& false_computation) {
-  return builder_.Conditional(predicate, true_operand,
-                              true_computation.computation(), false_operand,
-                              false_computation.computation());
+  return builder_.Conditional(
+      predicate.op(), true_operand.op(), true_computation.computation(),
+      false_operand.op(), false_computation.computation());
 }
 
-StatusOr<bool> LocalComputationBuilder::IsConstant(
-    const ComputationDataHandle& operand, int64 num_parameters) {
-  return builder_.IsConstant(operand, num_parameters);
+StatusOr<bool> LocalComputationBuilder::IsConstant(const LocalOp& operand) {
+  return builder_.IsConstant(operand.op());
 }
 
-StatusOr<std::unique_ptr<Literal>> LocalComputationBuilder::ComputeConstant(
-    const ComputationDataHandle& operand, const Layout* output_layout,
-    tensorflow::gtl::ArraySlice<Literal> parameters) {
-  return builder_.ComputeConstant(operand, output_layout, parameters);
+StatusOr<LocalComputation*> LocalComputationBuilder::BuildConstantSubGraph(
+    const LocalOp& operand) {
+  TF_ASSIGN_OR_RETURN(XlaComputation computation,
+                      builder_.BuildConstantSubGraph(operand.op()));
+  return new LocalComputation(std::move(computation));
 }
 
 #define _FORWARD(method_name, return_sig, args_sig, args)    \
@@ -534,23 +555,19 @@ StatusOr<std::unique_ptr<Literal>> LocalComputationBuilder::ComputeConstant(
     return builder_.method_name args;                        \
   }
 
-#define _FORWARD_UNOP(method_name)             \
-  _FORWARD(method_name, ComputationDataHandle, \
-           (const ComputationDataHandle& operand), (operand))
-
-#define _FORWARD_BINOP(method_name)                                        \
-  _FORWARD(                                                                \
-      method_name, ComputationDataHandle,                                  \
-      (const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, \
-       tensorflow::gtl::ArraySlice<int64> broadcast_dimensions),           \
-      (lhs, rhs, broadcast_dimensions))
-
-#define _FORWARD_TRIOP(method_name)                                        \
-  _FORWARD(                                                                \
-      method_name, ComputationDataHandle,                                  \
-      (const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, \
-       const ComputationDataHandle& ehs),                                  \
-      (lhs, rhs, ehs))
+#define _FORWARD_UNOP(method_name) \
+  _FORWARD(method_name, LocalOp, (const LocalOp& operand), (operand.op()))
+
+#define _FORWARD_BINOP(method_name)                                   \
+  _FORWARD(method_name, LocalOp,                                      \
+           (const LocalOp& lhs, const LocalOp& rhs,                   \
+            tensorflow::gtl::ArraySlice<int64> broadcast_dimensions), \
+           (lhs.op(), rhs.op(), broadcast_dimensions))
+
+#define _FORWARD_TRIOP(method_name)                                      \
+  _FORWARD(method_name, LocalOp,                                         \
+           (const LocalOp& lhs, const LocalOp& rhs, const LocalOp& ehs), \
+           (lhs.op(), rhs.op(), ehs.op()))
 
 _FORWARD_TRIOP(Select)
 _FORWARD_TRIOP(Clamp)
diff --git a/tensorflow/compiler/xla/python/local_computation_builder.h b/tensorflow/compiler/xla/python/local_computation_builder.h
index 5ec097846a..a06b85b4ea 100644
--- a/tensorflow/compiler/xla/python/local_computation_builder.h
+++ b/tensorflow/compiler/xla/python/local_computation_builder.h
@@ -17,9 +17,10 @@ limitations under the License.
 #define TENSORFLOW_COMPILER_XLA_PYTHON_LOCAL_COMPUTATION_BUILDER_H_
 
 #include "tensorflow/compiler/xla/client/client_library.h"
-#include "tensorflow/compiler/xla/client/computation_builder.h"
 #include "tensorflow/compiler/xla/client/executable_build_options.h"
 #include "tensorflow/compiler/xla/client/local_client.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_builder.h"
+#include "tensorflow/compiler/xla/client/xla_client/xla_computation.h"
 #include "tensorflow/compiler/xla/service/shaped_buffer.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
 #include "tensorflow/core/lib/gtl/array_slice.h"
@@ -97,25 +98,37 @@ class CompiledLocalComputation {
   std::unique_ptr<LocalExecutable> executable_;
 };
 
-// Wraps a Computation produced by a LocalComputationBuilder. The
+// Wraps a XlaComputation produced by a LocalComputationBuilder. The
 // Compile method compiles the computation to a (local) executable via
 // the client library's local client. This class is intended to be
 // made available to Python via SWIG.
 class LocalComputation {
  public:
-  LocalComputation(Computation computation);
+  LocalComputation(XlaComputation computation);
 
   StatusOr<CompiledLocalComputation*> Compile(
       const std::vector<Shape>& argument_shapes,
       const ExecutableBuildOptions* build_options);
 
-  const Computation& computation() const;
+  const XlaComputation& computation() const;
 
   // Returns the return-value shape for this computation.
   StatusOr<Shape> GetReturnValueShape() const;
 
  private:
-  Computation computation_;
+  XlaComputation computation_;
+};
+
+// Wraps a XlaOp produced by a LocalComputationBuilder. This class is intended
+// to be made available to Python via SWIG.
+class LocalOp {
+ public:
+  LocalOp(const XlaOp& op);
+
+  const XlaOp& op() const;
+
+ private:
+  XlaOp op_;
 };
 
 // Wraps the ComputationBuilder API in order to:
@@ -135,166 +148,137 @@ class LocalComputationBuilder {
   // Returns an owned LocalComputation to the caller on success.
   StatusOr<LocalComputation*> Build();
 
-  ComputationDataHandle Parameter(int64 parameter_number, const Shape& shape,
-                                  const string& name);
+  LocalOp Parameter(int64 parameter_number, const Shape& shape,
+                    const string& name);
 
-  std::unique_ptr<Shape> GetShape(const ComputationDataHandle& operand);
+  std::unique_ptr<Shape> GetShape(const LocalOp& operand);
 
   // Returns the shape of the current return value for the computation.
   StatusOr<Shape> GetReturnValueShape();
 
-  ComputationDataHandle Infeed(const Shape& shape);
+  LocalOp Infeed(const Shape& shape);
 
-  void Outfeed(const ComputationDataHandle& operand, const Shape& shape,
+  void Outfeed(const LocalOp& operand, const Shape& shape,
                const string& outfeed_config);
 
-  ComputationDataHandle ConstantLiteral(const Literal& literal);
+  LocalOp ConstantLiteral(const Literal& literal);
 
-  ComputationDataHandle Broadcast(
-      const ComputationDataHandle& operand,
-      tensorflow::gtl::ArraySlice<int64> broadcast_sizes);
+  LocalOp Broadcast(const LocalOp& operand,
+                    tensorflow::gtl::ArraySlice<int64> broadcast_sizes);
 
-  ComputationDataHandle Pad(const ComputationDataHandle& operand,
-                            const ComputationDataHandle& padding_value,
-                            const PaddingConfig& padding_config);
+  LocalOp Pad(const LocalOp& operand, const LocalOp& padding_value,
+              const PaddingConfig& padding_config);
 
-  ComputationDataHandle Reshape(const ComputationDataHandle& operand,
-                                tensorflow::gtl::ArraySlice<int64> dimensions,
-                                tensorflow::gtl::ArraySlice<int64> new_sizes);
+  LocalOp Reshape(const LocalOp& operand,
+                  tensorflow::gtl::ArraySlice<int64> dimensions,
+                  tensorflow::gtl::ArraySlice<int64> new_sizes);
 
-  ComputationDataHandle Collapse(const ComputationDataHandle& operand,
-                                 tensorflow::gtl::ArraySlice<int64> dimensions);
+  LocalOp Collapse(const LocalOp& operand,
+                   tensorflow::gtl::ArraySlice<int64> dimensions);
 
-  ComputationDataHandle CrossReplicaSum(const ComputationDataHandle& operand);
+  LocalOp CrossReplicaSum(const LocalOp& operand);
 
-  ComputationDataHandle Slice(const ComputationDataHandle& operand,
-                              tensorflow::gtl::ArraySlice<int64> start_indices,
-                              tensorflow::gtl::ArraySlice<int64> limit_indices,
-                              tensorflow::gtl::ArraySlice<int64> strides);
+  LocalOp Slice(const LocalOp& operand,
+                tensorflow::gtl::ArraySlice<int64> start_indices,
+                tensorflow::gtl::ArraySlice<int64> limit_indices,
+                tensorflow::gtl::ArraySlice<int64> strides);
 
-  ComputationDataHandle SliceInDim(const ComputationDataHandle& operand,
-                                   int64 start_index, int64 limit_index,
-                                   int64 stride, int64 dimno);
+  LocalOp SliceInDim(const LocalOp& operand, int64 start_index,
+                     int64 limit_index, int64 stride, int64 dimno);
 
-  ComputationDataHandle DynamicSlice(
-      const ComputationDataHandle& operand,
-      const ComputationDataHandle& start_indices,
-      tensorflow::gtl::ArraySlice<int64> slice_sizes);
+  LocalOp DynamicSlice(const LocalOp& operand, const LocalOp& start_indices,
+                       tensorflow::gtl::ArraySlice<int64> slice_sizes);
 
-  ComputationDataHandle DynamicUpdateSlice(
-      const ComputationDataHandle& operand, const ComputationDataHandle& update,
-      const ComputationDataHandle& start_indices);
+  LocalOp DynamicUpdateSlice(const LocalOp& operand, const LocalOp& update,
+                             const LocalOp& start_indices);
 
-  ComputationDataHandle ConcatInDim(
-      tensorflow::gtl::ArraySlice<ComputationDataHandle> operands,
-      int64 dimension);
+  LocalOp ConcatInDim(tensorflow::gtl::ArraySlice<LocalOp> operands,
+                      int64 dimension);
 
-  ComputationDataHandle SelectAndScatterWithGeneralPadding(
-      const ComputationDataHandle& operand, const LocalComputation& select,
+  LocalOp SelectAndScatterWithGeneralPadding(
+      const LocalOp& operand, const LocalComputation& select,
       tensorflow::gtl::ArraySlice<int64> window_dimensions,
       tensorflow::gtl::ArraySlice<int64> window_strides,
       tensorflow::gtl::ArraySlice<std::pair<int64, int64> > padding,
-      const ComputationDataHandle& source,
-      const ComputationDataHandle& init_value, const LocalComputation& scatter);
+      const LocalOp& source, const LocalOp& init_value,
+      const LocalComputation& scatter);
 
-  ComputationDataHandle Tuple(
-      tensorflow::gtl::ArraySlice<ComputationDataHandle> elements);
+  LocalOp Tuple(tensorflow::gtl::ArraySlice<LocalOp> elements);
 
-  ComputationDataHandle GetTupleElement(const ComputationDataHandle& tuple_data,
-                                        int64 index);
+  LocalOp GetTupleElement(const LocalOp& tuple_data, int64 index);
 
-  ComputationDataHandle Dot(const ComputationDataHandle& lhs,
-                            const ComputationDataHandle& rhs);
+  LocalOp Dot(const LocalOp& lhs, const LocalOp& rhs);
 
-  ComputationDataHandle DotGeneral(
-      const ComputationDataHandle& lhs, const ComputationDataHandle& rhs,
-      const DotDimensionNumbers& dimension_numbers);
+  LocalOp DotGeneral(const LocalOp& lhs, const LocalOp& rhs,
+                     const DotDimensionNumbers& dimension_numbers);
 
-  ComputationDataHandle ConvGeneralDilated(
-      const ComputationDataHandle& lhs, const ComputationDataHandle& rhs,
+  LocalOp ConvGeneralDilated(
+      const LocalOp& lhs, const LocalOp& rhs,
       tensorflow::gtl::ArraySlice<int64> window_strides,
       tensorflow::gtl::ArraySlice<std::pair<int64, int64> > padding,
       tensorflow::gtl::ArraySlice<int64> lhs_dilation,
       tensorflow::gtl::ArraySlice<int64> rhs_dilation,
       const ConvolutionDimensionNumbers& dimension_numbers);
 
-  ComputationDataHandle ConvertElementType(const ComputationDataHandle& operand,
-                                           PrimitiveType new_element_type);
+  LocalOp ConvertElementType(const LocalOp& operand,
+                             PrimitiveType new_element_type);
 
-  ComputationDataHandle Call(
-      const LocalComputation& local_computation,
-      tensorflow::gtl::ArraySlice<ComputationDataHandle> operands);
+  LocalOp Call(const LocalComputation& local_computation,
+               tensorflow::gtl::ArraySlice<LocalOp> operands);
 
-  ComputationDataHandle Transpose(
-      const ComputationDataHandle& operand,
-      tensorflow::gtl::ArraySlice<int64> permutation);
+  LocalOp Transpose(const LocalOp& operand,
+                    tensorflow::gtl::ArraySlice<int64> permutation);
 
-  ComputationDataHandle Rev(const ComputationDataHandle& operand,
-                            tensorflow::gtl::ArraySlice<int64> dimensions);
+  LocalOp Rev(const LocalOp& operand,
+              tensorflow::gtl::ArraySlice<int64> dimensions);
 
-  ComputationDataHandle Map(
-      tensorflow::gtl::ArraySlice<ComputationDataHandle> operands,
-      const LocalComputation& local_computation,
-      tensorflow::gtl::ArraySlice<int64> dimensions,
-      tensorflow::gtl::ArraySlice<ComputationDataHandle> static_operands);
+  LocalOp Map(tensorflow::gtl::ArraySlice<LocalOp> operands,
+              const LocalComputation& local_computation,
+              tensorflow::gtl::ArraySlice<int64> dimensions,
+              tensorflow::gtl::ArraySlice<LocalOp> static_operands);
 
-  ComputationDataHandle Reduce(
-      const ComputationDataHandle& operand,
-      const ComputationDataHandle& init_value,
-      const LocalComputation& local_computation,
-      tensorflow::gtl::ArraySlice<int64> dimensions_to_reduce);
+  LocalOp Reduce(const LocalOp& operand, const LocalOp& init_value,
+                 const LocalComputation& local_computation,
+                 tensorflow::gtl::ArraySlice<int64> dimensions_to_reduce);
 
-  ComputationDataHandle ReduceWindowWithGeneralPadding(
-      const ComputationDataHandle& operand,
-      const ComputationDataHandle& init_value,
+  LocalOp ReduceWindowWithGeneralPadding(
+      const LocalOp& operand, const LocalOp& init_value,
       const LocalComputation& local_computation,
       tensorflow::gtl::ArraySlice<int64> window_dimensions,
       tensorflow::gtl::ArraySlice<int64> window_strides,
       tensorflow::gtl::ArraySlice<std::pair<int64, int64> > padding);
 
-  ComputationDataHandle RngNormal(const ComputationDataHandle& mu,
-                                  const ComputationDataHandle& sigma,
-                                  const Shape& shape);
+  LocalOp RngNormal(const LocalOp& mu, const LocalOp& sigma,
+                    const Shape& shape);
 
-  ComputationDataHandle RngUniform(const ComputationDataHandle& a,
-                                   const ComputationDataHandle& b,
-                                   const Shape& shape);
+  LocalOp RngUniform(const LocalOp& a, const LocalOp& b, const Shape& shape);
 
-  ComputationDataHandle While(const LocalComputation& condition,
-                              const LocalComputation& body,
-                              const ComputationDataHandle& init);
+  LocalOp While(const LocalComputation& condition, const LocalComputation& body,
+                const LocalOp& init);
 
-  ComputationDataHandle Conditional(const ComputationDataHandle& predicate,
-                                    const ComputationDataHandle& true_operand,
-                                    const LocalComputation& true_computation,
-                                    const ComputationDataHandle& false_operand,
-                                    const LocalComputation& false_computation);
+  LocalOp Conditional(const LocalOp& predicate, const LocalOp& true_operand,
+                      const LocalComputation& true_computation,
+                      const LocalOp& false_operand,
+                      const LocalComputation& false_computation);
 
-  StatusOr<bool> IsConstant(const ComputationDataHandle& operand,
-                            int64 num_parameters);
+  StatusOr<bool> IsConstant(const LocalOp& operand);
 
-  StatusOr<std::unique_ptr<Literal> > ComputeConstant(
-      const ComputationDataHandle& operand, const Layout* output_layout,
-      tensorflow::gtl::ArraySlice<Literal> parameters);
+  StatusOr<LocalComputation*> BuildConstantSubGraph(const LocalOp& operand);
 
 #define _FORWARD(method_name, return_sig, args_sig) \
   return_sig method_name args_sig;
 
-#define _FORWARD_UNOP(method_name)             \
-  _FORWARD(method_name, ComputationDataHandle, \
-           (const ComputationDataHandle& operand))
+#define _FORWARD_UNOP(method_name) \
+  _FORWARD(method_name, LocalOp, (const LocalOp& operand))
 
-#define _FORWARD_BINOP(method_name)                                        \
-  _FORWARD(                                                                \
-      method_name, ComputationDataHandle,                                  \
-      (const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, \
-       tensorflow::gtl::ArraySlice<int64> broadcast_dimensions))
+#define _FORWARD_BINOP(method_name)                 \
+  _FORWARD(method_name, LocalOp,                    \
+           (const LocalOp& lhs, const LocalOp& rhs, \
+            tensorflow::gtl::ArraySlice<int64> broadcast_dimensions))
 
-#define _FORWARD_TRIOP(method_name)                                        \
-  _FORWARD(                                                                \
-      method_name, ComputationDataHandle,                                  \
-      (const ComputationDataHandle& lhs, const ComputationDataHandle& rhs, \
-       const ComputationDataHandle& ehs))
+#define _FORWARD_TRIOP(method_name) \
+  _FORWARD(method_name, LocalOp,    \
+           (const LocalOp& lhs, const LocalOp& rhs, const LocalOp& ehs))
 
   _FORWARD_TRIOP(Select)
   _FORWARD_TRIOP(Clamp)
@@ -338,7 +322,7 @@ class LocalComputationBuilder {
 #undef _FORWARD_TRIOP
 
  private:
-  ComputationBuilder builder_;
+  XlaBuilder builder_;
 };
 
 // Functions for freeing resources from the Python side.
diff --git a/tensorflow/compiler/xla/python/local_computation_builder.i b/tensorflow/compiler/xla/python/local_computation_builder.i
index b8cce5a5f7..04c56bbba9 100644
--- a/tensorflow/compiler/xla/python/local_computation_builder.i
+++ b/tensorflow/compiler/xla/python/local_computation_builder.i
@@ -22,9 +22,8 @@ limitations under the License.
 //
 //    C++                                  Python
 // -------------------------------------+---------------------------------------
-//  ComputationDataHandle              <-> int
 //  ArraySlice<int64>                  <-  sequence of int
-//  ArraySlice<ComputationDataHandle>  <-  sequence of int
+//  ArraySlice<LocalOp>                <-  sequence of LocalOp
 //  Literal                            <-> (nested tuple of) numpy ndarray
 //  std::vector<Literal>               <-  sequence of (nested tuple of) ndarray
 //  Shape                               -> pair holding (dtype, dimensions)
@@ -91,12 +90,9 @@ limitations under the License.
 // One central reason for the Python-side indirection is that the
 // Python-side objects produced by the typemaps in this file are
 // further packaged up by xla_client before being passed on. For
-// instance, xla_client wraps the long produced for a C++
-// ComputationDataHandle in a Python ComputationDataHandle proto,
-// rather than exposing a raw long outside of the client. Similarly,
-// the Python pair produced for a C++ Shape is further wrapped in a
-// Python class (xla_client.Shape) so as not to expose the raw pair
-// externally.
+// instance, the Python pair produced for a C++ Shape is further
+// wrapped in a Python class (xla_client.Shape) so as not to expose
+// the raw pair externally.
 //
 // Other SWIG object wrappers (e.g. of LocalComputation) are further
 // wrapped by xla_client in order to set up a custom destructor that
@@ -124,6 +120,7 @@ using namespace xla;
 using namespace xla::swig;
 
 namespace xla {
+
 namespace swig {
 
 bool GetIntAttr(PyObject* o, const char* field, int64* result) {
@@ -177,21 +174,6 @@ bool HandleStringAttribute(PyObject* o,
 tensorflow::ImportNumpy();
 %}
 
-// ComputationDataHandle
-
-%typemap(in) const ComputationDataHandle& (ComputationDataHandle temp) {
-  const int64 handle = numpy::PyIntOrPyLongToLong($input);
-  if (handle == -1 && PyErr_Occurred()) {
-    SWIG_fail;
-  }
-  temp.set_handle(handle);
-  $1 = &temp;
-}
-
-%typemap(out) ComputationDataHandle {
-  $result = numpy::LongToPyIntOrPyLong($1.handle());
-}
-
 %typemap(out) StatusOr<xla::swig::CompiledLocalComputation*> {
   if ($1.ok()) {
     auto* value = $1.ValueOrDie();
@@ -301,33 +283,23 @@ tensorflow::ImportNumpy();
   $1 = temps;
 }
 
-// ComputationDataHandle
+// ArraySlice<LocalOp>
 
-%typemap(in) tensorflow::gtl::ArraySlice<ComputationDataHandle>
-    (std::vector<ComputationDataHandle> temps) {
+%typemap(in) tensorflow::gtl::ArraySlice<xla::swig::LocalOp>(
+      std::vector<LocalOp> temps) {
   if (!PySequence_Check($input)) {
     PyErr_SetString(PyExc_TypeError, "Argument is not a sequence");
     SWIG_fail;
   }
   const int size = PySequence_Size($input);
-  temps.resize(size);
   for (int i = 0; i < size; ++i) {
     PyObject* o = PySequence_GetItem($input, i);
-    PyObject* py_int = numpy::PyNumberToPyInt(o);
-    if (!py_int) {
-      PyErr_SetString(
-          PyExc_TypeError,
-          "Argument sequence element cannot be converted to int");
-      SWIG_fail;
-    }
-    const int64 handle = numpy::PyIntOrPyLongToLong(py_int);
-    if (handle == -1 && PyErr_Occurred()) {
-      Py_DECREF(py_int);
-      Py_DECREF(o);
+    LocalOp* op;
+    if ((SWIG_ConvertPtr(o, (void**)&op, $descriptor(xla::swig::LocalOp*),
+                         SWIG_POINTER_EXCEPTION)) == -1) {
       SWIG_fail;
     }
-    temps[i].set_handle(handle);
-    Py_DECREF(py_int);
+    temps.push_back(*op);
     Py_DECREF(o);
   }
   $1 = temps;
@@ -934,6 +906,7 @@ tensorflow::ImportNumpy();
 %unignore xla::swig::LocalComputation;
 %unignore xla::swig::LocalComputation::Compile;
 %unignore xla::swig::LocalComputation::GetReturnValueShape;
+%unignore xla::swig::LocalOp;
 %unignore xla::swig::LocalComputationBuilder;
 %unignore xla::swig::LocalComputationBuilder::LocalComputationBuilder;
 %unignore xla::swig::LocalComputationBuilder::Build;
diff --git a/tensorflow/compiler/xla/python/xla_client.py b/tensorflow/compiler/xla/python/xla_client.py
index f6809b6b87..1d5b75d1be 100644
--- a/tensorflow/compiler/xla/python/xla_client.py
+++ b/tensorflow/compiler/xla/python/xla_client.py
@@ -335,20 +335,6 @@ def _wrap_shape(shape_info):
     return Shape.array_shape(dtype, dims)
 
 
-def _wrap_data_handle(handle):
-  cdh = xla_data_pb2.ComputationDataHandle()
-  cdh.handle = handle
-  return cdh
-
-
-def _unwrap_data_handle(handle_proto):
-  return handle_proto.handle
-
-
-def _unwrap_data_handles(handle_protos):
-  return [_unwrap_data_handle(cdh) for cdh in handle_protos]
-
-
 def require_numpy_array_layout(value):
   if isinstance(value, tuple):
     return tuple(require_numpy_array_layout(x) for x in value)
@@ -535,9 +521,9 @@ class ComputationBuilder(object):
     queue for subsequent use in the computation.
 
     Returns:
-      A  ComputationDataHandle message.
+      A LocalOp.
     """
-    return _wrap_data_handle(self._client.Infeed(shape))
+    return self._client.Infeed(shape)
 
   def Outfeed(self, operand):
     """Enqueues an outfeed op onto the computation.
@@ -545,9 +531,7 @@ class ComputationBuilder(object):
     Outfeed operations enqueue data, using the given operand, onto the XLA
     outfeed queue for subsequent dequeue via the client API.
     """
-    self._client.Outfeed(
-        _unwrap_data_handle(operand), self.GetShape(operand),
-        ''.encode('utf-8'))
+    self._client.Outfeed(operand, self.GetShape(operand), ''.encode('utf-8'))
 
   def Constant(self, value):
     """Enqueues a constant op onto the computation.
@@ -557,10 +541,10 @@ class ComputationBuilder(object):
              to one of the supported types.
 
     Returns:
-      A ComputationDataHandle message.
+      A LocalOp.
     """
     value = require_numpy_array_layout(value)
-    return _wrap_data_handle(self._client.ConstantLiteral(value))
+    return self._client.ConstantLiteral(value)
 
   def ConstantF32Scalar(self, value):
     """Convenience method to enqueue a scalar F32 constant op.
@@ -569,7 +553,7 @@ class ComputationBuilder(object):
       value: a floating-point number.
 
     Returns:
-      A ComputationDataHandle message.
+      A LocalOp.
     """
     return self.Constant(np.array(value, dtype=np.float32))
 
@@ -580,7 +564,7 @@ class ComputationBuilder(object):
       value: a floating-point number.
 
     Returns:
-      A ComputationDataHandle message.
+      A LocalOp.
     """
     return self.Constant(np.array(value, dtype=np.float64))
 
@@ -591,7 +575,7 @@ class ComputationBuilder(object):
       value: a floating-point number.
 
     Returns:
-      A ComputationDataHandle message.
+      A LocalOp.
     """
     return self.Constant(np.array(value, dtype=np.int32))
 
@@ -602,7 +586,7 @@ class ComputationBuilder(object):
       value: a floating-point number.
 
     Returns:
-      A ComputationDataHandle message.
+      A LocalOp.
     """
     return self.Constant(np.array(value, dtype=np.int64))
 
@@ -613,7 +597,7 @@ class ComputationBuilder(object):
       value: a boolean value.
 
     Returns:
-      A ComputationDataHandle message.
+      A LocalOp.
     """
     return self.Constant(np.array(value, dtype=np.bool))
 
@@ -629,15 +613,14 @@ class ComputationBuilder(object):
         parameters, use it for *all* parameters to avoid clashes.
 
     Returns:
-      A ComputationDataHandle message.
+      A LocalOp.
     """
     if name is None:
       name = ''
     if parameter_num is None:
       parameter_num = next(self._parameter_numbering)
 
-    return _wrap_data_handle(
-        self._client.Parameter(parameter_num, shape, name.encode('utf8')))
+    return self._client.Parameter(parameter_num, shape, name.encode('utf8'))
 
   def ParameterFromNumpy(self, value, name=None, parameter_num=None):
     """Enqueues a Parameter op onto the computation.
@@ -649,7 +632,7 @@ class ComputationBuilder(object):
       parameter_num: as in ParameterWithShape.
 
     Returns:
-      A ComputationDataHandle message.
+      A LocalOp.
     """
     return self.ParameterWithShape(
         Shape.from_pyval(value), name=name, parameter_num=parameter_num)
@@ -658,14 +641,13 @@ class ComputationBuilder(object):
     """Enqueues a broadcast operation onto the computation.
 
     Args:
-      operand: the operand ComputationDataHandle to broadcast.
+      operand: the operand LocalOp to broadcast.
       sizes: an iterable of broadcast sizes.
 
     Returns:
-      A ComputationDataHandle representing the added broadcast op.
+      A LocalOp representing the added broadcast op.
     """
-    return _wrap_data_handle(
-        self._client.Broadcast(_unwrap_data_handle(operand), sizes))
+    return self._client.Broadcast(operand, sizes)
 
   def Concatenate(self, operands, dimension):
     """Enqueues a concatenate operation onto the computation.
@@ -675,10 +657,9 @@ class ComputationBuilder(object):
       dimension: the dimension in which to perform the concatenation.
 
     Returns:
-      A ComputationDataHandle representing the added concatenate op.
+      A LocalOp representing the added concatenate op.
     """
-    return _wrap_data_handle(
-        self._client.ConcatInDim(_unwrap_data_handles(operands), dimension))
+    return self._client.ConcatInDim(operands, dimension)
 
   def ConvertElementType(self, operand, new_element_type):
     """Enqueues an element type conversion operation onto the computation.
@@ -688,14 +669,12 @@ class ComputationBuilder(object):
       new_element_type: the target primitive type.
 
     Returns:
-      A ComputationDataHandle representing the added conversion op.
+      A LocalOp representing the added conversion op.
     """
-    return _wrap_data_handle(
-        self._client.ConvertElementType(
-            _unwrap_data_handle(operand), new_element_type))
+    return self._client.ConvertElementType(operand, new_element_type)
 
   def GetShape(self, operand):
-    return _wrap_shape(self._client.GetShape(_unwrap_data_handle(operand)))
+    return _wrap_shape(self._client.GetShape(operand))
 
   def GetReturnValueShape(self):
     return _wrap_shape(self._client.GetReturnValueShape())
@@ -707,40 +686,35 @@ class ComputationBuilder(object):
     """Enqueues a Pad operation onto the computation.
 
     Args:
-      operand: ComputationDataHandle representing the array to pad.
-      padding_value: ComputationDataHandle representing the scalar pad value.
+      operand: LocalOp representing the array to pad.
+      padding_value: LocalOp representing the scalar pad value.
       padding_config: either an xla_data_pb2.PaddingConfig or a list of integer
         triples (edge_padding_low, edge_padding_high, interior_padding)
         representing the configuration of the padding operation.
 
     Returns:
-      A ComputationDataHandle representing the added Pad op.
+      A LocalOp representing the added Pad op.
     """
     if not isinstance(padding_config, xla_data_pb2.PaddingConfig):
       padding_config = GetPaddingConfigFromTriples(padding_config)
-    return _wrap_data_handle(
-        self._client.Pad(_unwrap_data_handle(operand),
-                         _unwrap_data_handle(padding_value),
-                         padding_config))
+    return self._client.Pad(operand, padding_value, padding_config)
 
   def Reshape(self, operand, dimensions, new_sizes):
     """Enqueues a reshape op onto the computation.
 
     Args:
-      operand: ComputationDataHandle representing the array to be reshaped.
+      operand: LocalOp representing the array to be reshaped.
       dimensions: sequence of integers encoding the order in which dimensions
         are collapsed or None, in which case dimensions are flattened in order.
       new_sizes: sequence of integers encoding the new dimension sizes (shape).
 
     Returns:
-      A ComputationDataHandle representing the added Reshape op.
+      A LocalOp representing the added Reshape op.
     """
     if dimensions is None:
       ndim = len(self.GetShape(operand).dimensions())
       dimensions = tuple(range(ndim))
-    return _wrap_data_handle(
-        self._client.Reshape(
-            _unwrap_data_handle(operand), dimensions, new_sizes))
+    return self._client.Reshape(operand, dimensions, new_sizes)
 
   def CrossReplicaSum(self, operand):
     """CrossReplicaSum op.
@@ -749,67 +723,56 @@ class ComputationBuilder(object):
       operand: the operand to sum across replica instances.
 
     Returns:
-      A ComputationDataHandle that has the sum of the value among all replicas.
+      A LocalOp that has the sum of the value among all replicas.
     """
-    return _wrap_data_handle(
-        self._client.CrossReplicaSum(_unwrap_data_handle(operand)))
+    return self._client.CrossReplicaSum(operand)
 
   def Collapse(self, operand, dimensions):
     """Collapse op."""
-    return _wrap_data_handle(
-        self._client.Collapse(_unwrap_data_handle(operand), dimensions))
+    return self._client.Collapse(operand, dimensions)
 
   def Trans(self, operand):
     """Specialized matrix transpose op."""
-    return _wrap_data_handle(
-        self._client.Transpose(_unwrap_data_handle(operand), [1, 0]))
+    return self._client.Transpose(operand, [1, 0])
 
   def Transpose(self, operand, permutation):
     """Transpose op."""
-    return _wrap_data_handle(
-        self._client.Transpose(_unwrap_data_handle(operand), permutation))
+    return self._client.Transpose(operand, permutation)
 
   def Rev(self, operand, dimensions):
     """Rev op."""
-    return _wrap_data_handle(
-        self._client.Rev(_unwrap_data_handle(operand), dimensions))
+    return self._client.Rev(operand, dimensions)
 
   def Clamp(self, min, operand, max):  # pylint: disable=redefined-builtin
     """Clamp op."""
-    return _wrap_data_handle(
-        self._client.Clamp(_unwrap_data_handle(min),
-                           _unwrap_data_handle(operand),
-                           _unwrap_data_handle(max)))
+    return self._client.Clamp(min, operand, max)
 
   def SelectAndScatter(self, operand, select, window_dimensions, window_strides,
                        padding, source, init_value, scatter):
     """Select and scatter op, used by the gradient of ReduceWindow.
 
     Args:
-      operand: ComputationDataHandle for array of dimension N and type T over
+      operand: LocalOp for array of dimension N and type T over
         which the windows slide.
       select: Computation of type (T, T) -> Pred to apply to the elements of
         each window to indicate which element is selected.
       window_dimensions: sequence of N integers for dimensions of the window.
       window_strides: sequence of N integers for the strides of the window.
       padding: PaddingType representing either 'SAME' or 'VALID ' padding.
-      source: ComputationDataHandle for array of type T with values to scatter.
-      init_value: ComputationDataHandle of scalar type T for initial out value.
+      source: LocalOp for array of type T with values to scatter.
+      init_value: LocalOp of scalar type T for initial out value.
       scatter: Computation of type (T, T) -> T to apply to each scatter source
         element with its destination element.
 
     Returns:
-      A ComputationDataHandle representing the added SelectAndScatter op.
+      A LocalOp representing the added SelectAndScatter op.
     """
     pads = _convert_padding_type_to_pad_values(
         padding, self.GetShape(operand).dimensions(),
         window_dimensions, window_strides)
-    return _wrap_data_handle(
-        self._client.SelectAndScatterWithGeneralPadding(
-            _unwrap_data_handle(operand), select.c_local_computation,
-            window_dimensions, window_strides, pads,
-            _unwrap_data_handle(source), _unwrap_data_handle(init_value),
-            scatter.c_local_computation))
+    return self._client.SelectAndScatterWithGeneralPadding(
+        operand, select.c_local_computation, window_dimensions, window_strides,
+        pads, source, init_value, scatter.c_local_computation)
 
   def Select(self, pred, on_true, on_false):
     """Element-wise selection op.
@@ -817,17 +780,13 @@ class ComputationBuilder(object):
     Constructs an output array from elements of two input arrays, based on the
     values of a predicate array.
     """
-    return _wrap_data_handle(
-        self._client.Select(
-            _unwrap_data_handle(pred),
-            _unwrap_data_handle(on_true),
-            _unwrap_data_handle(on_false)))
+    return self._client.Select(pred, on_true, on_false)
 
   def Slice(self, operand, start_indices, limit_indices, strides=None):
     """Enqueues a slice operation onto the computation.
 
     Args:
-      operand: ComputationDataHandle for the N dimensional array to be sliced.
+      operand: LocalOp for the N dimensional array to be sliced.
       start_indices: iterable of N integers containing the starting indices of
         the slice for each dimension.
       limit_indices: iterable of N integers containing the ending indices
@@ -836,207 +795,177 @@ class ComputationBuilder(object):
         each dimension.
 
     Returns:
-      A ComputationDataHandle representing the added Slice op.
+      A LocalOp representing the added Slice op.
     """
     if strides is None:
       start_indices = list(start_indices)
       strides = [1] * len(start_indices)
-    return _wrap_data_handle(
-        self._client.Slice(
-            _unwrap_data_handle(operand), start_indices, limit_indices,
-            strides))
+    return self._client.Slice(operand, start_indices, limit_indices, strides)
 
   def SliceInDim(self, operand, start_index, limit_index, stride, dimno):
     """Enqueues a slice-in-dimension operation onto the computation.
 
     Args:
-      operand: ComputationDataHandle for the N dimensional array to be sliced.
+      operand: LocalOp for the N dimensional array to be sliced.
       start_index: an integer containing the start index of the slice.
       limit_index: an integer containing the end index of the slice.
       stride: an integer containing the stride size for the slice.
       dimno: an integer indicating the dimension along which to slice.
 
     Returns:
-      A ComputationDataHandle representing the added Slice op.
+      A LocalOp representing the added Slice op.
     """
-    return _wrap_data_handle(
-        self._client.SliceInDim(
-            _unwrap_data_handle(operand), start_index, limit_index, stride,
-            dimno))
+    return self._client.SliceInDim(operand, start_index, limit_index, stride,
+                                   dimno)
 
   def DynamicSlice(self, operand, start_indices, slice_sizes):
     """Enqueues a slice op with dynamic start indices onto the computation.
 
     Args:
-      operand: ComputationDataHandle for the N dimensional array to be sliced.
-      start_indices: ComputationDataHandle for the 1D array of N integers
+      operand: LocalOp for the N dimensional array to be sliced.
+      start_indices: LocalOp for the 1D array of N integers
         containing the starting indices of the slice.
       slice_sizes: iterable of N integers containing the slice sizes in each
         dimension.
 
     Returns:
-      A ComputationDataHandle representing the added DynamicSlice op.
+      A LocalOp representing the added DynamicSlice op.
     """
-    return _wrap_data_handle(
-        self._client.DynamicSlice(
-            _unwrap_data_handle(operand),
-            _unwrap_data_handle(start_indices),
-            slice_sizes))
+    return self._client.DynamicSlice(operand, start_indices, slice_sizes)
 
   def DynamicUpdateSlice(self, operand, update, start_indices):
     """Enqueues a dynamic update slice operation onto the computation.
 
     Args:
-      operand: ComputationDataHandle for the N dimensional array to be updated.
+      operand: LocalOp for the N dimensional array to be updated.
       update: N dimensional array comprising the slice update.
       start_indices: Rank-1 array of N integers comprising the starting indices
         of the slice along each dimension.
     Returns:
-      A ComputationDataHandle representing the added DynamicUpdateSlice op.
+      A LocalOp representing the added DynamicUpdateSlice op.
     """
-    return _wrap_data_handle(
-        self._client.DynamicUpdateSlice(
-            _unwrap_data_handle(operand),
-            _unwrap_data_handle(update),
-            _unwrap_data_handle(start_indices)))
+    return self._client.DynamicUpdateSlice(operand, update, start_indices)
 
   def Tuple(self, *ops):
     """Enqueues a tuple operation onto the computation.
 
     Args:
-      ops: a sequence of tuple operands (each a ComputationDataHandle).
+      ops: a sequence of tuple operands (each a LocalOp).
 
     Returns:
-      A ComputationDataHandle representing the added Tuple op.
+      A LocalOp representing the added Tuple op.
     """
-    return _wrap_data_handle(self._client.Tuple(_unwrap_data_handles(ops)))
+    return self._client.Tuple(ops)
 
   def GetTupleElement(self, tup, index):
     """Enqueues a 'get tuple element' operation onto the computation.
 
     Args:
-      tup: the tuple operand (a ComputationDataHandle).
+      tup: the tuple operand (a LocalOp).
       index: numeric index to select from the tuple.
 
     Returns:
-      A ComputationDataHandle representing the added GetTupleElement op.
+      A LocalOp representing the added GetTupleElement op.
     """
-    return _wrap_data_handle(
-        self._client.GetTupleElement(_unwrap_data_handle(tup), index))
+    return self._client.GetTupleElement(tup, index)
 
   def Call(self, computation_to_apply, operands):
     """Enqueues a call operation onto the computation.
 
     Args:
       computation_to_apply: a Computation object.
-      operands: an iterable of ComputationDataHandle. The number and types of
+      operands: an iterable of LocalOp. The number and types of
         operands must match the arity of computation_to_apply.
 
     Returns:
-      A ComputationDataHandle representing the added call op.
+      A LocalOp representing the added call op.
     """
-    return _wrap_data_handle(
-        self._client.Call(computation_to_apply.c_local_computation,
-                          _unwrap_data_handles(operands)))
+    return self._client.Call(computation_to_apply.c_local_computation, operands)
 
   def Map(self, operands, computation_to_apply, dimensions, static_operands=()):
     """Enqueues a map operation onto the computation.
 
     Args:
-      operands: an iterable of ComputationDataHandle.
+      operands: an iterable of LocalOp.
       computation_to_apply: a Computation object.
       dimensions: dimensions over which to apply map the function.
       static_operands: auxiliary arguments passed to the applied computation.
 
     Returns:
-      A ComputationDataHandle representing the added Map op.
+      A LocalOp representing the added Map op.
     """
-    return _wrap_data_handle(
-        self._client.Map(
-            _unwrap_data_handles(operands),
-            computation_to_apply.c_local_computation,
-            dimensions,
-            _unwrap_data_handles(static_operands)))
+    return self._client.Map(operands, computation_to_apply.c_local_computation,
+                            dimensions, static_operands)
 
   def Reduce(self, operand, init_value, computation_to_apply, dimensions):
     """Enqueues a reduction operation onto the computation.
 
     Args:
-      operand: reduction operand (ComputationDataHandle).
-      init_value: reduction initial value (ComputationDataHandle).
+      operand: reduction operand (LocalOp).
+      init_value: reduction initial value (LocalOp).
       computation_to_apply: a Computation object - binary reduction function.
       dimensions: sequence of dimensions (integers) to reduce on.
 
     Returns:
-      A ComputationDataHandle representing the added Reduce op.
+      A LocalOp representing the added Reduce op.
     """
-    return _wrap_data_handle(
-        self._client.Reduce(
-            _unwrap_data_handle(operand),
-            _unwrap_data_handle(init_value),
-            computation_to_apply.c_local_computation,
-            dimensions))
+    return self._client.Reduce(operand, init_value,
+                               computation_to_apply.c_local_computation,
+                               dimensions)
 
   def ReduceWindow(self, operand, init_value, computation_to_apply,
                    window_dimensions, window_strides, padding):
     """Enqueues a windowed reduction operation onto the computation.
 
     Args:
-      operand: reduction operand (ComputationDataHandle).
-      init_value: reduction initial value (ComputationDataHandle).
+      operand: reduction operand (LocalOp).
+      init_value: reduction initial value (LocalOp).
       computation_to_apply: a binary reduction function (Computation).
       window_dimensions: dimensions of window (sequence of integers).
       window_strides: strides for window (sequence of integers).
       padding: PaddingType representing either 'SAME' or 'VALID' padding.
 
     Returns:
-      A ComputationDataHandle representing the added ReduceWindow op.
+      A LocalOp representing the added ReduceWindow op.
     """
     pads = _convert_padding_type_to_pad_values(
         padding, self.GetShape(operand).dimensions(), window_dimensions,
         window_strides)
-    return _wrap_data_handle(
-        self._client.ReduceWindowWithGeneralPadding(
-            _unwrap_data_handle(operand),
-            _unwrap_data_handle(init_value),
-            computation_to_apply.c_local_computation,
-            window_dimensions, window_strides, pads))
+    return self._client.ReduceWindowWithGeneralPadding(
+        operand, init_value, computation_to_apply.c_local_computation,
+        window_dimensions, window_strides, pads)
 
   def RngNormal(self, mu, sigma, dims):
     """Enqueues an RngNormal operation onto the computation.
 
     Args:
-      mu: A ComputationDataHandle to an F32 scalar specifying the mean.
-      sigma: A ComputationDataHandle to an F32 scalar specifying the standard
+      mu: A LocalOp to an F32 scalar specifying the mean.
+      sigma: A LocalOp to an F32 scalar specifying the standard
         deviation.
       dims: A 1D array-like of nonnegative integers specifying the dimensions.
 
-    Returns: a ComputationDataHandle to the generated array of F32 values.
+    Returns: a LocalOp to the generated array of F32 values.
     """
     shape = Shape.array_shape(self.GetShape(mu).element_type(), dims)
-    return _wrap_data_handle(
-        self._client.RngNormal(
-            _unwrap_data_handle(mu), _unwrap_data_handle(sigma), shape))
+    return self._client.RngNormal(mu, sigma, shape)
 
   def RngUniform(self, a, b, dims):
     """Enqueues an RngUniform operation onto the computation.
 
     Args:
-      a: a ComputationDataHandle to an F32, S32, or U32 scalar (consistent with
+      a: a LocalOp to an F32, S32, or U32 scalar (consistent with
         the type of b) specifying the low end of the interval [a, b) over which
         values are generated.
-      b: a ComputationDataHandle to an F32, S32, or U32 scalar (consistent with
+      b: a LocalOp to an F32, S32, or U32 scalar (consistent with
         the type of a) specifying the high end of the interval [a, b) over which
         values are generated.
       dims: A 1D array-like of nonnegative integers specifying the dimensions.
 
-    Returns: a ComputationDataHandle to the generated array of values with the
+    Returns: a LocalOp to the generated array of values with the
       same numeric type (F32, S32, or U32) as the arguments a and b.
     """
     shape = Shape.array_shape(self.GetShape(a).element_type(), dims)
-    return _wrap_data_handle(
-        self._client.RngUniform(
-            _unwrap_data_handle(a), _unwrap_data_handle(b), shape))
+    return self._client.RngUniform(a, b, shape)
 
   def While(self, cond, body, init):
     """Enqueues a While operation onto the computation.
@@ -1044,112 +973,105 @@ class ComputationBuilder(object):
     Args:
       cond: a Computation for the loop condition, which has type T -> PRED
       body: a Computation for the loop body, which has type T -> T
-      init: a ComputationDataHandle for the initial parameter, which has type T
+      init: a LocalOp for the initial parameter, which has type T
 
-    Returns: a ComputationDataHandle representing the While operation.
+    Returns: a LocalOp representing the While operation.
     """
-    return _wrap_data_handle(
-        self._client.While(cond.c_local_computation,
-                           body.c_local_computation,
-                           _unwrap_data_handle(init)))
+    return self._client.While(cond.c_local_computation,
+                              body.c_local_computation, init)
 
   def Conditional(self, pred, true_operand, true_computation, false_operand,
                   false_computation):
     """Enqueues a Conditional operation onto the computation.
 
     Args:
-      predicate: a ComputationDataHandle to test, which has scalar type PRED
-      true_operand: a ComputationDataHandle of type T_0
+      predicate: a LocalOp to test, which has scalar type PRED
+      true_operand: a LocalOp of type T_0
       true_computation: a Computation to apply to true_operand, type T_0 -> S
       false_operand: a ComputationDatahandle of type T_1
       false_computation: a Computation to apply to false_operand, type T_1 -> S
 
-    Returns: a ComputationDataHandle representing the Conditional operation.
+    Returns: a LocalOp representing the Conditional operation.
     """
-    return _wrap_data_handle(
-        self._client.Conditional(
-            _unwrap_data_handle(pred), _unwrap_data_handle(true_operand),
-            true_computation.c_local_computation,
-            _unwrap_data_handle(false_operand),
-            false_computation.c_local_computation))
+    return self._client.Conditional(
+        pred, true_operand, true_computation.c_local_computation, false_operand,
+        false_computation.c_local_computation)
 
-  def IsConstant(self, operand, num_parameters=0):
-    """Enqueues an IsConstant operation onto the computation.
+  def IsConstant(self, operand):
+    """Checks whether the given operand is a compile-time constant.
 
     Args:
       operand: a ComputationDataHandle to test.
-      num_parameters: optional int, number of computation parameters to treat as
-        constant (default 0).
 
     Returns: bool indicating whether `operand` is a compile-time constant,
-      meaning its value does not depend on parameters with index greater than or
-      equal to `num_parameters`.
+      meaning its value does not depend on any parametersor, or on stateful
+      operators such as `RngNormal` or `Infeed`.
+    """
+    return self._client.IsConstant(operand)
+
+  def BuildConstantSubGraph(self, operand):
+    """Builds a constant sub graph.
+
+    Args:
+      operand: a LocalOp to test.
+    Returns: a LocalComputation that is rooted on the given `operand` which is a
+      compile-time constant.
     """
-    return self._client.IsConstant(_unwrap_data_handle(operand), num_parameters)
+    return self._client.BuildConstantSubGraph(operand)
 
   def Dot(self, lhs, rhs):
     """Enqueues a dot operation onto the computation.
 
     Args:
-      lhs: ComputationDataHandle for the rank 1 or rank 2 left-hand-side array.
-      rhs: ComputationDataHandle for the rank 1 or rank 2 right-hand-side array.
+      lhs: LocalOp for the rank 1 or rank 2 left-hand-side array.
+      rhs: LocalOp for the rank 1 or rank 2 right-hand-side array.
 
-    Returns: a ComputationDataHandle representing the Dot operation.
+    Returns: a LocalOp representing the Dot operation.
     """
-    return _wrap_data_handle(
-        self._client.Dot(_unwrap_data_handle(lhs), _unwrap_data_handle(rhs)))
+    return self._client.Dot(lhs, rhs)
 
   def DotGeneral(self, lhs, rhs, dimension_numbers):
     """Enqueues a general dot operation onto the computation.
 
     Args:
-      lhs: ComputationDataHandle for the left-hand-side array.
-      rhs: ComputationDataHandle for the right-hand-side array.
+      lhs: LocalOp for the left-hand-side array.
+      rhs: LocalOp for the right-hand-side array.
       dimension_numbers: either an xla_data_pb2.DotDimensionNumbers or a nested
         tuple ((lhs_contract, rhs_contract), (lhs_batch, rhs_batch)) of lists of
         integers representing the dimensions to treat as contracting dimensions
         and batch dimensions on each input operand.
 
-    Returns: a ComputationDataHandle representing the DotGeneral operation.
+    Returns: a LocalOp representing the DotGeneral operation.
     """
     if not isinstance(dimension_numbers, xla_data_pb2.DotDimensionNumbers):
       dimension_numbers = GetDotDimensionsFromLists(dimension_numbers)
-    return _wrap_data_handle(
-        self._client.DotGeneral(
-            _unwrap_data_handle(lhs), _unwrap_data_handle(rhs),
-            dimension_numbers))
+    return self._client.DotGeneral(lhs, rhs, dimension_numbers)
 
   def Conv(self, lhs, rhs, window_strides, padding):
     """Enqueues a Conv operation onto the computation.
 
     Args:
-      lhs: ComputationDataHandle for the rank N+2 array of inputs.
-      rhs: ComputationDataHandle for the rank N+2 array of kernel weights.
+      lhs: LocalOp for the rank N+2 array of inputs.
+      rhs: LocalOp for the rank N+2 array of kernel weights.
       window_strides: length-N array-like of integer kernel strides.
       padding: PaddingType representing either 'SAME' or 'VALID' padding.
 
-    Returns: a ComputationDataHandle representing the Conv operation.
+    Returns: a LocalOp representing the Conv operation.
     """
     pads = _convert_padding_type_to_pad_values(
         padding, self.GetShape(lhs).dimensions()[2:],
         self.GetShape(rhs).dimensions()[2:], window_strides)
     dimension_numbers = self._GetConvDimensionNumbers(len(window_strides))
-    return _wrap_data_handle(
-        self._client.ConvGeneralDilated(_unwrap_data_handle(lhs),
-                                        _unwrap_data_handle(rhs),
-                                        window_strides,
-                                        pads,
-                                        (),
-                                        (),
-                                        dimension_numbers))
+    return self._client.ConvGeneralDilated(lhs, rhs, window_strides, pads, (),
+                                           (), dimension_numbers)
 
   def ConvWithGeneralPadding(self, lhs, rhs, window_strides, padding,
                              lhs_dilation, rhs_dilation):
     """Enqueues a ConvWithGeneralPadding operation onto the computation.
 
     Args:
-      lhs: ComputationDataHandle for the rank N+2 array of inputs.
-      rhs: ComputationDataHandle for the rank N+2 array of kernel weights.
+      lhs: LocalOp for the rank N+2 array of inputs.
+      rhs: LocalOp for the rank N+2 array of kernel weights.
       window_strides: length-N array-like of kernel strides.
       padding: length-N array-like of pairs of integers of (low, high) padding.
       lhs_dilation: length-N array-like of dilation factors.
@@ -1159,14 +1081,9 @@ class ComputationBuilder(object):
       A ComputationdataHandle representing the added ConvWithGeneralPadding op.
     """
     dimension_numbers = self._GetConvDimensionNumbers(len(window_strides))
-    return _wrap_data_handle(
-        self._client.ConvGeneralDilated(_unwrap_data_handle(lhs),
-                                        _unwrap_data_handle(rhs),
-                                        window_strides,
-                                        padding,
-                                        lhs_dilation,
-                                        rhs_dilation,
-                                        dimension_numbers))
+    return self._client.ConvGeneralDilated(lhs, rhs, window_strides, padding,
+                                           lhs_dilation, rhs_dilation,
+                                           dimension_numbers)
 
   def _GetConvDimensionNumbers(self, num_spatial_dims):
     """Create ConvolutionDimensionNumbers proto for convolutions."""
@@ -1196,15 +1113,14 @@ def _forward_methods_to_local_builder():
     """Generate a forwarding method that wraps/unwraps data handles."""
 
     def forward(self, *args, **kwargs):
-      unwrapped_args = [_unwrap_data_handle(arg) for arg in args]
+      arg_list = list(args)
 
-      if is_binop and len(unwrapped_args) < 3:
-        unwrapped_args.append(kwargs.get('broadcast_dimensions', ()))
+      if is_binop and len(arg_list) < 3:
+        arg_list.append(kwargs.get('broadcast_dimensions', ()))
 
-      return _wrap_data_handle(
-          target_method(
-              self._client,  # pylint: disable=protected-access
-              *unwrapped_args))
+      return target_method(
+          self._client,  # pylint: disable=protected-access
+          *arg_list)
 
     return forward
 
-- 
GitLab


From 5ca373b4b64167f8b0fcab96d7d2e7886ea31b6a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 12:28:42 -0700
Subject: [PATCH 264/395] Some fixes to support another TF graph: 1. Fix
 ResolveBatchNormalization to avoid deleting arrays that may still be used. 2.
 Correctly count the number of ops using a given array, even when some ops use
 the same array as more than one of their inputs. 3. In PropagateFixedSizes
 for Concatenation ops, when resolving a -1 wildcard to a fixed value, we were
 doing so in a local 'axis' variable without actually updating op->axis! The
 resulting -1 value still in op->axis tripped runtime code, causing the
 concatenation to misbehave during inference.

PiperOrigin-RevId: 195454037
---
 .../graph_transformations/propagate_fixed_sizes.cc    | 11 +++++------
 .../resolve_batch_normalization.cc                    |  6 +++---
 tensorflow/contrib/lite/toco/tooling_util.cc          |  4 ++++
 3 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
index 4923f83d91..b02b02c5be 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
@@ -670,8 +670,7 @@ void ProcessConcatenationOperator(Model* model, ConcatenationOperator* op) {
   const auto& first_input_array = model->GetArray(op->inputs[0]);
   output_array.copy_shape(first_input_array.shape());
   // Negative axis means the count starts at the back of the dims().
-  int axis = op->axis;
-  if (axis < 0) axis += first_input_array.shape().dims().size();
+  if (op->axis < 0) op->axis += first_input_array.shape().dims().size();
   // Determine the concat size, and enfore that all inputs have
   // the same dimensions count.
   int concat_size = 0;
@@ -684,14 +683,14 @@ void ProcessConcatenationOperator(Model* model, ConcatenationOperator* op) {
     CHECK_EQ(input_array.shape().dimensions_count(),
              output_array.shape().dimensions_count());
     const std::vector<int>& input_dims = input_array.shape().dims();
-    CHECK_LT(axis, input_dims.size());
-    concat_size += input_dims[axis];
+    CHECK_LT(op->axis, input_dims.size());
+    concat_size += input_dims[op->axis];
   }
   // Write out the concat_size on the output array shape.
   auto& output_shape = *output_array.mutable_shape();
   auto& output_dims = *output_shape.mutable_dims();
-  CHECK_LT(axis, output_shape.dimensions_count());
-  output_dims[axis] = concat_size;
+  CHECK_LT(op->axis, output_shape.dimensions_count());
+  output_dims[op->axis] = concat_size;
 }
 
 void ProcessRangeOperator(Model* model, RangeOperator* op) {
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_normalization.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_normalization.cc
index 2b3ee36ad1..8f2c1f8162 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_normalization.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_batch_normalization.cc
@@ -134,9 +134,9 @@ bool ResolveBatchNormalization::Run(Model* model, std::size_t op_index) {
   }
 
   // Remove the old param arrays
-  model->EraseArray(bn_op->inputs[1]);
-  model->EraseArray(bn_op->inputs[2]);
-  model->EraseArray(bn_op->inputs[3]);
+  DeleteArrayIfUsedOnce(bn_op->inputs[1], model);
+  DeleteArrayIfUsedOnce(bn_op->inputs[2], model);
+  DeleteArrayIfUsedOnce(bn_op->inputs[3], model);
 
   // Remove the old operator
   DCHECK_EQ(bn_it->get(), bn_op);
diff --git a/tensorflow/contrib/lite/toco/tooling_util.cc b/tensorflow/contrib/lite/toco/tooling_util.cc
index 86ee1f3761..341d45e753 100644
--- a/tensorflow/contrib/lite/toco/tooling_util.cc
+++ b/tensorflow/contrib/lite/toco/tooling_util.cc
@@ -143,6 +143,10 @@ int CountOpsWithInput(const Model& model, const string& array_name) {
     for (auto& input : op->inputs) {
       if (input == array_name) {
         count++;
+        // Breaking here is important: some graphs have ops that use the
+        // same array as more than one of their inputs, and in that case
+        // we want it counted only once.
+        break;
       }
     }
   }
-- 
GitLab


From 67b5e724121c5874425936fe01318642508d9975 Mon Sep 17 00:00:00 2001
From: Justin Lebar <jlebar@google.com>
Date: Fri, 4 May 2018 14:40:02 -0700
Subject: [PATCH 265/395] [XLA:GPU] Mark floating-point division as an
 inexpensive op.

"Expensive" really means "so expensive you'd choose not to fuse in order
to avoid doing it twice".  FP division definitely isn't that expensive.

PiperOrigin-RevId: 195473524
---
 .../xla/service/gpu/instruction_fusion.cc     | 13 +++++
 .../xla/service/gpu/instruction_fusion.h      |  2 +
 .../service/gpu/instruction_fusion_test.cc    | 56 +++++++++++++++++++
 3 files changed, 71 insertions(+)

diff --git a/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc b/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc
index 85ecbe8fdb..c5eb721185 100644
--- a/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc
+++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc
@@ -48,6 +48,19 @@ bool IsFusile(const HloInstruction& hlo) {
 
 }  // namespace
 
+/*static*/ bool GpuInstructionFusion::IsExpensive(
+    const HloInstruction& instruction) {
+  switch (instruction.opcode()) {
+    // We say that floating-point division is cheap on the GPU.
+    case HloOpcode::kDivide:
+      return !ShapeUtil::ElementIsFloating(instruction.shape()) &&
+             InstructionFusion::IsExpensive(instruction);
+
+    default:
+      return InstructionFusion::IsExpensive(instruction);
+  }
+}
+
 bool GpuInstructionFusion::ShouldFuse(HloInstruction* consumer,
                                       int64 operand_index) {
   HloInstruction* producer = consumer->mutable_operand(operand_index);
diff --git a/tensorflow/compiler/xla/service/gpu/instruction_fusion.h b/tensorflow/compiler/xla/service/gpu/instruction_fusion.h
index bb2990e6df..9fb06b0a24 100644
--- a/tensorflow/compiler/xla/service/gpu/instruction_fusion.h
+++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion.h
@@ -27,6 +27,8 @@ class GpuInstructionFusion : public InstructionFusion {
   explicit GpuInstructionFusion(bool may_duplicate)
       : InstructionFusion(GpuInstructionFusion::IsExpensive, may_duplicate) {}
 
+  static bool IsExpensive(const HloInstruction& instruction);
+
   bool ShouldFuse(HloInstruction* consumer, int64 operand_index) override;
 
   HloInstruction::FusionKind ChooseKind(
diff --git a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc
index 4b231c449f..6c9a805ad6 100644
--- a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc
+++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc
@@ -253,5 +253,61 @@ TEST_F(InstructionFusionTest, DotOutputFusion) {
                    op::Dot(op::Parameter(), op::Transpose(op::Parameter()))));
 }
 
+// Compute sum(1/p0), where p0 has type f32, twice.  Check that the division is
+// duplicated and fused into both reduces.
+TEST_F(InstructionFusionTest, FloatingPointDivIsCheap) {
+  auto module = tools::Parse(R"(
+  HloModule test_module
+  Add {
+    lhs = f32[] parameter(0)
+    rhs = f32[] parameter(1)
+    ROOT add = f32[] add(lhs, rhs)
+  }
+  ENTRY TestComputation {
+    zero = f32[] constant(0)
+    one = f32[] constant(1)
+    p0 = f32[100] parameter(0)
+    recip = f32[100] divide(one, p0)
+    sum1 = f32[] reduce(recip, zero), dimensions={0}, to_apply=Add
+    sum2 = f32[] reduce(recip, zero), dimensions={0}, to_apply=Add
+    ROOT root = (f32[], f32[]) tuple(sum1, sum2)
+  })")
+                    .ValueOrDie();
+
+  EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true)
+                  .Run(module.get())
+                  .ValueOrDie());
+
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root, op::Tuple(op::Fusion(), op::Fusion()));
+}
+
+// Compute sum(100/p0), where p0 has type s32, twice.  Check that the division
+// is *not* duplicated and fused into both reduces, because we say that integer
+// division is not cheap.
+TEST_F(InstructionFusionTest, IntegerDivIsNotCheap) {
+  auto module = tools::Parse(R"(
+  HloModule test_module
+  Add {
+    lhs = s32[] parameter(0)
+    rhs = s32[] parameter(1)
+    ROOT add = s32[] add(lhs, rhs)
+  }
+  ENTRY TestComputation {
+    zero = s32[] constant(0)
+    one_hundred = s32[] constant(100)
+    p0 = s32[100] parameter(0)
+    recip = s32[100] divide(one_hundred, p0)
+    sum1 = s32[] reduce(recip, zero), dimensions={0}, to_apply=Add
+    sum2 = s32[] reduce(recip, zero), dimensions={0}, to_apply=Add
+    ROOT mul = (s32[], s32[]) tuple(sum1, sum2)
+  })")
+                    .ValueOrDie();
+
+  EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true)
+                   .Run(module.get())
+                   .ValueOrDie());
+}
+
 }  // namespace gpu
 }  // namespace xla
-- 
GitLab


From 4d0388d22060a61f40965127c153c681b2412c50 Mon Sep 17 00:00:00 2001
From: James Qin <jamesqin@google.com>
Date: Fri, 4 May 2018 14:53:58 -0700
Subject: [PATCH 266/395] Fix build failure for macos py3

PiperOrigin-RevId: 195475780
---
 tensorflow/python/debug/examples/debug_tflearn_iris.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/debug/examples/debug_tflearn_iris.py b/tensorflow/python/debug/examples/debug_tflearn_iris.py
index 00090b21fe..7cbaae46b4 100644
--- a/tensorflow/python/debug/examples/debug_tflearn_iris.py
+++ b/tensorflow/python/debug/examples/debug_tflearn_iris.py
@@ -140,7 +140,7 @@ def main(_):
 
   # Make predictions, using tfdbg hook.
   predict_results = classifier.predict(test_input_fn, hooks=hooks)
-  print("A prediction result: %s" % predict_results.next())
+  print("A prediction result: %s" % next(predict_results))
 
 
 if __name__ == "__main__":
-- 
GitLab


From cb1775e9525ae621d23708a3d64a6cad897be95e Mon Sep 17 00:00:00 2001
From: Benoit Steiner <bsteiner@google.com>
Date: Fri, 4 May 2018 15:14:00 -0700
Subject: [PATCH 267/395] Identify and prune nodes that can never be executed

PiperOrigin-RevId: 195478951
---
 tensorflow/core/grappler/optimizers/BUILD     |   1 +
 .../grappler/optimizers/loop_optimizer.cc     | 140 ++++++++++++++++++
 .../core/grappler/optimizers/loop_optimizer.h |   1 +
 .../optimizers/loop_optimizer_test.cc         | 107 +++++++++++++
 tensorflow/core/grappler/utils.h              |   4 +-
 5 files changed, 251 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/BUILD b/tensorflow/core/grappler/optimizers/BUILD
index 5b5e1e024e..900dfa95c5 100644
--- a/tensorflow/core/grappler/optimizers/BUILD
+++ b/tensorflow/core/grappler/optimizers/BUILD
@@ -604,6 +604,7 @@ cc_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/grappler:graph_view",
         "//tensorflow/core/grappler:grappler_item",
         "//tensorflow/core/grappler:op_types",
         "//tensorflow/core/grappler:utils",
diff --git a/tensorflow/core/grappler/optimizers/loop_optimizer.cc b/tensorflow/core/grappler/optimizers/loop_optimizer.cc
index 5adc5b9227..7d3520febc 100644
--- a/tensorflow/core/grappler/optimizers/loop_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/loop_optimizer.cc
@@ -27,6 +27,7 @@ limitations under the License.
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/tensor_shape.pb.h"
 #include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/grappler/graph_view.h"
 #include "tensorflow/core/grappler/grappler_item.h"
 #include "tensorflow/core/grappler/op_types.h"
 #include "tensorflow/core/grappler/optimizers/constant_folding.h"
@@ -504,6 +505,140 @@ Status RemoveStackOps(const std::unordered_set<string>& nodes_to_preserve,
   return Status::OK();
 }
 
+Status RemoveDeadBranches(const std::unordered_set<string>& nodes_to_preserve,
+                          GraphDef* optimized_graph) {
+  std::unordered_set<const NodeDef*> dead_nodes;
+  std::unordered_map<NodeDef*, std::set<int>> dead_merge_inputs;
+  // TODO(bsteiner): also rewrite switches as identity. For now we just record
+  // them
+  std::unordered_set<GraphView::OutputPort, GraphView::HashPort>
+      identity_switches;
+
+  GraphView view(optimized_graph);
+  for (const NodeDef& node : optimized_graph->node()) {
+    if (!IsSwitch(node)) {
+      continue;
+    }
+    if (nodes_to_preserve.find(node.name()) != nodes_to_preserve.end()) {
+      continue;
+    }
+    GraphView::InputPort ctrl_port(&node, 1);
+    GraphView::OutputPort ctrl_node = view.GetRegularFanin(ctrl_port);
+    if (!IsConstant(*ctrl_node.node)) {
+      continue;
+    }
+    Tensor selector;
+    CHECK(selector.FromProto(ctrl_node.node->attr().at("value").tensor()));
+    const int dead_fanout = selector.scalar<bool>()() ? 0 : 1;
+    GraphView::OutputPort dead(const_cast<NodeDef*>(&node), dead_fanout);
+    identity_switches.insert(dead);
+
+    SetVector<GraphView::InputPort, GraphView::HashPort> zombie_inputs;
+    for (const GraphView::InputPort& port : view.GetFanout(dead)) {
+      if (dead_nodes.find(port.node) == dead_nodes.end()) {
+        zombie_inputs.PushBack(port);
+      }
+    }
+    // If we encounter a single node that must be preserved in the fanout of the
+    // switch node we need to preserve the entire switch fanout: we therefore
+    // work on a local copy that only gets committed to the master copy once the
+    // whole fanout has been explored.
+    std::unordered_set<const NodeDef*> local_dead_nodes = dead_nodes;
+    std::unordered_map<NodeDef*, std::set<int>> local_dead_merge_inputs =
+        dead_merge_inputs;
+    bool found_node_to_preserve = false;
+    while (!found_node_to_preserve && !zombie_inputs.Empty()) {
+      GraphView::InputPort dead = zombie_inputs.PopBack();
+      if (nodes_to_preserve.find(dead.node->name()) !=
+          nodes_to_preserve.end()) {
+        found_node_to_preserve = true;
+        break;
+      }
+
+      if (local_dead_nodes.find(dead.node) != local_dead_nodes.end()) {
+        continue;
+      }
+
+      if (IsMerge(*dead.node)) {
+        const int fanout = dead.node->attr().at("N").i();
+        if (fanout > 2) {
+          // This never happens in practice, so we'll just skip these to
+          // simplify the code for now.
+          found_node_to_preserve = true;
+          break;
+        }
+        GraphView::OutputPort value_index(dead.node, 1);
+        const std::unordered_set<GraphView::InputPort, GraphView::HashPort>&
+            index_fanout = view.GetFanout(value_index);
+        if (!index_fanout.empty()) {
+          // The 2nd output (that indicates which input is propagated) is
+          // connected. This never happens in practice, so we'll just skip this
+          // case to simplify the code for now.
+          found_node_to_preserve = true;
+          break;
+        }
+
+        bool fully_dead = false;
+        if (dead.port_id < 0) {
+          // If the control dependency never gets triggered the merge will also
+          // never get triggered.
+          local_dead_nodes.insert(dead.node);
+          fully_dead = true;
+        } else {
+          local_dead_merge_inputs[dead.node].insert(dead.port_id);
+          if (local_dead_merge_inputs[dead.node].size() ==
+              dead.node->attr().at("N").i()) {
+            fully_dead = true;
+          }
+          if (fully_dead) {
+            local_dead_nodes.insert(dead.node);
+            for (const GraphView::InputPort& port :
+                 view.GetFanouts(*dead.node, true)) {
+              zombie_inputs.PushBack(port);
+            }
+          }
+        }
+      } else {
+        if (local_dead_nodes.insert(dead.node).second) {
+          for (const GraphView::InputPort& dead_fanout :
+               view.GetFanouts(*dead.node, true)) {
+            zombie_inputs.PushBack(dead_fanout);
+          }
+        }
+      }
+    }
+    if (!found_node_to_preserve) {
+      std::swap(dead_nodes, local_dead_nodes);
+      std::swap(dead_merge_inputs, local_dead_merge_inputs);
+    }
+  }
+
+  int last = optimized_graph->node_size() - 1;
+  for (int i = optimized_graph->node_size() - 1; i >= 0; --i) {
+    NodeDef* node = optimized_graph->mutable_node(i);
+    if (dead_nodes.find(node) != dead_nodes.end()) {
+      optimized_graph->mutable_node()->SwapElements(i, last);
+      last--;
+    }
+  }
+  optimized_graph->mutable_node()->DeleteSubrange(last + 1, dead_nodes.size());
+
+  for (const auto& itr : dead_merge_inputs) {
+    NodeDef* dead_node = itr.first;
+    if (dead_nodes.find(dead_node) != dead_nodes.end()) {
+      // The node has been pruned since all its inputs are dead.
+      continue;
+    }
+    const std::set<int>& dead_inputs = itr.second;
+    for (int index : dead_inputs) {
+      dead_node->mutable_input()->DeleteSubrange(index, 1);
+    }
+    dead_node->set_op("Identity");
+    dead_node->mutable_attr()->erase("N");
+  }
+  return Status::OK();
+}
+
 }  // namespace
 
 Status LoopOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
@@ -517,6 +652,11 @@ Status LoopOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
   if (options_.enable_stack_push_removal) {
     TF_RETURN_IF_ERROR(RemoveStackOps(item.NodesToPreserve(), optimized_graph));
   }
+  if (opt_level_ == RewriterConfig::AGGRESSIVE &&
+      options_.enable_dead_branch_removal) {
+    TF_RETURN_IF_ERROR(
+        RemoveDeadBranches(item.NodesToPreserve(), optimized_graph));
+  }
 
   return Status::OK();
 }
diff --git a/tensorflow/core/grappler/optimizers/loop_optimizer.h b/tensorflow/core/grappler/optimizers/loop_optimizer.h
index 764506f7c1..85b8e65543 100644
--- a/tensorflow/core/grappler/optimizers/loop_optimizer.h
+++ b/tensorflow/core/grappler/optimizers/loop_optimizer.h
@@ -54,6 +54,7 @@ class LoopOptimizer : public GraphOptimizer {
   struct LoopOptimizerOptions {
     bool enable_loop_invariant_node_motion = false;
     bool enable_stack_push_removal = true;
+    bool enable_dead_branch_removal = true;
 
     static LoopOptimizerOptions Default(RewriterConfig::Toggle opt_level) {
       LoopOptimizerOptions options;
diff --git a/tensorflow/core/grappler/optimizers/loop_optimizer_test.cc b/tensorflow/core/grappler/optimizers/loop_optimizer_test.cc
index 10ec544424..6fd177b710 100644
--- a/tensorflow/core/grappler/optimizers/loop_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/loop_optimizer_test.cc
@@ -589,5 +589,112 @@ TEST_F(LoopOptimizerTest, RemovePushWithoutMatchingPop) {
   }
 }
 
+TEST_F(LoopOptimizerTest, RemoveDeadBranches) {
+  Scope scope = Scope::NewRootScope();
+  Output v_in = ops::Variable(scope.WithOpName("v_in"), {3}, DT_FLOAT);
+
+  Output ctrl1 = ops::Const(scope.WithOpName("ctrl1"), false, TensorShape({}));
+  ops::Switch s1(scope.WithOpName("switch1"), v_in, ctrl1);
+  Output square1 = ops::Square(scope.WithOpName("square1"), s1.output_false);
+  Output sqrt1 = ops::Sqrt(scope.WithOpName("sqrt1"), s1.output_true);
+
+  Output ctrl2 = ops::Const(scope.WithOpName("ctrl2"), true, TensorShape({}));
+  ops::Switch s2(scope.WithOpName("switch2"), v_in, ctrl2);
+  Output square2 = ops::Square(scope.WithOpName("square2"), s2.output_false);
+  Output sqrt2 = ops::Sqrt(scope.WithOpName("sqrt2"), s2.output_true);
+
+  Output ctrl3 = ops::Const(scope.WithOpName("ctrl3"), false, TensorShape({}));
+  ops::Switch s3(scope.WithOpName("switch3"), v_in, ctrl3);
+  Output square3 = ops::Square(scope.WithOpName("square3"), s3.output_false);
+  Output sqrt3 = ops::Sqrt(scope.WithOpName("sqrt3"), s3.output_true);
+
+  Output ctrl4 = ops::Const(scope.WithOpName("ctrl4"), false, TensorShape({}));
+  ops::Switch s4(scope.WithOpName("switch4"), v_in, ctrl4);
+  Output square4 = ops::Square(scope.WithOpName("square4"), s4.output_false);
+  Output sqrt4 = ops::Sqrt(scope.WithOpName("sqrt4"), s4.output_true);
+
+  ops::Merge m1(scope.WithOpName("m1"), {square1, sqrt1});
+  ops::Merge m2(scope.WithOpName("m2"), {v_in, square1});
+  ops::Merge m3(scope.WithOpName("m3"), {v_in, sqrt1});
+  ops::Merge m4(scope.WithOpName("m4"), {square1, sqrt2});
+  ops::Merge m5(scope.WithOpName("m5"), {square2, sqrt1});
+  ops::Merge m6(scope.WithOpName("m6").WithControlDependencies(sqrt2),
+                {v_in, square1});
+  ops::Merge m7(scope.WithOpName("m7").WithControlDependencies(sqrt1),
+                {v_in, square1});
+
+  ops::Switch s5(scope.WithOpName("switch5"), v_in, ctrl1);
+  Output id1 = ops::Identity(scope.WithOpName("id1"), s5.output_false);
+  Output id2 = ops::Identity(scope.WithOpName("id2"), s5.output_true);
+  ops::Merge m8(scope.WithOpName("m8"), {id1, id2});
+
+  ops::Switch s6(scope.WithOpName("switch6"), v_in, ctrl1);
+  Output id3 = ops::Identity(scope.WithOpName("id3"), s6.output_false);
+  Output id4 = ops::Identity(scope.WithOpName("id4"), s6.output_true);
+  ops::Merge m9(scope.WithOpName("m9"), {id3, id4});
+
+  GrapplerItem item;
+  item.fetch.push_back("m8");
+  item.fetch.push_back("id4");
+
+  TF_CHECK_OK(scope.ToGraphDef(&item.graph));
+
+  LoopOptimizer optimizer(RewriterConfig::AGGRESSIVE);
+  GraphDef output;
+  Status status = optimizer.Optimize(nullptr, item, &output);
+  TF_CHECK_OK(status);
+
+  for (const NodeDef& node : output.node()) {
+    // These nodes should have been pruned
+    EXPECT_NE("Square1", node.name());
+    EXPECT_NE("Sqrt2", node.name());
+    EXPECT_NE("m5", node.name());
+    EXPECT_NE("m7", node.name());
+
+    if (node.name() == "m1") {
+      // sqrt1 is dead
+      EXPECT_EQ("Identity", node.op());
+      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ("square1", node.input(0));
+    } else if (node.name() == "m2") {
+      // both inputs are alive
+      EXPECT_EQ("Merge", node.op());
+      EXPECT_EQ(2, node.input_size());
+      EXPECT_EQ("v_in", node.input(0));
+      EXPECT_EQ("square1", node.input(1));
+    } else if (node.name() == "m3") {
+      // sqrt1 is dead
+      EXPECT_EQ("Identity", node.op());
+      EXPECT_EQ(1, node.input_size());
+      EXPECT_EQ("v_in", node.input(0));
+    } else if (node.name() == "m4") {
+      // both inputs are alive
+      EXPECT_EQ("Merge", node.op());
+      EXPECT_EQ(2, node.input_size());
+      EXPECT_EQ("square1", node.input(0));
+      EXPECT_EQ("sqrt2", node.input(1));
+    } else if (node.name() == "m6") {
+      // both inputs are alive and the control dependency can get triggered
+      EXPECT_EQ("Merge", node.op());
+      EXPECT_EQ(3, node.input_size());
+      EXPECT_EQ("v_in", node.input(0));
+      EXPECT_EQ("square1", node.input(1));
+      EXPECT_EQ("^sqrt2", node.input(2));
+    } else if (node.name() == "m8") {
+      // The node is to be preserved because of a fetch
+      EXPECT_EQ("Merge", node.op());
+      EXPECT_EQ(2, node.input_size());
+      EXPECT_EQ("id1", node.input(0));
+      EXPECT_EQ("id2", node.input(1));
+    } else if (node.name() == "m9") {
+      // The node is to be preserved because of a fetch
+      EXPECT_EQ("Merge", node.op());
+      EXPECT_EQ(2, node.input_size());
+      EXPECT_EQ("id3", node.input(0));
+      EXPECT_EQ("id4", node.input(1));
+    }
+  }
+}
+
 }  // namespace grappler
 }  // namespace tensorflow
diff --git a/tensorflow/core/grappler/utils.h b/tensorflow/core/grappler/utils.h
index b87ae05546..1c6fef59ea 100644
--- a/tensorflow/core/grappler/utils.h
+++ b/tensorflow/core/grappler/utils.h
@@ -65,7 +65,7 @@ class NodeMap {
 // A vector with a set. The set stores the same elements as the vector, and
 // quickly answers whether a value is in the vector. Duplicated elements are not
 // allowed for now.
-template <class T>
+template <class T, class Hash = std::hash<T>>
 class SetVector {
  public:
   // Returns false if value already existed in the set, true otherwise.
@@ -91,7 +91,7 @@ class SetVector {
   void Reserve(int64 size) { vector_.reserve(size); }
 
  private:
-  std::unordered_set<T> set_;
+  std::unordered_set<T, Hash> set_;
   std::vector<T> vector_;
 };
 
-- 
GitLab


From c92de2f3fc81c701ab29408a8a84cd6e41e96fe5 Mon Sep 17 00:00:00 2001
From: "karl@kubx.ca" <karl@kubx.ca>
Date: Sat, 5 May 2018 10:44:20 -0400
Subject: [PATCH 268/395] Skip all ops with function attribute by default

---
 tensorflow/core/api_def/BUILD                      |  6 ------
 .../api_def/java_api/api_def_FilterDataset.pbtxt   |  4 ----
 .../api_def/java_api/api_def_FlatMapDataset.pbtxt  |  4 ----
 tensorflow/core/api_def/java_api/api_def_For.pbtxt |  4 ----
 .../java_api/api_def_GeneratorDataset.pbtxt        |  4 ----
 .../java_api/api_def_GroupByWindowDataset.pbtxt    |  4 ----
 tensorflow/core/api_def/java_api/api_def_If.pbtxt  |  4 ----
 .../java_api/api_def_InterleaveDataset.pbtxt       |  4 ----
 .../java_api/api_def_MapAndBatchDataset.pbtxt      |  4 ----
 .../core/api_def/java_api/api_def_MapDataset.pbtxt |  4 ----
 .../api_def/java_api/api_def_OneShotIterator.pbtxt |  4 ----
 .../api_def_ParallelInterleaveDataset.pbtxt        |  4 ----
 .../java_api/api_def_ParallelMapDataset.pbtxt      |  4 ----
 .../core/api_def/java_api/api_def_RemoteCall.pbtxt |  4 ----
 .../api_def/java_api/api_def_ScanDataset.pbtxt     |  4 ----
 .../java_api/api_def_SymbolicGradient.pbtxt        |  4 ----
 .../core/api_def/java_api/api_def_While.pbtxt      |  4 ----
 tensorflow/java/BUILD                              |  1 -
 tensorflow/java/src/gen/cc/op_generator.cc         | 14 +++++++++++++-
 19 files changed, 13 insertions(+), 72 deletions(-)
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_FilterDataset.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_FlatMapDataset.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_For.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_GeneratorDataset.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_GroupByWindowDataset.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_If.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_InterleaveDataset.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_MapAndBatchDataset.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_MapDataset.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_OneShotIterator.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_ParallelInterleaveDataset.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_ParallelMapDataset.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_RemoteCall.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_ScanDataset.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_SymbolicGradient.pbtxt
 delete mode 100644 tensorflow/core/api_def/java_api/api_def_While.pbtxt

diff --git a/tensorflow/core/api_def/BUILD b/tensorflow/core/api_def/BUILD
index 06b797e32e..1454a1d9b2 100644
--- a/tensorflow/core/api_def/BUILD
+++ b/tensorflow/core/api_def/BUILD
@@ -30,12 +30,6 @@ filegroup(
     visibility = ["//tensorflow:internal"],
 )
 
-filegroup(
-    name = "java_api_def",
-    srcs = glob(["java_api/*"]),
-    visibility = ["//tensorflow:internal"],
-)
-
 cc_library(
     name = "excluded_ops_lib",
     srcs = ["excluded_ops.cc"],
diff --git a/tensorflow/core/api_def/java_api/api_def_FilterDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_FilterDataset.pbtxt
deleted file mode 100644
index debd7e5709..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_FilterDataset.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "FilterDataset"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_FlatMapDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_FlatMapDataset.pbtxt
deleted file mode 100644
index 329ab15ef5..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_FlatMapDataset.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "FlatMapDataset"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_For.pbtxt b/tensorflow/core/api_def/java_api/api_def_For.pbtxt
deleted file mode 100644
index caabc947bb..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_For.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "For"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_GeneratorDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_GeneratorDataset.pbtxt
deleted file mode 100644
index a6e5167c30..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_GeneratorDataset.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "GeneratorDataset"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_GroupByWindowDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_GroupByWindowDataset.pbtxt
deleted file mode 100644
index 4c0b2084a8..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_GroupByWindowDataset.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "GroupByWindowDataset"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_If.pbtxt b/tensorflow/core/api_def/java_api/api_def_If.pbtxt
deleted file mode 100644
index 13b8635ca7..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_If.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "If"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_InterleaveDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_InterleaveDataset.pbtxt
deleted file mode 100644
index ed748d4d2a..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_InterleaveDataset.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "InterleaveDataset"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_MapAndBatchDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_MapAndBatchDataset.pbtxt
deleted file mode 100644
index cb96bf63d8..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_MapAndBatchDataset.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "MapAndBatchDataset"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_MapDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_MapDataset.pbtxt
deleted file mode 100644
index e0ab8dd9db..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_MapDataset.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "MapDataset"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_OneShotIterator.pbtxt b/tensorflow/core/api_def/java_api/api_def_OneShotIterator.pbtxt
deleted file mode 100644
index 13130e6882..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_OneShotIterator.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "OneShotIterator"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_ParallelInterleaveDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_ParallelInterleaveDataset.pbtxt
deleted file mode 100644
index 6a985d24fa..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_ParallelInterleaveDataset.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "ParallelInterleaveDataset"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_ParallelMapDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_ParallelMapDataset.pbtxt
deleted file mode 100644
index 64f25b9e5e..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_ParallelMapDataset.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "ParallelMapDataset"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_RemoteCall.pbtxt b/tensorflow/core/api_def/java_api/api_def_RemoteCall.pbtxt
deleted file mode 100644
index 2ccb5c8cf3..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_RemoteCall.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "RemoteCall"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_ScanDataset.pbtxt b/tensorflow/core/api_def/java_api/api_def_ScanDataset.pbtxt
deleted file mode 100644
index 3463e60049..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_ScanDataset.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "ScanDataset"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_SymbolicGradient.pbtxt b/tensorflow/core/api_def/java_api/api_def_SymbolicGradient.pbtxt
deleted file mode 100644
index 88c3acea74..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_SymbolicGradient.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "SymbolicGradient"
-  visibility: SKIP
-}
diff --git a/tensorflow/core/api_def/java_api/api_def_While.pbtxt b/tensorflow/core/api_def/java_api/api_def_While.pbtxt
deleted file mode 100644
index 33756682c3..0000000000
--- a/tensorflow/core/api_def/java_api/api_def_While.pbtxt
+++ /dev/null
@@ -1,4 +0,0 @@
-op {
-  graph_op_name: "While"
-  visibility: SKIP
-}
\ No newline at end of file
diff --git a/tensorflow/java/BUILD b/tensorflow/java/BUILD
index 7cd0208dbf..0cc8e7c3e2 100644
--- a/tensorflow/java/BUILD
+++ b/tensorflow/java/BUILD
@@ -72,7 +72,6 @@ tf_java_op_gen_srcjar(
     name = "java_op_gen_sources",
     api_def_srcs = [
         "//tensorflow/core/api_def:base_api_def",
-        "//tensorflow/core/api_def:java_api_def",
     ],
     base_package = "org.tensorflow.op",
     gen_tool = ":java_op_gen_tool",
diff --git a/tensorflow/java/src/gen/cc/op_generator.cc b/tensorflow/java/src/gen/cc/op_generator.cc
index 7355b3a395..f4cefbe933 100644
--- a/tensorflow/java/src/gen/cc/op_generator.cc
+++ b/tensorflow/java/src/gen/cc/op_generator.cc
@@ -420,6 +420,18 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
   writer.EndType();
 }
 
+bool CanGenerateOp(const OpDef& op_def, const ApiDef& api_def) {
+  if (api_def.visibility() == ApiDef::SKIP) {
+    return false;
+  }
+  for (const auto& attr : op_def.attr()) {
+    if (attr.type() == "func") {
+      return false;  // TODO(karllessard) add support for function attributes
+    }
+  }
+  return true;
+}
+
 }  // namespace
 
 Status OpGenerator::Run(const OpList& op_list, const string& base_package,
@@ -441,7 +453,7 @@ Status OpGenerator::Run(const OpList& op_list, const string& base_package,
   api_map.UpdateDocs();
   for (const auto& op_def : op_list.op()) {
     const ApiDef* api_def = api_map.GetApiDef(op_def.name());
-    if (api_def->visibility() != ApiDef::SKIP) {
+    if (CanGenerateOp(op_def, *api_def)) {
       OpSpec op(OpSpec::Create(op_def, *api_def));
       for (const EndpointSpec& endpoint : op.endpoints()) {
         GenerateOp(op, endpoint, base_package, output_dir, env_);
-- 
GitLab


From 90bbbdcc42a67c93ba8dcbc66f9c1d06909c48cb Mon Sep 17 00:00:00 2001
From: Karl Lessard <karllessard@users.noreply.github.com>
Date: Sat, 5 May 2018 10:48:22 -0400
Subject: [PATCH 269/395] Remove comment left-over

---
 tensorflow/core/api_def/BUILD | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/core/api_def/BUILD b/tensorflow/core/api_def/BUILD
index 1454a1d9b2..19d6438809 100644
--- a/tensorflow/core/api_def/BUILD
+++ b/tensorflow/core/api_def/BUILD
@@ -4,7 +4,6 @@
 # The following targets can be used to access ApiDefs:
 #   :base_api_def
 #   :python_api_def
-#   :java_api_def
 
 package(
     default_visibility = ["//visibility:private"],
-- 
GitLab


From ab48fb528221152299fb08da8116d2eca54b8423 Mon Sep 17 00:00:00 2001
From: Justin Lebar <jlebar@google.com>
Date: Fri, 4 May 2018 15:40:07 -0700
Subject: [PATCH 270/395] [XLA] Print allowed attributes when the user
 specifies an invalid attr.

PiperOrigin-RevId: 195482974
---
 .../compiler/xla/tools/parser/hlo_parser.cc   | 30 +++++++++++++------
 .../xla/tools/parser/hlo_parser_test.cc       |  2 +-
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/tensorflow/compiler/xla/tools/parser/hlo_parser.cc b/tensorflow/compiler/xla/tools/parser/hlo_parser.cc
index 3a945fb3b1..40dc0730ce 100644
--- a/tensorflow/compiler/xla/tools/parser/hlo_parser.cc
+++ b/tensorflow/compiler/xla/tools/parser/hlo_parser.cc
@@ -30,6 +30,7 @@ namespace {
 
 using tensorflow::StringPiece;
 using tensorflow::gtl::optional;
+using tensorflow::str_util::Join;
 using tensorflow::str_util::Split;
 using tensorflow::str_util::SplitAndParseAsInts;
 using tensorflow::strings::Printf;
@@ -53,7 +54,7 @@ class HloParser {
   std::unique_ptr<HloModule> ConsumeHloModule() { return std::move(module_); }
 
   // Returns the error information.
-  string GetError() const { return tensorflow::str_util::Join(error_, "\n"); }
+  string GetError() const { return Join(error_, "\n"); }
 
  private:
   // ParseXXX returns false if an error occurred.
@@ -245,7 +246,7 @@ bool HloParser::Error(LocTy loc, StringPiece msg) {
   error_lines.push_back(std::string(lexer_.GetLine(loc)));
   error_lines.push_back(col == 0 ? "" : StrCat(string(col - 1, ' '), "^"));
 
-  error_.push_back(tensorflow::str_util::Join(error_lines, "\n"));
+  error_.push_back(Join(error_lines, "\n"));
   VLOG(1) << "Error: " << error_.back();
   return false;
 }
@@ -1488,11 +1489,10 @@ bool HloParser::ParseDenseLiteral(std::unique_ptr<Literal>* literal,
     std::vector<int64> elems_seen_until_dim(elems_seen_per_dim.begin(),
                                             elems_seen_per_dim.begin() + dim);
     return StrCat("[",
-                  tensorflow::str_util::Join(
-                      elems_seen_until_dim, ",",
-                      [](string* out, const int64& num_elems) {
-                        tensorflow::strings::StrAppend(out, num_elems - 1);
-                      }),
+                  Join(elems_seen_until_dim, ",",
+                       [](string* out, const int64& num_elems) {
+                         tensorflow::strings::StrAppend(out, num_elems - 1);
+                       }),
                   "]");
   };
   do {
@@ -1680,7 +1680,7 @@ bool HloParser::ParseSparseLiteralHelper(std::unique_ptr<Literal>* literal,
         return Error(
             index_loc,
             StrCat("invalid multi-dimension index for shape with rank ", rank,
-                   ": [", tensorflow::str_util::Join(index, ", "), "]"));
+                   ": [", Join(index, ", "), "]"));
       }
     }
     if (!ParseToken(TokKind::kColon,
@@ -1848,7 +1848,19 @@ bool HloParser::ParseAttributeHelper(
   }
   auto attr_it = attrs.find(name);
   if (attr_it == attrs.end()) {
-    return Error(loc, Printf("unexpected attribute %s", name.c_str()));
+    string allowed_attrs;
+    if (attrs.empty()) {
+      allowed_attrs = "No attributes are allowed here.";
+    } else {
+      allowed_attrs = StrCat(
+          "Allowed attributes: ",
+          Join(attrs, ", ",
+               [&](string* out, const std::pair<string, AttrConfig>& kv) {
+                 StrAppend(out, kv.first);
+               }));
+    }
+    return Error(loc, Printf("unexpected attribute \"%s\".  %s", name.c_str(),
+                             allowed_attrs.c_str()));
   }
   AttrTy attr_type = attr_it->second.attr_type;
   void* attr_out_ptr = attr_it->second.result;
diff --git a/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc b/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc
index 4e085bc89c..d38d8907a6 100644
--- a/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc
+++ b/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc
@@ -1138,7 +1138,7 @@ ENTRY %TwoSendRecvBothWayRecvFist.v3 () -> f32[] {
 
 )";
   ExpectHasSubstr(Parse(original).status().error_message(),
-                  "unexpected attribute calls");
+                  "unexpected attribute \"calls\"");
 }
 
 TEST_F(HloParserTest, MissingAttribute) {
-- 
GitLab


From 008a3b69a601dc68fd940eb8a03b0c445714a339 Mon Sep 17 00:00:00 2001
From: Karmel Allison <karmel@google.com>
Date: Fri, 4 May 2018 16:01:02 -0700
Subject: [PATCH 271/395] Add the ability to export separate SavedModels for
 train and eval mode to Estimator with two new methods, available in
 tf.contrib: export_all_saved_models and export_saved_model_for_mode.

PiperOrigin-RevId: 195485922
---
 tensorflow/contrib/estimator/BUILD            |  38 ++
 tensorflow/contrib/estimator/__init__.py      |   3 +
 .../estimator/python/estimator/export.py      | 216 ++++++++++
 .../estimator/python/estimator/export_test.py | 391 ++++++++++++++++++
 tensorflow/python/estimator/BUILD             |   1 +
 tensorflow/python/estimator/estimator.py      | 346 +++++++++++++---
 tensorflow/python/estimator/estimator_test.py | 336 ++++++++++++++-
 tensorflow/python/estimator/export/export.py  | 325 +++++++++++----
 .../python/estimator/export/export_output.py  | 223 +++++++++-
 .../estimator/export/export_output_test.py    | 110 +++++
 .../python/estimator/export/export_test.py    | 253 +++++++++++-
 tensorflow/python/estimator/model_fn.py       |   8 +
 tensorflow/python/saved_model/builder_impl.py |  54 ++-
 tensorflow/python/saved_model/constants.py    |   6 +
 .../python/saved_model/saved_model_test.py    |  90 ++++
 .../python/saved_model/signature_constants.py |   6 +
 .../python/saved_model/signature_def_utils.py |   2 +
 .../saved_model/signature_def_utils_impl.py   |  56 +++
 .../saved_model/signature_def_utils_test.py   |  95 +++++
 .../python/saved_model/tag_constants.py       |   5 +
 20 files changed, 2373 insertions(+), 191 deletions(-)
 create mode 100644 tensorflow/contrib/estimator/python/estimator/export.py
 create mode 100644 tensorflow/contrib/estimator/python/estimator/export_test.py

diff --git a/tensorflow/contrib/estimator/BUILD b/tensorflow/contrib/estimator/BUILD
index 571e2e3a5d..e9a68801ef 100644
--- a/tensorflow/contrib/estimator/BUILD
+++ b/tensorflow/contrib/estimator/BUILD
@@ -17,6 +17,7 @@ py_library(
         ":boosted_trees",
         ":dnn",
         ":dnn_linear_combined",
+        ":export",
         ":extenders",
         ":head",
         ":linear",
@@ -180,6 +181,43 @@ py_test(
     ],
 )
 
+py_library(
+    name = "export",
+    srcs = [
+        "python/estimator/export.py",
+    ],
+    srcs_version = "PY2AND3",
+    deps = [
+        "//tensorflow/python/estimator:model_fn",
+    ],
+)
+
+py_test(
+    name = "export_test",
+    size = "medium",
+    srcs = ["python/estimator/export_test.py"],
+    srcs_version = "PY2AND3",
+    tags = ["notsan"],  # b/62863147
+    deps = [
+        ":export",
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:metrics",
+        "//tensorflow/python:parsing_ops",
+        "//tensorflow/python:session",
+        "//tensorflow/python:state_ops",
+        "//tensorflow/python:training",
+        "//tensorflow/python:util",
+        "//tensorflow/python:variables",
+        "//tensorflow/python/estimator",
+        "//tensorflow/python/estimator:export_export",
+        "//tensorflow/python/estimator:export_output",
+        "//tensorflow/python/estimator:model_fn",
+        "//tensorflow/python/saved_model:loader",
+        "//tensorflow/python/saved_model:tag_constants",
+    ],
+)
+
 py_library(
     name = "head",
     srcs = [
diff --git a/tensorflow/contrib/estimator/__init__.py b/tensorflow/contrib/estimator/__init__.py
index d43b3ea6bf..ec502f86dd 100644
--- a/tensorflow/contrib/estimator/__init__.py
+++ b/tensorflow/contrib/estimator/__init__.py
@@ -22,6 +22,7 @@ from __future__ import print_function
 from tensorflow.contrib.estimator.python.estimator.boosted_trees import *
 from tensorflow.contrib.estimator.python.estimator.dnn import *
 from tensorflow.contrib.estimator.python.estimator.dnn_linear_combined import *
+from tensorflow.contrib.estimator.python.estimator.export import *
 from tensorflow.contrib.estimator.python.estimator.extenders import *
 from tensorflow.contrib.estimator.python.estimator.head import *
 from tensorflow.contrib.estimator.python.estimator.linear import *
@@ -56,6 +57,8 @@ _allowed_symbols = [
     'TowerOptimizer',
     'RNNClassifier',
     'RNNEstimator',
+    'export_saved_model_for_mode',
+    'export_all_saved_models',
 ]
 
 remove_undocumented(__name__, allowed_exception_list=_allowed_symbols)
diff --git a/tensorflow/contrib/estimator/python/estimator/export.py b/tensorflow/contrib/estimator/python/estimator/export.py
new file mode 100644
index 0000000000..e7e366a3f2
--- /dev/null
+++ b/tensorflow/contrib/estimator/python/estimator/export.py
@@ -0,0 +1,216 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Wrapper for methods to export train/eval graphs from Estimator."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.estimator import model_fn as model_fn_lib
+
+
+def export_saved_model_for_mode(
+    estimator, export_dir_base, input_receiver_fn,
+    assets_extra=None,
+    as_text=False,
+    checkpoint_path=None,
+    strip_default_attrs=False,
+    mode=model_fn_lib.ModeKeys.PREDICT):
+  # pylint: disable=line-too-long
+  """Exports a single train/eval/predict graph as a SavedModel.
+
+  For a detailed guide, see
+  @{$saved_model#using_savedmodel_with_estimators$Using SavedModel with Estimators}.
+
+  Sample usage:
+  ```python
+  classifier = tf.estimator.LinearClassifier(
+      feature_columns=[age, language])
+  classifier.train(input_fn=input_fn, steps=1000)
+
+  feature_spec = {
+      'age': tf.placeholder(dtype=tf.int64),
+      'language': array_ops.placeholder(dtype=tf.string)
+  }
+  label_spec = tf.placeholder(dtype=dtypes.int64)
+
+  train_rcvr_fn = tf.contrib.estimator.build_raw_supervised_input_receiver_fn(
+      feature_spec, label_spec)
+
+  export_dir = tf.contrib.estimator.export_saved_model_for_mode(
+      classifier,
+      export_dir_base='my_model/',
+      input_receiver_fn=train_rcvr_fn,
+      mode=model_fn_lib.ModeKeys.TRAIN)
+
+  # export_dir is a timestamped directory with the SavedModel, which
+  # can be used for serving, analysis with TFMA, or directly loaded in.
+  with ops.Graph().as_default() as graph:
+    with session.Session(graph=graph) as sess:
+      loader.load(sess, [tag_constants.TRAINING], export_dir)
+      ...
+  ```
+
+  This method takes an input_receiver_fn and mode. For the mode passed in,
+  this method builds a new graph by calling the input_receiver_fn to obtain
+  feature and label `Tensor`s. Next, this method calls the `Estimator`'s
+  model_fn in the passed mode to generate the model graph based on
+  those features and labels, and restores the given checkpoint
+  (or, lacking that, the most recent checkpoint) into the graph.
+  Finally, it creates a timestamped export directory below the
+  export_dir_base, and writes a `SavedModel` into it containing
+  the `MetaGraphDef` for the given mode and its associated signatures.
+
+  For prediction, the exported `MetaGraphDef` will provide one `SignatureDef`
+  for each element of the export_outputs dict returned from the model_fn,
+  named using the same keys.  One of these keys is always
+  signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY, indicating which
+  signature will be served when a serving request does not specify one.
+  For each signature, the outputs are provided by the corresponding
+  `ExportOutput`s, and the inputs are always the input receivers provided by
+  the serving_input_receiver_fn.
+
+  For training and evaluation, the train_op is stored in an extra collection,
+  and loss, metrics, and predictions are included in a SignatureDef for the
+  mode in question.
+
+  Extra assets may be written into the SavedModel via the assets_extra
+  argument.  This should be a dict, where each key gives a destination path
+  (including the filename) relative to the assets.extra directory.  The
+  corresponding value gives the full path of the source file to be copied.
+  For example, the simple case of copying a single file without renaming it
+  is specified as `{'my_asset_file.txt': '/path/to/my_asset_file.txt'}`.
+
+  Args:
+    estimator: an instance of tf.estimator.Estimator
+    export_dir_base: A string containing a directory in which to create
+      timestamped subdirectories containing exported SavedModels.
+    input_receiver_fn: a function that takes no argument and
+      returns the appropriate subclass of `InputReceiver`.
+    assets_extra: A dict specifying how to populate the assets.extra directory
+      within the exported SavedModel, or `None` if no extra assets are needed.
+    as_text: whether to write the SavedModel proto in text format.
+    checkpoint_path: The checkpoint path to export.  If `None` (the default),
+      the most recent checkpoint found within the model directory is chosen.
+    strip_default_attrs: Boolean. If `True`, default-valued attributes will be
+      removed from the NodeDefs. For a detailed guide, see
+      [Stripping Default-Valued Attributes](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/README.md#stripping-default-valued-attributes).
+    mode: tf.estimator.ModeKeys value indicating with mode will be exported.
+
+  Returns:
+    The string path to the exported directory.
+
+  Raises:
+    ValueError: if input_receiver_fn is None, no export_outputs
+      are provided, or no checkpoint can be found.
+  """
+  # pylint: enable=line-too-long
+
+  # pylint: disable=protected-access
+  return estimator._export_saved_model_for_mode(
+      export_dir_base, input_receiver_fn,
+      assets_extra=assets_extra,
+      as_text=as_text,
+      checkpoint_path=checkpoint_path,
+      strip_default_attrs=strip_default_attrs,
+      mode=mode)
+  # pylint: enable=protected-access
+
+
+def export_all_saved_models(
+    estimator, export_dir_base, input_receiver_fn_map,
+    assets_extra=None,
+    as_text=False,
+    checkpoint_path=None,
+    strip_default_attrs=False):
+  # pylint: disable=line-too-long
+  """Exports requested train/eval/predict graphs as separate SavedModels.
+
+  This is a wrapper around export_saved_model_for_mode that accepts
+  multiple modes simultaneously and creates directories for each under
+  export_dir_base. See `Estimator.export_saved_model_for_mode` for
+  further details as to how the export works for each mode.
+
+  Sample usage:
+  ```python
+  classifier = tf.estimator.LinearClassifier(
+      feature_columns=[age, language])
+  classifier.train(input_fn=input_fn)
+
+  feature_spec = {
+      'age': tf.placeholder(dtype=tf.int64),
+      'language': array_ops.placeholder(dtype=tf.string)
+  }
+  label_spec = tf.placeholder(dtype=dtypes.int64)
+
+  train_rcvr_fn = tf.contrib.estimator.build_raw_supervised_input_receiver_fn(
+      feature_spec, label_spec)
+
+  serve_rcvr_fn = tf.estimator.export.build_parsing_serving_input_receiver_fn(
+      feature_spec)
+
+  rcvr_fn_map = {
+      model_fn_lib.ModeKeys.TRAIN: train_rcvr_fn,
+      model_fn_lib.ModeKeys.PREDICT: serve_rcvr_fn,
+  }
+
+  export_dirs = tf.contrib.estimator.export_all_saved_models(
+      classifier,
+      export_dir_base='my_model/',
+      input_receiver_fn_map=rcvr_fn_map)
+
+  # export_dirs is a dict of directories with SavedModels, which
+  # can be used for serving, analysis with TFMA, or directly loaded in.
+  with ops.Graph().as_default() as graph:
+    with session.Session(graph=graph) as sess:
+      loader.load(sess, [tag_constants.TRAINING],
+                  export_dirs[tf.estimator.ModeKeys.TRAIN])
+      ...
+  ```
+
+  Args:
+    estimator: an instance of tf.estimator.Estimator
+    export_dir_base: A string containing a directory in which to create
+      timestamped subdirectories containing exported SavedModels.
+    input_receiver_fn_map: dict of tf.estimator.ModeKeys to input_receiver_fn
+      mappings, where the input_receiver_fn is a function that takes no
+      argument and returns the appropriate subclass of `InputReceiver`.
+    assets_extra: A dict specifying how to populate the assets.extra directory
+      within the exported SavedModel, or `None` if no extra assets are needed.
+    as_text: whether to write the SavedModel proto in text format.
+    checkpoint_path: The checkpoint path to export.  If `None` (the default),
+      the most recent checkpoint found within the model directory is chosen.
+    strip_default_attrs: Boolean. If `True`, default-valued attributes will be
+      removed from the NodeDefs. For a detailed guide, see
+      [Stripping Default-Valued Attributes](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/README.md#stripping-default-valued-attributes).
+
+  Returns:
+    A dict of tf.estimator.ModeKeys value to string path for each exported
+    directory.
+
+  Raises:
+    ValueError: if any input_receiver_fn is None, no export_outputs
+      are provided, or no checkpoint can be found.
+  """
+  # pylint: enable=line-too-long
+
+  # pylint: disable=protected-access
+  return estimator._export_all_saved_models(
+      export_dir_base, input_receiver_fn_map,
+      assets_extra=assets_extra,
+      as_text=as_text,
+      checkpoint_path=checkpoint_path,
+      strip_default_attrs=strip_default_attrs)
+  # pylint: enable=protected-access
diff --git a/tensorflow/contrib/estimator/python/estimator/export_test.py b/tensorflow/contrib/estimator/python/estimator/export_test.py
new file mode 100644
index 0000000000..89d02582e1
--- /dev/null
+++ b/tensorflow/contrib/estimator/python/estimator/export_test.py
@@ -0,0 +1,391 @@
+# Copyright 2016 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 contrib wrapping of export_saved_model_for_mode functionality.
+
+These are direct copies of the tests included in core, with import locations
+changed. These should be removed when the functionality in core is part of the
+public API.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import tempfile
+
+from tensorflow.contrib.estimator.python.estimator import export as contrib_export
+from tensorflow.python.client import session
+from tensorflow.python.estimator import estimator
+from tensorflow.python.estimator import model_fn as model_fn_lib
+from tensorflow.python.estimator.export import export
+from tensorflow.python.estimator.export import export_output
+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
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import metrics as metrics_lib
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import loader
+from tensorflow.python.saved_model import tag_constants
+from tensorflow.python.training import training
+from tensorflow.python.util import compat
+
+
+def _model_fn_for_export_tests(features, labels, mode):
+  _, _ = features, labels
+  variables.Variable(1., name='weight')
+  scores = constant_op.constant([3.])
+  classes = constant_op.constant(['wumpus'])
+  update_global_step = state_ops.assign_add(training.get_global_step(), 1)
+  with ops.control_dependencies([update_global_step]):
+    train_op = constant_op.constant(2.)
+  return model_fn_lib.EstimatorSpec(
+      mode,
+      predictions=constant_op.constant(10.),
+      loss=constant_op.constant(1.),
+      train_op=train_op,
+      export_outputs={
+          'test': export_output.ClassificationOutput(scores, classes)})
+
+
+def _x_y_input_fn():
+  return ({'x': constant_op.constant([[1], [1]]),
+           'y': constant_op.constant([[2], [2]])},
+          constant_op.constant([[1], [1]]))
+
+
+def _model_fn_with_x_y(features, labels, mode):
+  _ = labels
+  variables.Variable(1., name='weight')
+  scores = constant_op.constant([3.])
+  classes = constant_op.constant(['wumpus'])
+  if mode == model_fn_lib.ModeKeys.PREDICT:
+    variables.Variable(36., name='name_collision')
+    return model_fn_lib.EstimatorSpec(
+        mode,
+        predictions=constant_op.constant(10.),
+        export_outputs={
+            'test': export_output.ClassificationOutput(scores, classes)})
+  else:
+    prefix = 'eval_' if mode == model_fn_lib.ModeKeys.EVAL else ''
+
+    multiplied = math_ops.multiply(
+        features['x'], features['y'], name='{}multiplied'.format(prefix))
+    metrics = {'mean': metrics_lib.mean(features['x'] - features['y'],
+                                        name='{}mean'.format(prefix))}
+    variables.Variable(1., name='later_var')
+    variables.Variable(3., name='name_collision')
+    return model_fn_lib.EstimatorSpec(
+        mode,
+        predictions=multiplied,
+        loss=constant_op.constant(1.),
+        train_op=state_ops.assign_add(training.get_global_step(), 1),
+        eval_metric_ops=metrics)
+
+
+def _get_serving_input_receiver_fn():
+  feature_spec = {'x': parsing_ops.VarLenFeature(dtype=dtypes.int64),
+                  'y': parsing_ops.VarLenFeature(dtype=dtypes.int64)}
+  return export.build_parsing_serving_input_receiver_fn(feature_spec)
+
+
+def _get_supervised_input_receiver_fn():
+  feature_spec = {
+      'x': array_ops.placeholder(
+          dtype=dtypes.int64, shape=(2, 1), name='feature_x'),
+      'y': array_ops.placeholder(
+          dtype=dtypes.int64, shape=(2, 1), name='feature_y')
+      }
+  label_spec = array_ops.placeholder(
+      dtype=dtypes.float32, shape=[1], name='truth')
+
+  return export.build_raw_supervised_input_receiver_fn(
+      feature_spec, label_spec)
+
+
+class EstimatorExportTest(test.TestCase):
+
+  def test_export_saved_model_train(self):
+    self._test_export_saved_model_for_mode(
+        _get_supervised_input_receiver_fn(), model_fn_lib.ModeKeys.TRAIN)
+
+  def test_export_saved_model_eval(self):
+    self._test_export_saved_model_for_mode(
+        _get_supervised_input_receiver_fn(), model_fn_lib.ModeKeys.EVAL)
+
+  def test_export_saved_model_predict(self):
+    self._test_export_saved_model_for_mode(
+        _get_serving_input_receiver_fn(), model_fn_lib.ModeKeys.PREDICT)
+
+  def _test_export_saved_model_for_mode(self, input_receiver_fn, mode):
+    tmpdir = tempfile.mkdtemp()
+    est = estimator.Estimator(model_fn=_model_fn_for_export_tests)
+    est.train(input_fn=_x_y_input_fn, steps=1)
+
+    # Perform the export.
+    export_dir_base = os.path.join(
+        compat.as_bytes(tmpdir), compat.as_bytes('export'))
+    export_dir = contrib_export.export_saved_model_for_mode(
+        est, export_dir_base, input_receiver_fn, mode=mode)
+
+    # Check that all the files are in the right places.
+    self.assertTrue(gfile.Exists(export_dir_base))
+    self._validate_exported_files(export_dir)
+
+    # Restore, to validate that the export was well-formed.
+    tag_set = model_fn_lib.EXPORT_TAG_MAP[mode]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, tag_set, export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertFalse('name_collision_1' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def test_export_all_saved_models_proto_roundtrip_receiver_map(self):
+    input_receiver_fn_map = {
+        model_fn_lib.ModeKeys.PREDICT: _get_serving_input_receiver_fn()
+    }
+    export_dirs, tmpdir = self._test_export_all_saved_models(
+        input_receiver_fn_map)
+
+    self.assertEqual(len(export_dirs), 1)
+    # Restore, to validate that the export was well-formed.
+    export_dir = export_dirs[model_fn_lib.ModeKeys.PREDICT]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.SERVING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('input_example_tensor' in graph_ops)
+        self.assertTrue('ParseExample/ParseExample' in graph_ops)
+        self.assertFalse('feature_x' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def test_export_all_saved_models_proto_roundtrip_train_only(self):
+    input_receiver_fn_map = {
+        model_fn_lib.ModeKeys.TRAIN: _get_supervised_input_receiver_fn(),
+    }
+    export_dirs, tmpdir = self._test_export_all_saved_models(
+        input_receiver_fn_map)
+
+    self.assertEqual(len(export_dirs), 1)
+    # Restore, to validate that the export was well-formed.
+    export_dir = export_dirs[model_fn_lib.ModeKeys.TRAIN]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.TRAINING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('multiplied' in graph_ops)
+        self.assertTrue('mean/update_op' in graph_ops)
+        self.assertFalse('eval_multiplied' in graph_ops)
+        self.assertTrue('feature_x' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def test_export_all_saved_models_proto_roundtrip_eval_only(self):
+    input_receiver_fn_map = {
+        model_fn_lib.ModeKeys.EVAL: _get_supervised_input_receiver_fn()
+    }
+    export_dirs, tmpdir = self._test_export_all_saved_models(
+        input_receiver_fn_map)
+
+    self.assertEqual(len(export_dirs), 1)
+    # Restore, to validate that the export was well-formed.
+    export_dir = export_dirs[model_fn_lib.ModeKeys.EVAL]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.EVAL], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('eval_multiplied' in graph_ops)
+        self.assertTrue('eval_mean/value' in graph_ops)
+        self.assertFalse('multiplied' in graph_ops)
+        self.assertTrue('feature_x' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def test_export_all_saved_models_proto_roundtrip_no_serving(self):
+    input_receiver_fn_map = {
+        model_fn_lib.ModeKeys.TRAIN: _get_supervised_input_receiver_fn(),
+        model_fn_lib.ModeKeys.EVAL: _get_supervised_input_receiver_fn()
+    }
+    export_dirs, tmpdir = self._test_export_all_saved_models(
+        input_receiver_fn_map)
+
+    self.assertEqual(len(export_dirs), 2)
+    # Restore, to validate that the export was well-formed.
+    export_dir = export_dirs[model_fn_lib.ModeKeys.TRAIN]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.TRAINING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('multiplied' in graph_ops)
+        self.assertFalse('eval_multiplied' in graph_ops)
+        self.assertTrue('feature_x' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+    export_dir = export_dirs[model_fn_lib.ModeKeys.EVAL]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.EVAL], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('eval_multiplied' in graph_ops)
+        self.assertFalse('multiplied' in graph_ops)
+        # TODO(karmel): is this the desired behavior when names are shared?
+        self.assertTrue('feature_x_1' in graph_ops)
+        self.assertTrue('feature_y_1' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def test_export_all_saved_models_proto_roundtrip_three_defs(self):
+    input_receiver_fn_map = {
+        model_fn_lib.ModeKeys.TRAIN: _get_supervised_input_receiver_fn(),
+        model_fn_lib.ModeKeys.EVAL: _get_supervised_input_receiver_fn(),
+        model_fn_lib.ModeKeys.PREDICT: _get_serving_input_receiver_fn()
+    }
+    export_dirs, tmpdir = self._test_export_all_saved_models(
+        input_receiver_fn_map)
+
+    # Restore, to validate that the export was well-formed.
+    for mode, tag_set in model_fn_lib.EXPORT_TAG_MAP.items():
+      export_dir = export_dirs[mode]
+      with ops.Graph().as_default() as graph:
+        with session.Session(graph=graph) as sess:
+          loader.load(sess, tag_set, export_dir)
+          graph_ops = [x.name for x in graph.get_operations()]
+          self.assertTrue('global_step/Assign' in graph_ops)
+          self.assertTrue('global_step/Initializer/zeros' in graph_ops)
+          self.assertTrue('weight' in graph_ops)
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def test_export_all_saved_models_proto_roundtrip_all_vars(self):
+    input_receiver_fn_map = {
+        model_fn_lib.ModeKeys.TRAIN: _get_supervised_input_receiver_fn(),
+        model_fn_lib.ModeKeys.PREDICT: _get_serving_input_receiver_fn()
+    }
+    export_dirs, tmpdir = self._test_export_all_saved_models(
+        input_receiver_fn_map)
+
+    export_dir = export_dirs[model_fn_lib.ModeKeys.TRAIN]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.TRAINING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('later_var' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+    export_dir = export_dirs[model_fn_lib.ModeKeys.PREDICT]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.SERVING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertFalse('later_var' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def test_export_all_saved_models_name_collision(self):
+    input_receiver_fn_map = {
+        model_fn_lib.ModeKeys.TRAIN: _get_supervised_input_receiver_fn(),
+        model_fn_lib.ModeKeys.PREDICT: _get_serving_input_receiver_fn()
+    }
+    export_dirs, tmpdir = self._test_export_all_saved_models(
+        input_receiver_fn_map)
+
+    export_dir = export_dirs[model_fn_lib.ModeKeys.TRAIN]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.TRAINING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('name_collision' in graph_ops)
+        self.assertFalse('name_collision_1' in graph_ops)
+        collection_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
+        self.assertEqual(3, collection_vars[-1].eval())
+
+    export_dir = export_dirs[model_fn_lib.ModeKeys.PREDICT]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.SERVING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('name_collision' in graph_ops)
+        self.assertFalse('name_collision_1' in graph_ops)
+        collection_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
+        # This is a non-obvious detail: when we load the estimator spec
+        # for predict, name_collision gets set to 36. However, we then restore
+        # from checkpoint, which should overwrite that var and make it the 3
+        # from training. In practice, this would not be a good way to write
+        # a model_fn, but leaving this check in for now to ensure consistency
+        # with what would happen given our current order of spec, then
+        # checkpoint.
+        self.assertEqual(3, collection_vars[-1].eval())
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def _test_export_all_saved_models(self, input_receiver_fn_map):
+    tmpdir = tempfile.mkdtemp()
+    est = estimator.Estimator(model_fn=_model_fn_with_x_y)
+    est.train(input_fn=_x_y_input_fn, steps=1)
+
+    # Perform the export.
+    export_dir_base = os.path.join(
+        compat.as_bytes(tmpdir), compat.as_bytes('export'))
+    export_dirs = contrib_export.export_all_saved_models(
+        est, export_dir_base, input_receiver_fn_map)
+
+    # Check that all the files are in the right places.
+    self.assertTrue(gfile.Exists(export_dir_base))
+
+    for _, export_dir in export_dirs.items():
+      self._validate_exported_files(export_dir)
+
+    return export_dirs, tmpdir
+
+  def _validate_exported_files(self, export_dir):
+    self.assertTrue(gfile.Exists(export_dir))
+    self.assertTrue(gfile.Exists(os.path.join(
+        compat.as_bytes(export_dir),
+        compat.as_bytes('saved_model.pb'))))
+    self.assertTrue(gfile.Exists(os.path.join(
+        compat.as_bytes(export_dir),
+        compat.as_bytes('variables'))))
+    self.assertTrue(gfile.Exists(os.path.join(
+        compat.as_bytes(export_dir),
+        compat.as_bytes('variables/variables.index'))))
+    self.assertTrue(gfile.Exists(os.path.join(
+        compat.as_bytes(export_dir),
+        compat.as_bytes('variables/variables.data-00000-of-00001'))))
+
+
+if __name__ == '__main__':
+  test.main()
diff --git a/tensorflow/python/estimator/BUILD b/tensorflow/python/estimator/BUILD
index 56dec1eaa1..b25cc7aa26 100644
--- a/tensorflow/python/estimator/BUILD
+++ b/tensorflow/python/estimator/BUILD
@@ -91,6 +91,7 @@ py_library(
         "//tensorflow/python:training",
         "//tensorflow/python:util",
         "//tensorflow/python/saved_model:signature_constants",
+        "//tensorflow/python/saved_model:tag_constants",
         "@six_archive//:six",
     ],
 )
diff --git a/tensorflow/python/estimator/estimator.py b/tensorflow/python/estimator/estimator.py
index 530a4a24ef..9ae64d230e 100644
--- a/tensorflow/python/estimator/estimator.py
+++ b/tensorflow/python/estimator/estimator.py
@@ -37,9 +37,8 @@ from tensorflow.python.eager import context
 from tensorflow.python.estimator import model_fn as model_fn_lib
 from tensorflow.python.estimator import run_config
 from tensorflow.python.estimator import util
-from tensorflow.python.estimator.export.export import build_all_signature_defs
-from tensorflow.python.estimator.export.export import get_temp_export_dir
-from tensorflow.python.estimator.export.export import get_timestamped_export_dir
+from tensorflow.python.estimator.export import export as export_helpers
+from tensorflow.python.estimator.export import export_output
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import random_seed
 from tensorflow.python.ops import array_ops
@@ -51,7 +50,6 @@ from tensorflow.python.platform import gfile
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.saved_model import builder as saved_model_builder
 from tensorflow.python.saved_model import constants
-from tensorflow.python.saved_model import tag_constants
 from tensorflow.python.summary import summary
 from tensorflow.python.summary.writer import writer_cache
 from tensorflow.python.training import device_setter
@@ -609,73 +607,283 @@ class Estimator(object):
           are provided, or no checkpoint can be found.
     """
     # pylint: enable=line-too-long
+    return self._export_saved_model_for_mode(
+        export_dir_base,
+        serving_input_receiver_fn,
+        assets_extra=assets_extra,
+        as_text=as_text,
+        checkpoint_path=checkpoint_path,
+        strip_default_attrs=strip_default_attrs,
+        mode=model_fn_lib.ModeKeys.PREDICT)
+
+  def _export_all_saved_models(
+      self, export_dir_base, input_receiver_fn_map,
+      assets_extra=None,
+      as_text=False,
+      checkpoint_path=None,
+      strip_default_attrs=False):
+    # pylint: disable=line-too-long
+    """Exports requested train/eval/predict graphs as separate SavedModels.
+
+    This is a wrapper around export_saved_model_for_mode that accepts
+    multiple modes simultaneously and creates directories for each under
+    export_dir_base. See `Estimator.export_saved_model_for_mode` for
+    further details as to how the export works for each mode.
+
+    See tf.contrib.estimator.export_all_saved_models for the currently
+    exposed version of this function.
+
+    Args:
+      export_dir_base: A string containing a directory in which to create
+        timestamped subdirectories containing exported SavedModels.
+      input_receiver_fn_map: dict of tf.estimator.ModeKeys to input_receiver_fn
+        mappings, where the input_receiver_fn is a function that takes no
+        argument and returns the appropriate subclass of `InputReceiver`.
+      assets_extra: A dict specifying how to populate the assets.extra directory
+        within the exported SavedModel, or `None` if no extra assets are needed.
+      as_text: whether to write the SavedModel proto in text format.
+      checkpoint_path: The checkpoint path to export.  If `None` (the default),
+        the most recent checkpoint found within the model directory is chosen.
+      strip_default_attrs: Boolean. If `True`, default-valued attributes will be
+        removed from the NodeDefs. For a detailed guide, see
+        [Stripping Default-Valued Attributes](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/README.md#stripping-default-valued-attributes).
+
+    Returns:
+      A dict of tf.estimator.ModeKeys value to string path for each exported
+      directory.
+
+    Raises:
+      ValueError: if any input_receiver_fn is None, no export_outputs
+        are provided, or no checkpoint can be found.
+    """
+    # pylint: enable=line-too-long
+    # TODO(b/65561022): Consider allowing multiple input_receiver_fns per mode.
+    exported = {}
+    for mode, input_receiver_fn in input_receiver_fn_map.items():
+      export_mode_dir = os.path.join(
+          compat.as_bytes(export_dir_base),
+          compat.as_bytes(mode))
+      gfile.MakeDirs(export_mode_dir)
+
+      exported_path = self._export_saved_model_for_mode(
+          export_mode_dir,
+          input_receiver_fn,
+          assets_extra=assets_extra,
+          as_text=as_text,
+          checkpoint_path=checkpoint_path,
+          strip_default_attrs=strip_default_attrs,
+          mode=mode)
+
+      exported[mode] = exported_path
+
+    return exported
+
+  def _export_saved_model_for_mode(
+      self, export_dir_base, input_receiver_fn,
+      assets_extra=None,
+      as_text=False,
+      checkpoint_path=None,
+      strip_default_attrs=False,
+      mode=model_fn_lib.ModeKeys.PREDICT):
+    # pylint: disable=line-too-long
+    """Exports a single train/eval/predict graph as a SavedModel.
+
+    For a detailed guide, see
+    @{$saved_model#using_savedmodel_with_estimators$Using SavedModel with Estimators}.
+
+    See tf.contrib.estimator.export_saved_model_for_mode for the currently
+    exposed version of this function.
+
+    This method takes an input_receiver_fn and mode. For the mode passed in,
+    this method builds a new graph by calling the input_receiver_fn to obtain
+    feature and label `Tensor`s. Next, this method calls the `Estimator`'s
+    model_fn in the passed mode to generate the model graph based on
+    those features and labels, and restores the given checkpoint
+    (or, lacking that, the most recent checkpoint) into the graph.
+    Finally, it creates a timestamped export directory below the
+    export_dir_base, and writes a `SavedModel` into it containing
+    the `MetaGraphDef` for the given mode and its associated signatures.
+
+    For prediction, the exported `MetaGraphDef` will provide one `SignatureDef`
+    for each element of the export_outputs dict returned from the model_fn,
+    named using the same keys.  One of these keys is always
+    signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY, indicating which
+    signature will be served when a serving request does not specify one.
+    For each signature, the outputs are provided by the corresponding
+    `ExportOutput`s, and the inputs are always the input receivers provided by
+    the serving_input_receiver_fn.
+
+    For training and evaluation, the train_op is stored in an extra collection,
+    and loss, metrics, and predictions are included in a SignatureDef for the
+    mode in question.
+
+    Extra assets may be written into the SavedModel via the assets_extra
+    argument.  This should be a dict, where each key gives a destination path
+    (including the filename) relative to the assets.extra directory.  The
+    corresponding value gives the full path of the source file to be copied.
+    For example, the simple case of copying a single file without renaming it
+    is specified as `{'my_asset_file.txt': '/path/to/my_asset_file.txt'}`.
+
+    Args:
+      export_dir_base: A string containing a directory in which to create
+        timestamped subdirectories containing exported SavedModels.
+      input_receiver_fn: a function that takes no argument and
+        returns the appropriate subclass of `InputReceiver`.
+      assets_extra: A dict specifying how to populate the assets.extra directory
+        within the exported SavedModel, or `None` if no extra assets are needed.
+      as_text: whether to write the SavedModel proto in text format.
+      checkpoint_path: The checkpoint path to export.  If `None` (the default),
+        the most recent checkpoint found within the model directory is chosen.
+      strip_default_attrs: Boolean. If `True`, default-valued attributes will be
+        removed from the NodeDefs. For a detailed guide, see
+        [Stripping Default-Valued Attributes](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/README.md#stripping-default-valued-attributes).
+      mode: tf.estimator.ModeKeys value indicating with mode will be exported.
+
+    Returns:
+      The string path to the exported directory.
+
+    Raises:
+      ValueError: if input_receiver_fn is None, no export_outputs
+        are provided, or no checkpoint can be found.
+    """
+    # pylint: enable=line-too-long
     with context.graph_mode():
-      if serving_input_receiver_fn is None:
-        raise ValueError('serving_input_receiver_fn must be defined.')
+      if not input_receiver_fn:
+        raise ValueError('An input_receiver_fn must be defined.')
 
-      with ops.Graph().as_default() as g:
-        self._create_and_assert_global_step(g)
-        random_seed.set_random_seed(self._config.tf_random_seed)
-        serving_input_receiver = serving_input_receiver_fn()
+      if not checkpoint_path:
+        # Locate the latest checkpoint
+        checkpoint_path = saver.latest_checkpoint(self._model_dir)
+      if not checkpoint_path:
+        raise ValueError("Couldn't find trained model at %s." % self._model_dir)
 
-        # Call the model_fn and collect the export_outputs.
-        estimator_spec = self._call_model_fn(
-            features=serving_input_receiver.features,
-            labels=None,
-            mode=model_fn_lib.ModeKeys.PREDICT,
-            config=self.config)
-
-        # Build the SignatureDefs from receivers and all outputs
-        signature_def_map = build_all_signature_defs(
-            serving_input_receiver.receiver_tensors,
-            estimator_spec.export_outputs,
-            serving_input_receiver.receiver_tensors_alternatives)
-
-        if not checkpoint_path:
-          # Locate the latest checkpoint
-          checkpoint_path = saver.latest_checkpoint(self._model_dir)
-        if not checkpoint_path:
-          raise ValueError(
-              "Couldn't find trained model at %s." % self._model_dir)
-
-        export_dir = get_timestamped_export_dir(export_dir_base)
-        temp_export_dir = get_temp_export_dir(export_dir)
-
-        # TODO(soergel): Consider whether MonitoredSession makes sense here
-        with tf_session.Session(config=self._session_config) as session:
-
-          saver_for_restore = estimator_spec.scaffold.saver or saver.Saver(
-              sharded=True)
-          saver_for_restore.restore(session, checkpoint_path)
-
-          local_init_op = (
-              estimator_spec.scaffold.local_init_op or
-              monitored_session.Scaffold.default_local_init_op())
-
-          # Perform the export
-          builder = saved_model_builder.SavedModelBuilder(temp_export_dir)
-          builder.add_meta_graph_and_variables(
-              session, [tag_constants.SERVING],
-              signature_def_map=signature_def_map,
-              assets_collection=ops.get_collection(
-                  ops.GraphKeys.ASSET_FILEPATHS),
-              legacy_init_op=local_init_op,
-              strip_default_attrs=strip_default_attrs)
-          builder.save(as_text)
-
-        # Add the extra assets
-        if assets_extra:
-          assets_extra_path = os.path.join(compat.as_bytes(temp_export_dir),
-                                           compat.as_bytes('assets.extra'))
-          for dest_relative, source in assets_extra.items():
-            dest_absolute = os.path.join(compat.as_bytes(assets_extra_path),
-                                         compat.as_bytes(dest_relative))
-            dest_path = os.path.dirname(dest_absolute)
-            gfile.MakeDirs(dest_path)
-            gfile.Copy(source, dest_absolute)
-
-        gfile.Rename(temp_export_dir, export_dir)
-        return export_dir
+      export_dir = export_helpers.get_timestamped_export_dir(export_dir_base)
+      temp_export_dir = export_helpers.get_temp_export_dir(export_dir)
+
+      builder = saved_model_builder.SavedModelBuilder(temp_export_dir)
+
+      self._add_meta_graph_and_variables_for_mode(
+          builder, input_receiver_fn, checkpoint_path,
+          strip_default_attrs, mode)
+
+      builder.save(as_text)
+
+      # Add the extra assets
+      if assets_extra:
+        assets_extra_path = os.path.join(compat.as_bytes(temp_export_dir),
+                                         compat.as_bytes('assets.extra'))
+        for dest_relative, source in assets_extra.items():
+          dest_absolute = os.path.join(compat.as_bytes(assets_extra_path),
+                                       compat.as_bytes(dest_relative))
+          dest_path = os.path.dirname(dest_absolute)
+          gfile.MakeDirs(dest_path)
+          gfile.Copy(source, dest_absolute)
+
+      gfile.Rename(temp_export_dir, export_dir)
+      return export_dir
+
+  def _add_meta_graph_and_variables_for_mode(
+      self, builder, input_receiver_fn, checkpoint_path, strip_default_attrs,
+      mode=model_fn_lib.ModeKeys.PREDICT):
+    # pylint: disable=line-too-long
+    """Loads variables and adds them along with a MetaGraphDef for saving.
+
+    Args:
+      builder: instance of SavedModelBuilder that will be used for saving.
+      input_receiver_fn: a function that takes no argument and
+        returns the appropriate subclass of `InputReceiver`.
+      checkpoint_path: The checkpoint path to export.  If `None` (the default),
+        the most recent checkpoint found within the model directory is chosen.
+      strip_default_attrs: Boolean. If `True`, default-valued attributes will be
+        removed from the NodeDefs. For a detailed guide, see
+        [Stripping Default-Valued Attributes](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/README.md#stripping-default-valued-attributes).
+      mode: tf.estimator.ModeKeys value indicating which mode will be exported.
+    """
+    # pylint: enable=line-too-long
+    with ops.Graph().as_default() as g:
+      self._create_and_assert_global_step(g)
+      random_seed.set_random_seed(self._config.tf_random_seed)
+
+      input_receiver = input_receiver_fn()
+
+      # Call the model_fn and collect the export_outputs.
+      estimator_spec = self._call_model_fn(
+          features=input_receiver.features,
+          labels=getattr(input_receiver, 'labels', None),
+          mode=mode,
+          config=self.config)
+
+      export_outputs = self._get_export_outputs_for_spec(estimator_spec)
+
+      # Build the SignatureDefs from receivers and all outputs
+      signature_def_map = export_helpers.build_all_signature_defs(
+          input_receiver.receiver_tensors,
+          export_outputs,
+          getattr(input_receiver, 'receiver_tensors_alternatives', None),
+          serving_only=(mode == model_fn_lib.ModeKeys.PREDICT))
+
+      with tf_session.Session(config=self._session_config) as session:
+
+        export_tags = model_fn_lib.EXPORT_TAG_MAP[mode]
+
+        local_init_op = (
+            estimator_spec.scaffold.local_init_op or
+            monitored_session.Scaffold.default_local_init_op())
+
+        saver_for_restore = estimator_spec.scaffold.saver or saver.Saver(
+            sharded=True)
+        saver_for_restore.restore(session, checkpoint_path)
+
+        # We add the train op explicitly for now, so that we don't have to
+        # change the Builder public interface. Note that this is a no-op
+        # for prediction, where train_op is None.
+        builder._add_train_op(estimator_spec.train_op)  # pylint: disable=protected-access
+
+        builder.add_meta_graph_and_variables(
+            session,
+            tags=export_tags,
+            signature_def_map=signature_def_map,
+            assets_collection=ops.get_collection(
+                ops.GraphKeys.ASSET_FILEPATHS),
+            strip_default_attrs=strip_default_attrs,
+            legacy_init_op=local_init_op)
+
+  def _get_export_outputs_for_spec(self, estimator_spec):
+    """Given an EstimatorSpec, determine what our export outputs should be.
+
+    EstimatorSpecs contain export_outputs that are used for serving, but for
+    training and eval graphs, we must wrap the tensors of interest in
+    appropriate ExportOutput objects.
+
+    Args:
+      estimator_spec: EstimatorSpec object that will be exported.
+
+    Returns:
+      a dict mapping export_output_name to ExportOutput object.
+
+    Raises:
+      ValueError: if an appropriate ExportOutput cannot be found for the
+        passed EstimatorSpec.mode
+    """
+    mode = estimator_spec.mode
+    if mode == model_fn_lib.ModeKeys.PREDICT:
+      outputs = estimator_spec.export_outputs
+    else:
+      if mode == model_fn_lib.ModeKeys.TRAIN:
+        output_class = export_output.TrainOutput
+      elif mode == model_fn_lib.ModeKeys.EVAL:
+        output_class = export_output.EvalOutput
+      else:
+        raise ValueError(
+            'Export output type not found for mode: {}'.format(mode))
+
+      export_out = output_class(
+          loss=estimator_spec.loss,
+          predictions=estimator_spec.predictions,
+          metrics=estimator_spec.eval_metric_ops)
+      outputs = {mode: export_out}
+
+    return outputs
 
   def _get_features_from_input_fn(self, input_fn, mode):
     """Extracts the `features` from return values of `input_fn`."""
@@ -1544,3 +1752,5 @@ def _get_default_warm_start_settings(warm_start_from):
   else:
     raise ValueError('warm_start_from must be a string or a WarmStartSettings, '
                      'instead got {}'.format(type(warm_start_from)))
+
+
diff --git a/tensorflow/python/estimator/estimator_test.py b/tensorflow/python/estimator/estimator_test.py
index 76b45b7f57..02088e5134 100644
--- a/tensorflow/python/estimator/estimator_test.py
+++ b/tensorflow/python/estimator/estimator_test.py
@@ -1865,6 +1865,41 @@ def _model_fn_for_export_tests(features, labels, mode):
           'test': export_output.ClassificationOutput(scores, classes)})
 
 
+def _x_y_input_fn():
+  return ({'x': constant_op.constant([[1], [1]]),
+           'y': constant_op.constant([[2], [2]])},
+          constant_op.constant([[1], [1]]))
+
+
+def _model_fn_with_x_y(features, labels, mode):
+  _ = labels
+  variables.Variable(1., name='weight')
+  scores = constant_op.constant([3.])
+  classes = constant_op.constant(['wumpus'])
+  if mode == model_fn_lib.ModeKeys.PREDICT:
+    variables.Variable(36., name='name_collision')
+    return model_fn_lib.EstimatorSpec(
+        mode,
+        predictions=constant_op.constant(10.),
+        export_outputs={
+            'test': export_output.ClassificationOutput(scores, classes)})
+  else:
+    prefix = 'eval_' if mode == model_fn_lib.ModeKeys.EVAL else ''
+
+    multiplied = math_ops.multiply(
+        features['x'], features['y'], name='{}multiplied'.format(prefix))
+    metrics = {'mean': metrics_lib.mean(features['x'] - features['y'],
+                                        name='{}mean'.format(prefix))}
+    variables.Variable(1., name='later_var')
+    variables.Variable(3., name='name_collision')
+    return model_fn_lib.EstimatorSpec(
+        mode,
+        predictions=multiplied,
+        loss=constant_op.constant(1.),
+        train_op=state_ops.assign_add(training.get_global_step(), 1),
+        eval_metric_ops=metrics)
+
+
 def _model_fn_with_saveables_for_export_tests(features, labels, mode):
   _, _ = features, labels
   table = saver_test_utils.CheckpointedOp(name='v2')
@@ -1881,21 +1916,41 @@ def _model_fn_with_saveables_for_export_tests(features, labels, mode):
           'test': export_output.PredictOutput({'prediction': prediction})})
 
 
+def _get_serving_input_receiver_fn():
+  feature_spec = {'x': parsing_ops.VarLenFeature(dtype=dtypes.int64),
+                  'y': parsing_ops.VarLenFeature(dtype=dtypes.int64)}
+  return export.build_parsing_serving_input_receiver_fn(feature_spec)
+
+
+def _get_supervised_input_receiver_fn():
+  feature_spec = {
+      'x': array_ops.placeholder(
+          dtype=dtypes.int64, shape=(2, 1), name='feature_x'),
+      'y': array_ops.placeholder(
+          dtype=dtypes.int64, shape=(2, 1), name='feature_y')
+      }
+  label_spec = array_ops.placeholder(
+      dtype=dtypes.float32, shape=[1], name='truth')
+
+  return export.build_raw_supervised_input_receiver_fn(feature_spec, label_spec)
+
+
 _VOCAB_FILE_CONTENT = 'emerson\nlake\npalmer\n'
 _EXTRA_FILE_CONTENT = 'kermit\npiggy\nralph\n'
 
 
 class EstimatorExportTest(test.TestCase):
 
-  def test_export_savedmodel_proto_roundtrip(self):
-    tmpdir = tempfile.mkdtemp()
-    est = estimator.Estimator(model_fn=_model_fn_for_export_tests)
-    est.train(input_fn=dummy_input_fn, steps=1)
+  def test_export_savedmodel_proto_roundtrip_raw_receiver(self):
     feature_spec = {'x': parsing_ops.VarLenFeature(dtype=dtypes.int64),
                     'y': parsing_ops.VarLenFeature(dtype=dtypes.int64)}
     serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
         feature_spec)
 
+    tmpdir = tempfile.mkdtemp()
+    est = estimator.Estimator(model_fn=_model_fn_for_export_tests)
+    est.train(input_fn=dummy_input_fn, steps=1)
+
     # Perform the export.
     export_dir_base = os.path.join(
         compat.as_bytes(tmpdir), compat.as_bytes('export'))
@@ -1904,6 +1959,266 @@ class EstimatorExportTest(test.TestCase):
 
     # Check that all the files are in the right places.
     self.assertTrue(gfile.Exists(export_dir_base))
+    self._validate_exported_files(export_dir)
+
+    # Restore, to validate that the export was well-formed.
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.SERVING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('input_example_tensor' in graph_ops)
+        self.assertTrue('ParseExample/ParseExample' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+  def test_export_saved_model_train(self):
+    self._test_export_saved_model_for_mode(
+        _get_supervised_input_receiver_fn(), model_fn_lib.ModeKeys.TRAIN)
+
+  def test_export_saved_model_eval(self):
+    self._test_export_saved_model_for_mode(
+        _get_supervised_input_receiver_fn(), model_fn_lib.ModeKeys.EVAL)
+
+  def test_export_saved_model_predict(self):
+    self._test_export_saved_model_for_mode(
+        _get_serving_input_receiver_fn(), model_fn_lib.ModeKeys.PREDICT)
+
+  def _test_export_saved_model_for_mode(self, input_receiver_fn, mode):
+    tmpdir = tempfile.mkdtemp()
+    est = estimator.Estimator(model_fn=_model_fn_for_export_tests)
+    est.train(input_fn=_x_y_input_fn, steps=1)
+
+    # Perform the export.
+    export_dir_base = os.path.join(
+        compat.as_bytes(tmpdir), compat.as_bytes('export'))
+    export_dir = est._export_saved_model_for_mode(
+        export_dir_base, input_receiver_fn, mode=mode)
+
+    # Check that all the files are in the right places.
+    self.assertTrue(gfile.Exists(export_dir_base))
+    self._validate_exported_files(export_dir)
+
+    # Restore, to validate that the export was well-formed.
+    tag_set = model_fn_lib.EXPORT_TAG_MAP[mode]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, tag_set, export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertFalse('name_collision_1' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def test_export_all_saved_models_proto_roundtrip_receiver_map(self):
+    input_receiver_fn_map = {
+        model_fn_lib.ModeKeys.PREDICT: _get_serving_input_receiver_fn()
+    }
+    export_dirs, tmpdir = self._test_export_all_saved_models(
+        input_receiver_fn_map)
+
+    self.assertEqual(len(export_dirs), 1)
+    # Restore, to validate that the export was well-formed.
+    export_dir = export_dirs[model_fn_lib.ModeKeys.PREDICT]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.SERVING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('input_example_tensor' in graph_ops)
+        self.assertTrue('ParseExample/ParseExample' in graph_ops)
+        self.assertFalse('feature_x' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def test_export_all_saved_models_proto_roundtrip_train_only(self):
+    input_receiver_fn_map = {
+        model_fn_lib.ModeKeys.TRAIN: _get_supervised_input_receiver_fn(),
+    }
+    export_dirs, tmpdir = self._test_export_all_saved_models(
+        input_receiver_fn_map)
+
+    self.assertEqual(len(export_dirs), 1)
+    # Restore, to validate that the export was well-formed.
+    export_dir = export_dirs[model_fn_lib.ModeKeys.TRAIN]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.TRAINING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('multiplied' in graph_ops)
+        self.assertTrue('mean/update_op' in graph_ops)
+        self.assertFalse('eval_multiplied' in graph_ops)
+        self.assertTrue('feature_x' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def test_export_all_saved_models_proto_roundtrip_eval_only(self):
+    input_receiver_fn_map = {
+        model_fn_lib.ModeKeys.EVAL: _get_supervised_input_receiver_fn()
+    }
+    export_dirs, tmpdir = self._test_export_all_saved_models(
+        input_receiver_fn_map)
+
+    self.assertEqual(len(export_dirs), 1)
+    # Restore, to validate that the export was well-formed.
+    export_dir = export_dirs[model_fn_lib.ModeKeys.EVAL]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.EVAL], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('eval_multiplied' in graph_ops)
+        self.assertTrue('eval_mean/value' in graph_ops)
+        self.assertFalse('multiplied' in graph_ops)
+        self.assertTrue('feature_x' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def test_export_all_saved_models_proto_roundtrip_no_serving(self):
+    input_receiver_fn_map = {
+        model_fn_lib.ModeKeys.TRAIN: _get_supervised_input_receiver_fn(),
+        model_fn_lib.ModeKeys.EVAL: _get_supervised_input_receiver_fn()
+    }
+    export_dirs, tmpdir = self._test_export_all_saved_models(
+        input_receiver_fn_map)
+
+    self.assertEqual(len(export_dirs), 2)
+    # Restore, to validate that the export was well-formed.
+    export_dir = export_dirs[model_fn_lib.ModeKeys.TRAIN]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.TRAINING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('multiplied' in graph_ops)
+        self.assertFalse('eval_multiplied' in graph_ops)
+        self.assertTrue('feature_x' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+    export_dir = export_dirs[model_fn_lib.ModeKeys.EVAL]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.EVAL], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('eval_multiplied' in graph_ops)
+        self.assertFalse('multiplied' in graph_ops)
+        # TODO(karmel): is this the desired behavior when names are shared?
+        self.assertTrue('feature_x_1' in graph_ops)
+        self.assertTrue('feature_y_1' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def test_export_all_saved_models_proto_roundtrip_three_defs(self):
+    input_receiver_fn_map = {
+        model_fn_lib.ModeKeys.TRAIN: _get_supervised_input_receiver_fn(),
+        model_fn_lib.ModeKeys.EVAL: _get_supervised_input_receiver_fn(),
+        model_fn_lib.ModeKeys.PREDICT: _get_serving_input_receiver_fn()
+    }
+    export_dirs, tmpdir = self._test_export_all_saved_models(
+        input_receiver_fn_map)
+
+    # Restore, to validate that the export was well-formed.
+    for mode, tag_set in model_fn_lib.EXPORT_TAG_MAP.items():
+      export_dir = export_dirs[mode]
+      with ops.Graph().as_default() as graph:
+        with session.Session(graph=graph) as sess:
+          loader.load(sess, tag_set, export_dir)
+          graph_ops = [x.name for x in graph.get_operations()]
+          self.assertTrue('global_step/Assign' in graph_ops)
+          self.assertTrue('global_step/Initializer/zeros' in graph_ops)
+          self.assertTrue('weight' in graph_ops)
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def test_export_all_saved_models_proto_roundtrip_all_vars(self):
+    input_receiver_fn_map = {
+        model_fn_lib.ModeKeys.TRAIN: _get_supervised_input_receiver_fn(),
+        model_fn_lib.ModeKeys.PREDICT: _get_serving_input_receiver_fn()
+    }
+    export_dirs, tmpdir = self._test_export_all_saved_models(
+        input_receiver_fn_map)
+
+    export_dir = export_dirs[model_fn_lib.ModeKeys.TRAIN]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.TRAINING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('later_var' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+    export_dir = export_dirs[model_fn_lib.ModeKeys.PREDICT]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.SERVING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertFalse('later_var' in graph_ops)
+        self.assertTrue('weight' in graph_ops)
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def test_export_all_saved_models_name_collision(self):
+    input_receiver_fn_map = {
+        model_fn_lib.ModeKeys.TRAIN: _get_supervised_input_receiver_fn(),
+        model_fn_lib.ModeKeys.PREDICT: _get_serving_input_receiver_fn()
+    }
+    export_dirs, tmpdir = self._test_export_all_saved_models(
+        input_receiver_fn_map)
+
+    export_dir = export_dirs[model_fn_lib.ModeKeys.TRAIN]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.TRAINING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('name_collision' in graph_ops)
+        self.assertFalse('name_collision_1' in graph_ops)
+        collection_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
+        self.assertEqual(3, collection_vars[-1].eval())
+
+    export_dir = export_dirs[model_fn_lib.ModeKeys.PREDICT]
+    with ops.Graph().as_default() as graph:
+      with session.Session(graph=graph) as sess:
+        loader.load(sess, [tag_constants.SERVING], export_dir)
+        graph_ops = [x.name for x in graph.get_operations()]
+        self.assertTrue('name_collision' in graph_ops)
+        self.assertFalse('name_collision_1' in graph_ops)
+        collection_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
+        # This is a non-obvious detail: when we load the estimator spec
+        # for predict, name_collision gets set to 36. However, we then restore
+        # from checkpoint, which should overwrite that var and make it the 3
+        # from training. In practice, this would not be a good way to write
+        # a model_fn, but leaving this check in for now to ensure consistency
+        # with what would happen given our current order of spec, then
+        # checkpoint.
+        self.assertEqual(3, collection_vars[-1].eval())
+
+    # Clean up.
+    gfile.DeleteRecursively(tmpdir)
+
+  def _test_export_all_saved_models(self, input_receiver_fn_map):
+    tmpdir = tempfile.mkdtemp()
+    est = estimator.Estimator(model_fn=_model_fn_with_x_y)
+    est.train(input_fn=_x_y_input_fn, steps=1)
+
+    # Perform the export.
+    export_dir_base = os.path.join(
+        compat.as_bytes(tmpdir), compat.as_bytes('export'))
+    export_dirs = est._export_all_saved_models(
+        export_dir_base, input_receiver_fn_map)
+
+    # Check that all the files are in the right places.
+    self.assertTrue(gfile.Exists(export_dir_base))
+
+    for _, export_dir in export_dirs.items():
+      self._validate_exported_files(export_dir)
+
+    return export_dirs, tmpdir
+
+  def _validate_exported_files(self, export_dir):
     self.assertTrue(gfile.Exists(export_dir))
     self.assertTrue(gfile.Exists(os.path.join(
         compat.as_bytes(export_dir),
@@ -1918,18 +2233,6 @@ class EstimatorExportTest(test.TestCase):
         compat.as_bytes(export_dir),
         compat.as_bytes('variables/variables.data-00000-of-00001'))))
 
-    # Restore, to validate that the export was well-formed.
-    with ops.Graph().as_default() as graph:
-      with session.Session(graph=graph) as sess:
-        loader.load(sess, [tag_constants.SERVING], export_dir)
-        graph_ops = [x.name for x in graph.get_operations()]
-        self.assertTrue('input_example_tensor' in graph_ops)
-        self.assertTrue('ParseExample/ParseExample' in graph_ops)
-        self.assertTrue('weight' in graph_ops)
-
-    # Clean up.
-    gfile.DeleteRecursively(tmpdir)
-
   def test_export_savedmodel_with_saveables_proto_roundtrip(self):
     tmpdir = tempfile.mkdtemp()
     est = estimator.Estimator(
@@ -2485,5 +2788,6 @@ class EstimatorIntegrationTest(test.TestCase):
                                        serving_input_receiver_fn)
     self.assertTrue(gfile.Exists(export_dir))
 
+
 if __name__ == '__main__':
   test.main()
diff --git a/tensorflow/python/estimator/export/export.py b/tensorflow/python/estimator/export/export.py
index 41c1f5a2e2..9aafb56679 100644
--- a/tensorflow/python/estimator/export/export.py
+++ b/tensorflow/python/estimator/export/export.py
@@ -40,6 +40,60 @@ from tensorflow.python.util.tf_export import tf_export
 
 _SINGLE_FEATURE_DEFAULT_NAME = 'feature'
 _SINGLE_RECEIVER_DEFAULT_NAME = 'input'
+_SINGLE_LABEL_DEFAULT_NAME = 'label'
+
+
+def _wrap_and_check_receiver_tensors(receiver_tensors):
+  """Ensure that receiver_tensors is a dict of str to Tensor mappings.
+
+  Args:
+    receiver_tensors: dict of str to Tensors, or a single Tensor.
+
+  Returns:
+    dict of str to Tensors; this is the original dict if one was passed, or
+    the original tensor wrapped in a dictionary.
+
+  Raises:
+    ValueError: if receiver_tensors is None, or has non-string keys,
+      or non-Tensor values
+  """
+  if receiver_tensors is None:
+    raise ValueError('receiver_tensors must be defined.')
+  if not isinstance(receiver_tensors, dict):
+    receiver_tensors = {_SINGLE_RECEIVER_DEFAULT_NAME: receiver_tensors}
+  for name, tensor in receiver_tensors.items():
+    _check_tensor_key(name, error_label='receiver_tensors')
+    _check_tensor(tensor, name, error_label='receiver_tensor')
+  return receiver_tensors
+
+
+def _check_tensor(tensor, name, error_label='feature'):
+  """Check that passed `tensor` is a Tensor or SparseTensor."""
+  if not (isinstance(tensor, ops.Tensor)
+          or isinstance(tensor, sparse_tensor.SparseTensor)):
+    fmt_name = ' {}'.format(name) if name else ''
+    value_error = ValueError(
+        '{}{} must be a Tensor or SparseTensor.'.format(error_label, fmt_name))
+    # NOTE(ericmc): This if-else block is a specific carve-out for
+    # LabeledTensor, which has a `.tensor` attribute and which is
+    # convertible to tf.Tensor via ops.convert_to_tensor.
+    # Allowing all types convertible to tf.Tensor is considered by soergel@
+    # to be too permissive.
+    # TODO(soergel): accept any type convertible to Tensor,
+    # as in cl/193238295 snapshot #6.
+    if hasattr(tensor, 'tensor'):
+      try:
+        ops.convert_to_tensor(tensor)
+      except TypeError:
+        raise value_error
+    else:
+      raise value_error
+
+
+def _check_tensor_key(name, error_label='feature'):
+  if not isinstance(name, six.string_types):
+    raise ValueError(
+        '{} keys must be strings: {}.'.format(error_label, name))
 
 
 @tf_export('estimator.export.ServingInputReceiver')
@@ -51,16 +105,18 @@ class ServingInputReceiver(collections.namedtuple(
   The expected return values are:
     features: A `Tensor`, `SparseTensor`, or dict of string to `Tensor` or
       `SparseTensor`, specifying the features to be passed to the model.
-    receiver_tensors: a `Tensor`, or a dict of string to `Tensor`, specifying
-      input nodes where this receiver expects to be fed by default.  Typically,
-      this is a single placeholder expecting serialized `tf.Example` protos.
+    receiver_tensors: A `Tensor`, `SparseTensor`, or dict of string to `Tensor`
+      or `SparseTensor`, specifying input nodes where this receiver expects to
+      be fed by default.  Typically, this is a single placeholder expecting
+      serialized `tf.Example` protos.
     receiver_tensors_alternatives: a dict of string to additional
-      groups of receiver tensors, each of which may be a `Tensor` or a dict of
-      string to `Tensor`.  These named receiver tensor alternatives generate
-      additional serving signatures, which may be used to feed inputs at
-      different points within the input receiver subgraph.  A typical usage is
-      to allow feeding raw feature `Tensor`s *downstream* of the
-      tf.parse_example() op.  Defaults to None.
+      groups of receiver tensors, each of which may be a `Tensor`,
+      `SparseTensor`, or dict of string to `Tensor` or`SparseTensor`.
+      These named receiver tensor alternatives generate additional serving
+      signatures, which may be used to feed inputs at different points within
+      the input receiver subgraph.  A typical usage is to allow feeding raw
+      feature `Tensor`s *downstream* of the tf.parse_example() op.
+      Defaults to None.
   """
 
   def __new__(cls, features, receiver_tensors,
@@ -70,36 +126,10 @@ class ServingInputReceiver(collections.namedtuple(
     if not isinstance(features, dict):
       features = {_SINGLE_FEATURE_DEFAULT_NAME: features}
     for name, tensor in features.items():
-      if not isinstance(name, six.string_types):
-        raise ValueError('feature keys must be strings: {}.'.format(name))
-      if not (isinstance(tensor, ops.Tensor)
-              or isinstance(tensor, sparse_tensor.SparseTensor)):
-        value_error = ValueError(
-            'feature {} must be a Tensor or SparseTensor.'.format(name))
-        # NOTE(ericmc): This if-else block is a specific carve-out for
-        # LabeledTensor, which has a `.tensor` attribute and which is
-        # convertible to tf.Tensor via ops.convert_to_tensor.
-        # Allowing all types convertible to tf.Tensor is considered by soergel@
-        # to be too permissive.
-        if hasattr(tensor, 'tensor'):
-          try:
-            ops.convert_to_tensor(tensor)
-          except TypeError:
-            raise value_error
-        else:
-          raise value_error
-
-    if receiver_tensors is None:
-      raise ValueError('receiver_tensors must be defined.')
-    if not isinstance(receiver_tensors, dict):
-      receiver_tensors = {_SINGLE_RECEIVER_DEFAULT_NAME: receiver_tensors}
-    for name, tensor in receiver_tensors.items():
-      if not isinstance(name, six.string_types):
-        raise ValueError(
-            'receiver_tensors keys must be strings: {}.'.format(name))
-      if not isinstance(tensor, ops.Tensor):
-        raise ValueError(
-            'receiver_tensor {} must be a Tensor.'.format(name))
+      _check_tensor_key(name)
+      _check_tensor(tensor, name)
+
+    receiver_tensors = _wrap_and_check_receiver_tensors(receiver_tensors)
 
     if receiver_tensors_alternatives is not None:
       if not isinstance(receiver_tensors_alternatives, dict):
@@ -115,14 +145,9 @@ class ServingInputReceiver(collections.namedtuple(
           receiver_tensors_alternatives[alternative_name] = (
               receiver_tensors_alt)
         for name, tensor in receiver_tensors_alt.items():
-          if not isinstance(name, six.string_types):
-            raise ValueError(
-                'receiver_tensors keys must be strings: {}.'.format(name))
-          if not (isinstance(tensor, ops.Tensor)
-                  or isinstance(tensor, sparse_tensor.SparseTensor)):
-            raise ValueError(
-                'receiver_tensor {} must be a Tensor or SparseTensor.'.format(
-                    name))
+          _check_tensor_key(name, error_label='receiver_tensors_alternative')
+          _check_tensor(
+              tensor, name, error_label='receiver_tensors_alternative')
 
     return super(ServingInputReceiver, cls).__new__(
         cls,
@@ -155,25 +180,25 @@ class TensorServingInputReceiver(collections.namedtuple(
   The expected return values are:
     features: A single `Tensor` or `SparseTensor`, representing the feature
       to be passed to the model.
-    receiver_tensors: a `Tensor`, or a dict of string to `Tensor`, specifying
-      input nodes where this receiver expects to be fed by default.  Typically,
-      this is a single placeholder expecting serialized `tf.Example` protos.
+    receiver_tensors: A `Tensor`, `SparseTensor`, or dict of string to `Tensor`
+      or `SparseTensor`, specifying input nodes where this receiver expects to
+      be fed by default.  Typically, this is a single placeholder expecting
+      serialized `tf.Example` protos.
     receiver_tensors_alternatives: a dict of string to additional
-      groups of receiver tensors, each of which may be a `Tensor` or a dict of
-      string to `Tensor`.  These named receiver tensor alternatives generate
-      additional serving signatures, which may be used to feed inputs at
-      different points within the input receiver subgraph.  A typical usage is
-      to allow feeding raw feature `Tensor`s *downstream* of the
-      tf.parse_example() op.  Defaults to None.
+      groups of receiver tensors, each of which may be a `Tensor`,
+      `SparseTensor`, or dict of string to `Tensor` or`SparseTensor`.
+      These named receiver tensor alternatives generate additional serving
+      signatures, which may be used to feed inputs at different points within
+      the input receiver subgraph.  A typical usage is to allow feeding raw
+      feature `Tensor`s *downstream* of the tf.parse_example() op.
+      Defaults to None.
   """
 
   def __new__(cls, features, receiver_tensors,
               receiver_tensors_alternatives=None):
     if features is None:
       raise ValueError('features must be defined.')
-    if not (isinstance(features, ops.Tensor)
-            or isinstance(features, sparse_tensor.SparseTensor)):
-      raise ValueError('feature must be a Tensor or SparseTensor.')
+    _check_tensor(features, None)
 
     receiver = ServingInputReceiver(
         features=features,
@@ -187,6 +212,49 @@ class TensorServingInputReceiver(collections.namedtuple(
         receiver_tensors_alternatives=receiver.receiver_tensors_alternatives)
 
 
+class SupervisedInputReceiver(collections.namedtuple(
+    'SupervisedInputReceiver',
+    ['features', 'labels', 'receiver_tensors'])):
+  """A return type for a training_input_receiver_fn or eval_input_receiver_fn.
+
+  This differs from a ServingInputReceiver in that (1) this receiver expects
+  a set of labels to be passed in with features, and (2) this receiver does
+  not support receiver_tensors_alternatives, which are primarily used for
+  serving.
+
+  The expected return values are:
+    features: A `Tensor`, `SparseTensor`, or dict of string to `Tensor` or
+      `SparseTensor`, specifying the features to be passed to the model.
+    labels: A `Tensor`, `SparseTensor`, or dict of string to `Tensor` or
+      `SparseTensor`, specifying the labels to be passed to the model.
+    receiver_tensors: A `Tensor`, `SparseTensor`, or dict of string to `Tensor`
+      or `SparseTensor`, specifying input nodes where this receiver expects to
+      be fed by default.  Typically, this is a single placeholder expecting
+      serialized `tf.Example` protos.
+
+  """
+
+  def __new__(cls, features, labels, receiver_tensors):
+    # Both features and labels can be dicts or raw tensors.
+    for input_vals, error_label in ((features, 'feature'), (labels, 'label')):
+      if input_vals is None:
+        raise ValueError('{}s must be defined.'.format(error_label))
+      if isinstance(input_vals, dict):
+        for name, tensor in input_vals.items():
+          _check_tensor_key(name, error_label=error_label)
+          _check_tensor(tensor, name, error_label=error_label)
+      else:
+        _check_tensor(input_vals, None, error_label=error_label)
+
+    receiver_tensors = _wrap_and_check_receiver_tensors(receiver_tensors)
+
+    return super(SupervisedInputReceiver, cls).__new__(
+        cls,
+        features=features,
+        labels=labels,
+        receiver_tensors=receiver_tensors)
+
+
 @tf_export('estimator.export.build_parsing_serving_input_receiver_fn')
 def build_parsing_serving_input_receiver_fn(feature_spec,
                                             default_batch_size=None):
@@ -216,6 +284,23 @@ def build_parsing_serving_input_receiver_fn(feature_spec,
   return serving_input_receiver_fn
 
 
+def _placeholder_from_tensor(t, default_batch_size=None):
+  shape_list = t.get_shape().as_list()
+  shape_list[0] = default_batch_size
+  shape = tensor_shape.TensorShape(shape_list)
+
+  # Reuse the feature tensor's op name (t.op.name) for the placeholder,
+  # excluding the index from the tensor's name (t.name):
+  # t.name = "%s:%d" % (t.op.name, t._value_index)
+  return array_ops.placeholder(dtype=t.dtype, shape=shape, name=t.op.name)
+
+
+def _placeholders_from_receiver_tensors_dict(
+    input_vals, default_batch_size=None):
+  return {name: _placeholder_from_tensor(t, default_batch_size)
+          for name, t in input_vals.items()}
+
+
 @tf_export('estimator.export.build_raw_serving_input_receiver_fn')
 def build_raw_serving_input_receiver_fn(features, default_batch_size=None):
   """Build a serving_input_receiver_fn expecting feature Tensors.
@@ -233,17 +318,9 @@ def build_raw_serving_input_receiver_fn(features, default_batch_size=None):
   """
   def serving_input_receiver_fn():
     """A serving_input_receiver_fn that expects features to be fed directly."""
-    receiver_tensors = {}
-    for name, t in features.items():
-      shape_list = t.get_shape().as_list()
-      shape_list[0] = default_batch_size
-      shape = tensor_shape.TensorShape(shape_list)
-
-      # Reuse the feature tensor's op name (t.op.name) for the placeholder,
-      # excluding the index from the tensor's name (t.name):
-      # t.name = "%s:%d" % (t.op.name, t._value_index)
-      receiver_tensors[name] = array_ops.placeholder(
-          dtype=t.dtype, shape=shape, name=t.op.name)
+    receiver_tensors = _placeholders_from_receiver_tensors_dict(
+        features, default_batch_size)
+
     # TODO(b/34885899): remove the unnecessary copy
     # The features provided are simply the placeholders, but we defensively copy
     # the dict because it may be mutated.
@@ -252,13 +329,100 @@ def build_raw_serving_input_receiver_fn(features, default_batch_size=None):
   return serving_input_receiver_fn
 
 
+def build_raw_supervised_input_receiver_fn(
+    features, labels, default_batch_size=None):
+  """Build a supervised_input_receiver_fn for raw features and labels.
+
+  This function wraps tensor placeholders in a supervised_receiver_fn
+  with the expectation that the features and labels appear precisely as
+  the model_fn expects them. Features and labels can therefore be dicts of
+  tensors, or raw tensors.
+
+  Args:
+    features: a dict of string to `Tensor` or `Tensor`.
+    labels: a dict of string to `Tensor` or `Tensor`.
+    default_batch_size: the number of query examples expected per batch.
+        Leave unset for variable batch size (recommended).
+
+  Returns:
+    A supervised_input_receiver_fn.
+
+  Raises:
+    ValueError: if features and labels have overlapping keys.
+  """
+  # Check for overlapping keys before beginning.
+  try:
+    feat_keys = features.keys()
+  except AttributeError:
+    feat_keys = [_SINGLE_RECEIVER_DEFAULT_NAME]
+  try:
+    label_keys = labels.keys()
+  except AttributeError:
+    label_keys = [_SINGLE_LABEL_DEFAULT_NAME]
+
+  overlap_keys = set(feat_keys) & set(label_keys)
+  if overlap_keys:
+    raise ValueError('Features and labels must have distinct keys. '
+                     'Found overlapping keys: {}'.format(overlap_keys))
+
+  def supervised_input_receiver_fn():
+    """A receiver_fn that expects pass-through features and labels."""
+    if not isinstance(features, dict):
+      features_cp = _placeholder_from_tensor(features, default_batch_size)
+      receiver_features = {_SINGLE_RECEIVER_DEFAULT_NAME: features_cp}
+    else:
+      receiver_features = _placeholders_from_receiver_tensors_dict(
+          features, default_batch_size)
+      features_cp = receiver_features
+
+    if not isinstance(labels, dict):
+      labels_cp = _placeholder_from_tensor(labels, default_batch_size)
+      receiver_labels = {_SINGLE_LABEL_DEFAULT_NAME: labels_cp}
+    else:
+      receiver_labels = _placeholders_from_receiver_tensors_dict(
+          labels, default_batch_size)
+      labels_cp = receiver_labels
+
+    receiver_tensors = dict(receiver_features)
+    receiver_tensors.update(receiver_labels)
+    return SupervisedInputReceiver(features_cp, labels_cp, receiver_tensors)
+
+  return supervised_input_receiver_fn
+
+
 ### Below utilities are specific to SavedModel exports.
 
 
 def build_all_signature_defs(receiver_tensors,
                              export_outputs,
-                             receiver_tensors_alternatives=None):
-  """Build `SignatureDef`s for all export outputs."""
+                             receiver_tensors_alternatives=None,
+                             serving_only=True):
+  """Build `SignatureDef`s for all export outputs.
+
+  Args:
+    receiver_tensors: a `Tensor`, or a dict of string to `Tensor`, specifying
+      input nodes where this receiver expects to be fed by default.  Typically,
+      this is a single placeholder expecting serialized `tf.Example` protos.
+    export_outputs: a dict of ExportOutput instances, each of which has
+      an as_signature_def instance method that will be called to retrieve
+      the signature_def for all export output tensors.
+    receiver_tensors_alternatives: a dict of string to additional
+      groups of receiver tensors, each of which may be a `Tensor` or a dict of
+      string to `Tensor`.  These named receiver tensor alternatives generate
+      additional serving signatures, which may be used to feed inputs at
+      different points within the input receiver subgraph.  A typical usage is
+      to allow feeding raw feature `Tensor`s *downstream* of the
+      tf.parse_example() op.  Defaults to None.
+    serving_only: boolean; if true, resulting signature defs will only include
+      valid serving signatures. If false, all requested signatures will be
+      returned.
+
+  Returns:
+    signature_def representing all passed args.
+
+  Raises:
+    ValueError: if export_outputs is not a dict
+  """
   if not isinstance(receiver_tensors, dict):
     receiver_tensors = {_SINGLE_RECEIVER_DEFAULT_NAME: receiver_tensors}
   if export_outputs is None or not isinstance(export_outputs, dict):
@@ -293,17 +457,24 @@ def build_all_signature_defs(receiver_tensors,
   _log_signature_report(signature_def_map, excluded_signatures)
 
   # The above calls to export_output.as_signature_def should return only
-  # valid signatures; if there is a validity problem, they raise ValueError,
-  # which we ignore above. Consequently the call to is_valid_signature here
-  # should not remove anything else; it's just an extra sanity check.
-  return {k: v for k, v in signature_def_map.items()
-          if signature_def_utils.is_valid_signature(v)}
+  # valid signatures; if there is a validity problem, they raise a ValueError,
+  # in which case we exclude that signature from signature_def_map above.
+  # The is_valid_signature check ensures that the signatures produced are
+  # valid for serving, and acts as an additional sanity check for export
+  # signatures produced for serving. We skip this check for training and eval
+  # signatures, which are not intended for serving.
+  if serving_only:
+    signature_def_map = {k: v for k, v in signature_def_map.items()
+                         if signature_def_utils.is_valid_signature(v)}
+  return signature_def_map
 
 
 _FRIENDLY_METHOD_NAMES = {
     signature_constants.CLASSIFY_METHOD_NAME: 'Classify',
     signature_constants.REGRESS_METHOD_NAME: 'Regress',
     signature_constants.PREDICT_METHOD_NAME: 'Predict',
+    signature_constants.SUPERVISED_TRAIN_METHOD_NAME: 'Train',
+    signature_constants.SUPERVISED_EVAL_METHOD_NAME: 'Eval',
 }
 
 
diff --git a/tensorflow/python/estimator/export/export_output.py b/tensorflow/python/estimator/export/export_output.py
index 87b964be37..d387ea2940 100644
--- a/tensorflow/python/estimator/export/export_output.py
+++ b/tensorflow/python/estimator/export/export_output.py
@@ -38,6 +38,8 @@ class ExportOutput(object):
 
   __metaclass__ = abc.ABCMeta
 
+  _SEPARATOR_CHAR = '/'
+
   @abc.abstractmethod
   def as_signature_def(self, receiver_tensors):
     """Generate a SignatureDef proto for inclusion in a MetaGraphDef.
@@ -51,6 +53,52 @@ class ExportOutput(object):
     """
     pass
 
+  def _check_output_key(self, key, error_label):
+    # For multi-head models, the key can be a tuple.
+    if isinstance(key, tuple):
+      key = self._SEPARATOR_CHAR.join(key)
+
+    if not isinstance(key, six.string_types):
+      raise ValueError(
+          '{} output key must be a string; got {}.'.format(error_label, key))
+    return key
+
+  def _wrap_and_check_outputs(
+      self, outputs, single_output_default_name, error_label=None):
+    """Wraps raw tensors as dicts and checks type.
+
+    Note that we create a new dict here so that we can overwrite the keys
+    if necessary.
+
+    Args:
+      outputs: A `Tensor` or a dict of string to `Tensor`.
+      single_output_default_name: A string key for use in the output dict
+        if the provided `outputs` is a raw tensor.
+      error_label: descriptive string for use in error messages. If none,
+        single_output_default_name will be used.
+
+    Returns:
+      A dict of tensors
+
+    Raises:
+      ValueError: if the outputs dict keys are not strings or tuples of strings
+        or the values are not Tensors.
+    """
+    if not isinstance(outputs, dict):
+      outputs = {single_output_default_name: outputs}
+
+    output_dict = {}
+    for key, value in outputs.items():
+      error_name = error_label or single_output_default_name
+      key = self._check_output_key(key, error_name)
+      if not isinstance(value, ops.Tensor):
+        raise ValueError(
+            '{} output value must be a Tensor; got {}.'.format(
+                error_name, value))
+
+      output_dict[key] = value
+    return output_dict
+
 
 @tf_export('estimator.export.ClassificationOutput')
 class ClassificationOutput(ExportOutput):
@@ -154,9 +202,6 @@ class RegressionOutput(ExportOutput):
     return signature_def_utils.regression_signature_def(examples, self.value)
 
 
-_SINGLE_OUTPUT_DEFAULT_NAME = 'output'
-
-
 @tf_export('estimator.export.PredictOutput')
 class PredictOutput(ExportOutput):
   """Represents the output of a generic prediction head.
@@ -165,6 +210,7 @@ class PredictOutput(ExportOutput):
 
   Named outputs must be provided as a dict from string to `Tensor`,
   """
+  _SINGLE_OUTPUT_DEFAULT_NAME = 'output'
 
   def __init__(self, outputs):
     """Constructor for PredictOutput.
@@ -177,16 +223,9 @@ class PredictOutput(ExportOutput):
       ValueError: if the outputs is not dict, or any of its keys are not
           strings, or any of its values are not `Tensor`s.
     """
-    if not isinstance(outputs, dict):
-      outputs = {_SINGLE_OUTPUT_DEFAULT_NAME: outputs}
-    for key, value in outputs.items():
-      if not isinstance(key, six.string_types):
-        raise ValueError(
-            'Prediction output key must be a string; got {}.'.format(key))
-      if not isinstance(value, ops.Tensor):
-        raise ValueError(
-            'Prediction output value must be a Tensor; got {}.'.format(value))
-    self._outputs = outputs
+
+    self._outputs = self._wrap_and_check_outputs(
+        outputs, self._SINGLE_OUTPUT_DEFAULT_NAME, error_label='Prediction')
 
   @property
   def outputs(self):
@@ -195,3 +234,161 @@ class PredictOutput(ExportOutput):
   def as_signature_def(self, receiver_tensors):
     return signature_def_utils.predict_signature_def(receiver_tensors,
                                                      self.outputs)
+
+
+class _SupervisedOutput(ExportOutput):
+  """Represents the output of a supervised training or eval process."""
+  __metaclass__ = abc.ABCMeta
+
+  LOSS_NAME = 'loss'
+  PREDICTIONS_NAME = 'predictions'
+  METRICS_NAME = 'metrics'
+
+  METRIC_VALUE_SUFFIX = 'value'
+  METRIC_UPDATE_SUFFIX = 'update_op'
+
+  _loss = None
+  _predictions = None
+  _metrics = None
+
+  def __init__(self, loss=None, predictions=None, metrics=None):
+    """Constructor for SupervisedOutput (ie, Train or Eval output).
+
+    Args:
+      loss: dict of Tensors or single Tensor representing calculated loss.
+      predictions: dict of Tensors or single Tensor representing model
+        predictions.
+      metrics: dict of (metric_value, update_op) tuples, or a single tuple.
+        metric_value must be a Tensor, and update_op must be a Tensor or Op.
+
+    Raises:
+      ValueError: if any of the outputs' dict keys are not strings or tuples of
+        strings or the values are not Tensors (or Operations in the case of
+        update_op).
+    """
+
+    if loss is not None:
+      loss_dict = self._wrap_and_check_outputs(loss, self.LOSS_NAME)
+      self._loss = self._prefix_output_keys(loss_dict, self.LOSS_NAME)
+    if predictions is not None:
+      pred_dict = self._wrap_and_check_outputs(
+          predictions, self.PREDICTIONS_NAME)
+      self._predictions = self._prefix_output_keys(
+          pred_dict, self.PREDICTIONS_NAME)
+    if metrics is not None:
+      self._metrics = self._wrap_and_check_metrics(metrics)
+
+  def _prefix_output_keys(self, output_dict, output_name):
+    """Prepend output_name to the output_dict keys if it doesn't exist.
+
+    This produces predictable prefixes for the pre-determined outputs
+    of SupervisedOutput.
+
+    Args:
+      output_dict: dict of string to Tensor, assumed valid.
+      output_name: prefix string to prepend to existing keys.
+
+    Returns:
+      dict with updated keys and existing values.
+    """
+
+    new_outputs = {}
+    for key, val in output_dict.items():
+      key = self._prefix_key(key, output_name)
+      new_outputs[key] = val
+    return new_outputs
+
+  def _prefix_key(self, key, output_name):
+    if key.find(output_name) != 0:
+      key = output_name + self._SEPARATOR_CHAR + key
+    return key
+
+  def _wrap_and_check_metrics(self, metrics):
+    """Handle the saving of metrics.
+
+    Metrics is either a tuple of (value, update_op), or a dict of such tuples.
+    Here, we separate out the tuples and create a dict with names to tensors.
+
+    Args:
+      metrics: dict of (metric_value, update_op) tuples, or a single tuple.
+
+    Returns:
+      dict of output_names to tensors
+
+    Raises:
+      ValueError: if the dict key is not a string, or the metric values or ops
+        are not tensors.
+    """
+    if not isinstance(metrics, dict):
+      metrics = {self.METRICS_NAME: metrics}
+
+    outputs = {}
+    for key, (metric_val, metric_op) in metrics.items():
+      key = self._check_output_key(key, self.METRICS_NAME)
+      key = self._prefix_key(key, self.METRICS_NAME)
+
+      val_name = key + self._SEPARATOR_CHAR + self.METRIC_VALUE_SUFFIX
+      op_name = key + self._SEPARATOR_CHAR + self.METRIC_UPDATE_SUFFIX
+      if not isinstance(metric_val, ops.Tensor):
+        raise ValueError(
+            '{} output value must be a Tensor; got {}.'.format(
+                key, metric_val))
+      if (not isinstance(metric_op, ops.Tensor) and
+          not isinstance(metric_op, ops.Operation)):
+        raise ValueError(
+            '{} update_op must be a Tensor or Operation; got {}.'.format(
+                key, metric_op))
+      outputs[val_name] = metric_val
+      outputs[op_name] = metric_op
+
+    return outputs
+
+  @property
+  def loss(self):
+    return self._loss
+
+  @property
+  def predictions(self):
+    return self._predictions
+
+  @property
+  def metrics(self):
+    return self._metrics
+
+  @abc.abstractmethod
+  def _get_signature_def_fn(self):
+    """Returns a function that produces a SignatureDef given desired outputs."""
+    pass
+
+  def as_signature_def(self, receiver_tensors):
+    signature_def_fn = self._get_signature_def_fn()
+    return signature_def_fn(
+        receiver_tensors, self.loss, self.predictions, self.metrics)
+
+
+class TrainOutput(_SupervisedOutput):
+  """Represents the output of a supervised training process.
+
+  This class generates the appropriate signature def for exporting
+  training output by type-checking and wrapping loss, predictions, and metrics
+  values.
+  """
+
+  def _get_signature_def_fn(self):
+    return signature_def_utils.supervised_train_signature_def
+
+
+class EvalOutput(_SupervisedOutput):
+  """Represents the output of a supervised eval process.
+
+  This class generates the appropriate signature def for exporting
+  eval output by type-checking and wrapping loss, predictions, and metrics
+  values.
+  """
+
+  def _get_signature_def_fn(self):
+    return signature_def_utils.supervised_eval_signature_def
+
+
+
+
diff --git a/tensorflow/python/estimator/export/export_output_test.py b/tensorflow/python/estimator/export/export_output_test.py
index 7090e53d80..b21ba91b0f 100644
--- a/tensorflow/python/estimator/export/export_output_test.py
+++ b/tensorflow/python/estimator/export/export_output_test.py
@@ -225,5 +225,115 @@ class ExportOutputTest(test.TestCase):
       })
 
 
+class MockSupervisedOutput(export_output_lib._SupervisedOutput):
+  """So that we can test the abstract class methods directly."""
+
+  def _get_signature_def_fn(self):
+    pass
+
+
+class SupervisedOutputTest(test.TestCase):
+
+  def test_supervised_outputs_valid(self):
+    """Tests that no errors are raised when provided outputs are valid."""
+    loss = {"my_loss": constant_op.constant([0])}
+    predictions = {u"output1": constant_op.constant(["foo"])}
+    metrics = {"metrics": (constant_op.constant([0]),
+                           constant_op.constant([10])),
+               "metrics2": (constant_op.constant([0]),
+                            constant_op.constant([10]))}
+
+    outputter = MockSupervisedOutput(loss, predictions, metrics)
+    self.assertEqual(outputter.loss["loss/my_loss"], loss["my_loss"])
+    self.assertEqual(
+        outputter.predictions["predictions/output1"], predictions["output1"])
+    self.assertEqual(outputter.metrics["metrics/value"], metrics["metrics"][0])
+    self.assertEqual(
+        outputter.metrics["metrics2/update_op"], metrics["metrics2"][1])
+
+    # Single Tensor is OK too
+    outputter = MockSupervisedOutput(
+        loss["my_loss"], predictions["output1"], metrics["metrics"])
+    self.assertEqual(outputter.loss, {"loss": loss["my_loss"]})
+    self.assertEqual(
+        outputter.predictions, {"predictions": predictions["output1"]})
+    self.assertEqual(outputter.metrics["metrics/value"], metrics["metrics"][0])
+
+  def test_supervised_outputs_none(self):
+    outputter = MockSupervisedOutput(
+        constant_op.constant([0]), None, None)
+    self.assertEqual(len(outputter.loss), 1)
+    self.assertEqual(outputter.predictions, None)
+    self.assertEqual(outputter.metrics, None)
+
+  def test_supervised_outputs_invalid(self):
+    with self.assertRaisesRegexp(ValueError, "predictions output value must"):
+      MockSupervisedOutput(constant_op.constant([0]), [3], None)
+    with self.assertRaisesRegexp(ValueError, "loss output value must"):
+      MockSupervisedOutput("str", None, None)
+    with self.assertRaisesRegexp(ValueError, "metrics output value must"):
+      MockSupervisedOutput(None, None, (15.3, 4))
+    with self.assertRaisesRegexp(ValueError, "loss output key must"):
+      MockSupervisedOutput({25: "Tensor"}, None, None)
+
+  def test_supervised_outputs_tuples(self):
+    """Tests that no errors are raised when provided outputs are valid."""
+    loss = {("my", "loss"): constant_op.constant([0])}
+    predictions = {(u"output1", "2"): constant_op.constant(["foo"])}
+    metrics = {("metrics", "twice"): (constant_op.constant([0]),
+                                      constant_op.constant([10]))}
+
+    outputter = MockSupervisedOutput(loss, predictions, metrics)
+    self.assertEqual(set(outputter.loss.keys()), set(["loss/my/loss"]))
+    self.assertEqual(set(outputter.predictions.keys()),
+                     set(["predictions/output1/2"]))
+    self.assertEqual(set(outputter.metrics.keys()),
+                     set(["metrics/twice/value", "metrics/twice/update_op"]))
+
+  def test_supervised_outputs_no_prepend(self):
+    """Tests that no errors are raised when provided outputs are valid."""
+    loss = {"loss": constant_op.constant([0])}
+    predictions = {u"predictions": constant_op.constant(["foo"])}
+    metrics = {u"metrics": (constant_op.constant([0]),
+                            constant_op.constant([10]))}
+
+    outputter = MockSupervisedOutput(loss, predictions, metrics)
+    self.assertEqual(set(outputter.loss.keys()), set(["loss"]))
+    self.assertEqual(set(outputter.predictions.keys()), set(["predictions"]))
+    self.assertEqual(set(outputter.metrics.keys()),
+                     set(["metrics/value", "metrics/update_op"]))
+
+  def test_train_signature_def(self):
+    loss = {"my_loss": constant_op.constant([0])}
+    predictions = {u"output1": constant_op.constant(["foo"])}
+    metrics = {"metrics": (constant_op.constant([0]),
+                           constant_op.constant([10]))}
+
+    outputter = export_output_lib.TrainOutput(loss, predictions, metrics)
+
+    receiver = {u"features": constant_op.constant(100, shape=(100, 2)),
+                "labels": constant_op.constant(100, shape=(100, 1))}
+    sig_def = outputter.as_signature_def(receiver)
+
+    self.assertTrue("loss/my_loss" in sig_def.outputs)
+    self.assertTrue("metrics/value" in sig_def.outputs)
+    self.assertTrue("predictions/output1" in sig_def.outputs)
+    self.assertTrue("features" in sig_def.inputs)
+
+  def test_eval_signature_def(self):
+    loss = {"my_loss": constant_op.constant([0])}
+    predictions = {u"output1": constant_op.constant(["foo"])}
+
+    outputter = export_output_lib.EvalOutput(loss, predictions, None)
+
+    receiver = {u"features": constant_op.constant(100, shape=(100, 2)),
+                "labels": constant_op.constant(100, shape=(100, 1))}
+    sig_def = outputter.as_signature_def(receiver)
+
+    self.assertTrue("loss/my_loss" in sig_def.outputs)
+    self.assertFalse("metrics/value" in sig_def.outputs)
+    self.assertTrue("predictions/output1" in sig_def.outputs)
+    self.assertTrue("features" in sig_def.inputs)
+
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/python/estimator/export/export_test.py b/tensorflow/python/estimator/export/export_test.py
index c203be7dac..0af587f2a8 100644
--- a/tensorflow/python/estimator/export/export_test.py
+++ b/tensorflow/python/estimator/export/export_test.py
@@ -54,7 +54,7 @@ ops.register_tensor_conversion_function(LabeledTensorMock,
                                         _convert_labeled_tensor_mock_to_tensor)
 
 
-class ExportTest(test_util.TensorFlowTestCase):
+class ServingInputReceiverTest(test_util.TensorFlowTestCase):
 
   def test_serving_input_receiver_constructor(self):
     """Tests that no errors are raised when input is expected."""
@@ -161,6 +161,165 @@ class ExportTest(test_util.TensorFlowTestCase):
     with self.assertRaises(ValueError):
       _ = export.ServingInputReceiver(feature, receiver_tensor)
 
+
+class SupervisedInputReceiverTest(test_util.TensorFlowTestCase):
+
+  def test_input_receiver_constructor(self):
+    """Tests that no errors are raised when input is expected."""
+    features = {
+        "feature0": constant_op.constant([0]),
+        u"feature1": constant_op.constant([1]),
+        "feature2": sparse_tensor.SparseTensor(
+            indices=[[0, 0]], values=[1], dense_shape=[1, 1]),
+    }
+    labels = {
+        "classes": constant_op.constant([0] * 100),
+    }
+
+    receiver_tensors = {
+        "example0": array_ops.placeholder(dtypes.string, name="example0"),
+        u"example1": array_ops.placeholder(dtypes.string, name="example1"),
+    }
+    export.SupervisedInputReceiver(features, labels, receiver_tensors)
+
+  def test_input_receiver_raw_values(self):
+    """Tests that no errors are raised when input is expected."""
+    features = {
+        "feature0": constant_op.constant([0]),
+        u"feature1": constant_op.constant([1]),
+        "feature2": sparse_tensor.SparseTensor(
+            indices=[[0, 0]], values=[1], dense_shape=[1, 1]),
+    }
+
+    labels = {
+        "classes": constant_op.constant([0] * 100),
+    }
+
+    receiver_tensors = {
+        "example0": array_ops.placeholder(dtypes.string, name="example0"),
+        u"example1": array_ops.placeholder(dtypes.string, name="example1"),
+    }
+    rec = export.SupervisedInputReceiver(
+        features["feature2"], labels, receiver_tensors)
+    self.assertIsInstance(rec.features, sparse_tensor.SparseTensor)
+
+    rec = export.SupervisedInputReceiver(
+        features, labels["classes"], receiver_tensors)
+    self.assertIsInstance(rec.labels, ops.Tensor)
+
+  def test_input_receiver_features_invalid(self):
+    features = constant_op.constant([0] * 100)
+    labels = constant_op.constant([0])
+    receiver_tensors = {
+        "example0": array_ops.placeholder(dtypes.string, name="example0"),
+        u"example1": array_ops.placeholder(dtypes.string, name="example1"),
+    }
+
+    with self.assertRaisesRegexp(ValueError, "features must be defined"):
+      export.SupervisedInputReceiver(
+          features=None,
+          labels=labels,
+          receiver_tensors=receiver_tensors)
+
+    with self.assertRaisesRegexp(ValueError, "feature keys must be strings"):
+      export.SupervisedInputReceiver(
+          features={1: constant_op.constant([1])},
+          labels=labels,
+          receiver_tensors=receiver_tensors)
+
+    with self.assertRaisesRegexp(ValueError, "label keys must be strings"):
+      export.SupervisedInputReceiver(
+          features=features,
+          labels={1: constant_op.constant([1])},
+          receiver_tensors=receiver_tensors)
+
+    with self.assertRaisesRegexp(
+        ValueError, "feature feature1 must be a Tensor or SparseTensor"):
+      export.SupervisedInputReceiver(
+          features={"feature1": [1]},
+          labels=labels,
+          receiver_tensors=receiver_tensors)
+
+    with self.assertRaisesRegexp(
+        ValueError, "feature must be a Tensor or SparseTensor"):
+      export.SupervisedInputReceiver(
+          features=[1],
+          labels=labels,
+          receiver_tensors=receiver_tensors)
+
+    with self.assertRaisesRegexp(
+        ValueError, "label must be a Tensor or SparseTensor"):
+      export.SupervisedInputReceiver(
+          features=features,
+          labels=100,
+          receiver_tensors=receiver_tensors)
+
+  def test_input_receiver_receiver_tensors_invalid(self):
+    features = {
+        "feature0": constant_op.constant([0]),
+        u"feature1": constant_op.constant([1]),
+        "feature2": sparse_tensor.SparseTensor(
+            indices=[[0, 0]], values=[1], dense_shape=[1, 1]),
+    }
+    labels = constant_op.constant([0])
+
+    with self.assertRaisesRegexp(
+        ValueError, "receiver_tensors must be defined"):
+      export.SupervisedInputReceiver(
+          features=features,
+          labels=labels,
+          receiver_tensors=None)
+
+    with self.assertRaisesRegexp(
+        ValueError, "receiver_tensors keys must be strings"):
+      export.SupervisedInputReceiver(
+          features=features,
+          labels=labels,
+          receiver_tensors={
+              1: array_ops.placeholder(dtypes.string, name="example0")})
+
+    with self.assertRaisesRegexp(
+        ValueError, "receiver_tensor example1 must be a Tensor"):
+      export.SupervisedInputReceiver(
+          features=features,
+          labels=labels,
+          receiver_tensors={"example1": [1]})
+
+  def test_single_feature_single_receiver(self):
+    feature = constant_op.constant(5)
+    label = constant_op.constant(5)
+    receiver_tensor = array_ops.placeholder(dtypes.string)
+    input_receiver = export.SupervisedInputReceiver(
+        feature, label, receiver_tensor)
+
+    # single receiver is automatically named
+    receiver_key, = input_receiver.receiver_tensors.keys()
+    self.assertEqual("input", receiver_key)
+
+  def test_multi_feature_single_receiver(self):
+    features = {"foo": constant_op.constant(5),
+                "bar": constant_op.constant(6)}
+    labels = {"value": constant_op.constant(5)}
+    receiver_tensor = array_ops.placeholder(dtypes.string)
+    _ = export.SupervisedInputReceiver(features, labels, receiver_tensor)
+
+  def test_multi_feature_multi_receiver(self):
+    features = {"foo": constant_op.constant(5),
+                "bar": constant_op.constant(6)}
+    labels = {"value": constant_op.constant(5)}
+    receiver_tensors = {"baz": array_ops.placeholder(dtypes.int64),
+                        "qux": array_ops.placeholder(dtypes.float32)}
+    _ = export.SupervisedInputReceiver(features, labels, receiver_tensors)
+
+  def test_feature_labeled_tensor(self):
+    feature = LabeledTensorMock()
+    label = constant_op.constant(5)
+    receiver_tensor = array_ops.placeholder(dtypes.string)
+    _ = export.SupervisedInputReceiver(feature, label, receiver_tensor)
+
+
+class ExportTest(test_util.TensorFlowTestCase):
+
   def test_build_parsing_serving_input_receiver_fn(self):
     feature_spec = {"int_feature": parsing_ops.VarLenFeature(dtypes.int64),
                     "float_feature": parsing_ops.VarLenFeature(dtypes.float32)}
@@ -237,6 +396,69 @@ class ExportTest(test_util.TensorFlowTestCase):
           dtypes.int32,
           serving_input_receiver.receiver_tensors["feature_2"].dtype)
 
+  def test_build_raw_supervised_input_receiver_fn(self):
+    features = {"feature_1": constant_op.constant(["hello"]),
+                "feature_2": constant_op.constant([42])}
+    labels = {"foo": constant_op.constant([5]),
+              "bar": constant_op.constant([6])}
+    input_receiver_fn = export.build_raw_supervised_input_receiver_fn(
+        features, labels)
+    with ops.Graph().as_default():
+      input_receiver = input_receiver_fn()
+      self.assertEqual(set(["feature_1", "feature_2"]),
+                       set(input_receiver.features.keys()))
+      self.assertEqual(set(["foo", "bar"]),
+                       set(input_receiver.labels.keys()))
+      self.assertEqual(set(["feature_1", "feature_2", "foo", "bar"]),
+                       set(input_receiver.receiver_tensors.keys()))
+      self.assertEqual(
+          dtypes.string, input_receiver.receiver_tensors["feature_1"].dtype)
+      self.assertEqual(
+          dtypes.int32, input_receiver.receiver_tensors["feature_2"].dtype)
+
+  def test_build_raw_supervised_input_receiver_fn_raw_tensors(self):
+    features = {"feature_1": constant_op.constant(["hello"]),
+                "feature_2": constant_op.constant([42])}
+    labels = {"foo": constant_op.constant([5]),
+              "bar": constant_op.constant([6])}
+    input_receiver_fn1 = export.build_raw_supervised_input_receiver_fn(
+        features["feature_1"], labels)
+    input_receiver_fn2 = export.build_raw_supervised_input_receiver_fn(
+        features["feature_1"], labels["foo"])
+    with ops.Graph().as_default():
+      input_receiver = input_receiver_fn1()
+      self.assertIsInstance(input_receiver.features, ops.Tensor)
+      self.assertEqual(set(["foo", "bar"]),
+                       set(input_receiver.labels.keys()))
+      self.assertEqual(set(["input", "foo", "bar"]),
+                       set(input_receiver.receiver_tensors.keys()))
+
+      input_receiver = input_receiver_fn2()
+      self.assertIsInstance(input_receiver.features, ops.Tensor)
+      self.assertIsInstance(input_receiver.labels, ops.Tensor)
+      self.assertEqual(set(["input", "label"]),
+                       set(input_receiver.receiver_tensors.keys()))
+
+  def test_build_raw_supervised_input_receiver_fn_batch_size(self):
+    features = {"feature_1": constant_op.constant(["hello"]),
+                "feature_2": constant_op.constant([42])}
+    labels = {"foo": constant_op.constant([5]),
+              "bar": constant_op.constant([6])}
+    input_receiver_fn = export.build_raw_supervised_input_receiver_fn(
+        features, labels, default_batch_size=10)
+    with ops.Graph().as_default():
+      input_receiver = input_receiver_fn()
+      self.assertEqual([10], input_receiver.receiver_tensors["feature_1"].shape)
+      self.assertEqual([10], input_receiver.features["feature_1"].shape)
+
+  def test_build_raw_supervised_input_receiver_fn_overlapping_keys(self):
+    features = {"feature_1": constant_op.constant(["hello"]),
+                "feature_2": constant_op.constant([42])}
+    labels = {"feature_1": constant_op.constant([5]),
+              "bar": constant_op.constant([6])}
+    with self.assertRaises(ValueError):
+      export.build_raw_supervised_input_receiver_fn(features, labels)
+
   def test_build_all_signature_defs_without_receiver_alternatives(self):
     receiver_tensor = array_ops.placeholder(dtypes.string)
     output_1 = constant_op.constant([1.])
@@ -404,6 +626,35 @@ class ExportTest(test_util.TensorFlowTestCase):
     self.assertTrue(int(time_1) < int(time_2))
     self.assertTrue(int(time_2) < int(time_3))
 
+  def test_build_all_signature_defs_serving_only(self):
+    receiver_tensor = {"input": array_ops.placeholder(dtypes.string)}
+    output_1 = constant_op.constant([1.])
+    export_outputs = {
+        signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
+            export_output.PredictOutput(outputs=output_1),
+        "train": export_output.TrainOutput(loss=output_1),
+    }
+
+    signature_defs = export.build_all_signature_defs(
+        receiver_tensor, export_outputs)
+
+    expected_signature_defs = {
+        "serving_default": signature_def_utils.predict_signature_def(
+            receiver_tensor, {"output": output_1})
+    }
+
+    self.assertDictEqual(expected_signature_defs, signature_defs)
+
+    signature_defs = export.build_all_signature_defs(
+        receiver_tensor, export_outputs, serving_only=False)
+
+    expected_signature_defs.update({
+        "train": signature_def_utils.supervised_train_signature_def(
+            receiver_tensor, loss={"loss": output_1})
+    })
+
+    self.assertDictEqual(expected_signature_defs, signature_defs)
+
 
 class TensorServingReceiverTest(test_util.TensorFlowTestCase):
 
diff --git a/tensorflow/python/estimator/model_fn.py b/tensorflow/python/estimator/model_fn.py
index 8111ab564c..4ab2578769 100644
--- a/tensorflow/python/estimator/model_fn.py
+++ b/tensorflow/python/estimator/model_fn.py
@@ -28,6 +28,7 @@ from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_shape
 from tensorflow.python.ops import array_ops
 from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model import tag_constants
 from tensorflow.python.training import monitored_session
 from tensorflow.python.training import session_run_hook
 from tensorflow.python.util import nest
@@ -53,6 +54,13 @@ class ModeKeys(object):
 LOSS_METRIC_KEY = 'loss'
 AVERAGE_LOSS_METRIC_KEY = 'average_loss'
 
+# Mapping of the modes to appropriate tag_constants that are used for saving.
+EXPORT_TAG_MAP = {
+    ModeKeys.PREDICT: [tag_constants.SERVING],
+    ModeKeys.TRAIN: [tag_constants.TRAINING],
+    ModeKeys.EVAL: [tag_constants.EVAL],
+}
+
 
 @tf_export('estimator.EstimatorSpec')
 class EstimatorSpec(
diff --git a/tensorflow/python/saved_model/builder_impl.py b/tensorflow/python/saved_model/builder_impl.py
index 3447d917e9..071033b066 100644
--- a/tensorflow/python/saved_model/builder_impl.py
+++ b/tensorflow/python/saved_model/builder_impl.py
@@ -168,6 +168,25 @@ class SavedModelBuilder(object):
         raise TypeError("main_op needs to be an Operation: %r" % main_op)
       ops.add_to_collection(constants.MAIN_OP_KEY, main_op)
 
+  def _add_train_op(self, train_op):
+    """Add train op to the SavedModel.
+
+    Note that this functionality is in development, and liable to be
+    moved elsewhere.
+
+    Args:
+      train_op: Op or group of ops that are used for training. These are
+        stored as a collection with key TRAIN_OP_KEY, but not executed.
+
+    Raises:
+      TypeError if Train op is not of type `Operation`.
+    """
+    if train_op is not None:
+      if (not isinstance(train_op, ops.Tensor) and
+          not isinstance(train_op, ops.Operation)):
+        raise TypeError("train_op needs to be a Tensor or Op: %r" % train_op)
+      ops.add_to_collection(constants.TRAIN_OP_KEY, train_op)
+
   def _tag_and_add_meta_graph(self, meta_graph_def, tags, signature_def_map):
     """Tags the meta graph def and adds it to the SavedModel.
 
@@ -238,6 +257,20 @@ class SavedModelBuilder(object):
         for outputs_key in outputs:
           self._validate_tensor_info(outputs[outputs_key])
 
+  def _add_collections(
+      self, assets_collection, legacy_init_op, main_op, train_op):
+    """Add asset and op collections to be saved."""
+    # Save asset files and write them to disk, if any.
+    self._save_and_write_assets(assets_collection)
+
+    if main_op is None:
+      # Add legacy init op to the SavedModel.
+      self._maybe_add_legacy_init_op(legacy_init_op)
+    else:
+      self._add_main_op(main_op)
+
+    self._add_train_op(train_op)
+
   def add_meta_graph(self,
                      tags,
                      signature_def_map=None,
@@ -285,14 +318,8 @@ class SavedModelBuilder(object):
     # properly populated.
     self._validate_signature_def_map(signature_def_map)
 
-    # Save asset files and write them to disk, if any.
-    self._save_and_write_assets(assets_collection)
-
-    if main_op is None:
-      # Add legacy init op to the SavedModel.
-      self._maybe_add_legacy_init_op(legacy_init_op)
-    else:
-      self._add_main_op(main_op)
+    # Add assets and ops
+    self._add_collections(assets_collection, legacy_init_op, main_op, None)
 
     # Initialize a saver to generate a sharded output for all saveables in the
     # current scope.
@@ -351,6 +378,7 @@ class SavedModelBuilder(object):
       strip_default_attrs: Boolean. If `True`, default-valued attributes will be
         removed from the NodeDefs. For a detailed guide, see
         [Stripping Default-Valued Attributes](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/README.md#stripping-default-valued-attributes).
+
     """
     # pylint: enable=line-too-long
     if self._has_saved_variables:
@@ -362,8 +390,8 @@ class SavedModelBuilder(object):
     # properly populated.
     self._validate_signature_def_map(signature_def_map)
 
-    # Save asset files and write them to disk, if any.
-    self._save_and_write_assets(assets_collection)
+    # Add assets and ops
+    self._add_collections(assets_collection, legacy_init_op, main_op, None)
 
     # Create the variables sub-directory, if it does not exist.
     variables_dir = os.path.join(
@@ -376,12 +404,6 @@ class SavedModelBuilder(object):
         compat.as_text(variables_dir),
         compat.as_text(constants.VARIABLES_FILENAME))
 
-    if main_op is None:
-      # Add legacy init op to the SavedModel.
-      self._maybe_add_legacy_init_op(legacy_init_op)
-    else:
-      self._add_main_op(main_op)
-
     # Initialize a saver to generate a sharded output for all saveables in the
     # current scope.
     saver = tf_saver.Saver(
diff --git a/tensorflow/python/saved_model/constants.py b/tensorflow/python/saved_model/constants.py
index 34206c6f6d..61c6ffbd0d 100644
--- a/tensorflow/python/saved_model/constants.py
+++ b/tensorflow/python/saved_model/constants.py
@@ -41,6 +41,10 @@ MAIN_OP_KEY = "saved_model_main_op"
 tf_export("saved_model.constants.MAIN_OP_KEY").export_constant(
     __name__, "MAIN_OP_KEY")
 
+# CollectionDef key for the SavedModel train op.
+# Not exported while export_all_saved_models is in contrib.
+TRAIN_OP_KEY = "saved_model_train_op"
+
 # Schema version for SavedModel.
 SAVED_MODEL_SCHEMA_VERSION = 1
 tf_export("saved_model.constants.SAVED_MODEL_SCHEMA_VERSION").export_constant(
@@ -65,3 +69,5 @@ tf_export("saved_model.constants.VARIABLES_DIRECTORY").export_constant(
 VARIABLES_FILENAME = "variables"
 tf_export("saved_model.constants.VARIABLES_FILENAME").export_constant(
     __name__, "VARIABLES_FILENAME")
+
+
diff --git a/tensorflow/python/saved_model/saved_model_test.py b/tensorflow/python/saved_model/saved_model_test.py
index 804255375e..a4d994fd43 100644
--- a/tensorflow/python/saved_model/saved_model_test.py
+++ b/tensorflow/python/saved_model/saved_model_test.py
@@ -734,6 +734,96 @@ class SavedModelTest(test.TestCase):
         builder.add_meta_graph_and_variables(
             sess, ["foo"], legacy_init_op=legacy_init_op)
 
+  def testTrainOp(self):
+    export_dir = self._get_export_dir("test_train_op")
+    builder = saved_model_builder.SavedModelBuilder(export_dir)
+
+    with self.test_session(graph=ops.Graph()) as sess:
+      # Add `v1` and `v2` variables to the graph.
+      v1 = variables.Variable(1, name="v1")
+      ops.add_to_collection("v", v1)
+      v2 = variables.Variable(2, name="v2")
+      ops.add_to_collection("v", v2)
+
+      sess.run(variables.global_variables_initializer())
+      train_op = state_ops.assign_add(v1, v2)
+
+      sess.run(train_op)
+      # TODO(karmel): remove explicit call when in the public method.
+      builder._add_train_op(train_op)
+      builder.add_meta_graph_and_variables(sess, ["foo"])
+
+    # Save the SavedModel to disk.
+    builder.save()
+
+    with self.test_session(graph=ops.Graph()) as sess:
+      loader.load(sess, ["foo"], export_dir)
+      self.assertEqual(3, ops.get_collection("v")[0].eval())
+      self.assertEqual(2, ops.get_collection("v")[1].eval())
+      self.assertIsInstance(
+          ops.get_collection(constants.TRAIN_OP_KEY)[0], ops.Tensor)
+
+  def testTrainOpGroup(self):
+    export_dir = self._get_export_dir("test_train_op_group")
+    builder = saved_model_builder.SavedModelBuilder(export_dir)
+
+    with self.test_session(graph=ops.Graph()) as sess:
+      # Add `v1` and `v2` variables to the graph.
+      v1 = variables.Variable(1, name="v1")
+      ops.add_to_collection("v", v1)
+      v2 = variables.Variable(2, name="v2")
+      ops.add_to_collection("v", v2)
+
+      sess.run(variables.global_variables_initializer())
+      train_op = control_flow_ops.group()
+
+      sess.run(train_op)
+      # TODO(karmel): remove explicit call when in the public method.
+      builder._add_train_op(train_op)
+      builder.add_meta_graph_and_variables(sess, ["foo"])
+
+    # Save the SavedModel to disk.
+    builder.save()
+
+    with self.test_session(graph=ops.Graph()) as sess:
+      loader.load(sess, ["foo"], export_dir)
+      self.assertEqual(1, ops.get_collection("v")[0].eval())
+      self.assertEqual(2, ops.get_collection("v")[1].eval())
+      self.assertIsInstance(
+          ops.get_collection(constants.TRAIN_OP_KEY)[0], ops.Operation)
+
+  def testTrainOpAfterVariables(self):
+    export_dir = self._get_export_dir("test_train_op_after_variables")
+    builder = saved_model_builder.SavedModelBuilder(export_dir)
+
+    with self.test_session(graph=ops.Graph()) as sess:
+      # Add `v1` and `v2` variables to the graph.
+      v1 = variables.Variable(1, name="v1")
+      ops.add_to_collection("v", v1)
+      v2 = variables.Variable(2, name="v2")
+      ops.add_to_collection("v", v2)
+
+      sess.run(variables.global_variables_initializer())
+      builder.add_meta_graph_and_variables(sess, ["pre_foo"])
+
+      train_op = state_ops.assign_add(v1, v2)
+      sess.run(train_op)
+      # TODO(karmel): remove explicit call when in the public method.
+      builder._add_train_op(train_op)
+      builder.add_meta_graph(["foo"])
+
+    # Save the SavedModel to disk.
+    builder.save()
+
+    with self.test_session(graph=ops.Graph()) as sess:
+      loader.load(sess, ["foo"], export_dir)
+      self.assertIsInstance(
+          ops.get_collection(constants.TRAIN_OP_KEY)[0], ops.Tensor)
+
+    with self.test_session(graph=ops.Graph()) as sess:
+      loader.load(sess, ["pre_foo"], export_dir)
+      self.assertFalse(ops.get_collection(constants.TRAIN_OP_KEY))
+
   def testMultipleAssets(self):
     export_dir = self._get_export_dir("test_multiple_assets")
     builder = saved_model_builder.SavedModelBuilder(export_dir)
diff --git a/tensorflow/python/saved_model/signature_constants.py b/tensorflow/python/saved_model/signature_constants.py
index 819f351291..99007a9634 100644
--- a/tensorflow/python/saved_model/signature_constants.py
+++ b/tensorflow/python/saved_model/signature_constants.py
@@ -94,3 +94,9 @@ tf_export("saved_model.signature_constants.REGRESS_OUTPUTS").export_constant(
     __name__, "REGRESS_OUTPUTS")
 
 ################################################################################
+# Train/Eval API constants.
+# Not exported while export_all_saved_models is in contrib.
+
+SUPERVISED_TRAIN_METHOD_NAME = "tensorflow/supervised/training"
+
+SUPERVISED_EVAL_METHOD_NAME = "tensorflow/supervised/eval"
diff --git a/tensorflow/python/saved_model/signature_def_utils.py b/tensorflow/python/saved_model/signature_def_utils.py
index ea0f52f17e..27d6b70e9d 100644
--- a/tensorflow/python/saved_model/signature_def_utils.py
+++ b/tensorflow/python/saved_model/signature_def_utils.py
@@ -26,6 +26,8 @@ from tensorflow.python.saved_model.signature_def_utils_impl import classificatio
 from tensorflow.python.saved_model.signature_def_utils_impl import is_valid_signature
 from tensorflow.python.saved_model.signature_def_utils_impl import predict_signature_def
 from tensorflow.python.saved_model.signature_def_utils_impl import regression_signature_def
+from tensorflow.python.saved_model.signature_def_utils_impl import supervised_eval_signature_def
+from tensorflow.python.saved_model.signature_def_utils_impl import supervised_train_signature_def
 # pylint: enable=unused-import
 
 del absolute_import
diff --git a/tensorflow/python/saved_model/signature_def_utils_impl.py b/tensorflow/python/saved_model/signature_def_utils_impl.py
index d033159188..f8ad788f77 100644
--- a/tensorflow/python/saved_model/signature_def_utils_impl.py
+++ b/tensorflow/python/saved_model/signature_def_utils_impl.py
@@ -185,6 +185,62 @@ def predict_signature_def(inputs, outputs):
   return signature_def
 
 
+def supervised_train_signature_def(
+    inputs, loss, predictions=None, metrics=None):
+  return _supervised_signature_def(
+      signature_constants.SUPERVISED_TRAIN_METHOD_NAME, inputs, loss=loss,
+      predictions=predictions, metrics=metrics)
+
+
+def supervised_eval_signature_def(
+    inputs, loss, predictions=None, metrics=None):
+  return _supervised_signature_def(
+      signature_constants.SUPERVISED_EVAL_METHOD_NAME, inputs, loss=loss,
+      predictions=predictions, metrics=metrics)
+
+
+def _supervised_signature_def(
+    method_name, inputs, loss=None, predictions=None,
+    metrics=None):
+  """Creates a signature for training and eval data.
+
+  This function produces signatures that describe the inputs and outputs
+  of a supervised process, such as training or evaluation, that
+  results in loss, metrics, and the like. Note that this function only requires
+  inputs to be not None.
+
+  Args:
+    method_name: Method name of the SignatureDef as a string.
+    inputs: dict of string to `Tensor`.
+    loss: dict of string to `Tensor` representing computed loss.
+    predictions: dict of string to `Tensor` representing the output predictions.
+    metrics: dict of string to `Tensor` representing metric ops.
+
+  Returns:
+    A train- or eval-flavored signature_def.
+
+  Raises:
+    ValueError: If inputs or outputs is `None`.
+  """
+  if inputs is None or not inputs:
+    raise ValueError('{} inputs cannot be None or empty.'.format(method_name))
+
+  signature_inputs = {key: utils.build_tensor_info(tensor)
+                      for key, tensor in inputs.items()}
+
+  signature_outputs = {}
+  for output_set in (loss, predictions, metrics):
+    if output_set is not None:
+      sig_out = {key: utils.build_tensor_info(tensor)
+                 for key, tensor in output_set.items()}
+      signature_outputs.update(sig_out)
+
+  signature_def = build_signature_def(
+      signature_inputs, signature_outputs, method_name)
+
+  return signature_def
+
+
 @tf_export('saved_model.signature_def_utils.is_valid_signature')
 def is_valid_signature(signature_def):
   """Determine whether a SignatureDef can be served by TensorFlow Serving."""
diff --git a/tensorflow/python/saved_model/signature_def_utils_test.py b/tensorflow/python/saved_model/signature_def_utils_test.py
index b2bd14db8c..ebc5450633 100644
--- a/tensorflow/python/saved_model/signature_def_utils_test.py
+++ b/tensorflow/python/saved_model/signature_def_utils_test.py
@@ -180,6 +180,101 @@ class SignatureDefUtilsTest(test.TestCase):
     self.assertEqual(types_pb2.DT_STRING, output2_tensor_info_actual.dtype)
     self.assertEqual(0, len(output2_tensor_info_actual.tensor_shape.dim))
 
+  def testTrainSignatureDef(self):
+    self._testSupervisedSignatureDef(
+        signature_def_utils_impl.supervised_train_signature_def,
+        signature_constants.SUPERVISED_TRAIN_METHOD_NAME)
+
+  def testEvalSignatureDef(self):
+    self._testSupervisedSignatureDef(
+        signature_def_utils_impl.supervised_eval_signature_def,
+        signature_constants.SUPERVISED_EVAL_METHOD_NAME)
+
+  def _testSupervisedSignatureDef(self, fn_to_test, method_name):
+    inputs = {
+        "input-1": constant_op.constant("a", name="input-1"),
+        "input-2": constant_op.constant("b", name="input-2"),
+    }
+    loss = {"loss-1": constant_op.constant(0.45, name="loss-1")}
+    predictions = {
+        "classes": constant_op.constant([100], name="classes"),
+    }
+    metrics_val = constant_op.constant(100.0, name="metrics_val")
+    metrics = {
+        "metrics/value": metrics_val,
+        "metrics/update_op": array_ops.identity(metrics_val, name="metrics_op"),
+    }
+
+    signature_def = fn_to_test(inputs, loss, predictions, metrics)
+
+    self.assertEqual(method_name, signature_def.method_name)
+
+    # Check inputs in signature def.
+    self.assertEqual(2, len(signature_def.inputs))
+    input1_tensor_info_actual = (signature_def.inputs["input-1"])
+    self.assertEqual("input-1:0", input1_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_STRING, input1_tensor_info_actual.dtype)
+    self.assertEqual(0, len(input1_tensor_info_actual.tensor_shape.dim))
+    input2_tensor_info_actual = (signature_def.inputs["input-2"])
+    self.assertEqual("input-2:0", input2_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_STRING, input2_tensor_info_actual.dtype)
+    self.assertEqual(0, len(input2_tensor_info_actual.tensor_shape.dim))
+
+    # Check outputs in signature def.
+    self.assertEqual(4, len(signature_def.outputs))
+    self.assertEqual("loss-1:0", signature_def.outputs["loss-1"].name)
+    self.assertEqual(types_pb2.DT_FLOAT, signature_def.outputs["loss-1"].dtype)
+
+    self.assertEqual("classes:0", signature_def.outputs["classes"].name)
+    self.assertEqual(1, len(signature_def.outputs["classes"].tensor_shape.dim))
+
+    self.assertEqual(
+        "metrics_val:0", signature_def.outputs["metrics/value"].name)
+    self.assertEqual(
+        types_pb2.DT_FLOAT, signature_def.outputs["metrics/value"].dtype)
+
+    self.assertEqual(
+        "metrics_op:0", signature_def.outputs["metrics/update_op"].name)
+    self.assertEqual(
+        types_pb2.DT_FLOAT, signature_def.outputs["metrics/value"].dtype)
+
+  def testTrainSignatureDefMissingInputs(self):
+    self._testSupervisedSignatureDefMissingInputs(
+        signature_def_utils_impl.supervised_train_signature_def,
+        signature_constants.SUPERVISED_TRAIN_METHOD_NAME)
+
+  def testEvalSignatureDefMissingInputs(self):
+    self._testSupervisedSignatureDefMissingInputs(
+        signature_def_utils_impl.supervised_eval_signature_def,
+        signature_constants.SUPERVISED_EVAL_METHOD_NAME)
+
+  def _testSupervisedSignatureDefMissingInputs(self, fn_to_test, method_name):
+    inputs = {
+        "input-1": constant_op.constant("a", name="input-1"),
+        "input-2": constant_op.constant("b", name="input-2"),
+    }
+    loss = {"loss-1": constant_op.constant(0.45, name="loss-1")}
+    predictions = {
+        "classes": constant_op.constant([100], name="classes"),
+    }
+    metrics_val = constant_op.constant(100, name="metrics_val")
+    metrics = {
+        "metrics/value": metrics_val,
+        "metrics/update_op": array_ops.identity(metrics_val, name="metrics_op"),
+    }
+
+    with self.assertRaises(ValueError):
+      signature_def = fn_to_test(
+          {}, loss=loss, predictions=predictions, metrics=metrics)
+
+    signature_def = fn_to_test(inputs, loss=loss)
+    self.assertEqual(method_name, signature_def.method_name)
+    self.assertEqual(1, len(signature_def.outputs))
+
+    signature_def = fn_to_test(inputs, metrics=metrics, loss=loss)
+    self.assertEqual(method_name, signature_def.method_name)
+    self.assertEqual(3, len(signature_def.outputs))
+
   def testGetShapeAndTypes(self):
     inputs = {
         "input-1": constant_op.constant(["a", "b"]),
diff --git a/tensorflow/python/saved_model/tag_constants.py b/tensorflow/python/saved_model/tag_constants.py
index 5a797da791..c82154e7b9 100644
--- a/tensorflow/python/saved_model/tag_constants.py
+++ b/tensorflow/python/saved_model/tag_constants.py
@@ -32,6 +32,9 @@ TRAINING = "train"
 tf_export("saved_model.tag_constants.TRAINING").export_constant(
     __name__, "TRAINING")
 
+# Tag for the `eval` graph. Not exported while the export logic is in contrib.
+EVAL = "eval"
+
 # Tag for the `gpu` graph.
 GPU = "gpu"
 tf_export("saved_model.tag_constants.GPU").export_constant(__name__, "GPU")
@@ -39,3 +42,5 @@ tf_export("saved_model.tag_constants.GPU").export_constant(__name__, "GPU")
 # Tag for the `tpu` graph.
 TPU = "tpu"
 tf_export("saved_model.tag_constants.TPU").export_constant(__name__, "TPU")
+
+
-- 
GitLab


From 037e52e20157985d3f385f8e0426cdde3f5aae2b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 16:37:27 -0700
Subject: [PATCH 272/395] Expose read-only versions of tensors in tflite.

PiperOrigin-RevId: 195491701
---
 tensorflow/contrib/lite/interpreter.h | 37 ++++++++++++++++++++++-----
 1 file changed, 31 insertions(+), 6 deletions(-)

diff --git a/tensorflow/contrib/lite/interpreter.h b/tensorflow/contrib/lite/interpreter.h
index 1074f64263..0450e86ae7 100644
--- a/tensorflow/contrib/lite/interpreter.h
+++ b/tensorflow/contrib/lite/interpreter.h
@@ -201,7 +201,7 @@ class Interpreter {
   // Overrides execution plan. This bounds checks indices sent in.
   TfLiteStatus SetExecutionPlan(const std::vector<int>& new_plan);
 
-  // Get a tensor data structure.
+  // Get a mutable tensor data structure.
   // TODO(aselle): Create a safe ArrayHandle interface to avoid exposing this
   // read/write access to structure
   TfLiteTensor* tensor(int tensor_index) {
@@ -210,9 +210,14 @@ class Interpreter {
     return &context_.tensors[tensor_index];
   }
 
+  // Get an immutable tensor data structure.
+  const TfLiteTensor* tensor(int tensor_index) const {
+    if (tensor_index >= context_.tensors_size || tensor_index < 0)
+      return nullptr;
+    return &context_.tensors[tensor_index];
+  }
+
   // Get a pointer to an operation and registration data structure if in bounds.
-  // TODO(aselle): Create a safe ArrayHandle interface to avoid exposing this
-  // read/write access to structure
   const std::pair<TfLiteNode, TfLiteRegistration>* node_and_registration(
       int node_index) const {
     if (node_index >= nodes_and_registration_.size() || node_index < 0)
@@ -220,7 +225,8 @@ class Interpreter {
     return &nodes_and_registration_[node_index];
   }
 
-  // Perform a checked cast to the appropriate tensor type.
+  // Perform a checked cast to the appropriate tensor type (mutable pointer
+  // version).
   template <class T>
   T* typed_tensor(int tensor_index) {
     if (TfLiteTensor* tensor_ptr = tensor(tensor_index)) {
@@ -231,6 +237,18 @@ class Interpreter {
     return nullptr;
   }
 
+  // Perform a checked cast to the appropriate tensor type (immutable pointer
+  // version).
+  template <class T>
+  const T* typed_tensor(int tensor_index) const {
+    if (const TfLiteTensor* tensor_ptr = tensor(tensor_index)) {
+      if (tensor_ptr->type == typeToTfLiteType<T>()) {
+        return reinterpret_cast<const T*>(tensor_ptr->data.raw);
+      }
+    }
+    return nullptr;
+  }
+
   // Return a pointer into the data of a given input tensor. The given index
   // must be between 0 and inputs().size().
   template <class T>
@@ -238,13 +256,20 @@ class Interpreter {
     return typed_tensor<T>(inputs_[index]);
   }
 
-  // Return a pointer into the data of a given output tensor. The given index
-  // must be between 0 and outputs().size().
+  // Return a mutable pointer into the data of a given output tensor. The given
+  // index must be between 0 and outputs().size().
   template <class T>
   T* typed_output_tensor(int index) {
     return typed_tensor<T>(outputs_[index]);
   }
 
+  // Return an immutable pointer into the data of a given output tensor. The
+  // given index must be between 0 and outputs().size().
+  template <class T>
+  const T* typed_output_tensor(int index) const {
+    return typed_tensor<T>(outputs_[index]);
+  }
+
   // Change the dimensionality of a given tensor. Note, this is only acceptable
   // for tensor indices that are inputs.
   // Returns status of failure or success.
-- 
GitLab


From fa1d92f70adf52d9258384e8528f9a7203a141dd Mon Sep 17 00:00:00 2001
From: Bjarke Hammersholt Roune <broune@google.com>
Date: Fri, 4 May 2018 16:51:06 -0700
Subject: [PATCH 273/395] Add infrastructure for a backend-specific
 configuration for each op. This is intentionally not exposed in
 ComputationBuilder and is not intended for use or to be set at all prior to
 the last backend-specific part of compilation.

PiperOrigin-RevId: 195493500
---
 tensorflow/compiler/xla/service/hlo.proto     |   3 +
 .../compiler/xla/service/hlo_computation.cc   |  52 ++++-----
 .../compiler/xla/service/hlo_computation.h    |  20 ++--
 .../compiler/xla/service/hlo_graph_dumper.cc  |  43 +++++---
 .../compiler/xla/service/hlo_graph_dumper.h   |   5 +-
 .../compiler/xla/service/hlo_instruction.cc   | 100 +++++-------------
 .../compiler/xla/service/hlo_instruction.h    |  59 +++++++++--
 tensorflow/compiler/xla/service/hlo_module.cc |  12 +++
 tensorflow/compiler/xla/service/hlo_module.h  |  19 ++++
 .../compiler/xla/service/hlo_verifier.cc      |  71 ++++++-------
 tensorflow/compiler/xla/statusor.h            |  11 +-
 tensorflow/compiler/xla/statusor_test.cc      |   8 ++
 .../compiler/xla/tools/parser/hlo_parser.cc   |  10 +-
 .../xla/tools/parser/hlo_parser_test.cc       |  22 +++-
 14 files changed, 259 insertions(+), 176 deletions(-)

diff --git a/tensorflow/compiler/xla/service/hlo.proto b/tensorflow/compiler/xla/service/hlo.proto
index aa6860880b..1f7c1cffd3 100644
--- a/tensorflow/compiler/xla/service/hlo.proto
+++ b/tensorflow/compiler/xla/service/hlo.proto
@@ -147,6 +147,9 @@ message HloInstructionProto {
   repeated int64 called_computation_ids = 38;
 
   xla.OpSharding sharding = 40;
+
+  // Backend configuration for the instruction. Has backend-specific meaning.
+  string backend_config = 43;
 }
 
 // Serialization of HloComputation.
diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc
index 594413e88f..17e43c3cb8 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.cc
+++ b/tensorflow/compiler/xla/service/hlo_computation.cc
@@ -347,6 +347,11 @@ std::list<HloComputation*> HloComputation::MakeEmbeddedComputationsList()
   // To avoid special handling of this computation, cast away const of
   // 'this'. 'this' is immediately removed from the post order after
   // construction.
+  //
+  // TODO(b/78350259): This violates const-correctness, since while the original
+  // computation is not returned, we still retrieve non-const computations from
+  // a const one. Consider also avoiding const for HloComputation, or review XLA
+  // for const-correctness of non-HloInstruction* types like this.
   ComputeComputationPostOrder(const_cast<HloComputation*>(this), &visited,
                               &post_order);
 
@@ -723,18 +728,25 @@ Status HloComputation::Accept(
   return this->Accept(&visitor);
 }
 
-std::unique_ptr<HloComputation> HloComputation::Clone(const string& suffix,
-                                                      HloModule* module) {
+std::unique_ptr<HloComputation> HloComputation::Clone(
+    const string& suffix, HloModule* module,
+    HloInstruction::CloneMap* clone_map) {
   return CloneWithReplacements(
       /*replacements=*/std::unordered_map<const HloInstruction*,
                                           std::unique_ptr<HloInstruction>>(),
-      module, suffix);
+      module, clone_map, suffix);
 }
 
 std::unique_ptr<HloComputation> HloComputation::CloneWithReplacements(
     std::unordered_map<const HloInstruction*, std::unique_ptr<HloInstruction>>
         replacements,
-    HloModule* module, const string& suffix) {
+    HloModule* module, HloInstruction::CloneMap* clone_map,
+    const string& suffix) {
+  HloInstruction::CloneMap local_clone_map;
+  if (clone_map == nullptr) {
+    clone_map = &local_clone_map;
+  }
+
   // Look up instr in the replacements map, and return either the replacement,
   // or instr, if the replacement isn't present.
   //
@@ -756,24 +768,19 @@ std::unique_ptr<HloComputation> HloComputation::CloneWithReplacements(
     }
   }
 
-  std::unordered_map<HloInstruction*, HloInstruction*> clone_map;
   std::vector<std::unique_ptr<HloInstruction>> instructions;
   std::unique_ptr<HloInstruction> new_instr = nullptr;
   for (auto instr : postorder) {
     std::vector<HloInstruction*> new_operands;
     for (auto operand : instr->operands()) {
       auto replaced_operand = replace(operand);
-      // If replaced_operand is null, that means 'replacements' asked us not to
-      // include operand in the new computation.  But we can't do that, because
-      // operand is used by instr.
       CHECK_NE(replaced_operand, nullptr)
-          << "replacements map tried to eliminate a used instruction "
-          << operand->ToString() << ", used by " << instr->ToString();
-      new_operands.push_back(FindOrDie(clone_map, replaced_operand));
+          << "Replacements map specifies to leave out " << operand->ToString()
+          << ", but it is used by " << instr->ToString() << ".";
+      new_operands.push_back(FindOrDie(*clone_map, replaced_operand));
     }
-    new_instr =
-        instr->CloneWithNewOperands(instr->shape(), new_operands, module);
-    InsertOrDie(&clone_map, instr, new_instr.get());
+    new_instr = instr->CloneWithNewOperands(instr->shape(), new_operands,
+                                            module, clone_map);
     instructions.push_back(std::move(new_instr));
   }
   Builder builder(name() + "." + suffix);
@@ -781,27 +788,24 @@ std::unique_ptr<HloComputation> HloComputation::CloneWithReplacements(
     builder.AddInstruction(std::move(instr));
   }
   auto result = builder.Build(
-      /*root_instruction=*/FindOrDie(clone_map, replace(root_instruction())));
+      /*root_instruction=*/FindOrDie(*clone_map, replace(root_instruction())));
 
   // Clone control dependencies.
   for (auto instr : postorder) {
-    HloInstruction* new_instr = FindOrDie(clone_map, instr);
+    HloInstruction* new_instr = FindOrDie(*clone_map, instr);
     for (auto successor : instr->control_successors()) {
       auto replaced_successor = replace(successor);
-
-      // successor may not be in clone_map, because it might have been
-      // removed by the replacements map.
-      if (replaced_successor == nullptr) {
-        continue;
-      }
+      CHECK_NE(replaced_successor, nullptr)
+          << "Replacements map specifies to leave out " << successor->ToString()
+          << ", but it is control-depended-on by " << instr->ToString() << ".";
 
       TF_CHECK_OK(new_instr->AddControlDependencyTo(
-          FindOrDie(clone_map, replaced_successor)));
+          FindOrDie(*clone_map, replaced_successor)));
     }
   }
 
   // We cloned the elements of 'replacements', so they're all going to be
-  // destroyed.  HloInstructions need to be detached from their operands before
+  // destroyed. HloInstructions need to be detached from their operands before
   // they're destroyed, otherwise they stick around in the operands' users lists
   // and cause use-after-frees.
   for (auto& kv : replacements) {
diff --git a/tensorflow/compiler/xla/service/hlo_computation.h b/tensorflow/compiler/xla/service/hlo_computation.h
index 9d3f6e9a2c..9898355625 100644
--- a/tensorflow/compiler/xla/service/hlo_computation.h
+++ b/tensorflow/compiler/xla/service/hlo_computation.h
@@ -291,11 +291,17 @@ class HloComputation {
       const std::function<Status(const HloInstruction*)>& visitor_func) const;
 
   // Returns a deep copy of this computation including all instructions.
-  // If the module pointer is not nullptr, it will be the module where
-  // the cloned computations will be added to (in order to support deep
-  // cloning).
-  std::unique_ptr<HloComputation> Clone(const string& suffix = "clone",
-                                        HloModule* module = nullptr);
+  //
+  // If the module pointer is not nullptr, then the cloned computations will be
+  // added to this module in order to support deep cloning. Otherwise the module
+  // of the computation is used.
+  //
+  // If clone_map is not nullptr, then each original instruction that is cloned
+  // will be inserted and map to its clone. clone_map should not already contain
+  // any of the instructions to clone.
+  std::unique_ptr<HloComputation> Clone(
+      const string& suffix = "clone", HloModule* module = nullptr,
+      HloInstruction::CloneMap* clone_map = nullptr);
 
   // Like Clone(), but if an instruction is present in replacement_map, we use
   // the map's value to replace that instruction in the cloned computation.
@@ -305,7 +311,9 @@ class HloComputation {
   std::unique_ptr<HloComputation> CloneWithReplacements(
       std::unordered_map<const HloInstruction*, std::unique_ptr<HloInstruction>>
           replacements,
-      HloModule* module = nullptr, const string& suffix = "clone");
+      HloModule* module = nullptr,
+      HloInstruction::CloneMap* clone_map = nullptr,
+      const string& suffix = "clone");
 
   // Returns true if the given instruction can be removed from the computation.
   // Parameter instructions cannot be removed without violating invariants of
diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc
index bb4db89f0a..794f1b4682 100644
--- a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc
+++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc
@@ -322,11 +322,13 @@ class HloDotDumper {
  public:
   HloDotDumper(const HloComputation* computation, tensorflow::StringPiece label,
                const DebugOptions& debug_options, bool show_metadata,
-               const HloExecutionProfile* profile, NodeFilter filter)
+               bool show_backend_config, const HloExecutionProfile* profile,
+               NodeFilter filter)
       : computation_(computation),
         label_(label.ToString()),
         debug_options_(debug_options),
         show_metadata_(show_metadata),
+        show_backend_config_(show_backend_config),
         profile_(profile),
         filter_(std::move(filter)) {}
 
@@ -365,6 +367,7 @@ class HloDotDumper {
   string GetInstructionNodeShape(const HloInstruction* instr);
   string GetInstructionNodeLabel(const HloInstruction* instr);
   string GetInstructionNodeMetadata(const HloInstruction* instr);
+  string GetInstructionNodeBackendConfig(const HloInstruction* instr);
   string GetInstructionNodeExtraInfo(const HloInstruction* instr);
   string GetInstructionNodeInlinedOperands(const HloInstruction* instr);
   void AddInstructionIncomingEdges(const HloInstruction* instr);
@@ -393,6 +396,7 @@ class HloDotDumper {
   const string label_;                 // overall name for the graph
   const DebugOptions& debug_options_;
   const bool show_metadata_;
+  const bool show_backend_config_;
   const HloExecutionProfile* profile_;  // may be null
   const NodeFilter filter_;
 
@@ -611,6 +615,10 @@ tooltip = " ";
     if (!extra_info.empty()) {
       StrAppend(&subcomp_label, "<br/>", extra_info);
     }
+    string node_backend_config = GetInstructionNodeBackendConfig(parent_instr);
+    if (!node_backend_config.empty()) {
+      StrAppend(&subcomp_label, "<br/>", node_backend_config);
+    }
 
     bool highlight = filter_.Highlight(parent_instr);
     const char* fillcolor;
@@ -765,6 +773,7 @@ string HloDotDumper::DumpInstruction(const HloInstruction* instr) {
   string node_shape = GetInstructionNodeShape(instr);
   string node_label = GetInstructionNodeLabel(instr);
   string node_metadata = GetInstructionNodeMetadata(instr);
+  string node_backend_config = GetInstructionNodeBackendConfig(instr);
   string extra_info = GetInstructionNodeExtraInfo(instr);
   string inlined_constants = GetInstructionNodeInlinedOperands(instr);
   string trivial_subcomputation = GetInstructionTrivialComputationStr(instr);
@@ -782,8 +791,8 @@ string HloDotDumper::DumpInstruction(const HloInstruction* instr) {
   }
   // Build the text that will be displayed inside the node.
   string node_body = node_label;
-  for (const string& s :
-       {trivial_subcomputation, node_metadata, extra_info, inlined_constants}) {
+  for (const string& s : {trivial_subcomputation, node_metadata,
+                          node_backend_config, extra_info, inlined_constants}) {
     if (!s.empty()) {
       StrAppend(&node_body, "<br/>", s);
     }
@@ -1078,6 +1087,15 @@ string HloDotDumper::GetInstructionNodeMetadata(const HloInstruction* instr) {
   return Join(lines, "<br/>");
 }
 
+string HloDotDumper::GetInstructionNodeBackendConfig(
+    const HloInstruction* instr) {
+  if (!show_backend_config_ || instr->backend_config().empty()) {
+    return "";
+  }
+
+  return StrCat("backend_config=\"", instr->backend_config(), "\"");
+}
+
 string HloDotDumper::GetInstructionNodeExtraInfo(const HloInstruction* instr) {
   std::vector<string> lines;
 
@@ -1404,7 +1422,7 @@ string ExportGraph(const string& graph,
 string DumpGraph(const HloComputation& computation, const string& label,
                  const DebugOptions& debug_options,
                  const HloExecutionProfile* hlo_execution_profile,
-                 bool show_metadata) {
+                 bool show_metadata, bool show_backend_config) {
   GraphRendererInterface::GraphKind graph_kind;
   string graph;
   if (debug_options.xla_hlo_dump_as_graphdef()) {
@@ -1414,9 +1432,10 @@ string DumpGraph(const HloComputation& computation, const string& label,
                                                           &graph));
     graph_kind = GraphRendererInterface::TF_GRAPHDEF;
   } else {
-    graph = HloDotDumper(&computation, label, debug_options, show_metadata,
-                         hlo_execution_profile, NodeFilter())
-                .Dump();
+    graph =
+        HloDotDumper(&computation, label, debug_options, show_metadata,
+                     show_backend_config, hlo_execution_profile, NodeFilter())
+            .Dump();
     graph_kind = GraphRendererInterface::DOT_GRAPH;
   }
 
@@ -1427,15 +1446,15 @@ string DumpGraph(const HloComputation& computation, const string& label,
 }
 
 string DumpNeighborhoodAround(const HloInstruction& node, int radius,
-                              bool show_metadata) {
+                              bool show_metadata, bool show_backend_config) {
   auto debug_options = node.GetModule()->config().debug_options();
   string label =
       StrCat("Neighborhood of ", radius, " nodes around ", node.name());
   NodeFilter filter = MakeNodeFilter(&node, radius);
-  string graph =
-      HloDotDumper(node.parent(), label, debug_options, show_metadata,
-                   /*profile=*/nullptr, filter)
-          .Dump();
+  string graph = HloDotDumper(node.parent(), label, debug_options,
+                              show_metadata, show_backend_config,
+                              /*profile=*/nullptr, filter)
+                     .Dump();
   return ExportGraph(graph, GraphRendererInterface::DOT_GRAPH, debug_options);
 }
 
diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.h b/tensorflow/compiler/xla/service/hlo_graph_dumper.h
index 2704aae1e3..fc8e1468ac 100644
--- a/tensorflow/compiler/xla/service/hlo_graph_dumper.h
+++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.h
@@ -56,7 +56,7 @@ string MaybeDumpHloModule(const HloModule& module, const string& label,
 string DumpGraph(const HloComputation& computation, const string& label,
                  const DebugOptions& debug_options,
                  const HloExecutionProfile* hlo_execution_profile = nullptr,
-                 bool show_metadata = false);
+                 bool show_metadata = false, bool show_backend_config = false);
 
 // Like DumpGraph, but renders only nodes "near" the given node in the graph.
 //
@@ -64,7 +64,8 @@ string DumpGraph(const HloComputation& computation, const string& label,
 // (roughly) corresponds to the max distance a node may be from the primary node
 // before it's omitted from the graph.
 string DumpNeighborhoodAround(const HloInstruction& node, int radius,
-                              bool show_metadata = false);
+                              bool show_metadata = false,
+                              bool show_backend_config = false);
 
 // Dumps the HloModule::ToString() as a file into the provided directory path
 // suffixed with the provided label.
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.cc b/tensorflow/compiler/xla/service/hlo_instruction.cc
index a714d0e114..2c733726a6 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.cc
+++ b/tensorflow/compiler/xla/service/hlo_instruction.cc
@@ -109,6 +109,7 @@ StatusOr<std::unique_ptr<HloInstruction>> HloInstruction::CreateFromProto(
   instruction->name_ = proto.name();
 
   instruction->metadata_ = proto.metadata();
+  instruction->set_backend_config(proto.backend_config());
   if (proto.has_literal()) {
     TF_ASSIGN_OR_RETURN(instruction->literal_,
                         Literal::CreateFromProto(proto.literal()));
@@ -1231,12 +1232,15 @@ bool HloInstruction::HasSideEffect() const {
 std::unique_ptr<HloInstruction> HloInstruction::CloneWithNewOperands(
     const Shape& shape,
     tensorflow::gtl::ArraySlice<HloInstruction*> new_operands,
-    HloModule* module) const {
+    HloModule* module, CloneMap* clone_map) const {
   VLOG(3) << "CloneWithNewOperands:\n  " << ToString();
   VLOG(3) << "  new operands:";
   for (const HloInstruction* new_operand : new_operands) {
     VLOG(3) << "    %" << new_operand->name();
   }
+  if (module == nullptr) {
+    module = GetModule();
+  }
 
   std::unique_ptr<HloInstruction> clone;
 
@@ -1342,7 +1346,8 @@ std::unique_ptr<HloInstruction> HloInstruction::CloneWithNewOperands(
       break;
     case HloOpcode::kFft:
       CHECK_EQ(new_operands.size(), 1);
-      return CreateFft(shape, new_operands[0], fft_type_, fft_length_);
+      clone = CreateFft(shape, new_operands[0], fft_type_, fft_length_);
+      break;
     case HloOpcode::kCrossReplicaSum:
       clone = CreateCrossReplicaSum(shape, new_operands);
       break;
@@ -1415,9 +1420,15 @@ std::unique_ptr<HloInstruction> HloInstruction::CloneWithNewOperands(
     case HloOpcode::kConstant:
       clone = CreateConstant(literal_->CloneToUnique());
       break;
-    case HloOpcode::kFusion:
-      clone = CloneFusionWithNewOperands(shape, new_operands, module);
+    case HloOpcode::kFusion: {
+      CHECK_NE(module, nullptr);
+      auto new_fused_computation = module->AddEmbeddedComputation(
+          fused_instructions_computation()->Clone("clone", module, clone_map));
+      clone = CreateFusion(/*shape=*/shape, /*fusion_kind=*/fusion_kind(),
+                           /*operands=*/new_operands,
+                           /*fusion_computation=*/new_fused_computation);
       break;
+    }
     case HloOpcode::kParameter:
       clone = CreateParameter(parameter_number_, shape, name_);
       break;
@@ -1481,15 +1492,19 @@ std::unique_ptr<HloInstruction> HloInstruction::CloneWithNewOperands(
   }
   SetupDerivedInstruction(clone.get());
   clone->set_parent(parent_);
+  clone->set_backend_config(backend_config());
+  if (clone_map != nullptr) {
+    InsertOrDie(clone_map, this, clone.get());
+  }
   return clone;
 }
 
 HloInstruction::~HloInstruction() {}
 
-std::unique_ptr<HloInstruction> HloInstruction::Clone(const string& suffix,
-                                                      HloModule* module) const {
+std::unique_ptr<HloInstruction> HloInstruction::Clone(
+    const string& suffix, HloModule* module, CloneMap* clone_map) const {
   std::unique_ptr<HloInstruction> clone =
-      CloneWithNewOperands(shape_, operands_, module);
+      CloneWithNewOperands(shape_, operands_, module, clone_map);
   if (suffix.empty()) {
     clone->name_ = name();
   } else {
@@ -1526,71 +1541,6 @@ std::unique_ptr<HloInstruction> HloInstruction::Clone(const string& suffix,
   return clone;
 }
 
-std::unique_ptr<HloInstruction> HloInstruction::CloneFusionWithNewOperands(
-    const Shape& shape, tensorflow::gtl::ArraySlice<HloInstruction*> operands,
-    HloModule* module) const {
-  CHECK_EQ(opcode_, HloOpcode::kFusion);
-  CHECK(parent() != nullptr);
-
-  auto new_instruction =
-      WrapUnique(new HloInstruction(HloOpcode::kFusion, shape));
-  // Add the operands to our new fusion instruction.
-  for (HloInstruction* new_operand : operands) {
-    new_instruction->AppendOperand(new_operand);
-  }
-  // Clone all the fused instructions for the new fusion instruction.
-  HloInstructionMap<HloInstruction*> old_to_new;
-  std::list<std::unique_ptr<HloInstruction>> new_fused_instructions;
-  // Create the list of fused parameters by mapping through the cloned,
-  // fused instructions.
-  for (HloInstruction* old_fused_parameter :
-       fused_instructions_computation()->parameter_instructions()) {
-    new_fused_instructions.push_back(
-        old_fused_parameter->Clone("clone", module));
-    HloInstruction* new_fusion_parameter = new_fused_instructions.back().get();
-    InsertOrDie(&old_to_new, old_fused_parameter, new_fusion_parameter);
-  }
-  for (auto old_fused_instruction :
-       fused_instructions_computation()->MakeInstructionPostOrder()) {
-    if (old_fused_instruction->opcode() == HloOpcode::kParameter) {
-      FindOrDie(old_to_new, old_fused_instruction);
-      continue;
-    }
-    std::vector<HloInstruction*> new_operands;
-    for (int64 operand_idx = 0;
-         operand_idx < old_fused_instruction->operand_count(); ++operand_idx) {
-      HloInstruction* old_operand =
-          old_fused_instruction->mutable_operand(operand_idx);
-      new_operands.push_back(FindOrDie(old_to_new, old_operand));
-    }
-    new_fused_instructions.push_back(
-        old_fused_instruction->CloneWithNewOperands(
-            old_fused_instruction->shape(), new_operands, module));
-    HloInstruction* new_fused_instruction = new_fused_instructions.back().get();
-    new_fused_instruction->set_parent(parent_);
-    InsertOrDie(&old_to_new, old_fused_instruction, new_fused_instruction);
-  }
-  new_instruction->fusion_kind_ = fusion_kind_;
-  auto computation_builder = HloComputation::Builder(
-      fused_instructions_computation()->name() + ".clone",
-      new_instruction.get());
-  // We iterated the fusion instructions in reverse post order which means
-  // that we must reverse our new list of fusion instructions.
-  for (auto new_fused_instruction_iter = new_fused_instructions.rbegin();
-       new_fused_instruction_iter != new_fused_instructions.rend();
-       ++new_fused_instruction_iter) {
-    computation_builder.AddInstruction(std::move(*new_fused_instruction_iter));
-  }
-  if (module == nullptr) {
-    module = GetModule();
-  }
-  auto fused_root_ = fused_expression_root();
-  new_instruction->called_computations_.push_back(
-      CHECK_NOTNULL(module)->AddEmbeddedComputation(
-          computation_builder.Build(FindOrDie(old_to_new, fused_root_))));
-  return new_instruction;
-}
-
 std::pair<const HloInstruction*, ShapeIndex>
 HloInstruction::LatestNonGteAncestorAndIndex() const {
   const HloInstruction* hlo = this;
@@ -2172,6 +2122,9 @@ string HloInstruction::ToString(const HloPrintOptions& options) const {
        !metadata_.source_file().empty())) {
     StrAppend(&result, ", metadata={", xla::OpMetadataToString(metadata_), "}");
   }
+  if (options.print_backend_config() && !backend_config().empty()) {
+    StrAppend(&result, ", backend_config=\"", CEscape(backend_config()), "\"");
+  }
   return result;
 }
 
@@ -2357,6 +2310,7 @@ std::vector<string> HloInstruction::ExtraAttributesToString(
     extra.push_back(
         StrCat("custom_call_target=\"", CEscape(custom_call_target_), "\""));
   }
+
   return extra;
 }
 
@@ -2386,6 +2340,7 @@ HloInstructionProto HloInstruction::ToProto() const {
   }
 
   *proto.mutable_metadata() = metadata_;
+  proto.set_backend_config(backend_config());
   if (literal_ != nullptr) {
     *proto.mutable_literal() = literal_->ToProto();
   }
@@ -2971,6 +2926,7 @@ Status HloInstruction::AcceptOrdered(
       continue;
     }
 
+    // TODO(b/78350259): Eliminate const laundering.
     HloInstruction* instruction =
         const_cast<HloInstruction*>(const_instruction);
 
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h
index a5e9aecb9e..19c8c11453 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.h
+++ b/tensorflow/compiler/xla/service/hlo_instruction.h
@@ -66,6 +66,7 @@ class HloPrintOptions {
       : print_large_constants_(false),
         print_subcomputation_references_(true),
         print_metadata_(true),
+        print_backend_config_(true),
         compact_operands_(false),
         print_operand_shape_(true),
         print_program_shape_(true),
@@ -77,6 +78,7 @@ class HloPrintOptions {
         .set_print_large_constants(true)
         .set_print_subcomputation_references(true)
         .set_print_metadata(false)
+        .set_print_backend_config(false)
         .set_print_operand_shape(false)
         .set_print_program_shape(false)
         .set_print_percent(false);
@@ -99,12 +101,18 @@ class HloPrintOptions {
     return *this;
   }
 
-  // If true, metatdata will be printed.
+  // If true, metadata will be printed.
   HloPrintOptions& set_print_metadata(bool value) {
     print_metadata_ = value;
     return *this;
   }
 
+  // If true, backend_config will be printed.
+  HloPrintOptions& set_print_backend_config(bool value) {
+    print_backend_config_ = value;
+    return *this;
+  }
+
   // If true, operands' shapes will be printed.
   HloPrintOptions& set_print_operand_shape(bool value) {
     print_operand_shape_ = value;
@@ -141,6 +149,7 @@ class HloPrintOptions {
     return print_subcomputation_references_;
   }
   bool print_metadata() const { return print_metadata_; }
+  bool print_backend_config() const { return print_metadata_; }
   bool compact_operands() const { return compact_operands_; }
   bool print_operand_shape() const { return print_operand_shape_; }
   bool print_program_shape() const { return print_program_shape_; }
@@ -151,6 +160,7 @@ class HloPrintOptions {
   bool print_large_constants_;
   bool print_subcomputation_references_;
   bool print_metadata_;
+  bool print_backend_config_;
   bool compact_operands_;
   bool print_operand_shape_;
   bool print_program_shape_;
@@ -643,6 +653,8 @@ class HloInstruction {
   // Detaches an instruction from its operands. That is, remove the instruction
   // from each operand's user set. This should only be called prior to
   // deallocating the instruction.
+  //
+  // TODO(b/78305363): Make this automatic when deleting an instruction.
   void DetachFromOperands();
 
   // Performs a postorder DFS visit using this node as the root. If
@@ -1157,23 +1169,30 @@ class HloInstruction {
   // Precondition: opcode() == HloOpcode::kRng
   RandomDistribution random_distribution() const;
 
+  // See documentation for Clone().
+  using CloneMap = std::unordered_map<const HloInstruction*, HloInstruction*>;
+
   // Clones the HLO instruction. The clone will have the same opcode, shape, and
   // operands. After creation the clone has no uses. "this" (the instruction
   // cloned from) is not changed. Suffix is the string to append to the name of
-  // the instruction to form the name of the cloned instruction.  If the module
-  // pointer is not nullptr, it will be the module where the cloned computations
-  // will be added to (in order to support deep cloning).  Ignores the control
-  // predecessors and successors of this HLO instruction.
+  // the instruction to form the name of the cloned instruction. Ignores the
+  // control predecessors and successors of this HLO instruction.
+  //
+  // If the module pointer is not nullptr, then any cloned computations will be
+  // added to this module in order to support deep cloning. Otherwise the module
+  // of the instruction is used.
+  //
+  // If clone_map is not nullptr, then each original instruction that is cloned
+  // will be inserted and map to its clone. clone_map should not already contain
+  // any of the instructions to clone.
   std::unique_ptr<HloInstruction> Clone(const string& suffix = "clone",
-                                        HloModule* module = nullptr) const;
+                                        HloModule* module = nullptr,
+                                        CloneMap* clone_map = nullptr) const;
 
-  // Clones the HLO instruction as above but with new shape and operands.  If
-  // the module pointer is not nullptr, it will be the module where the cloned
-  // computations will be added to (in order to support deep cloning).  Ignores
-  // the control predecessors and successors of this HLO instruction.
+  // Clones the HLO instruction as above but with new shape and operands.
   std::unique_ptr<HloInstruction> CloneWithNewOperands(
       const Shape& shape, tensorflow::gtl::ArraySlice<HloInstruction*> operands,
-      HloModule* module = nullptr) const;
+      HloModule* module = nullptr, CloneMap* clone_map = nullptr) const;
 
   // Returns the computations this instruction directly calls (if any).
   const std::vector<HloComputation*>& called_computations() const {
@@ -1262,6 +1281,19 @@ class HloInstruction {
   // if no id has been assigned yet).
   int unique_id() const { return unique_id_; }
 
+  // Returns the backend-specific configuration for how a backend should compile
+  // this HLO. The meaning of the field is backend specific. Not for use before
+  // or during general HLO optimization, since HLO optimizations do not preserve
+  // this field and they cannot interpret it due to its meaning being backend
+  // specific.
+  //
+  // TODO(b/78194644): Introduce structured configuration format as per
+  // go/xla-heuristics.
+  const string& backend_config() const { return backend_config_; }
+  void set_backend_config(string backend_config) {
+    backend_config_ = std::move(backend_config);
+  }
+
   // Sets the debug metadata for this instruction.
   void set_metadata(const OpMetadata& metadata) { metadata_ = metadata; }
   const OpMetadata& metadata() const { return metadata_; }
@@ -1283,6 +1315,7 @@ class HloInstruction {
   // Get/Set the number of partitions per outer dimension (in order, starting
   // with outer-most dimension first). Currently used by the parallel cpu
   // backend to partition HLOs into parallel tasks.
+  //
   // TODO(b/62783254) Replace these methods with a more general way to
   // annotate HLOs with backend-specific information.
   const std::vector<int64>& outer_dimension_partitions() const {
@@ -1510,6 +1543,10 @@ class HloInstruction {
   // The string representation of the infeed configuration.
   string infeed_config_;
 
+  // The backend-specific configuration for how a backend should compile this
+  // HLO. See the documentation on backend_config().
+  string backend_config_;
+
   // String identifier for instruction.
   string name_;
 
diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc
index c7a7192867..5308fb5848 100644
--- a/tensorflow/compiler/xla/service/hlo_module.cc
+++ b/tensorflow/compiler/xla/service/hlo_module.cc
@@ -46,6 +46,18 @@ HloModule::HloModule(const string& name, const HloModuleConfig& config)
       config_(config),
       unique_id_(next_unique_module_id_++) {}
 
+StatusOr<HloInstruction*> HloModule::LaunderConstInstructionFromModule(
+    const HloInstruction* hlo) {
+  if (hlo == nullptr) {
+    return nullptr;
+  }
+
+  TF_RET_CHECK(hlo->GetModule() == this);
+
+  // TODO(b/78350259): Eliminate const laundering.
+  return const_cast<HloInstruction*>(hlo);
+}
+
 HloComputation* HloModule::AddComputationInternal(
     std::unique_ptr<HloComputation> computation, bool is_entry,
     bool uniquify_names) {
diff --git a/tensorflow/compiler/xla/service/hlo_module.h b/tensorflow/compiler/xla/service/hlo_module.h
index f9674df812..1604a72612 100644
--- a/tensorflow/compiler/xla/service/hlo_module.h
+++ b/tensorflow/compiler/xla/service/hlo_module.h
@@ -217,6 +217,25 @@ class HloModule {
   // the lifetime of this process.
   int unique_id() const { return unique_id_; }
 
+  // Returns a non-const version of the passed-in const HloInstruction*. This is
+  // safe on the argument that if you have a non-const module, then you can
+  // access all instructions in the module as non-const.
+  //
+  // Returns an error if the passed-in instruction is not from this module,
+  // except that it is allowed to pass in a null pointer.
+  //
+  // TODO(b/78350259): Eliminate const laundering. The argument above is not
+  // reliable since at any time someone could add or discover a way for a
+  // non-const module to transitively contain a const HloInstruction. The
+  // reliable way to do this would be to create a const laundering map from a
+  // module, mapping each encountered HloInstruction to its non-const version
+  // and then look up each instruction in need of laundering in that map, but
+  // this is much more expensive and complicated. This returns a Status instead
+  // of doing a CHECK-failure in part to make it strongly apparent that this is
+  // something that can fail.
+  StatusOr<HloInstruction*> LaunderConstInstructionFromModule(
+      const HloInstruction* hlo);
+
  private:
   HloComputation* AddComputationInternal(
       std::unique_ptr<HloComputation> computation, bool is_entry,
diff --git a/tensorflow/compiler/xla/service/hlo_verifier.cc b/tensorflow/compiler/xla/service/hlo_verifier.cc
index 8a30cbf9cd..096ebb7946 100644
--- a/tensorflow/compiler/xla/service/hlo_verifier.cc
+++ b/tensorflow/compiler/xla/service/hlo_verifier.cc
@@ -116,7 +116,7 @@ Status ShapeVerifier::HandleOutfeed(HloInstruction* outfeed) {
   // produces no HLO value in the graph.
   if (!ShapeUtil::Compatible(outfeed->outfeed_shape(),
                              outfeed->operand(0)->shape())) {
-    return InvalidArgument(
+    return InternalError(
         "Expected outfeed to have shape compatible with operand's shape %s, "
         "actual shape is %s:\n%s",
         ShapeUtil::HumanString(outfeed->operand(0)->shape()).c_str(),
@@ -200,7 +200,7 @@ Status ShapeVerifier::HandleTranspose(HloInstruction* transpose) {
                      transpose->operand(0)->shape(), transpose->dimensions()));
 }
 
-Status ShapeVerifier::HandleParameter(HloInstruction*) {
+Status ShapeVerifier::HandleParameter(HloInstruction* hlo) {
   return tensorflow::Status::OK();
 }
 
@@ -410,7 +410,7 @@ Status CheckMixedPrecisionOperands(const HloInstruction* instruction) {
               if (fp_type == PRIMITIVE_TYPE_INVALID) {
                 fp_type = subshape.element_type();
               } else if (fp_type != subshape.element_type()) {
-                return FailedPrecondition(
+                return InternalError(
                     "Seen floating point types of different precisions in "
                     "%s, but mixed precision is disallowed.",
                     instruction->ToString().c_str());
@@ -490,7 +490,7 @@ Status ShapeVerifier::CheckShape(const HloInstruction* instruction,
       }
   }
   if (!compatible) {
-    return InvalidArgument(
+    return InternalError(
         "Expected instruction to have shape compatible with %s, actual "
         "shape is %s:\n%s",
         ShapeUtil::HumanString(inferred_shape).c_str(),
@@ -541,7 +541,7 @@ Status ShapeVerifier::CheckVariadicShape(const HloInstruction* instruction) {
 Status ShapeVerifier::CheckSameChannel(const HloInstruction* instr1,
                                        const HloInstruction* instr2) {
   if (instr1->channel_id() != instr2->channel_id()) {
-    return FailedPrecondition(
+    return InternalError(
         "Expected to have the same channel id, actual channel ids are: %s "
         "(%lld), %s (%lld)",
         instr1->ToString().c_str(), instr1->channel_id(),
@@ -571,22 +571,22 @@ string ComputationsToString(
 Status VerifyHloStructure(HloModule* module) {
   for (const HloComputation* computation : module->computations()) {
     if (computation->parent() == nullptr) {
-      return FailedPrecondition("Computation %s has a null parent pointer",
-                                computation->name().c_str());
+      return InternalError("Computation %s has a null parent pointer",
+                           computation->name().c_str());
     }
     if (computation->parent() != module) {
-      return FailedPrecondition(
+      return InternalError(
           "Computation %s parent() does not point to parent module",
           computation->name().c_str());
     }
 
     for (const HloInstruction* instruction : computation->instructions()) {
       if (instruction->parent() == nullptr) {
-        return FailedPrecondition("Instruction %s has a null parent pointer",
-                                  instruction->name().c_str());
+        return InternalError("Instruction %s has a null parent pointer",
+                             instruction->name().c_str());
       }
       if (instruction->parent() != computation) {
-        return FailedPrecondition(
+        return InternalError(
             "Instruction %s parent() does not point to parent computation",
             instruction->name().c_str());
       }
@@ -602,7 +602,7 @@ Status VerifyHloStructure(HloModule* module) {
       for (int i = 0; i < instruction->operand_count(); ++i) {
         const HloInstruction* operand = instruction->operand(i);
         if (operand->parent() != instruction->parent()) {
-          return FailedPrecondition(
+          return InternalError(
               "Operand %d (%s) of instruction %s is in a different "
               "computation: %s vs %s",
               i, operand->name().c_str(), instruction->name().c_str(),
@@ -619,7 +619,7 @@ Status HloVerifier::CheckFusionInstruction(HloInstruction* fusion) const {
   // The parent fusion instruction of the fusion computation must be 'fusion'.
   HloComputation* fused_computation = fusion->fused_instructions_computation();
   if (fusion != fused_computation->FusionInstruction()) {
-    return FailedPrecondition(
+    return InternalError(
         "Instruction of fused computation does not match expected instruction "
         "%s.",
         fusion->ToString().c_str());
@@ -635,37 +635,37 @@ Status HloVerifier::CheckFusionInstruction(HloInstruction* fusion) const {
   for (auto* instruction : fused_computation->instructions()) {
     if (fused_root == instruction) {
       if (root_owned) {
-        return FailedPrecondition("Root appears more than once in %s.",
-                                  fusion->ToString().c_str());
+        return InternalError("Root appears more than once in %s.",
+                             fusion->ToString().c_str());
       }
       root_owned = true;
     }
     for (int i = 0; i < fused_parameters.size(); ++i) {
       if (fused_parameters[i] == instruction) {
         if (parameter_owned[i]) {
-          return FailedPrecondition("Parameter appears more than once in %s.",
-                                    fusion->ToString().c_str());
+          return InternalError("Parameter appears more than once in %s.",
+                               fusion->ToString().c_str());
         }
         parameter_owned[i] = true;
       }
     }
   }
   if (!root_owned) {
-    return FailedPrecondition("Root not found in computation of %s.",
-                              fusion->ToString().c_str());
+    return InternalError("Root not found in computation of %s.",
+                         fusion->ToString().c_str());
   }
   // Make sure all the parameter_owned entries are set
   for (int i = 0; i < parameter_owned.size(); i++) {
     if (!parameter_owned[i]) {
-      return FailedPrecondition("Parameter %d not found in computation of %s.",
-                                i, fusion->ToString().c_str());
+      return InternalError("Parameter %d not found in computation of %s.", i,
+                           fusion->ToString().c_str());
     }
   }
 
   // Fused root must have no users.
   if (fused_root->user_count() != 0) {
-    return FailedPrecondition("Root of %s may not have users.",
-                              fusion->ToString().c_str());
+    return InternalError("Root of %s may not have users.",
+                         fusion->ToString().c_str());
   }
 
   // All uses of fused instructions must be in the fusion computation, and every
@@ -674,13 +674,13 @@ Status HloVerifier::CheckFusionInstruction(HloInstruction* fusion) const {
        fusion->fused_instructions_computation()->instructions()) {
     if (instruction != fused_root) {
       if (instruction->user_count() == 0) {
-        return FailedPrecondition(
-            "Non-root instruction %s in %s must have users.",
-            instruction->ToString().c_str(), fusion->ToString().c_str());
+        return InternalError("Non-root instruction %s in %s must have users.",
+                             instruction->ToString().c_str(),
+                             fusion->ToString().c_str());
       }
       for (auto& user : instruction->users()) {
         if (fused_computation != user->parent()) {
-          return FailedPrecondition(
+          return InternalError(
               "Non-root instruction %s in %s may not have external users.",
               instruction->ToString().c_str(), fusion->ToString().c_str());
         }
@@ -695,34 +695,33 @@ Status HloVerifier::CheckFusionInstruction(HloInstruction* fusion) const {
   for (auto fused_param : fused_parameters) {
     int64 param_no = fused_param->parameter_number();
     if (param_no < 0) {
-      return FailedPrecondition(
-          "Unexpected negative parameter number %lld in %s.", param_no,
-          fusion->ToString().c_str());
+      return InternalError("Unexpected negative parameter number %lld in %s.",
+                           param_no, fusion->ToString().c_str());
     }
     if (param_no >= fused_parameters.size()) {
-      return FailedPrecondition(
+      return InternalError(
           "Unexpected parameter number %lld in %s: higher then number of "
           "parameters %lu.",
           param_no, fusion->ToString().c_str(), fused_parameters.size());
     }
     if (parameter_numbers[param_no]) {
-      return FailedPrecondition(
+      return InternalError(
           "Did not expect parameter number %lld more than once in %s.",
           param_no, fusion->ToString().c_str());
     }
     parameter_numbers[param_no] = true;
     if (!ShapeUtil::Compatible(fused_param->shape(),
                                fusion->operand(param_no)->shape())) {
-      return FailedPrecondition(
+      return InternalError(
           "Shape mismatch between parameter number %lld and its operand in %s.",
           param_no, fusion->ToString().c_str());
     }
   }
-  // Make sure all the parameter_numbers entries were seen
+  // Make sure all the parameter_numbers entries were seen.
   for (int i = 0; i < parameter_numbers.size(); i++) {
     if (!parameter_numbers[i]) {
-      return FailedPrecondition("Did not see parameter number %d in %s.", i,
-                                fusion->ToString().c_str());
+      return InternalError("Did not see parameter number %d in %s.", i,
+                           fusion->ToString().c_str());
     }
   }
 
diff --git a/tensorflow/compiler/xla/statusor.h b/tensorflow/compiler/xla/statusor.h
index cccbce5fc8..0e1387c939 100644
--- a/tensorflow/compiler/xla/statusor.h
+++ b/tensorflow/compiler/xla/statusor.h
@@ -13,13 +13,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-// StatusOr<T> is the union of a Status object and a T
-// object. StatusOr models the concept of an object that is either a
-// usable value, or an error Status explaining why such a value is
-// not present. To this end, StatusOr<T> does not allow its Status
-// value to be Status::OK. Furthermore, the value of a StatusOr<T*>
-// must not be null. This is enforced by a debug check in most cases,
-// but even when it is not, clients must not set the value to null.
+// StatusOr<T> is the union of a Status object and a T object. StatusOr models
+// the concept of an object that is either a value, or an error Status
+// explaining why such a value is not present. To this end, StatusOr<T> does not
+// allow its Status value to be Status::OK.
 //
 // The primary use-case for StatusOr<T> is as the return value of a
 // function which may fail.
diff --git a/tensorflow/compiler/xla/statusor_test.cc b/tensorflow/compiler/xla/statusor_test.cc
index f9d25945bc..7d76370e85 100644
--- a/tensorflow/compiler/xla/statusor_test.cc
+++ b/tensorflow/compiler/xla/statusor_test.cc
@@ -75,6 +75,14 @@ TEST(StatusOr, ElementType) {
   static_assert(std::is_same<StatusOr<char>::element_type, char>(), "");
 }
 
+TEST(StatusOr, NullPointerStatusOr) {
+  // As a very special case, null-plain-pointer StatusOr used to be an
+  // error. Test that it no longer is.
+  StatusOr<int*> null_status(nullptr);
+  EXPECT_TRUE(null_status.ok());
+  EXPECT_EQ(null_status.ValueOrDie(), nullptr);
+}
+
 TEST(StatusOr, TestNoDefaultConstructorInitialization) {
   // Explicitly initialize it with an error code.
   StatusOr<NoDefaultConstructor> statusor(tensorflow::errors::Cancelled(""));
diff --git a/tensorflow/compiler/xla/tools/parser/hlo_parser.cc b/tensorflow/compiler/xla/tools/parser/hlo_parser.cc
index 40dc0730ce..156a06c596 100644
--- a/tensorflow/compiler/xla/tools/parser/hlo_parser.cc
+++ b/tensorflow/compiler/xla/tools/parser/hlo_parser.cc
@@ -440,6 +440,10 @@ bool HloParser::ParseInstruction(HloComputation::Builder* builder,
   optional<OpMetadata> metadata;
   attrs["metadata"] = {/*required=*/false, AttrTy::kMetadata, &metadata};
 
+  optional<string> backend_config;
+  attrs["backend_config"] = {/*required=*/false, AttrTy::kString,
+                             &backend_config};
+
   HloInstruction* instruction;
   switch (opcode) {
     case HloOpcode::kParameter: {
@@ -1094,8 +1098,7 @@ bool HloParser::ParseInstruction(HloComputation::Builder* builder,
 
   instruction->set_name(name);
 
-  // Add common attrs (sharding, control predecessors) to the instruction, if
-  // they were seen.
+  // Add shared attributes like metadata to the instruction, if they were seen.
   if (sharding) {
     instruction->set_sharding(
         HloSharding::FromProto(sharding.value()).ValueOrDie());
@@ -1112,6 +1115,9 @@ bool HloParser::ParseInstruction(HloComputation::Builder* builder,
   if (metadata) {
     instruction->set_metadata(*metadata);
   }
+  if (backend_config) {
+    instruction->set_backend_config(std::move(*backend_config));
+  }
   return AddInstruction(name, instruction, name_loc);
 }  // NOLINT(readability/fn_size)
 
diff --git a/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc b/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc
index d38d8907a6..e100d8cda1 100644
--- a/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc
+++ b/tensorflow/compiler/xla/tools/parser/hlo_parser_test.cc
@@ -65,7 +65,7 @@ ENTRY %axpy.v5 (alpha: f32[], x: f32[2,4], y: f32[2,4]) -> f32[2,4] {
 R"(HloModule constant_pred_module
 
 ENTRY %constant_pred () -> pred[] {
-  ROOT %constant = pred[] constant(true), metadata={op_type="const" op_name="\"it\'s not a problem\n" source_file="path/to/test.cc" source_line=68}
+  ROOT %constant = pred[] constant(true), metadata={op_type="const" op_name="\"it\'s not a problem\n" source_file="path/to/test.cc" source_line=68}, backend_config="foo\" bar"
 }
 
 )"
@@ -81,13 +81,14 @@ ENTRY %constant_s32 () -> s32[] {
 
 )"
 },
-// f32 constant, but the value is not a decimal
+// f32 constant, but the value is not a decimal and there is a backend
+// configuration
 {
 "ConstantF32",
 R"(HloModule ConstantF32_module
 
 ENTRY %ConstantF32.v4 () -> f32[] {
-  ROOT %constant = f32[] constant(42)
+  ROOT %constant = f32[] constant(42), backend_config="this is a configuration"
 }
 
 )"
@@ -1013,6 +1014,19 @@ ENTRY %SelectScalarS32True.v4 () -> s32[] {
   // but the constant names will not be exactly the same.
 }
 
+TEST_F(HloParserTest, ConfigurationField) {
+  const string original = R"(HloModule AModule
+ENTRY %configuration_test() -> s32[] {
+  %constant = s32[] constant(42), backend_config="foo bar"
+})";
+  auto result = Parse(original);
+  TF_ASSERT_OK(result.status());
+  EXPECT_EQ("foo bar", result.ValueOrDie()
+                           ->entry_computation()
+                           ->root_instruction()
+                           ->backend_config());
+}
+
 TEST_F(HloParserTest, LiteralDimensionsMismatch_1) {
   const string original = R"(HloModule some_2_module
 
@@ -1092,7 +1106,7 @@ ENTRY %Convolve1D1Window_0.v3 (input: f32[1,2,1], filter: f32[1,1,1]) -> f32[1,2
   %input = f32[1,2,1]{2,1,0} parameter(0)
   %copy = f32[1,2,1]{2,0,1} copy(f32[1,2,1]{2,1,0} %input)
   %filter = f32[1,1,1]{2,1,0} parameter(1)
-  ROOT %convolution = f32[1,2,1]{2,0,1} convolution(f32[1,2,1]{2,0,1} %copy, f32[1,1,1]{2,1,0} %filter), sharding={maximal device=1}, dim_labels=b0f_0io->b0f, window={pad=1_1 size=2}
+  ROOT %convolution = f32[1,2,1]{2,0,1} convolution(f32[1,2,1]{2,0,1} %copy, f32[1,1,1]{2,1,0} %filter), sharding={maximal device=1}, backend_config="foo", dim_labels=b0f_0io->b0f, window={pad=1_1 size=2}
 }
 
 )";
-- 
GitLab


From bf228e1435da0032d2529de93661b742ee8a7048 Mon Sep 17 00:00:00 2001
From: Jiri Simsa <jsimsa@google.com>
Date: Fri, 4 May 2018 17:03:52 -0700
Subject: [PATCH 274/395] [tf.data] Adding `num_parallel_calls` to
 `map_and_batch`.

PiperOrigin-RevId: 195495206
---
 .../kernel_tests/batch_dataset_op_test.py     |  44 +-
 .../contrib/data/python/ops/batching.py       |  47 +-
 .../base_api/api_def_MapAndBatchDataset.pbtxt |  35 +-
 .../api_def_MapAndBatchDatasetV2.pbtxt        |  54 ++
 .../kernels/data/map_and_batch_dataset_op.cc  | 773 +++++++++---------
 tensorflow/core/ops/dataset_ops.cc            |  13 +
 6 files changed, 538 insertions(+), 428 deletions(-)
 create mode 100644 tensorflow/core/api_def/base_api/api_def_MapAndBatchDatasetV2.pbtxt

diff --git a/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py b/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py
index 6588fd04ac..2568b899d7 100644
--- a/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py
+++ b/tensorflow/contrib/data/python/kernel_tests/batch_dataset_op_test.py
@@ -427,7 +427,9 @@ class BatchDatasetTest(test.TestCase):
     self.assertEqual([None], dataset.output_shapes[1][0].as_list())
     self.assertEqual([None, 30], dataset.output_shapes[1][1].as_list())
 
-  def _testMapAndBatchDatasetHelper(self, num_parallel_batches=1):
+  def _testMapAndBatchDatasetHelper(self,
+                                    num_parallel_calls=None,
+                                    num_parallel_batches=None):
     """Test a dataset that maps a TF function across its input elements."""
     # The pipeline is TensorSliceDataset ->
     # RepeatDataset(count) -> MapAndBatchDataset(square_3, batch_size).
@@ -446,6 +448,7 @@ class BatchDatasetTest(test.TestCase):
             batching.map_and_batch(
                 map_func=_map_fn,
                 batch_size=batch_size,
+                num_parallel_calls=num_parallel_calls,
                 num_parallel_batches=num_parallel_batches))
         .make_initializable_iterator())
     init_op = iterator.initializer
@@ -497,12 +500,18 @@ class BatchDatasetTest(test.TestCase):
       with self.assertRaises(errors.InvalidArgumentError):
         sess.run(init_op, feed_dict={count: 14, batch_size: 0})
 
-  def testMapAndBatchDataset(self):
+  def testMapAndBatch(self):
     return self._testMapAndBatchDatasetHelper()
 
-  def testMapAndBatchDatasetWithParallelBatching(self):
+  def testMapAndBatchWithParallelBatches(self):
     return self._testMapAndBatchDatasetHelper(num_parallel_batches=10)
 
+  def testMapAndBatchWithSequentialCalls(self):
+    return self._testMapAndBatchDatasetHelper(num_parallel_calls=1)
+
+  def testMapAndBatchWithParallelCalls(self):
+    return self._testMapAndBatchDatasetHelper(num_parallel_calls=2)
+
   def _testMapAndBatchPartialBatchHelper(self, drop_remainder=False):
     iterator = (
         dataset_ops.Dataset.range(10).apply(
@@ -682,7 +691,7 @@ class UnbatchDatasetSerializationTest(
 class MapAndBatchDatasetSerializationTest(
     dataset_serialization_test_base.DatasetSerializationTestBase):
 
-  def testSerializationCore(self):
+  def testNumParallelBatches(self):
     range_size = 11
     num_repeats = 2
     batch_size = 5
@@ -709,6 +718,33 @@ class MapAndBatchDatasetSerializationTest(
     self.run_core_tests(lambda: build_ds(10, True), lambda: build_ds(15, True),
                         num_outputs_drop_remainder)
 
+  def testNumParallelCalls(self):
+    range_size = 11
+    num_repeats = 2
+    batch_size = 5
+    total_outputs = range_size * num_repeats
+    num_outputs_drop_remainder = total_outputs // batch_size
+    num_outputs_keep_remainder = int(math.ceil(total_outputs / batch_size))
+    num_parallel_calls = 7
+
+    def build_ds(range_start, drop_remainder=False):
+
+      def _map_fn(x):
+        return math_ops.square(x)
+
+      return dataset_ops.Dataset.range(
+          range_start, range_start + range_size).repeat(num_repeats).apply(
+              batching.map_and_batch(
+                  map_func=_map_fn,
+                  batch_size=batch_size,
+                  num_parallel_calls=num_parallel_calls,
+                  drop_remainder=drop_remainder))
+
+    self.run_core_tests(lambda: build_ds(10), lambda: build_ds(15),
+                        num_outputs_keep_remainder)
+    self.run_core_tests(lambda: build_ds(10, True), lambda: build_ds(15, True),
+                        num_outputs_drop_remainder)
+
 
 class PaddedBatchDatasetSerializationTest(
     dataset_serialization_test_base.DatasetSerializationTestBase):
diff --git a/tensorflow/contrib/data/python/ops/batching.py b/tensorflow/contrib/data/python/ops/batching.py
index 42ec2b0b01..b9393de4e9 100644
--- a/tensorflow/contrib/data/python/ops/batching.py
+++ b/tensorflow/contrib/data/python/ops/batching.py
@@ -466,14 +466,14 @@ def assert_element_shape(expected_shapes):
 class _MapAndBatchDataset(dataset_ops.MapDataset):
   """A `Dataset` that maps a function over a batch of elements."""
 
-  def __init__(self, input_dataset, map_func, batch_size, num_parallel_batches,
+  def __init__(self, input_dataset, map_func, batch_size, num_parallel_calls,
                drop_remainder):
     """See `Dataset.map()` for details."""
     super(_MapAndBatchDataset, self).__init__(input_dataset, map_func)
     self._batch_size_t = ops.convert_to_tensor(
         batch_size, dtype=dtypes.int64, name="batch_size")
-    self._num_parallel_batches_t = ops.convert_to_tensor(
-        num_parallel_batches, dtype=dtypes.int64, name="num_parallel_batches")
+    self._num_parallel_calls_t = ops.convert_to_tensor(
+        num_parallel_calls, dtype=dtypes.int64, name="num_parallel_calls")
     self._drop_remainder_t = ops.convert_to_tensor(
         drop_remainder, dtype=dtypes.bool, name="drop_remainder")
 
@@ -483,12 +483,12 @@ class _MapAndBatchDataset(dataset_ops.MapDataset):
   def _as_variant_tensor(self):
     # pylint: disable=protected-access
     input_resource = self._input_dataset._as_variant_tensor()
-    return gen_dataset_ops.map_and_batch_dataset(
+    return gen_dataset_ops.map_and_batch_dataset_v2(
         input_resource,
         self._map_func.captured_inputs,
         f=self._map_func,
         batch_size=self._batch_size_t,
-        num_parallel_batches=self._num_parallel_batches_t,
+        num_parallel_calls=self._num_parallel_calls_t,
         drop_remainder=self._drop_remainder_t,
         output_types=nest.flatten(
             sparse.as_dense_types(self.output_types, self.output_classes)),
@@ -511,8 +511,9 @@ class _MapAndBatchDataset(dataset_ops.MapDataset):
 
 def map_and_batch(map_func,
                   batch_size,
-                  num_parallel_batches=1,
-                  drop_remainder=False):
+                  num_parallel_batches=None,
+                  drop_remainder=False,
+                  num_parallel_calls=None):
   """Fused implementation of `map` and `batch`.
 
   Maps `map_func` across `batch_size` consecutive elements of this dataset
@@ -528,21 +529,37 @@ def map_and_batch(map_func,
       nested structure of tensors.
     batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
       consecutive elements of this dataset to combine in a single batch.
-    num_parallel_batches: A `tf.int64` scalar `tf.Tensor`, representing the
-      number of batches to create in parallel. On one hand, higher values can
-      help mitigate the effect of stragglers. On the other hand, higher values
-      can increase contention if CPU is scarce.
-    drop_remainder: A `tf.bool` scalar `tf.Tensor`, representing whether the
-      last batch should be dropped in case its size is smaller than desired;
-      the default behavior is not to drop the smaller batch.
+    num_parallel_batches: (Optional.) A `tf.int64` scalar `tf.Tensor`,
+      representing the number of batches to create in parallel. On one hand,
+      higher values can help mitigate the effect of stragglers. On the other
+      hand, higher values can increase contention if CPU is scarce.
+    drop_remainder: (Optional.) A `tf.bool` scalar `tf.Tensor`, representing
+      whether the last batch should be dropped in case its size is smaller than
+      desired; the default behavior is not to drop the smaller batch.
+    num_parallel_calls: (Optional.) A `tf.int32` scalar `tf.Tensor`,
+        representing the number of elements to process in parallel. If not
+        specified, `batch_size * num_parallel_batches` elements will be
+        processed in parallel.
 
   Returns:
     A `Dataset` transformation function, which can be passed to
     @{tf.data.Dataset.apply}.
+
+  Raises:
+    ValueError: If both `num_parallel_batches` and `num_parallel_calls` are
+      specified.
   """
 
+  if num_parallel_batches is None and num_parallel_calls is None:
+    num_parallel_calls = batch_size
+  elif num_parallel_batches is not None and num_parallel_calls is None:
+    num_parallel_calls = batch_size * num_parallel_batches
+  elif num_parallel_batches is not None and num_parallel_calls is not None:
+    raise ValueError("The `num_parallel_batches` and `num_parallel_calls` "
+                     "arguments are mutually exclusive.")
+
   def _apply_fn(dataset):
     return _MapAndBatchDataset(dataset, map_func, batch_size,
-                               num_parallel_batches, drop_remainder)
+                               num_parallel_calls, drop_remainder)
 
   return _apply_fn
diff --git a/tensorflow/core/api_def/base_api/api_def_MapAndBatchDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_MapAndBatchDataset.pbtxt
index bf544703de..e230c51edf 100644
--- a/tensorflow/core/api_def/base_api/api_def_MapAndBatchDataset.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_MapAndBatchDataset.pbtxt
@@ -1,5 +1,19 @@
 op {
   graph_op_name: "MapAndBatchDataset"
+  visibility: HIDDEN
+  in_arg {
+    name: "input_dataset"
+    description: <<END
+A variant tensor representing the input dataset.
+END
+  }
+  in_arg {
+    name: "other_arguments"
+    description: <<END
+A list of tensors, typically values that were captured when building a closure
+for `f`.
+END
+  }
   in_arg {
     name: "batch_size"
     description: <<END
@@ -11,13 +25,26 @@ END
   in_arg {
     name: "num_parallel_batches"
     description: <<END
-A scalar representing the number of batches to create in
-parallel. Processing multiple batches in parallel benefits workloads prone to
-stragglers.
+A scalar representing the number of batches to create in parallel. Processing
+multiple batches in parallel benefits workloads prone to stragglers.
+END
+  }
+  in_arg {
+    name: "drop_remainder"
+    description: <<END
+A scalar representing whether the last batch should be dropped in case its size
+is smaller than desired.
+END
+  }
+  attr {
+    name: "f"
+    description: <<END
+A function to apply to the outputs of `input_dataset`.
 END
   }
-  summary: "Creates a dataset that applies `f` to the outputs of `input_dataset` and then"
+  summary: "Creates a dataset that fuses mapping with batching."
   description: <<END
+Creates a dataset that applies `f` to the outputs of `input_dataset` and then
 batches `batch_size` of them.
 
 Unlike a "MapDataset", which applies `f` sequentially, this dataset invokes up
diff --git a/tensorflow/core/api_def/base_api/api_def_MapAndBatchDatasetV2.pbtxt b/tensorflow/core/api_def/base_api/api_def_MapAndBatchDatasetV2.pbtxt
new file mode 100644
index 0000000000..81ef92cae0
--- /dev/null
+++ b/tensorflow/core/api_def/base_api/api_def_MapAndBatchDatasetV2.pbtxt
@@ -0,0 +1,54 @@
+op {
+  graph_op_name: "MapAndBatchDatasetV2"
+  visibility: HIDDEN
+  in_arg {
+    name: "input_dataset"
+    description: <<END
+A variant tensor representing the input dataset.
+END
+  }
+  in_arg {
+    name: "other_arguments"
+    description: <<END
+A list of tensors, typically values that were captured when building a closure
+for `f`.
+END
+  }
+  in_arg {
+    name: "batch_size"
+    description: <<END
+A scalar representing the number of elements to accumulate in a
+batch. It determines the number of concurrent invocations of `f` that process
+elements from `input_dataset` in parallel.
+END
+  }
+  in_arg {
+    name: "num_parallel_calls"
+    description: <<END
+A scalar representing the maximum number of parallel invocations of the `map_fn`
+function. Applying the `map_fn` on consecutive input elements in parallel has
+the potential to improve input pipeline throughput.
+END
+  }
+  in_arg {
+    name: "drop_remainder"
+    description: <<END
+A scalar representing whether the last batch should be dropped in case its size
+is smaller than desired.
+END
+  }
+  attr {
+    name: "f"
+    description: <<END
+A function to apply to the outputs of `input_dataset`.
+END
+  }
+  summary: "Creates a dataset that fuses mapping with batching."
+  description: <<END
+Creates a dataset that applies `f` to the outputs of `input_dataset` and then
+batches `batch_size` of them.
+
+Unlike a "MapDataset", which applies `f` sequentially, this dataset invokes up
+to `batch_size * num_parallel_batches` copies of `f` in parallel.
+END
+}
diff --git a/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc b/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc
index c9551fbf16..729b615e56 100644
--- a/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc
@@ -14,6 +14,8 @@ limitations under the License.
 ==============================================================================*/
 #define EIGEN_USE_THREADS
 
+#include <utility>
+
 #include "tensorflow/core/common_runtime/function.h"
 #include "tensorflow/core/framework/partial_tensor_shape.h"
 #include "tensorflow/core/framework/tensor.h"
@@ -21,6 +23,7 @@ limitations under the License.
 #include "tensorflow/core/kernels/data/dataset.h"
 #include "tensorflow/core/kernels/inplace_ops_functor.h"
 #include "tensorflow/core/lib/core/blocking_counter.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/lib/random/random.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/tracing.h"
@@ -36,7 +39,8 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
  public:
   explicit MapAndBatchDatasetOp(OpKernelConstruction* ctx)
       : UnaryDatasetOpKernel(ctx),
-        graph_def_version_(ctx->graph_def_version()) {
+        graph_def_version_(ctx->graph_def_version()),
+        op_version_(ctx->def().op() == "MapAndBatchDataset" ? 1 : 2) {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("f", &func_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("output_types", &output_types_));
     OP_REQUIRES_OK(ctx, ctx->GetAttr("output_shapes", &output_shapes_));
@@ -59,12 +63,29 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
         ctx, batch_size > 0,
         errors::InvalidArgument("batch_size must be greater than zero."));
 
-    int64 num_parallel_batches;
-    OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "num_parallel_batches",
-                                            &num_parallel_batches));
-    OP_REQUIRES(ctx, num_parallel_batches > 0,
-                errors::InvalidArgument(
-                    "num_parallel_batches must be greater than zero."));
+    int64 num_parallel_calls;
+    switch (op_version_) {
+      case 1:
+        int64 num_parallel_batches;
+        OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "num_parallel_batches",
+                                                &num_parallel_batches));
+        num_parallel_calls = num_parallel_batches * batch_size;
+        OP_REQUIRES(ctx, num_parallel_batches > 0,
+                    errors::InvalidArgument(
+                        "num_parallel_batches must be greater than zero."));
+        break;
+      case 2:
+        OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, "num_parallel_calls",
+                                                &num_parallel_calls));
+        OP_REQUIRES(ctx, num_parallel_calls > 0,
+                    errors::InvalidArgument(
+                        "num_parallel_calls must be greater than zero."));
+        break;
+      default:
+        OP_REQUIRES(ctx, false,
+                    errors::Unimplemented("Unsupported operation version %d.",
+                                          op_version_));
+    }
 
     bool drop_remainder;
     OP_REQUIRES_OK(ctx,
@@ -74,7 +95,7 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
     OP_REQUIRES_OK(ctx, CapturedFunction::Create(
                             func_, std::move(other_arguments), &captured_func));
 
-    *output = new Dataset(ctx, input, batch_size, num_parallel_batches,
+    *output = new Dataset(ctx, input, batch_size, num_parallel_calls,
                           drop_remainder, output_types_, output_shapes_, func_,
                           std::move(captured_func), &ctx->eigen_cpu_device());
   }
@@ -83,7 +104,7 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
   class Dataset : public GraphDatasetBase {
    public:
     Dataset(OpKernelContext* ctx, const DatasetBase* input, int64 batch_size,
-            int64 num_parallel_batches, bool drop_remainder,
+            int64 num_parallel_calls, bool drop_remainder,
             const DataTypeVector& output_types,
             const std::vector<PartialTensorShape>& output_shapes,
             const NameAttrList& func,
@@ -92,7 +113,7 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
         : GraphDatasetBase(ctx),
           input_(input),
           batch_size_(batch_size),
-          num_parallel_batches_(num_parallel_batches),
+          num_parallel_calls_(num_parallel_calls),
           drop_remainder_(drop_remainder),
           output_types_(output_types),
           output_shapes_(output_shapes),
@@ -128,9 +149,9 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
       TF_RETURN_IF_ERROR(b->AddParentDataset(ctx, input_, &input_graph_node));
       Node* batch_size_node;
       TF_RETURN_IF_ERROR(b->AddScalar(batch_size_, &batch_size_node));
-      Node* num_parallel_batches_node;
+      Node* num_parallel_calls_node;
       TF_RETURN_IF_ERROR(
-          b->AddScalar(num_parallel_batches_, &num_parallel_batches_node));
+          b->AddScalar(num_parallel_calls_, &num_parallel_calls_node));
       Node* drop_remainder_node;
       TF_RETURN_IF_ERROR(b->AddScalar(drop_remainder_, &drop_remainder_node));
 
@@ -153,7 +174,7 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
           this,
           {std::make_pair(0, input_graph_node),
            std::make_pair(2, batch_size_node),
-           std::make_pair(3, num_parallel_batches_node),
+           std::make_pair(3, num_parallel_calls_node),
            std::make_pair(4, drop_remainder_node)},  // Single tensor inputs.
           {std::make_pair(1, other_arguments)},      // Tensor list inputs.
           {std::make_pair("f", f),
@@ -168,129 +189,54 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
       explicit Iterator(const Params& params)
           : DatasetIterator<Dataset>(params),
             input_impl_(params.dataset->input_->MakeIterator(params.prefix)),
-            invocation_results_(params.dataset->batch_size_ *
-                                params.dataset->num_parallel_batches_),
-            batch_results_(params.dataset->num_parallel_batches_) {}
+            batch_results_((params.dataset->num_parallel_calls_ +
+                            params.dataset->batch_size_ - 1) /
+                           params.dataset->batch_size_) {
+        for (int i = 0; i < batch_results_.size(); ++i) {
+          batch_results_[i].Initialize(params.dataset->batch_size_);
+        }
+      }
 
       ~Iterator() override {
-        // TODO(mrry): Replace this cancellation logic with a
-        // CancellationManager. The syntax would be more heavyweight,
-        // but it would be possible to thread a cancellation manager
-        // through the IteratorContext to upstream,
-        // potentially-blocking iterators, when we add these.
         mutex_lock l(mu_);
-        if (current_batch_index_ != -1) {
-          for (size_t batch_index = 0;
-               batch_index < dataset()->num_parallel_batches_; ++batch_index) {
-            int64 num_elements;
-            WaitForBatch(batch_index, &num_elements).IgnoreError();
-            // Deallocate tensors allocated for the output.
-            batch_results_[batch_index].output.clear();
-          }
+        // Cancel the runner thread.
+        cancelled_ = true;
+        cond_var_.notify_all();
+        // Wait for all in-flight calls to complete.
+        while (num_calls_ > 0) {
+          cond_var_.wait(l);
         }
       }
 
-      // TODO(jsimsa): Implement and profile the following alternative design:
-      //
-      // 0. Set the number of in-flight batches and invocations independently
-      // (though obviously the max number of in-flight invocations must be <
-      // batch_size * num_parallel_batches). Maintain a current producing batch
-      // index and offset.
-      // 1. Issue invocations in order of batch and offset, as you do currently.
-      // 2. When an invocation finishes, increment the current producing batch
-      // and offset. If that invocation would start a new batch and give more
-      // than num_parallel_batches in-flight, block; else start the new
-      // invocation into that location.
-      // 3. When a GetNext() call arrives, block until there's a full batch.
-      // Before returning the batch, if the number of pending invocations is
-      // less than the max, issue that number of invocations.
       Status GetNextInternal(IteratorContext* ctx,
                              std::vector<Tensor>* out_tensors,
                              bool* end_of_sequence) override {
         mutex_lock l(mu_);
-
-        // One-time initialization.
-        if (current_batch_index_ == -1) {
-          current_batch_index_ = 0;
-          for (size_t i = 0; i < dataset()->num_parallel_batches_; ++i) {
-            StartInvocationBatch(ctx, i);
-          }
-        }
-
-        int64 num_elements = 0;
-        Status status = WaitForBatch(current_batch_index_, &num_elements);
-        if (num_elements == 0) {
-          *end_of_sequence = true;
-          return Status::OK();
-        }
-        if (!status.ok()) {
-          // Deallocate tensors allocated for the output.
-          batch_results_[current_batch_index_].output.clear();
-        } else {
-          if (num_elements < dataset()->batch_size_) {
-            if (dataset()->drop_remainder_) {
-              // Deallocate tensors allocated for the output.
-              batch_results_[current_batch_index_].output.clear();
-              *end_of_sequence = true;
-              return Status::OK();
-            }
-            const std::vector<Tensor>& output =
-                batch_results_[current_batch_index_].output;
-            for (size_t i = 0; i < output.size(); ++i) {
-              TensorShape component_shape(
-                  batch_results_[current_batch_index_].output[i].shape());
-              component_shape.set_dim(0, num_elements);
-              AllocatorAttributes attr;
-              attr.set_gpu_compatible(true);
-              Tensor component(ctx->allocator(attr), output[i].dtype(),
-                               component_shape);
-              TF_RETURN_IF_ERROR(
-                  CopyPartialBatch(&component, output[i], num_elements));
-              out_tensors->emplace_back(std::move(component));
-            }
-            // Deallocate tensors allocated for the output.
-            batch_results_[current_batch_index_].output.clear();
-          } else {
-            *out_tensors =
-                std::move(batch_results_[current_batch_index_].output);
-          }
-          *end_of_sequence = false;
-        }
-        StartInvocationBatch(ctx, current_batch_index_);
-        current_batch_index_ =
-            (current_batch_index_ + 1) % dataset()->num_parallel_batches_;
-        return status;
+        EnsureRunnerThreadStarted(ctx);
+        BatchResult* result = &batch_results_[ComputeIndex(input_batch_)];
+        WaitForBatch(result, &l);
+        return ProcessBatch(ctx, result, out_tensors, end_of_sequence);
       }
 
      protected:
       Status SaveInternal(IteratorStateWriter* writer) override {
         mutex_lock l(mu_);
-        if (current_batch_index_ == -1) {
-          // Iterator has not been used. Nothing to save.
-          return Status::OK();
+        // Wait for all in-flight calls to complete.
+        while (num_calls_ > 0) {
+          cond_var_.wait(l);
         }
-        TF_RETURN_IF_ERROR(writer->WriteScalar(full_name("current_batch_index"),
-                                               current_batch_index_));
+        CHECK_EQ(num_calls_, 0);
         TF_RETURN_IF_ERROR(SaveParent(writer, input_impl_));
-        // Wait for the map_fn dispatches made in `InvokeFunctionLocked` to
-        // finish. This may delay saving a checkpoint by a bit but keeps the
-        // code clean and also saves us from checkpointing the state of the
-        // `BlockingCounter`.
-        std::vector<int64> num_elements(batch_results_.size());
-        for (size_t i = 0; i < batch_results_.size(); i++) {
-          WaitForBatch(i, &num_elements[i]).IgnoreError();
-        }
-
-        TF_RETURN_IF_ERROR(writer->WriteScalar(
-            full_name("invocation_results_size"), invocation_results_.size()));
-        for (size_t i = 0; i < invocation_results_.size(); ++i) {
-          TF_RETURN_IF_ERROR(WriteInvocationResultLocked(writer, i));
-        }
+        TF_RETURN_IF_ERROR(
+            writer->WriteScalar(full_name("call_counter"), call_counter_));
+        TF_RETURN_IF_ERROR(
+            writer->WriteScalar(full_name("input_batch"), input_batch_));
+        TF_RETURN_IF_ERROR(
+            writer->WriteScalar(full_name("output_batch"), output_batch_));
         TF_RETURN_IF_ERROR(writer->WriteScalar(full_name("batch_results_size"),
                                                batch_results_.size()));
         for (size_t i = 0; i < batch_results_.size(); ++i) {
-          TF_RETURN_IF_ERROR(
-              WriteBatchResultLocked(writer, i, num_elements[i]));
+          TF_RETURN_IF_ERROR(WriteBatchResult(writer, i));
         }
         return Status::OK();
       }
@@ -298,70 +244,136 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
       Status RestoreInternal(IteratorContext* ctx,
                              IteratorStateReader* reader) override {
         mutex_lock l(mu_);
-        if (!reader->Contains(full_name("current_batch_index"))) {
-          // Iterator was never used so nothing to restore.
-          return Status::OK();
-        }
-        {
-          int64 temp;
-          TF_RETURN_IF_ERROR(
-              reader->ReadScalar(full_name("current_batch_index"), &temp));
-          current_batch_index_ = static_cast<int32>(temp);
-          if (current_batch_index_ != temp) {
-            return errors::Internal("Invalid value for current_batch_index ",
-                                    temp);
-          }
-        }
         TF_RETURN_IF_ERROR(RestoreParent(ctx, reader, input_impl_));
-        size_t invocation_results_size;
-        {
-          int64 temp;
-          TF_RETURN_IF_ERROR(
-              reader->ReadScalar(full_name("invocation_results_size"), &temp));
-          invocation_results_size = static_cast<size_t>(temp);
-          if (invocation_results_size != temp) {
-            return errors::Internal(
-                "Invalid value for invocation_results_size ", temp);
-          }
-        }
-        CHECK_EQ(invocation_results_.size(), invocation_results_size);
-        for (size_t i = 0; i < invocation_results_size; ++i) {
-          TF_RETURN_IF_ERROR(ReadInvocationResultLocked(reader, i));
-        }
-        size_t batch_results_size;
-        {
-          int64 temp;
-          TF_RETURN_IF_ERROR(
-              reader->ReadScalar(full_name("batch_results_size"), &temp));
-          batch_results_size = static_cast<size_t>(temp);
-          if (batch_results_size != temp) {
-            return errors::Internal("Invalid value for batch_results_size ",
-                                    temp);
-          }
-        }
+        TF_RETURN_IF_ERROR(
+            reader->ReadScalar(full_name("call_counter"), &call_counter_));
+        TF_RETURN_IF_ERROR(
+            reader->ReadScalar(full_name("input_batch"), &input_batch_));
+        TF_RETURN_IF_ERROR(
+            reader->ReadScalar(full_name("output_batch"), &output_batch_));
+        int64 batch_results_size;
+        TF_RETURN_IF_ERROR(reader->ReadScalar(full_name("batch_results_size"),
+                                              &batch_results_size));
         CHECK_EQ(batch_results_.size(), batch_results_size);
-        for (size_t i = 0; i < batch_results_size; ++i) {
-          TF_RETURN_IF_ERROR(ReadBatchResultLocked(ctx, reader, i));
+        for (int i = 0; i < batch_results_size; ++i) {
+          TF_RETURN_IF_ERROR(ReadBatchResult(ctx, reader, i));
         }
         return Status::OK();
       }
 
      private:
       struct BatchResult {
-        mutex mu ACQUIRED_AFTER(mu_);
-        bool output_allocated GUARDED_BY(mu);
+        mutex mu;
+        bool end_of_input GUARDED_BY(mu);
+        int64 num_elements GUARDED_BY(mu);
         std::vector<Tensor> output;
-        std::unique_ptr<BlockingCounter> counter;
+        bool output_allocated GUARDED_BY(mu);
+        Status status GUARDED_BY(mu);
+        // Used for coordination between the main thread and the callback
+        // threads. In particular, the main thread will wait for the value
+        // of `num_calls` to reach zero before processing the batch result.
+        condition_variable cond_var;  // access guarded by owner's mutex
+        // Counts the number of outstanding calls for this batch.
+        int64 num_calls;  // access guarded by owner's mutex
+
+        void Initialize(int64 batch_size) {
+          mutex_lock l(mu);
+          end_of_input = false;
+          num_calls = batch_size;
+          num_elements = 0;
+          output_allocated = false;
+          status = Status::OK();
+        }
+
+        void UpdateStatus(const Status& s) {
+          mutex_lock l(mu);
+          status.Update(s);
+        }
       };
 
-      struct InvocationResult {
-        Status status;
+      void Callback(const std::shared_ptr<IteratorContext>& ctx,
+                    BatchResult* result, std::vector<Tensor>* return_values,
+                    int64 offset, const Status& status) {
+        std::unique_ptr<std::vector<Tensor>> cleanup_retvals(return_values);
+        result->UpdateStatus(status);
+        if (status.ok()) {
+          EnsureOutputAllocated(ctx, result, return_values);
+          for (size_t i = 0; i < return_values->size(); ++i) {
+            const Tensor& tensor = return_values->at(i);
+            Tensor* batch = &(result->output)[i];
+            if (tensor.NumElements() !=
+                (batch->NumElements() / batch->dim_size(0))) {
+              TensorShape batch_shape = batch->shape();
+              batch_shape.RemoveDim(0);
+              result->UpdateStatus(errors::InvalidArgument(
+                  "Cannot add tensor to the batch: number of elements does not "
+                  "match. Shapes are: [tensor]: ",
+                  tensor.shape().DebugString(),
+                  ", [batch]: ", batch_shape.DebugString()));
+              break;
+            }
+            // TODO(mrry): Add a version of DoParallelConcat that allows us to
+            // move `tensor` where possible, to speed up string tensor batching.
+            Status copy_status = ::tensorflow::functor::DoParallelConcat(
+                *dataset()->device_, tensor, offset, batch);
+            if (!copy_status.ok()) {
+              result->UpdateStatus(copy_status);
+              break;
+            }
+          }
+        }
+        {
+          mutex_lock l(result->mu);
+          result->num_elements++;
+        }
+        {
+          mutex_lock l(mu_);
+          CallCompleted(result);
+        }
+      }
+
+      void CallCompleted(BatchResult* result) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+        num_calls_--;
+        cond_var_.notify_all();
+        result->num_calls--;
+        result->cond_var.notify_all();
+      }
+
+      void CallFunction(std::shared_ptr<IteratorContext> ctx,
+                        BatchResult* result, int64 offset) {
+        // Get the next input element.
+        std::vector<Tensor> input_element;
         bool end_of_input;
-        std::vector<Tensor> return_values;
-      };
+        Status status =
+            input_impl_->GetNext(ctx.get(), &input_element, &end_of_input);
+        {
+          mutex_lock l(mu_);
+          mutex_lock l2(result->mu);
+          result->end_of_input = result->end_of_input || end_of_input;
+          result->status.Update(status);
+          if (result->end_of_input || !result->status.ok()) {
+            CallCompleted(result);
+            return;
+          }
+        }
 
-      int64 ComputeInvocationIndex(int64 batch_index, int64 offset) {
-        return batch_index * dataset()->batch_size_ + offset;
+        // Call `captured_func_(input_element)`, using `Callback` to store the
+        // result in `result`.
+        (*ctx->runner())(std::bind(
+            [this, result, offset](std::shared_ptr<IteratorContext> ctx,
+                                   std::vector<Tensor> input_element) {
+              std::vector<Tensor>* return_values = new std::vector<Tensor>();
+              dataset()->captured_func_->RunAsync(
+                  ctx.get(), std::move(input_element), return_values,
+                  [this, ctx, result, return_values, offset](Status status) {
+                    Callback(ctx, result, return_values, offset, status);
+                  });
+            },
+            ctx, std::move(input_element)));
+      }
+
+      int64 ComputeIndex(int64 n) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+        return n % batch_results_.size();
       }
 
       Status CopyPartialBatch(Tensor* output, const Tensor& value,
@@ -387,253 +399,140 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
         return Status::OK();
       }
 
-      void EnsureOutputAllocated(IteratorContext* ctx,
-                                 BatchResult* batch_result,
-                                 const std::vector<Tensor>& return_values) {
-        mutex_lock l(batch_result->mu);
-        if (batch_result->output_allocated) {
+      void EnsureRunnerThreadStarted(IteratorContext* ctx)
+          EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+        if (!runner_thread_) {
+          std::shared_ptr<IteratorContext> ctx_copy(new IteratorContext(*ctx));
+          runner_thread_.reset(ctx->env()->StartThread(
+              {}, "runner_thread",
+              std::bind(&Iterator::RunnerThread, this, ctx_copy)));
+        }
+      }
+
+      void EnsureOutputAllocated(const std::shared_ptr<IteratorContext>& ctx,
+                                 BatchResult* result,
+                                 const std::vector<Tensor>* return_values) {
+        mutex_lock l(result->mu);
+        if (result->output_allocated) {
           return;
         }
-        const size_t num_components = return_values.size();
+        const size_t num_components = return_values->size();
         for (size_t i = 0; i < num_components; ++i) {
           TensorShape component_shape({dataset()->batch_size_});
-          component_shape.AppendShape(return_values[i].shape());
+          component_shape.AppendShape(return_values->at(i).shape());
           AllocatorAttributes attr;
           attr.set_gpu_compatible(true);
-          Tensor component(ctx->allocator(attr), return_values[i].dtype(),
+          Tensor component(ctx->allocator(attr), return_values->at(i).dtype(),
                            component_shape);
-          batch_result->output.emplace_back(std::move(component));
+          result->output.emplace_back(std::move(component));
         }
-        batch_result->output_allocated = true;
+        result->output_allocated = true;
       }
 
-      void InvokeFunctionLocked(IteratorContext* ctx, int64 batch_index,
-                                int64 offset) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
-        size_t index = ComputeInvocationIndex(batch_index, offset);
-        InvocationResult* result = &invocation_results_[index];
-        BatchResult* batch_result = &batch_results_[batch_index];
-
-        // Get the next input element.
-        std::vector<Tensor> input_element;
-        result->status =
-            input_impl_->GetNext(ctx, &input_element, &result->end_of_input);
-        if (result->end_of_input || !result->status.ok()) {
-          batch_result->counter->DecrementCount();
-          return;
+      Status ProcessBatch(IteratorContext* ctx, BatchResult* result,
+                          std::vector<Tensor>* out_tensors,
+                          bool* end_of_sequence) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+        auto cleanup =
+            gtl::MakeCleanup([this, result]() EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+              result->Initialize(dataset()->batch_size_);
+              input_batch_++;
+            });
+        mutex_lock l(result->mu);
+        if (result->num_elements == 0) {
+          *end_of_sequence = true;
+          return Status::OK();
         }
 
-        // Call `captured_func_(input_element)`, store the result in
-        // `result->return_values`, and notify `batch_result->counter`
-        // to unblock a consumer.
-        (*ctx->runner())(std::bind(
-            [this, result, batch_result, offset](
-                IteratorContext* ctx, std::vector<Tensor> input_element) {
-              dataset()->captured_func_->RunAsync(
-                  ctx, std::move(input_element), &result->return_values,
-                  [this, ctx, result, batch_result, offset](Status ret_status) {
-                    result->status.Update(ret_status);
-                    if (ret_status.ok()) {
-                      EnsureOutputAllocated(ctx, batch_result,
-                                            result->return_values);
-                      const size_t num_components =
-                          result->return_values.size();
-                      for (size_t i = 0; i < num_components; ++i) {
-                        const Tensor& tensor = result->return_values[i];
-                        Tensor* batch = &(batch_result->output)[i];
-                        if (tensor.NumElements() !=
-                            (batch->NumElements() / batch->dim_size(0))) {
-                          TensorShape batch_shape = batch->shape();
-                          batch_shape.RemoveDim(0);
-                          result->status.Update(errors::InvalidArgument(
-                              "Cannot add tensor to the batch: number of "
-                              "elements does not match. Shapes are: [tensor]: ",
-                              tensor.shape().DebugString(),
-                              ", [batch]: ", batch_shape.DebugString()));
-                          break;
-                        }
-                        // TODO(mrry): Add a version of DoParallelConcat that
-                        // allows us to move `tensor` where possible, to speed
-                        // up string tensor batching.
-                        Status copy_status =
-                            ::tensorflow::functor::DoParallelConcat(
-                                *dataset()->device_, tensor, offset, batch);
-                        if (!copy_status.ok()) {
-                          result->status.Update(copy_status);
-                          break;
-                        }
-                      }
-                    }
-                    delete ctx;
-                    // NOTE(mrry): We clear the return values here to release
-                    // any memory associated with them and to paralellize the
-                    // destruction of the tensors (which can be surprisingly
-                    // expensive for map functions with large numbers of return
-                    // values).
-                    result->return_values.clear();
-                    batch_result->counter->DecrementCount();
-                  });
-            },
-            new IteratorContext(*ctx), std::move(input_element)));
-      }
-
-      void StartInvocationBatch(IteratorContext* ctx, int64 batch_index)
-          EXCLUSIVE_LOCKS_REQUIRED(mu_) {
-        tracing::ScopedActivity activity(strings::StrCat(prefix(), "::Start"));
-        // Initialize batch result.
-        {
-          mutex_lock l(batch_results_[batch_index].mu);
-          batch_results_[batch_index].output_allocated = false;
-          batch_results_[batch_index].counter.reset(
-              new BlockingCounter(dataset()->batch_size_));
-        }
-        // Initialize invocation results.
-        for (size_t i = 0; i < dataset()->batch_size_; ++i) {
-          size_t index = ComputeInvocationIndex(batch_index, i);
-          InvocationResult* result = &invocation_results_[index];
-          // Reset the state of `result`; `result->return_values` was cleared
-          // when the previous invocation completed.
-          result->end_of_input = false;
-          result->status = Status::OK();
-        }
-        // Start individual invocations.
-        for (size_t i = 0; i < dataset()->batch_size_; ++i) {
-          InvokeFunctionLocked(ctx, batch_index, i);
+        if (!result->status.ok()) {
+          // Deallocate tensors allocated for the output.
+          result->output.clear();
+        } else {
+          if (result->num_elements < dataset()->batch_size_) {
+            if (dataset()->drop_remainder_) {
+              // Deallocate tensors allocated for the output.
+              result->output.clear();
+              *end_of_sequence = true;
+              return Status::OK();
+            }
+            const std::vector<Tensor>& output = result->output;
+            for (size_t i = 0; i < output.size(); ++i) {
+              TensorShape component_shape(result->output[i].shape());
+              component_shape.set_dim(0, result->num_elements);
+              AllocatorAttributes attr;
+              attr.set_gpu_compatible(true);
+              Tensor component(ctx->allocator(attr), output[i].dtype(),
+                               component_shape);
+              TF_RETURN_IF_ERROR(CopyPartialBatch(&component, output[i],
+                                                  result->num_elements));
+              out_tensors->emplace_back(std::move(component));
+            }
+            // Deallocate tensors allocated for the output.
+            result->output.clear();
+          } else {
+            *out_tensors = std::move(result->output);
+          }
+          *end_of_sequence = false;
         }
+        cond_var_.notify_all();
+        return result->status;
       }
 
-      Status WaitForBatch(int64 batch_index, int64* num_elements)
-          EXCLUSIVE_LOCKS_REQUIRED(mu_) {
-        tracing::ScopedActivity activity(strings::StrCat(prefix(), "::Wait"));
-        batch_results_[batch_index].counter->Wait();
-        Status status = Status::OK();
-        for (size_t i = 0; i < dataset()->batch_size_; ++i, ++*num_elements) {
-          size_t index = ComputeInvocationIndex(batch_index, i);
-          InvocationResult* result = &invocation_results_[index];
-          if (result->end_of_input) {
-            VLOG(3) << "end of input encountered at element[" << i << "]: ";
-            return Status::OK();
-          }
-          if (!result->status.ok()) {
-            VLOG(3) << "failed to process element[" << i
-                    << "]: " << result->status;
-            status.Update(result->status);
+      void RunnerThread(const std::shared_ptr<IteratorContext>& ctx) {
+        mutex_lock l(mu_);
+        while (true) {
+          while (!cancelled_ &&
+                 (num_calls_ == dataset()->num_parallel_calls_ ||
+                  (output_batch_ - input_batch_ == batch_results_.size()))) {
+            cond_var_.wait(l);
           }
-        }
-        return status;
-      }
 
-      Status WriteInvocationResultLocked(IteratorStateWriter* writer,
-                                         size_t index)
-          EXCLUSIVE_LOCKS_REQUIRED(mu_) {
-        const InvocationResult& result = invocation_results_[index];
-        string prefix = strings::StrCat("invocation_results_", index);
-        TF_RETURN_IF_ERROR(WriteStatusLocked(
-            writer, full_name(strings::StrCat(prefix, "_status")),
-            result.status));
-        if (result.end_of_input) {
-          TF_RETURN_IF_ERROR(writer->WriteScalar(
-              full_name(strings::StrCat(prefix, "_end_of_input")), ""));
-        }
-        TF_RETURN_IF_ERROR(writer->WriteScalar(
-            full_name(strings::StrCat(prefix, "_return_values_size")),
-            result.return_values.size()));
-        for (size_t i = 0; i < result.return_values.size(); i++) {
-          TF_RETURN_IF_ERROR(writer->WriteTensor(
-              full_name(strings::StrCat(prefix, "_return_values_", i)),
-              result.return_values[i]));
-        }
-        return Status::OK();
-      }
+          if (cancelled_) {
+            return;
+          }
 
-      Status ReadInvocationResultLocked(IteratorStateReader* reader,
-                                        size_t index)
-          EXCLUSIVE_LOCKS_REQUIRED(mu_) {
-        InvocationResult* result = &invocation_results_[index];
-        string prefix = strings::StrCat("invocation_results_", index);
-        TF_RETURN_IF_ERROR(ReadStatusLocked(
-            reader, full_name(strings::StrCat(prefix, "_status")),
-            &result->status));
-        result->end_of_input = reader->Contains(
-            full_name(strings::StrCat(prefix, "_end_of_input")));
-        size_t return_values_size;
-        {
-          int64 temp;
-          TF_RETURN_IF_ERROR(reader->ReadScalar(
-              full_name(strings::StrCat(prefix, "_return_values_size")),
-              &temp));
-          return_values_size = static_cast<size_t>(temp);
-          if (temp != return_values_size) {
-            return errors::Internal("Invalid value for return_values_size ",
-                                    return_values_size);
+          while (num_calls_ < dataset()->num_parallel_calls_ &&
+                 (output_batch_ - input_batch_ < batch_results_.size())) {
+            BatchResult* result = &batch_results_[ComputeIndex(output_batch_)];
+            int64 offset = call_counter_++ % dataset()->batch_size_;
+            num_calls_++;
+            mu_.unlock();
+            CallFunction(ctx, result, offset);
+            mu_.lock();
+            if (offset + 1 == dataset()->batch_size_) {
+              // Done scheduling calls for the current batch.
+              output_batch_++;
+            }
           }
         }
-        result->return_values.reserve(return_values_size);
-        for (size_t i = 0; i < return_values_size; i++) {
-          result->return_values.emplace_back();
-          TF_RETURN_IF_ERROR(reader->ReadTensor(
-              full_name(strings::StrCat(prefix, "_return_values_", i)),
-              &result->return_values.back()));
-        }
-        return Status::OK();
       }
 
-      Status WriteBatchResultLocked(IteratorStateWriter* writer, size_t index,
-                                    int64 num_elements)
+      void WaitForBatch(BatchResult* result, mutex_lock* l)
           EXCLUSIVE_LOCKS_REQUIRED(mu_) {
-        const BatchResult& result = batch_results_[index];
-        string prefix = strings::StrCat("batch_results_", index);
-        {
-          mutex_lock l(batch_results_[index].mu);
-          if (result.output_allocated) {
-            TF_RETURN_IF_ERROR(writer->WriteScalar(
-                full_name(strings::StrCat(prefix, "_output_allocated")), ""));
-          }
+        while (result->num_calls > 0) {
+          result->cond_var.wait(*l);
         }
-        TF_RETURN_IF_ERROR(writer->WriteScalar(
-            full_name(strings::StrCat(prefix, "_output_size")),
-            result.output.size()));
-        for (size_t i = 0; i < result.output.size(); i++) {
-          // If the batch is not full, we only store the first
-          // `num_elements` values. The rest of the batch tensor is
-          // *uninitialized* and accessing that will raise msan errors.
-          if (num_elements < dataset()->batch_size_) {
-            TF_RETURN_IF_ERROR(writer->WriteTensor(
-                full_name(strings::StrCat(prefix, "_output_", i)),
-                result.output[i].Slice(0, num_elements)));
-          } else {
-            TF_RETURN_IF_ERROR(writer->WriteTensor(
-                full_name(strings::StrCat(prefix, "_output_", i)),
-                result.output[i]));
-          }
-        }
-        return Status::OK();
       }
 
-      Status ReadBatchResultLocked(IteratorContext* ctx,
-                                   IteratorStateReader* reader, size_t index)
-          EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+      Status ReadBatchResult(IteratorContext* ctx, IteratorStateReader* reader,
+                             size_t index) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
         BatchResult* result = &batch_results_[index];
         string prefix = strings::StrCat("batch_results_", index);
-        {
-          mutex_lock l(batch_results_[index].mu);
-          result->output_allocated = reader->Contains(
-              full_name(strings::StrCat(prefix, "_output_allocated")));
-          // Simulate that the batch was fully generated.
-          batch_results_[index].counter.reset(new BlockingCounter(0));
-        }
-        size_t output_size;
-        {
-          int64 temp;
-          TF_RETURN_IF_ERROR(reader->ReadScalar(
-              full_name(strings::StrCat(prefix, "_output_size")), &temp));
-          output_size = static_cast<size_t>(temp);
-          if (temp != output_size) {
-            return errors::Internal("Invalid value for output_size ",
-                                    output_size);
-          }
-        }
+        mutex_lock l(result->mu);
+        result->end_of_input = reader->Contains(
+            full_name(strings::StrCat(prefix, "_end_of_input")));
+        TF_RETURN_IF_ERROR(
+            reader->ReadScalar(full_name(strings::StrCat(prefix, "_num_calls")),
+                               &result->num_calls));
+        TF_RETURN_IF_ERROR(reader->ReadScalar(
+            full_name(strings::StrCat(prefix, "_num_elements")),
+            &result->num_elements));
+        result->output_allocated = reader->Contains(
+            full_name(strings::StrCat(prefix, "_output_allocated")));
+        int64 output_size;
+        TF_RETURN_IF_ERROR(reader->ReadScalar(
+            full_name(strings::StrCat(prefix, "_output_size")), &output_size));
         result->output.reserve(output_size);
-        for (size_t i = 0; i < output_size; i++) {
+        for (int i = 0; i < output_size; i++) {
           Tensor t;
           TF_RETURN_IF_ERROR(reader->ReadTensor(
               full_name(strings::StrCat(prefix, "_output_", i)), &t));
@@ -653,25 +552,13 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
             result->output.emplace_back(std::move(t));
           }
         }
+        TF_RETURN_IF_ERROR(ReadStatus(
+            reader, strings::StrCat(prefix, "_status"), &result->status));
         return Status::OK();
       }
 
-      Status WriteStatusLocked(IteratorStateWriter* writer,
-                               const string& prefix, const Status& status)
-          EXCLUSIVE_LOCKS_REQUIRED(mu_) {
-        TF_RETURN_IF_ERROR(
-            writer->WriteScalar(full_name(strings::StrCat(prefix, "_code")),
-                                static_cast<int64>(status.code())));
-        if (!status.ok()) {
-          TF_RETURN_IF_ERROR(
-              writer->WriteScalar(full_name(strings::StrCat(prefix, "_msg")),
-                                  status.error_message()));
-        }
-        return Status::OK();
-      }
-
-      Status ReadStatusLocked(IteratorStateReader* reader, const string& prefix,
-                              Status* status) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+      Status ReadStatus(IteratorStateReader* reader, const string& prefix,
+                        Status* status) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
         int64 code_int;
         TF_RETURN_IF_ERROR(reader->ReadScalar(
             full_name(strings::StrCat(prefix, "_code")), &code_int));
@@ -687,17 +574,89 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
         }
         return Status::OK();
       }
+
+      Status WriteBatchResult(IteratorStateWriter* writer, size_t index)
+          EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+        BatchResult* result = &batch_results_[index];
+        string prefix = strings::StrCat("batch_results_", index);
+        mutex_lock l(result->mu);
+        if (result->end_of_input) {
+          TF_RETURN_IF_ERROR(writer->WriteScalar(
+              full_name(strings::StrCat(prefix, "_end_of_input")), ""));
+        }
+        TF_RETURN_IF_ERROR(writer->WriteScalar(
+            full_name(strings::StrCat(prefix, "_num_calls")),
+            result->num_calls));
+        TF_RETURN_IF_ERROR(writer->WriteScalar(
+            full_name(strings::StrCat(prefix, "_num_elements")),
+            result->num_elements));
+        if (result->output_allocated) {
+          TF_RETURN_IF_ERROR(writer->WriteScalar(
+              full_name(strings::StrCat(prefix, "_output_allocated")), ""));
+        }
+        TF_RETURN_IF_ERROR(writer->WriteScalar(
+            full_name(strings::StrCat(prefix, "_output_size")),
+            result->output.size()));
+        for (int i = 0; i < result->output.size(); i++) {
+          // If the batch is not full, we only store the first `num_elements`
+          // values. The rest of the batch tensor is *uninitialized* and
+          // accessing that will raise msan errors.
+          if (result->num_elements < dataset()->batch_size_) {
+            TF_RETURN_IF_ERROR(writer->WriteTensor(
+                full_name(strings::StrCat(prefix, "_output_", i)),
+                result->output[i].Slice(0, result->num_elements)));
+          } else {
+            TF_RETURN_IF_ERROR(writer->WriteTensor(
+                full_name(strings::StrCat(prefix, "_output_", i)),
+                result->output[i]));
+          }
+        }
+        TF_RETURN_IF_ERROR(WriteStatus(
+            writer, strings::StrCat(prefix, "_status"), result->status));
+        return Status::OK();
+      }
+
+      Status WriteStatus(IteratorStateWriter* writer, const string& prefix,
+                         const Status& status) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+        TF_RETURN_IF_ERROR(
+            writer->WriteScalar(full_name(strings::StrCat(prefix, "_code")),
+                                static_cast<int64>(status.code())));
+        if (!status.ok()) {
+          TF_RETURN_IF_ERROR(
+              writer->WriteScalar(full_name(strings::StrCat(prefix, "_msg")),
+                                  status.error_message()));
+        }
+        return Status::OK();
+      }
+
       mutex mu_;
-      int32 current_batch_index_ GUARDED_BY(mu_) = -1;
-      const std::unique_ptr<IteratorBase> input_impl_ GUARDED_BY(mu_);
-      std::vector<InvocationResult> invocation_results_ GUARDED_BY(mu_);
+      // Used for coordination between the main thread, the runner thread, and
+      // the callback threads. In particular, the runner thread should only
+      // schedule new calls when the number of in-flight calls is less than the
+      // user specified level of parallelism and there are slots available in
+      // the `batch_results_` buffer.
+      condition_variable cond_var_;
+      // Counts the number of outstanding calls for this batch.
+      int64 num_calls_ GUARDED_BY(mu_) = 0;
+      // Counts the total number of calls.
+      int64 call_counter_ GUARDED_BY(mu_) = 0;
+      const std::unique_ptr<IteratorBase> input_impl_;
+      // Identifies the next batch to be read by the caller.
+      int64 input_batch_ GUARDED_BY(mu_) = 0;
+      // Identifies the next batch to create.
+      int64 output_batch_ GUARDED_BY(mu_) = 0;
+      // Circular buffer for storing the (intermediate) batch results. When
+      // using `input_batch_` and `output_batch_` to index into the buffer,
+      // their value should be interpreted modulo the size of the buffer.
       std::vector<BatchResult> batch_results_ GUARDED_BY(mu_);
+      std::unique_ptr<Thread> runner_thread_ GUARDED_BY(mu_);
+      bool cancelled_ GUARDED_BY(mu_) = false;
     };
 
     const DatasetBase* const input_;
     const NameAttrList func_;
     const int64 batch_size_;
-    const int64 num_parallel_batches_;
+    const int64 num_parallel_calls_;
     const bool drop_remainder_;
     const DataTypeVector output_types_;
     const std::vector<PartialTensorShape> output_shapes_;
@@ -707,6 +666,7 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
   };
 
   const int graph_def_version_;
+  const int op_version_;
   DataTypeVector output_types_;
   std::vector<PartialTensorShape> output_shapes_;
   NameAttrList func_;
@@ -715,6 +675,9 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel {
 REGISTER_KERNEL_BUILDER(Name("MapAndBatchDataset").Device(DEVICE_CPU),
                         MapAndBatchDatasetOp);
 
+REGISTER_KERNEL_BUILDER(Name("MapAndBatchDatasetV2").Device(DEVICE_CPU),
+                        MapAndBatchDatasetOp);
+
 }  // namespace
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/ops/dataset_ops.cc b/tensorflow/core/ops/dataset_ops.cc
index 73174c184c..576946eddd 100644
--- a/tensorflow/core/ops/dataset_ops.cc
+++ b/tensorflow/core/ops/dataset_ops.cc
@@ -208,6 +208,19 @@ REGISTER_OP("MapAndBatchDataset")
     .Attr("output_shapes: list(shape) >= 1")
     .SetShapeFn(shape_inference::ScalarShape);
 
+REGISTER_OP("MapAndBatchDatasetV2")
+    .Input("input_dataset: variant")
+    .Input("other_arguments: Targuments")
+    .Input("batch_size: int64")
+    .Input("num_parallel_calls: int64")
+    .Input("drop_remainder: bool")
+    .Output("handle: variant")
+    .Attr("f: func")
+    .Attr("Targuments: list(type) >= 0")
+    .Attr("output_types: list(type) >= 1")
+    .Attr("output_shapes: list(shape) >= 1")
+    .SetShapeFn(shape_inference::ScalarShape);
+
 REGISTER_OP("PrefetchDataset")
     .Input("input_dataset: variant")
     .Input("buffer_size: int64")
-- 
GitLab


From dedcf1dc3f59416fb2eaebb4fc2e5470fc2ae647 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 17:18:35 -0700
Subject: [PATCH 275/395] Update ops-related pbtxt files.

PiperOrigin-RevId: 195497084
---
 .../core/ops/compat/ops_history.v1.pbtxt      | 48 +++++++++++++++++++
 tensorflow/core/ops/ops.pbtxt                 | 48 +++++++++++++++++++
 2 files changed, 96 insertions(+)

diff --git a/tensorflow/core/ops/compat/ops_history.v1.pbtxt b/tensorflow/core/ops/compat/ops_history.v1.pbtxt
index 3db00d8180..6880ceb505 100644
--- a/tensorflow/core/ops/compat/ops_history.v1.pbtxt
+++ b/tensorflow/core/ops/compat/ops_history.v1.pbtxt
@@ -28129,6 +28129,54 @@ op {
     minimum: 1
   }
 }
+op {
+  name: "MapAndBatchDatasetV2"
+  input_arg {
+    name: "input_dataset"
+    type: DT_VARIANT
+  }
+  input_arg {
+    name: "other_arguments"
+    type_list_attr: "Targuments"
+  }
+  input_arg {
+    name: "batch_size"
+    type: DT_INT64
+  }
+  input_arg {
+    name: "num_parallel_calls"
+    type: DT_INT64
+  }
+  input_arg {
+    name: "drop_remainder"
+    type: DT_BOOL
+  }
+  output_arg {
+    name: "handle"
+    type: DT_VARIANT
+  }
+  attr {
+    name: "f"
+    type: "func"
+  }
+  attr {
+    name: "Targuments"
+    type: "list(type)"
+    has_minimum: true
+  }
+  attr {
+    name: "output_types"
+    type: "list(type)"
+    has_minimum: true
+    minimum: 1
+  }
+  attr {
+    name: "output_shapes"
+    type: "list(shape)"
+    has_minimum: true
+    minimum: 1
+  }
+}
 op {
   name: "MapClear"
   attr {
diff --git a/tensorflow/core/ops/ops.pbtxt b/tensorflow/core/ops/ops.pbtxt
index 7156440b46..d741598b19 100644
--- a/tensorflow/core/ops/ops.pbtxt
+++ b/tensorflow/core/ops/ops.pbtxt
@@ -13939,6 +13939,54 @@ op {
     minimum: 1
   }
 }
+op {
+  name: "MapAndBatchDatasetV2"
+  input_arg {
+    name: "input_dataset"
+    type: DT_VARIANT
+  }
+  input_arg {
+    name: "other_arguments"
+    type_list_attr: "Targuments"
+  }
+  input_arg {
+    name: "batch_size"
+    type: DT_INT64
+  }
+  input_arg {
+    name: "num_parallel_calls"
+    type: DT_INT64
+  }
+  input_arg {
+    name: "drop_remainder"
+    type: DT_BOOL
+  }
+  output_arg {
+    name: "handle"
+    type: DT_VARIANT
+  }
+  attr {
+    name: "f"
+    type: "func"
+  }
+  attr {
+    name: "Targuments"
+    type: "list(type)"
+    has_minimum: true
+  }
+  attr {
+    name: "output_types"
+    type: "list(type)"
+    has_minimum: true
+    minimum: 1
+  }
+  attr {
+    name: "output_shapes"
+    type: "list(shape)"
+    has_minimum: true
+    minimum: 1
+  }
+}
 op {
   name: "MapClear"
   attr {
-- 
GitLab


From b81f2dcaaa2559880c7705fb0ea91d0d825fd59c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 17:46:13 -0700
Subject: [PATCH 276/395] Go: Update generated wrapper functions for TensorFlow
 ops. PiperOrigin-RevId: 195499636

---
 tensorflow/go/op/wrappers.go | 262 +++++++++++++++++------------------
 1 file changed, 131 insertions(+), 131 deletions(-)

diff --git a/tensorflow/go/op/wrappers.go b/tensorflow/go/op/wrappers.go
index c12ea51563..eed6dac071 100644
--- a/tensorflow/go/op/wrappers.go
+++ b/tensorflow/go/op/wrappers.go
@@ -2544,6 +2544,72 @@ func EditDistance(scope *Scope, hypothesis_indices tf.Output, hypothesis_values
 	return op.Output(0)
 }
 
+// Reverses specific dimensions of a tensor.
+//
+// Given a `tensor`, and a `bool` tensor `dims` representing the dimensions
+// of `tensor`, this operation reverses each dimension i of `tensor` where
+// `dims[i]` is `True`.
+//
+// `tensor` can have up to 8 dimensions. The number of dimensions
+// of `tensor` must equal the number of elements in `dims`. In other words:
+//
+// `rank(tensor) = size(dims)`
+//
+// For example:
+//
+// ```
+// # tensor 't' is [[[[ 0,  1,  2,  3],
+// #                  [ 4,  5,  6,  7],
+// #                  [ 8,  9, 10, 11]],
+// #                 [[12, 13, 14, 15],
+// #                  [16, 17, 18, 19],
+// #                  [20, 21, 22, 23]]]]
+// # tensor 't' shape is [1, 2, 3, 4]
+//
+// # 'dims' is [False, False, False, True]
+// reverse(t, dims) ==> [[[[ 3,  2,  1,  0],
+//                         [ 7,  6,  5,  4],
+//                         [ 11, 10, 9, 8]],
+//                        [[15, 14, 13, 12],
+//                         [19, 18, 17, 16],
+//                         [23, 22, 21, 20]]]]
+//
+// # 'dims' is [False, True, False, False]
+// reverse(t, dims) ==> [[[[12, 13, 14, 15],
+//                         [16, 17, 18, 19],
+//                         [20, 21, 22, 23]
+//                        [[ 0,  1,  2,  3],
+//                         [ 4,  5,  6,  7],
+//                         [ 8,  9, 10, 11]]]]
+//
+// # 'dims' is [False, False, True, False]
+// reverse(t, dims) ==> [[[[8, 9, 10, 11],
+//                         [4, 5, 6, 7],
+//                         [0, 1, 2, 3]]
+//                        [[20, 21, 22, 23],
+//                         [16, 17, 18, 19],
+//                         [12, 13, 14, 15]]]]
+// ```
+//
+// Arguments:
+//	tensor: Up to 8-D.
+//	dims: 1-D. The dimensions to reverse.
+//
+// Returns The same shape as `tensor`.
+func Reverse(scope *Scope, tensor tf.Output, dims tf.Output) (output tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	opspec := tf.OpSpec{
+		Type: "Reverse",
+		Input: []tf.Input{
+			tensor, dims,
+		},
+	}
+	op := scope.AddOperation(opspec)
+	return op.Output(0)
+}
+
 // Clips tensor values to a specified min and max.
 //
 // Given a tensor `t`, this operation returns a tensor of the same type and
@@ -2796,71 +2862,6 @@ func Asin(scope *Scope, x tf.Output) (y tf.Output) {
 	return op.Output(0)
 }
 
-// SparseToDenseAttr is an optional argument to SparseToDense.
-type SparseToDenseAttr func(optionalAttr)
-
-// SparseToDenseValidateIndices sets the optional validate_indices attribute to value.
-//
-// value: If true, indices are checked to make sure they are sorted in
-// lexicographic order and that there are no repeats.
-// If not specified, defaults to true
-func SparseToDenseValidateIndices(value bool) SparseToDenseAttr {
-	return func(m optionalAttr) {
-		m["validate_indices"] = value
-	}
-}
-
-// Converts a sparse representation into a dense tensor.
-//
-// Builds an array `dense` with shape `output_shape` such that
-//
-// ```
-// # If sparse_indices is scalar
-// dense[i] = (i == sparse_indices ? sparse_values : default_value)
-//
-// # If sparse_indices is a vector, then for each i
-// dense[sparse_indices[i]] = sparse_values[i]
-//
-// # If sparse_indices is an n by d matrix, then for each i in [0, n)
-// dense[sparse_indices[i][0], ..., sparse_indices[i][d-1]] = sparse_values[i]
-// ```
-//
-// All other values in `dense` are set to `default_value`.  If `sparse_values` is a
-// scalar, all sparse indices are set to this single value.
-//
-// Indices should be sorted in lexicographic order, and indices must not
-// contain any repeats. If `validate_indices` is true, these properties
-// are checked during execution.
-//
-// Arguments:
-//	sparse_indices: 0-D, 1-D, or 2-D.  `sparse_indices[i]` contains the complete
-// index where `sparse_values[i]` will be placed.
-//	output_shape: 1-D.  Shape of the dense output tensor.
-//	sparse_values: 1-D.  Values corresponding to each row of `sparse_indices`,
-// or a scalar value to be used for all sparse indices.
-//	default_value: Scalar value to set for indices not specified in
-// `sparse_indices`.
-//
-// Returns Dense output tensor of shape `output_shape`.
-func SparseToDense(scope *Scope, sparse_indices tf.Output, output_shape tf.Output, sparse_values tf.Output, default_value tf.Output, optional ...SparseToDenseAttr) (dense tf.Output) {
-	if scope.Err() != nil {
-		return
-	}
-	attrs := map[string]interface{}{}
-	for _, a := range optional {
-		a(attrs)
-	}
-	opspec := tf.OpSpec{
-		Type: "SparseToDense",
-		Input: []tf.Input{
-			sparse_indices, output_shape, sparse_values, default_value,
-		},
-		Attrs: attrs,
-	}
-	op := scope.AddOperation(opspec)
-	return op.Output(0)
-}
-
 // Computes the sum along sparse segments of a tensor.
 //
 // Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of
@@ -6469,72 +6470,6 @@ func SparseFillEmptyRows(scope *Scope, indices tf.Output, values tf.Output, dens
 	return op.Output(0), op.Output(1), op.Output(2), op.Output(3)
 }
 
-// Reverses specific dimensions of a tensor.
-//
-// Given a `tensor`, and a `bool` tensor `dims` representing the dimensions
-// of `tensor`, this operation reverses each dimension i of `tensor` where
-// `dims[i]` is `True`.
-//
-// `tensor` can have up to 8 dimensions. The number of dimensions
-// of `tensor` must equal the number of elements in `dims`. In other words:
-//
-// `rank(tensor) = size(dims)`
-//
-// For example:
-//
-// ```
-// # tensor 't' is [[[[ 0,  1,  2,  3],
-// #                  [ 4,  5,  6,  7],
-// #                  [ 8,  9, 10, 11]],
-// #                 [[12, 13, 14, 15],
-// #                  [16, 17, 18, 19],
-// #                  [20, 21, 22, 23]]]]
-// # tensor 't' shape is [1, 2, 3, 4]
-//
-// # 'dims' is [False, False, False, True]
-// reverse(t, dims) ==> [[[[ 3,  2,  1,  0],
-//                         [ 7,  6,  5,  4],
-//                         [ 11, 10, 9, 8]],
-//                        [[15, 14, 13, 12],
-//                         [19, 18, 17, 16],
-//                         [23, 22, 21, 20]]]]
-//
-// # 'dims' is [False, True, False, False]
-// reverse(t, dims) ==> [[[[12, 13, 14, 15],
-//                         [16, 17, 18, 19],
-//                         [20, 21, 22, 23]
-//                        [[ 0,  1,  2,  3],
-//                         [ 4,  5,  6,  7],
-//                         [ 8,  9, 10, 11]]]]
-//
-// # 'dims' is [False, False, True, False]
-// reverse(t, dims) ==> [[[[8, 9, 10, 11],
-//                         [4, 5, 6, 7],
-//                         [0, 1, 2, 3]]
-//                        [[20, 21, 22, 23],
-//                         [16, 17, 18, 19],
-//                         [12, 13, 14, 15]]]]
-// ```
-//
-// Arguments:
-//	tensor: Up to 8-D.
-//	dims: 1-D. The dimensions to reverse.
-//
-// Returns The same shape as `tensor`.
-func Reverse(scope *Scope, tensor tf.Output, dims tf.Output) (output tf.Output) {
-	if scope.Err() != nil {
-		return
-	}
-	opspec := tf.OpSpec{
-		Type: "Reverse",
-		Input: []tf.Input{
-			tensor, dims,
-		},
-	}
-	op := scope.AddOperation(opspec)
-	return op.Output(0)
-}
-
 // BiasAddGradAttr is an optional argument to BiasAddGrad.
 type BiasAddGradAttr func(optionalAttr)
 
@@ -24884,6 +24819,71 @@ func DecodeJSONExample(scope *Scope, json_examples tf.Output) (binary_examples t
 	return op.Output(0)
 }
 
+// SparseToDenseAttr is an optional argument to SparseToDense.
+type SparseToDenseAttr func(optionalAttr)
+
+// SparseToDenseValidateIndices sets the optional validate_indices attribute to value.
+//
+// value: If true, indices are checked to make sure they are sorted in
+// lexicographic order and that there are no repeats.
+// If not specified, defaults to true
+func SparseToDenseValidateIndices(value bool) SparseToDenseAttr {
+	return func(m optionalAttr) {
+		m["validate_indices"] = value
+	}
+}
+
+// Converts a sparse representation into a dense tensor.
+//
+// Builds an array `dense` with shape `output_shape` such that
+//
+// ```
+// # If sparse_indices is scalar
+// dense[i] = (i == sparse_indices ? sparse_values : default_value)
+//
+// # If sparse_indices is a vector, then for each i
+// dense[sparse_indices[i]] = sparse_values[i]
+//
+// # If sparse_indices is an n by d matrix, then for each i in [0, n)
+// dense[sparse_indices[i][0], ..., sparse_indices[i][d-1]] = sparse_values[i]
+// ```
+//
+// All other values in `dense` are set to `default_value`.  If `sparse_values` is a
+// scalar, all sparse indices are set to this single value.
+//
+// Indices should be sorted in lexicographic order, and indices must not
+// contain any repeats. If `validate_indices` is true, these properties
+// are checked during execution.
+//
+// Arguments:
+//	sparse_indices: 0-D, 1-D, or 2-D.  `sparse_indices[i]` contains the complete
+// index where `sparse_values[i]` will be placed.
+//	output_shape: 1-D.  Shape of the dense output tensor.
+//	sparse_values: 1-D.  Values corresponding to each row of `sparse_indices`,
+// or a scalar value to be used for all sparse indices.
+//	default_value: Scalar value to set for indices not specified in
+// `sparse_indices`.
+//
+// Returns Dense output tensor of shape `output_shape`.
+func SparseToDense(scope *Scope, sparse_indices tf.Output, output_shape tf.Output, sparse_values tf.Output, default_value tf.Output, optional ...SparseToDenseAttr) (dense tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	attrs := map[string]interface{}{}
+	for _, a := range optional {
+		a(attrs)
+	}
+	opspec := tf.OpSpec{
+		Type: "SparseToDense",
+		Input: []tf.Input{
+			sparse_indices, output_shape, sparse_values, default_value,
+		},
+		Attrs: attrs,
+	}
+	op := scope.AddOperation(opspec)
+	return op.Output(0)
+}
+
 // Computes the grayscale dilation of 4-D `input` and 3-D `filter` tensors.
 //
 // The `input` tensor has shape `[batch, in_height, in_width, depth]` and the
-- 
GitLab


From ad2e7f0503f3c9078124ed74ff5249c35cf7f267 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 18:07:41 -0700
Subject: [PATCH 277/395] Internal Change

PiperOrigin-RevId: 195501342
---
 tensorflow/contrib/distributions/BUILD | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/contrib/distributions/BUILD b/tensorflow/contrib/distributions/BUILD
index fad613155d..47f2ebca77 100644
--- a/tensorflow/contrib/distributions/BUILD
+++ b/tensorflow/contrib/distributions/BUILD
@@ -578,7 +578,7 @@ cuda_py_test(
 
 cuda_py_test(
     name = "wishart_test",
-    size = "small",
+    size = "medium",
     srcs = ["python/kernel_tests/wishart_test.py"],
     additional_deps = [
         ":distributions_py",
@@ -866,7 +866,7 @@ cuda_py_test(
 
 cuda_py_test(
     name = "batch_normalization_test",
-    size = "small",
+    size = "medium",
     srcs = ["python/kernel_tests/bijectors/batch_normalization_test.py"],
     additional_deps = [
         ":bijectors_py",
-- 
GitLab


From 59f0618cedb8fd6ea989471f64ec66619f26b83e Mon Sep 17 00:00:00 2001
From: Max Galkin <maxgalkin@google.com>
Date: Fri, 4 May 2018 18:17:45 -0700
Subject: [PATCH 278/395] GuaranteeConst is a NoOp for the
 op_level_cost_estiamtor.

PiperOrigin-RevId: 195501990
---
 tensorflow/core/grappler/costs/op_level_cost_estimator.cc | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/core/grappler/costs/op_level_cost_estimator.cc b/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
index b35873ce38..199b69452f 100644
--- a/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
+++ b/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
@@ -27,6 +27,7 @@ namespace grappler {
 
 constexpr int kOpsPerMac = 2;
 constexpr char kConst[] = "Const";
+constexpr char kGuaranteeConst[] = "GuaranteeConst";
 constexpr char kConv2d[] = "Conv2D";
 constexpr char kConv2dBackpropFilter[] = "Conv2DBackpropFilter";
 constexpr char kConv2dBackpropInput[] = "Conv2DBackpropInput";
@@ -205,6 +206,7 @@ OpLevelCostEstimator::OpLevelCostEstimator() {
       {kBatchMatMul, wrap(&OpLevelCostEstimator::PredictBatchMatMul)},
 
       {kNoOp, wrap(&OpLevelCostEstimator::PredictNoOp)},
+      {kGuaranteeConst, wrap(&OpLevelCostEstimator::PredictNoOp)},
 
       {kGather, wrap(&OpLevelCostEstimator::PredictGatherOrSlice)},
       {kGatherV2, wrap(&OpLevelCostEstimator::PredictGatherOrSlice)},
-- 
GitLab


From dd5ef1b9fc22b37e5eec87d659a3af064ca54b8b Mon Sep 17 00:00:00 2001
From: Allen Lavoie <allenl@google.com>
Date: Fri, 4 May 2018 18:25:18 -0700
Subject: [PATCH 279/395] Checkpointable: A small utility for exempting objects
 from __setattr__ tracking

Exposes it as tf.contrib.checkpoint.NoDependency. Objects wrapped in a
NoDependency object get unwrapped in __setattr__ and not tracked.

Removes the _save_counter dependency from tf.train.Checkpoint (the save counter
is still tracked as "save_counter" and always has been, so this is a
backwards-compatible dependency removal).

PiperOrigin-RevId: 195502562
---
 tensorflow/contrib/checkpoint/__init__.py     |  2 ++
 .../keras/_impl/keras/engine/network.py       |  6 +++-
 .../_impl/keras/model_subclassing_test.py     | 18 +++++++++++
 tensorflow/python/training/checkpointable.py  | 30 ++++++++++++++++++-
 .../python/training/checkpointable_test.py    | 10 +++++++
 .../python/training/checkpointable_utils.py   |  7 +++--
 6 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/tensorflow/contrib/checkpoint/__init__.py b/tensorflow/contrib/checkpoint/__init__.py
index d2c30f1215..e529b25b3c 100644
--- a/tensorflow/contrib/checkpoint/__init__.py
+++ b/tensorflow/contrib/checkpoint/__init__.py
@@ -19,6 +19,7 @@ For creating and managing dependencies:
 @@CheckpointableObjectGraph
 @@dot_graph_from_checkpoint
 @@object_metadata
+@@NoDependency
 @@split_dependency
 """
 
@@ -29,6 +30,7 @@ from __future__ import print_function
 from tensorflow.contrib.checkpoint.python.split_dependency import split_dependency
 from tensorflow.contrib.checkpoint.python.visualize import dot_graph_from_checkpoint
 from tensorflow.core.protobuf.checkpointable_object_graph_pb2 import CheckpointableObjectGraph
+from tensorflow.python.training.checkpointable import NoDependency
 from tensorflow.python.training.checkpointable_utils import object_metadata
 
 from tensorflow.python.util.all_util import remove_undocumented
diff --git a/tensorflow/python/keras/_impl/keras/engine/network.py b/tensorflow/python/keras/_impl/keras/engine/network.py
index a0229be346..9e75096249 100644
--- a/tensorflow/python/keras/_impl/keras/engine/network.py
+++ b/tensorflow/python/keras/_impl/keras/engine/network.py
@@ -318,6 +318,9 @@ class Network(base_layer.Layer):
           layer, name='layer-%d' % layer_index, overwrite=True)
 
   def __setattr__(self, name, value):
+    no_dependency = isinstance(value, checkpointable.NoDependency)
+    if no_dependency:
+      value = value.value
     if isinstance(value, (base_layer.Layer, Network)):
       try:
         is_graph_network = self._is_graph_network
@@ -332,7 +335,8 @@ class Network(base_layer.Layer):
             # In subclassed models, legacy layers (tf.layers) must always use
             # resource variables.
             value._use_resource_variables = True
-    if isinstance(value, checkpointable.CheckpointableBase):
+    if (not no_dependency
+        and isinstance(value, checkpointable.CheckpointableBase)):
       # Layer (and therefore Network/Model) inherit from CheckpointableBase
       # rather than Checkpointable, which means there is no Checkpointable
       # __setattr__ override (it would be a performance issue for functional
diff --git a/tensorflow/python/keras/_impl/keras/model_subclassing_test.py b/tensorflow/python/keras/_impl/keras/model_subclassing_test.py
index 295ad47f6b..3f850e57aa 100644
--- a/tensorflow/python/keras/_impl/keras/model_subclassing_test.py
+++ b/tensorflow/python/keras/_impl/keras/model_subclassing_test.py
@@ -28,7 +28,9 @@ from tensorflow.python.framework import tensor_shape
 from tensorflow.python.framework import test_util
 from tensorflow.python.keras._impl import keras
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.platform import test
+from tensorflow.python.training import checkpointable
 from tensorflow.python.training.rmsprop import RMSPropOptimizer
 
 try:
@@ -583,6 +585,22 @@ class ModelSubclassingTest(test.TestCase):
     loss = model.train_on_batch(x, y)
     self.assertGreater(loss, 0.1)
 
+  def test_no_dependency(self):
+    class Foo(keras.Model):
+
+      def __init__(self):
+        super(Foo, self).__init__()
+        self.isdep = keras.layers.Dense(1)
+        self.notdep = checkpointable.NoDependency(keras.layers.Dense(2))
+        self.notdep_var = checkpointable.NoDependency(
+            resource_variable_ops.ResourceVariable(1., name='notdep_var'))
+
+    m = Foo()
+    self.assertEqual([m.isdep, m.notdep], m.layers)
+    self.assertEqual(1, len(m._checkpoint_dependencies))
+    self.assertIs(m.isdep, m._checkpoint_dependencies[0].ref)
+    self.assertEqual('notdep_var:0', m.notdep_var.name)
+
 
 class CustomCallModel(keras.Model):
 
diff --git a/tensorflow/python/training/checkpointable.py b/tensorflow/python/training/checkpointable.py
index 05afd37ccd..d00312a1f3 100644
--- a/tensorflow/python/training/checkpointable.py
+++ b/tensorflow/python/training/checkpointable.py
@@ -659,6 +659,31 @@ class CheckpointableBase(object):
     return {}
 
 
+class NoDependency(object):
+  """Allows attribute assignment to `Checkpointable` objects with no dependency.
+
+  Example usage:
+  ```python
+  obj = Checkpointable()
+  obj.has_dependency = tf.Variable(0., name="dep")
+  obj.no_dependency = NoDependency(tf.Variable(1., name="nodep"))
+  assert obj.no_dependency.name == "nodep:0"
+  ```
+
+  `obj` in this example has a dependency on the variable "dep", and both
+  attributes contain un-wrapped `Variable` objects.
+
+  `NoDependency` also works with `tf.keras.Model`, but only for checkpoint
+  dependencies: wrapping a `Layer` in `NoDependency` will assign the (unwrapped)
+  `Layer` to the attribute without a checkpoint dependency, but the `Model` will
+  still track the `Layer` (so it will appear in `Model.layers`, and its
+  variables will appear in `Model.variables`).
+  """
+
+  def __init__(self, value):
+    self.value = value
+
+
 class Checkpointable(CheckpointableBase):
   """Manages dependencies on other objects.
 
@@ -691,8 +716,11 @@ class Checkpointable(CheckpointableBase):
     """Support self.foo = checkpointable syntax."""
     # Perform the attribute assignment, and potentially call other __setattr__
     # overrides such as that for tf.keras.Model.
+    no_dependency = isinstance(value, NoDependency)
+    if no_dependency:
+      value = value.value
     super(Checkpointable, self).__setattr__(name, value)
-    if isinstance(value, CheckpointableBase):
+    if not no_dependency and isinstance(value, CheckpointableBase):
       self._track_checkpointable(
           value, name=name,
           # Allow the user to switch the Checkpointable which is tracked by this
diff --git a/tensorflow/python/training/checkpointable_test.py b/tensorflow/python/training/checkpointable_test.py
index e79acb4975..85802cb661 100644
--- a/tensorflow/python/training/checkpointable_test.py
+++ b/tensorflow/python/training/checkpointable_test.py
@@ -34,6 +34,16 @@ class InterfaceTests(test.TestCase):
     root.leaf = duplicate_name_dep
     root._track_checkpointable(duplicate_name_dep, name="leaf", overwrite=True)
 
+  def testNoDependency(self):
+    root = checkpointable.Checkpointable()
+    hasdep = checkpointable.Checkpointable()
+    root.hasdep = hasdep
+    nodep = checkpointable.Checkpointable()
+    root.nodep = checkpointable.NoDependency(nodep)
+    self.assertEqual(1, len(root._checkpoint_dependencies))
+    self.assertIs(root._checkpoint_dependencies[0].ref, root.hasdep)
+    self.assertIs(root.hasdep, hasdep)
+    self.assertIs(root.nodep, nodep)
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/python/training/checkpointable_utils.py b/tensorflow/python/training/checkpointable_utils.py
index cf4112ff99..f2a2b411fd 100644
--- a/tensorflow/python/training/checkpointable_utils.py
+++ b/tensorflow/python/training/checkpointable_utils.py
@@ -1044,8 +1044,11 @@ class Checkpoint(checkpointable_lib.Checkpointable):
     if self._save_counter is None:
       # Initialized to 0 and incremented before saving.
       with ops.device("/cpu:0"):
-        self._save_counter = add_variable(
-            self, name="save_counter", initializer=0, dtype=dtypes.int64)
+        # add_variable creates a dependency named "save_counter"; NoDependency
+        # prevents creating a second dependency named "_save_counter".
+        self._save_counter = checkpointable_lib.NoDependency(
+            add_variable(self, name="save_counter", initializer=0,
+                         dtype=dtypes.int64))
 
   @property
   def save_counter(self):
-- 
GitLab


From 5fb53fe69afe7f9106a8bcb5632cea23cf227d78 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 18:49:08 -0700
Subject: [PATCH 280/395] add support for PadV2

PiperOrigin-RevId: 195503894
---
 tensorflow/contrib/lite/builtin_op_data.h     |   3 +
 tensorflow/contrib/lite/builtin_ops.h         |   1 +
 .../internal/optimized/optimized_ops.h        |  80 ++--
 .../internal/reference/reference_ops.h        |  21 +-
 tensorflow/contrib/lite/kernels/pad.cc        | 110 ++++--
 tensorflow/contrib/lite/kernels/pad_test.cc   | 368 +++++++++++++++++-
 tensorflow/contrib/lite/kernels/register.cc   |   2 +
 tensorflow/contrib/lite/kernels/test_util.cc  |  82 +---
 tensorflow/contrib/lite/kernels/test_util.h   |  85 +++-
 tensorflow/contrib/lite/model.cc              |   3 +
 tensorflow/contrib/lite/nnapi_delegate.cc     |   1 +
 tensorflow/contrib/lite/schema/schema.fbs     |   5 +
 .../contrib/lite/schema/schema_generated.h    | 160 ++++++--
 tensorflow/contrib/lite/testing/BUILD         |   1 +
 .../contrib/lite/testing/generate_examples.py |  54 +++
 .../testing/generated_examples_zip_test.cc    |   5 +-
 tensorflow/contrib/lite/toco/BUILD            |   1 +
 .../contrib/lite/toco/export_tensorflow.cc    |  34 ++
 .../graph_transformations.h                   |   1 +
 .../propagate_fixed_sizes.cc                  |  29 ++
 .../toco/graph_transformations/quantize.cc    |   1 +
 .../resolve_padv2_attributes.cc               |  55 +++
 .../contrib/lite/toco/import_tensorflow.cc    |  15 +
 tensorflow/contrib/lite/toco/model.h          |  24 ++
 .../contrib/lite/toco/tflite/operator.cc      |  17 +
 tensorflow/contrib/lite/toco/toco_tooling.cc  |   1 +
 tensorflow/contrib/lite/toco/tooling_util.cc  |   1 +
 27 files changed, 986 insertions(+), 174 deletions(-)
 create mode 100644 tensorflow/contrib/lite/toco/graph_transformations/resolve_padv2_attributes.cc

diff --git a/tensorflow/contrib/lite/builtin_op_data.h b/tensorflow/contrib/lite/builtin_op_data.h
index 4910c89eae..35cf43dd32 100644
--- a/tensorflow/contrib/lite/builtin_op_data.h
+++ b/tensorflow/contrib/lite/builtin_op_data.h
@@ -161,6 +161,9 @@ typedef struct {
 typedef struct {
 } TfLitePadParams;
 
+typedef struct {
+} TfLitePadV2Params;
+
 typedef struct {
   // TODO(ahentz): We can't have dynamic data in this struct, at least not yet.
   // For now we will fix the maximum possible number of dimensions.
diff --git a/tensorflow/contrib/lite/builtin_ops.h b/tensorflow/contrib/lite/builtin_ops.h
index 962a7a8970..d66b72843a 100644
--- a/tensorflow/contrib/lite/builtin_ops.h
+++ b/tensorflow/contrib/lite/builtin_ops.h
@@ -85,6 +85,7 @@ typedef enum {
   kTfLiteBuiltinMinimum = 57,
   kTfLiteBuiltinLess = 58,
   kTfLiteBuiltinNeg = 59,
+  kTfLiteBuiltinPadv2 = 60,
 } TfLiteBuiltinOperator;
 
 #ifdef __cplusplus
diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
index 4776726972..e2a1a6996d 100644
--- a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
@@ -5851,10 +5851,26 @@ inline void BatchToSpaceND(const T* input_data, const Dims<4>& input_dims,
 }
 
 template <typename T>
-inline void Pad(const T* input_data, const Dims<4>& input_dims,
-                const std::vector<int>& left_paddings,
-                const std::vector<int>& right_paddings, T* output_data,
-                const Dims<4>& output_dims, const int32_t pad_value) {
+void TypedMemset(void* ptr, T value, size_t num) {
+  // Optimization for common cases where memset() will suffice.
+  if (value == 0 || std::is_same<T, uint8_t>::value) {
+    memset(ptr, value, num * sizeof(T));
+  } else {
+    // Default implementation for cases where memset() will not preserve the
+    // bytes, e.g., typically when sizeof(T) > sizeof(uint8_t).
+    char* pos = static_cast<char*>(ptr);
+    for (size_t i = 0; i < num; ++i) {
+      memcpy(pos, &value, sizeof(T));
+      pos = pos + sizeof(T);
+    }
+  }
+}
+
+template <typename T>
+inline void PadV2(const T* input_data, const Dims<4>& input_dims,
+                  const std::vector<int>& left_paddings,
+                  const std::vector<int>& right_paddings, T* output_data,
+                  const Dims<4>& output_dims, const T pad_value) {
   gemmlowp::ScopedProfilingLabel label("Pad");
   TFLITE_DCHECK_EQ(left_paddings.size(), 4);
   TFLITE_DCHECK_EQ(right_paddings.size(), 4);
@@ -5877,27 +5893,28 @@ inline void Pad(const T* input_data, const Dims<4>& input_dims,
   const int input_depth = ArraySize(input_dims, 0);
 
   if (left_b_padding != 0) {
-    memset(output_data, pad_value,
-           left_b_padding * output_height * output_width * output_depth *
-               sizeof(T));
+    TypedMemset<T>(
+        output_data, pad_value,
+        left_b_padding * output_height * output_width * output_depth);
   }
   for (int out_b = left_b_padding; out_b < output_batch - right_b_padding;
        ++out_b) {
     if (left_h_padding != 0) {
-      memset(output_data + Offset(output_dims, 0, 0, 0, out_b), pad_value,
-             left_h_padding * output_width * output_depth * sizeof(T));
+      TypedMemset<T>(output_data + Offset(output_dims, 0, 0, 0, out_b),
+                     pad_value, left_h_padding * output_width * output_depth);
     }
     for (int out_h = left_h_padding; out_h < output_height - right_h_padding;
          ++out_h) {
       if (left_w_padding != 0) {
-        memset(output_data + Offset(output_dims, 0, 0, out_h, out_b), pad_value,
-               left_w_padding * output_depth * sizeof(T));
+        TypedMemset<T>(output_data + Offset(output_dims, 0, 0, out_h, out_b),
+                       pad_value, left_w_padding * output_depth);
       }
       for (int out_w = left_w_padding; out_w < output_width - right_w_padding;
            ++out_w) {
         if (left_d_padding != 0) {
-          memset(output_data + Offset(output_dims, 0, out_w, out_h, out_b),
-                 pad_value, left_d_padding * sizeof(T));
+          TypedMemset<T>(
+              output_data + Offset(output_dims, 0, out_w, out_h, out_b),
+              pad_value, left_d_padding);
         }
 
         T* out = output_data +
@@ -5908,35 +5925,46 @@ inline void Pad(const T* input_data, const Dims<4>& input_dims,
         memcpy(out, in, input_depth * sizeof(T));
 
         if (right_d_padding != 0) {
-          memset(
+          TypedMemset<T>(
               output_data + Offset(output_dims, output_depth - right_d_padding,
                                    out_w, out_h, out_b),
-              pad_value, right_d_padding * sizeof(T));
+              pad_value, right_d_padding);
         }
       }
       if (right_w_padding != 0) {
-        memset(
+        TypedMemset<T>(
             output_data + Offset(output_dims, 0, output_width - right_w_padding,
                                  out_h, out_b),
-            pad_value, right_w_padding * output_depth * sizeof(T));
+            pad_value, right_w_padding * output_depth);
       }
     }
     if (right_h_padding != 0) {
-      memset(output_data + Offset(output_dims, 0, 0,
-                                  output_height - right_h_padding, out_b),
-             pad_value,
-             right_h_padding * output_width * output_depth * sizeof(T));
+      TypedMemset<T>(
+          output_data +
+              Offset(output_dims, 0, 0, output_height - right_h_padding, out_b),
+          pad_value, right_h_padding * output_width * output_depth);
     }
   }
   if (right_b_padding != 0) {
-    memset(output_data +
-               Offset(output_dims, 0, 0, 0, output_batch - right_b_padding),
-           0,
-           right_b_padding * output_height * output_width * output_depth *
-               sizeof(T));
+    TypedMemset<T>(
+        output_data +
+            Offset(output_dims, 0, 0, 0, output_batch - right_b_padding),
+        pad_value,
+        right_b_padding * output_height * output_width * output_depth);
   }
 }
 
+// Legacy Pad() method that casts an int32_t to T before padding.
+template <typename T>
+inline void Pad(const T* input_data, const Dims<4>& input_dims,
+                const std::vector<int>& left_paddings,
+                const std::vector<int>& right_paddings, T* output_data,
+                const Dims<4>& output_dims, const int32_t pad_value) {
+  const T converted_pad_value = static_cast<T>(pad_value);
+  PadV2<T>(input_data, input_dims, left_paddings, right_paddings, output_data,
+           output_dims, converted_pad_value);
+}
+
 template <typename T>
 inline void Pad(const T* input_data, const Dims<4>& input_dims,
                 const std::vector<int>& left_paddings,
diff --git a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
index c6ed614593..05e6ca8e7e 100644
--- a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
@@ -3158,10 +3158,10 @@ inline void BatchToSpaceND(const T* input_data, const Dims<4>& input_dims,
 }
 
 template <typename T>
-inline void Pad(const T* input_data, const Dims<4>& input_dims,
-                const std::vector<int>& left_paddings,
-                const std::vector<int>& right_paddings, T* output_data,
-                const Dims<4>& output_dims, const int32_t pad_value) {
+inline void PadV2(const T* input_data, const Dims<4>& input_dims,
+                  const std::vector<int>& left_paddings,
+                  const std::vector<int>& right_paddings, T* output_data,
+                  const Dims<4>& output_dims, const T pad_value) {
   TFLITE_DCHECK_EQ(left_paddings.size(), 4);
   TFLITE_DCHECK_EQ(right_paddings.size(), 4);
 
@@ -3194,7 +3194,7 @@ inline void Pad(const T* input_data, const Dims<4>& input_dims,
               out_w >= output_width - right_w_padding ||
               out_d < left_d_padding ||
               out_d >= output_depth - right_d_padding) {
-            *out_ptr++ = static_cast<T>(pad_value);
+            *out_ptr++ = pad_value;
           } else {
             *out_ptr++ = *in_ptr++;
           }
@@ -3204,6 +3204,17 @@ inline void Pad(const T* input_data, const Dims<4>& input_dims,
   }
 }
 
+// Legacy Pad() method that casts an int32_t to T before padding.
+template <typename T>
+inline void Pad(const T* input_data, const Dims<4>& input_dims,
+                const std::vector<int>& left_paddings,
+                const std::vector<int>& right_paddings, T* output_data,
+                const Dims<4>& output_dims, const int32_t pad_value) {
+  const T converted_pad_value = static_cast<T>(pad_value);
+  PadV2<T>(input_data, input_dims, left_paddings, right_paddings, output_data,
+           output_dims, converted_pad_value);
+}
+
 template <typename T>
 inline void Pad(const T* input_data, const Dims<4>& input_dims,
                 const std::vector<int>& left_paddings,
diff --git a/tensorflow/contrib/lite/kernels/pad.cc b/tensorflow/contrib/lite/kernels/pad.cc
index 4f9449a225..9e1e4658e9 100644
--- a/tensorflow/contrib/lite/kernels/pad.cc
+++ b/tensorflow/contrib/lite/kernels/pad.cc
@@ -37,9 +37,15 @@ struct PadContext {
   PadContext(TfLiteContext* context, TfLiteNode* node) {
     input = GetInput(context, node, 0);
     paddings = GetInput(context, node, 1);
+    if (NumInputs(node) == 3) {
+      constant_values = GetOptionalInputTensor(context, node, 2);
+    } else {
+      constant_values = nullptr;
+    }
     output = GetOutput(context, node, 0);
     dims = NumDimensions(input);
   }
+  TfLiteTensor* constant_values;
   TfLiteTensor* input;
   TfLiteTensor* paddings;
   TfLiteTensor* output;
@@ -76,11 +82,15 @@ TfLiteStatus ResizeOutputTensor(TfLiteContext* context,
 }
 
 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
-  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
+  TF_LITE_ENSURE(context, NumInputs(node) == 2 || NumInputs(node) == 3);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
 
   PadContext op_context(context, node);
   TF_LITE_ENSURE_EQ(context, op_context.input->type, op_context.output->type);
+  if (op_context.constant_values != nullptr) {
+    TF_LITE_ENSURE_EQ(context, op_context.input->type,
+                      op_context.constant_values->type);
+  }
 
   // TODO(nupurgarg): Our current implementations rely on the inputs being 4D.
   TF_LITE_ENSURE_EQ(context, op_context.dims, 4);
@@ -98,6 +108,11 @@ template <KernelType kernel_type>
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   PadContext op_context(context, node);
 
+  if (op_context.constant_values != nullptr) {
+    // Ensure that constant_values is a scalar.
+    TF_LITE_ENSURE_EQ(context, NumElements(op_context.constant_values), 1);
+  }
+
   // Resize the output tensor if the output tensor is dynamic.
   if (IsDynamicTensor(op_context.output)) {
     TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context));
@@ -119,48 +134,70 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
     after_padding.push_back(paddings_data[idx * 2 + 1]);
   }
 
-#define TF_LITE_PAD(type, scalar, pad_value)                                \
-  type::Pad(GetTensorData<scalar>(op_context.input),                        \
-            GetTensorDims(op_context.input), before_padding, after_padding, \
-            GetTensorData<scalar>(op_context.output),                       \
-            GetTensorDims(op_context.output), pad_value)
+#define TF_LITE_PAD(type, scalar, pad_value)                                  \
+  type::PadV2(GetTensorData<scalar>(op_context.input),                        \
+              GetTensorDims(op_context.input), before_padding, after_padding, \
+              GetTensorData<scalar>(op_context.output),                       \
+              GetTensorDims(op_context.output), pad_value)
 
   switch (op_context.input->type) {
-    case kTfLiteFloat32:
+    case kTfLiteFloat32: {
+      float pad_value = op_context.constant_values == nullptr
+                            ? 0.f
+                            : *GetTensorData<float>(op_context.constant_values);
       if (kernel_type == kReference) {
-        TF_LITE_PAD(reference_ops, float, 0);
+        TF_LITE_PAD(reference_ops, float, pad_value);
       } else if (kernel_type == kGenericOptimized) {
-        TF_LITE_PAD(optimized_ops, float, 0);
+        TF_LITE_PAD(optimized_ops, float, pad_value);
+      }
+    } break;
+    case kTfLiteUInt8: {
+      uint8_t pad_value;
+      if (op_context.constant_values == nullptr) {
+        // Quantized Pad requires that 0 is represented in the quantized
+        // range.
+        TF_LITE_ENSURE(context, op_context.output->params.zero_point >=
+                                    std::numeric_limits<uint8_t>::min());
+        TF_LITE_ENSURE(context, op_context.output->params.zero_point <=
+                                    std::numeric_limits<uint8_t>::max());
+        pad_value = static_cast<uint8_t>(op_context.output->params.zero_point);
+      } else {
+        // Quantized Pad requires that 'constant_values' is represented in the
+        // same quantized range as the input and output tensors.
+        TF_LITE_ENSURE_EQ(context, op_context.output->params.zero_point,
+                          op_context.constant_values->params.zero_point);
+        TF_LITE_ENSURE_EQ(context, op_context.output->params.scale,
+                          op_context.constant_values->params.scale);
+        pad_value = *GetTensorData<uint8_t>(op_context.constant_values);
       }
-      break;
-    case kTfLiteUInt8:
-      // Quantized Pad requires that 0 is represented in the quantized range.
-      TF_LITE_ENSURE(context, op_context.output->params.zero_point >=
-                                  std::numeric_limits<uint8_t>::min());
-      TF_LITE_ENSURE(context, op_context.output->params.zero_point <=
-                                  std::numeric_limits<uint8_t>::max());
       if (kernel_type == kReference) {
-        TF_LITE_PAD(reference_ops, uint8_t,
-                    op_context.output->params.zero_point);
+        TF_LITE_PAD(reference_ops, uint8_t, pad_value);
       } else if (kernel_type == kGenericOptimized) {
-        TF_LITE_PAD(optimized_ops, uint8_t,
-                    op_context.output->params.zero_point);
+        TF_LITE_PAD(optimized_ops, uint8_t, pad_value);
       }
-      break;
-    case kTfLiteInt32:
+    } break;
+    case kTfLiteInt32: {
+      int32_t pad_value =
+          op_context.constant_values == nullptr
+              ? 0
+              : *GetTensorData<int32_t>(op_context.constant_values);
       if (kernel_type == kReference) {
-        TF_LITE_PAD(reference_ops, int32_t, 0);
+        TF_LITE_PAD(reference_ops, int32_t, pad_value);
       } else if (kernel_type == kGenericOptimized) {
-        TF_LITE_PAD(optimized_ops, int32_t, 0);
+        TF_LITE_PAD(optimized_ops, int32_t, pad_value);
       }
-      break;
-    case kTfLiteInt64:
+    } break;
+    case kTfLiteInt64: {
+      int64_t pad_value =
+          op_context.constant_values == nullptr
+              ? 0L
+              : *GetTensorData<int64_t>(op_context.constant_values);
       if (kernel_type == kReference) {
-        TF_LITE_PAD(reference_ops, int64_t, 0);
+        TF_LITE_PAD(reference_ops, int64_t, pad_value);
       } else if (kernel_type == kGenericOptimized) {
-        TF_LITE_PAD(optimized_ops, int64_t, 0);
+        TF_LITE_PAD(optimized_ops, int64_t, pad_value);
       }
-      break;
+    } break;
     default:
       context->ReportError(context, "Type is currently not supported by Pad.");
       return kTfLiteError;
@@ -185,6 +222,21 @@ TfLiteRegistration* Register_PAD_GENERIC_OPT() {
 
 TfLiteRegistration* Register_PAD() { return Register_PAD_GENERIC_OPT(); }
 
+// Also register Pad as PadV2.
+TfLiteRegistration* Register_PADV2_REF() {
+  static TfLiteRegistration r = {nullptr, nullptr, pad::Prepare,
+                                 pad::Eval<pad::kReference>};
+  return &r;
+}
+
+TfLiteRegistration* Register_PADV2_GENERIC_OPT() {
+  static TfLiteRegistration r = {nullptr, nullptr, pad::Prepare,
+                                 pad::Eval<pad::kGenericOptimized>};
+  return &r;
+}
+
+TfLiteRegistration* Register_PADV2() { return Register_PADV2_GENERIC_OPT(); }
+
 }  // namespace builtin
 }  // namespace ops
 }  // namespace tflite
diff --git a/tensorflow/contrib/lite/kernels/pad_test.cc b/tensorflow/contrib/lite/kernels/pad_test.cc
index c06237e572..f8b9064fbb 100644
--- a/tensorflow/contrib/lite/kernels/pad_test.cc
+++ b/tensorflow/contrib/lite/kernels/pad_test.cc
@@ -24,21 +24,26 @@ namespace {
 using ::testing::ElementsAreArray;
 using ::testing::Matcher;
 
+template <typename T>
 class PadOpModel : public SingleOpModel {
  public:
-  void SetInput(std::initializer_list<float> data) {
-    PopulateTensor<float>(input_, data);
+  void SetInput(std::initializer_list<T> data) {
+    PopulateTensor<T>(input_, data);
   }
 
   void SetQuantizedInput(std::initializer_list<float> data) {
     QuantizeAndPopulate<uint8_t>(input_, data);
   }
 
+  void SetQuantizedPadValue(float data) {
+    QuantizeAndPopulate<uint8_t>(constant_values_, {data});
+  }
+
   void SetPaddings(std::initializer_list<int> paddings) {
     PopulateTensor<int>(paddings_, paddings);
   }
 
-  std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
+  std::vector<T> GetOutput() { return ExtractVector<T>(output_); }
   std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
 
   std::vector<float> GetDequantizedOutput() {
@@ -50,6 +55,59 @@ class PadOpModel : public SingleOpModel {
   int input_;
   int output_;
   int paddings_;
+  int constant_values_;
+};
+
+namespace {
+
+// Returns the corresponding TensorType given the type T.
+template <typename T>
+TensorType GetTensorType() {
+  if (std::is_same<T, float>::value) return TensorType_FLOAT32;
+  if (std::is_same<T, int32_t>::value) return TensorType_INT32;
+  if (std::is_same<T, uint8_t>::value) return TensorType_UINT8;
+  return TensorType_MIN;  // default value
+}
+
+}  // namespace
+
+// Tests case where paddings is a const tensor. Type T is the dtype.
+template <typename T>
+class PadV2OpConstModel : public PadOpModel<T> {
+ public:
+  PadV2OpConstModel(const TensorData& input,
+                    std::initializer_list<int> paddings_shape,
+                    std::initializer_list<int> paddings, T constant_values,
+                    const TensorData& output) {
+    this->input_ = this->AddInput(input);
+    this->paddings_ =
+        this->AddConstInput(TensorType_INT32, paddings, paddings_shape);
+    this->constant_values_ =
+        this->AddConstInput(GetTensorType<T>(), {constant_values}, {1});
+
+    this->output_ = this->AddOutput(output);
+
+    this->SetBuiltinOp(BuiltinOperator_PADV2, BuiltinOptions_PadV2Options,
+                       CreatePadV2Options(this->builder_).Union());
+    this->BuildInterpreter({input.shape});
+  }
+
+  PadV2OpConstModel(const TensorData& input,
+                    std::initializer_list<int> paddings_shape,
+                    std::initializer_list<int> paddings,
+                    const TensorData& constant_values,
+                    const TensorData& output) {
+    this->input_ = this->AddInput(input);
+    this->paddings_ =
+        this->AddConstInput(TensorType_INT32, paddings, paddings_shape);
+    this->constant_values_ = this->AddInput(constant_values);
+
+    this->output_ = this->AddOutput(output);
+
+    this->SetBuiltinOp(BuiltinOperator_PADV2, BuiltinOptions_PadV2Options,
+                       CreatePadV2Options(this->builder_).Union());
+    this->BuildInterpreter({input.shape});
+  }
 };
 
 // Tests case where paddings is a const tensor.
@@ -58,7 +116,7 @@ class PadOpModel : public SingleOpModel {
 //    PadOpDynamicModel m(input_shape, paddings_shape, paddings_data);
 //    m.SetInput(input_data);
 //    m.Invoke();
-class PadOpConstModel : public PadOpModel {
+class PadOpConstModel : public PadOpModel<float> {
  public:
   PadOpConstModel(const TensorData& input,
                   std::initializer_list<int> paddings_shape,
@@ -66,6 +124,7 @@ class PadOpConstModel : public PadOpModel {
                   const TensorData& output) {
     input_ = AddInput(input);
     paddings_ = AddConstInput(TensorType_INT32, paddings, paddings_shape);
+    constant_values_ = AddNullInput();
     output_ = AddOutput(output);
 
     SetBuiltinOp(BuiltinOperator_PAD, BuiltinOptions_PadOptions,
@@ -74,6 +133,38 @@ class PadOpConstModel : public PadOpModel {
   }
 };
 
+// Test case where paddings is a non-const tensor.
+template <typename T>
+class PadV2OpDynamicModel : public PadOpModel<T> {
+ public:
+  PadV2OpDynamicModel(const TensorData& input,
+                      std::initializer_list<int> paddings_shape,
+                      T constant_values, const TensorData& output) {
+    this->input_ = this->AddInput(input);
+    this->paddings_ = this->AddInput(TensorType_INT32);
+    this->constant_values_ =
+        this->AddConstInput(GetTensorType<T>(), {constant_values}, {1});
+    this->output_ = this->AddOutput(output);
+
+    this->SetBuiltinOp(BuiltinOperator_PADV2, BuiltinOptions_PadV2Options,
+                       CreatePadV2Options(this->builder_).Union());
+    this->BuildInterpreter({input.shape, paddings_shape});
+  }
+  PadV2OpDynamicModel(const TensorData& input,
+                      std::initializer_list<int> paddings_shape,
+                      const TensorData& constant_values,
+                      const TensorData& output) {
+    this->input_ = this->AddInput(input);
+    this->paddings_ = this->AddInput(TensorType_INT32);
+    this->constant_values_ = this->AddInput(constant_values);
+    this->output_ = this->AddOutput(output);
+
+    this->SetBuiltinOp(BuiltinOperator_PADV2, BuiltinOptions_PadV2Options,
+                       CreatePadV2Options(this->builder_).Union());
+    this->BuildInterpreter({input.shape, paddings_shape});
+  }
+};
+
 // Test case where paddings is a non-const tensor.
 //
 // Example usage is as follows:
@@ -81,13 +172,14 @@ class PadOpConstModel : public PadOpModel {
 //    m.SetInput(input_data);
 //    m.SetPaddings(paddings_data);
 //    m.Invoke();
-class PadOpDynamicModel : public PadOpModel {
+class PadOpDynamicModel : public PadOpModel<float> {
  public:
   PadOpDynamicModel(const TensorData& input,
                     std::initializer_list<int> paddings_shape,
                     const TensorData& output) {
     input_ = AddInput(input);
     paddings_ = AddInput(TensorType_INT32);
+    constant_values_ = AddNullInput();
     output_ = AddOutput(output);
 
     SetBuiltinOp(BuiltinOperator_PAD, BuiltinOptions_PadOptions,
@@ -237,6 +329,272 @@ TEST_F(QuantizedPadOpTest, AdvancedDynamicTest) {
   EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 7, 1}));
 }
 
+TEST(PadV2OpTest, TooManyDimensions) {
+  EXPECT_DEATH(PadV2OpConstModel<float>(
+                   {TensorType_FLOAT32, {1, 2, 3, 4, 5, 6, 7, 8, 9}}, {9, 2},
+                   {1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9}, 0.0,
+                   {TensorType_FLOAT32}),
+               "dims != 4");
+}
+
+TEST(PadV2OpTest, UnequalDimensions) {
+  EXPECT_DEATH(
+      PadV2OpConstModel<float>({TensorType_FLOAT32, {1, 1, 2, 1}}, {3, 2},
+                               {1, 1, 2, 2, 3, 3}, 0.0, {TensorType_FLOAT32}),
+      "3 != 4");
+}
+
+TEST(PadV2OpTest, InvalidPadValue) {
+  EXPECT_DEATH(PadV2OpConstModel<float>({TensorType_FLOAT32, {1, 1, 2, 1}},
+                                        {4, 2}, {0, 0, 1, -1, 2, -1, 0, 0}, 0.0,
+                                        {TensorType_FLOAT32}),
+               "Pad value has to be greater than equal to 0.");
+}
+
+TEST(PadV2OpTest, SimpleConstTest) {
+  // Padding is represented as four 2-D lists representing above padding and
+  // below padding (i.e. {{0, 0}, {1, 1}, {1, 1}, {0, 0}}).
+  PadV2OpConstModel<float> m({TensorType_FLOAT32, {1, 2, 2, 1}}, {4, 2},
+                             {0, 0, 1, 1, 1, 1, 0, 0}, 0.0,
+                             {TensorType_FLOAT32});
+  m.SetInput({1, 2, 3, 4});
+  m.Invoke();
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({0, 0, 0, 0, 0, 1, 2, 0, 0, 3, 4,
+                                               0, 0, 0, 0, 0}));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1}));
+}
+
+TEST(PadV2OpTest, SimpleConstFloat32ValuedTest) {
+  // Padding is represented as four 2-D lists representing above padding and
+  // below padding (i.e. {{0, 0}, {1, 1}, {1, 1}, {0, 0}}).
+  PadV2OpConstModel<float> m({TensorType_FLOAT32, {1, 2, 2, 1}}, {4, 2},
+                             {0, 0, 1, 1, 1, 1, 0, 0}, 5, {TensorType_FLOAT32});
+  m.SetInput({1, 2, 3, 4});
+  m.Invoke();
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 5, 5, 5, 5, 1, 2, 5, 5, 3, 4,
+                                               5, 5, 5, 5, 5}));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1}));
+}
+
+TEST(PadV2OpTest, Simple4DConstFloat32ValuedTest) {
+  // Padding is represented as four 2-D lists representing above padding and
+  // below padding (i.e. {{0, 0}, {1, 1}, {1, 1}, {0, 0}}).
+  PadV2OpConstModel<float> m({TensorType_FLOAT32, {1, 1, 2, 1}}, {4, 2},
+                             {0, 1, 0, 0, 0, 0, 0, 1}, 5, {TensorType_FLOAT32});
+  m.SetInput({3, 3});
+  m.Invoke();
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 5, 3, 5, 5, 5, 5, 5}));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 1, 2, 2}));
+}
+
+TEST(PadV2OpTest, SimpleConstInt32ValuedTest) {
+  // Padding is represented as four 2-D lists representing above padding and
+  // below padding (i.e. {{0, 0}, {1, 1}, {1, 1}, {0, 0}}).
+  PadV2OpConstModel<int32_t> m({TensorType_INT32, {1, 2, 2, 1}}, {4, 2},
+                               {0, 0, 1, 1, 1, 1, 0, 0}, 5, {TensorType_INT32});
+  m.SetInput({1, 2, 3, 4});
+  m.Invoke();
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 5, 5, 5, 5, 1, 2, 5, 5, 3, 4,
+                                               5, 5, 5, 5, 5}));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1}));
+}
+
+TEST(PadV2OpTest, SimpleDynamicTest) {
+  PadV2OpDynamicModel<float> m({TensorType_FLOAT32, {1, 2, 2, 1}}, {4, 2}, 0.0,
+                               {TensorType_FLOAT32});
+  m.SetInput({1, 2, 3, 4});
+  m.SetPaddings({0, 0, 1, 1, 1, 1, 0, 0});
+  m.Invoke();
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({0, 0, 0, 0, 0, 1, 2, 0, 0, 3, 4,
+                                               0, 0, 0, 0, 0}));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1}));
+}
+
+TEST(PadV2OpTest, SimpleDynamicValuedTest) {
+  PadV2OpDynamicModel<float> m({TensorType_FLOAT32, {1, 2, 2, 1}}, {4, 2}, 5,
+                               {TensorType_FLOAT32});
+  m.SetInput({1, 2, 3, 4});
+  m.SetPaddings({0, 0, 1, 1, 1, 1, 0, 0});
+  m.Invoke();
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 5, 5, 5, 5, 1, 2, 5, 5, 3, 4,
+                                               5, 5, 5, 5, 5}));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1}));
+}
+
+TEST(PadV2OpTest, AdvancedConstTest) {
+  PadV2OpConstModel<float> m({TensorType_FLOAT32, {1, 2, 3, 1}}, {4, 2},
+                             {0, 0, 0, 2, 1, 3, 0, 0}, 0, {TensorType_FLOAT32});
+  m.SetInput({1, 2, 3, 4, 5, 6});
+  m.Invoke();
+  EXPECT_THAT(m.GetOutput(),
+              ElementsAreArray({0, 1, 2, 3, 0, 0, 0, 0, 4, 5, 6, 0, 0, 0,
+                                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 7, 1}));
+}
+
+TEST(PadV2OpTest, AdvancedDynamicTest) {
+  PadV2OpDynamicModel<float> m({TensorType_FLOAT32, {1, 2, 3, 1}}, {4, 2}, 0,
+                               {TensorType_FLOAT32});
+  m.SetInput({1, 2, 3, 4, 5, 6});
+  m.SetPaddings({0, 0, 0, 2, 1, 3, 0, 0});
+  m.Invoke();
+  EXPECT_THAT(m.GetOutput(),
+              ElementsAreArray({0, 1, 2, 3, 0, 0, 0, 0, 4, 5, 6, 0, 0, 0,
+                                0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 7, 1}));
+}
+
+class QuantizedPadV2OpTest : public ::testing::Test {
+ protected:
+  std::vector<Matcher<float>> DequantizedArrayNear(
+      const std::vector<float>& values, const float min, const float max) {
+    const float quantization_tolerance = (max - min) / 255.0;
+    return ArrayFloatNear(values, quantization_tolerance);
+  }
+};
+
+TEST_F(QuantizedPadV2OpTest, ZeroNotInQuantizationRange) {
+  // The test_util and actual quantization code currently ensure that the range
+  // must include zero, but if that ever changes, this test will catch it.
+  EXPECT_DEATH(
+      PadV2OpConstModel<float> m({TensorType_UINT8, {1, 2, 2, 1}, 1.0, 2.0},
+                                 {4, 2}, {0, 0, 1, 1, 1, 1, 0, 0}, 0,
+                                 {TensorType_UINT8, {}, 1.0, 2.0}),
+      ".*Check failed: f_min <= 0.*");
+}
+
+TEST_F(QuantizedPadV2OpTest, SimpleConstTest) {
+  // Padding is represented as four 2-D lists representing above padding and
+  // below padding (i.e. {{0, 0}, {1, 1}, {1, 1}, {0, 0}}).
+  PadV2OpConstModel<uint8_t> m({TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0},
+                               {4, 2}, {0, 0, 1, 1, 1, 1, 0, 0},
+                               {TensorType_UINT8, {1}, -1.0, 1.0},
+                               {TensorType_UINT8, {}, -1.0, 1.0});
+  m.SetQuantizedInput({-0.8, 0.2, 0.9, 0.7});
+  m.SetQuantizedPadValue(0);
+  m.Invoke();
+  EXPECT_THAT(m.GetDequantizedOutput(),
+              ElementsAreArray(DequantizedArrayNear(
+                  {0, 0, 0, 0, 0, -0.8, 0.2, 0, 0, 0.9, 0.7, 0, 0, 0, 0, 0},
+                  -1.0, 1.0)));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1}));
+}
+
+TEST_F(QuantizedPadV2OpTest, SimpleDynamicTest) {
+  PadV2OpDynamicModel<uint8_t> m({TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0},
+                                 {4, 2}, {TensorType_UINT8, {1}, -1.0, 1.0},
+                                 {TensorType_UINT8, {}, -1.0, 1.0});
+  m.SetQuantizedInput({-0.8, 0.2, 0.9, 0.7});
+  m.SetQuantizedPadValue(0);
+  m.SetPaddings({0, 0, 1, 1, 1, 1, 0, 0});
+  m.Invoke();
+  EXPECT_THAT(m.GetDequantizedOutput(),
+              ElementsAreArray(DequantizedArrayNear(
+                  {0, 0, 0, 0, 0, -0.8, 0.2, 0, 0, 0.9, 0.7, 0, 0, 0, 0, 0},
+                  -1.0, 1.0)));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1}));
+}
+
+TEST_F(QuantizedPadV2OpTest, AdvancedConstTest) {
+  PadV2OpConstModel<uint8_t> m({TensorType_UINT8, {1, 2, 3, 1}, -1.0, 1.0},
+                               {4, 2}, {0, 0, 0, 2, 1, 3, 0, 0},
+                               {TensorType_UINT8, {1}, -1.0, 1.0},
+                               {TensorType_UINT8, {}, -1.0, 1.0});
+  m.SetQuantizedInput({-0.8, 0.2, 0.9, 0.7, 0.1, -0.3});
+  m.SetQuantizedPadValue(0);
+  m.Invoke();
+  EXPECT_THAT(m.GetDequantizedOutput(),
+              ElementsAreArray(DequantizedArrayNear(
+                  {0, -0.8, 0.2, 0.9, 0, 0, 0, 0, 0.7, 0.1, -0.3, 0, 0, 0,
+                   0, 0,    0,   0,   0, 0, 0, 0, 0,   0,   0,    0, 0, 0},
+                  -1.0, 1.0)));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 7, 1}));
+}
+
+TEST_F(QuantizedPadV2OpTest, AdvancedDynamicTest) {
+  PadV2OpDynamicModel<uint8_t> m({TensorType_UINT8, {1, 2, 3, 1}, -1.0, 1.0},
+                                 {4, 2}, {TensorType_UINT8, {1}, -1.0, 1.0},
+                                 {TensorType_UINT8, {}, -1.0, 1.0});
+  m.SetQuantizedInput({-0.8, 0.2, 0.9, 0.7, 0.1, -0.3});
+  m.SetQuantizedPadValue(0);
+  m.SetPaddings({0, 0, 0, 2, 1, 3, 0, 0});
+  m.Invoke();
+  EXPECT_THAT(m.GetDequantizedOutput(),
+              ElementsAreArray(DequantizedArrayNear(
+                  {0, -0.8, 0.2, 0.9, 0, 0, 0, 0, 0.7, 0.1, -0.3, 0, 0, 0,
+                   0, 0,    0,   0,   0, 0, 0, 0, 0,   0,   0,    0, 0, 0},
+                  -1.0, 1.0)));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 7, 1}));
+}
+
+TEST_F(QuantizedPadV2OpTest, SimpleConstValuedTest) {
+  // Padding is represented as four 2-D lists representing above padding and
+  // below padding (i.e. {{0, 0}, {1, 1}, {1, 1}, {0, 0}}).
+  PadV2OpConstModel<uint8_t> m({TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0},
+                               {4, 2}, {0, 0, 1, 1, 1, 1, 0, 0},
+                               {TensorType_UINT8, {1}, -1.0, 1.0},
+                               {TensorType_UINT8, {}, -1.0, 1.0});
+  m.SetQuantizedInput({-0.8, 0.2, 0.9, 0.7});
+  m.SetQuantizedPadValue(-0.5);
+  m.Invoke();
+  EXPECT_THAT(m.GetDequantizedOutput(),
+              ElementsAreArray(DequantizedArrayNear(
+                  {-0.5, -0.5, -0.5, -0.5, -0.5, -0.8, 0.2, -0.5, -0.5, 0.9,
+                   0.7, -0.5, -0.5, -0.5, -0.5, -0.5},
+                  -1.0, 1.0)));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1}));
+}
+
+TEST_F(QuantizedPadV2OpTest, SimpleDynamicValuedTest) {
+  PadV2OpDynamicModel<uint8_t> m({TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0},
+                                 {4, 2}, {TensorType_UINT8, {1}, -1.0, 1.0},
+                                 {TensorType_UINT8, {}, -1.0, 1.0});
+  m.SetQuantizedInput({-0.8, 0.2, 0.9, 0.7});
+  m.SetQuantizedPadValue(-0.5);
+  m.SetPaddings({0, 0, 1, 1, 1, 1, 0, 0});
+  m.Invoke();
+  EXPECT_THAT(m.GetDequantizedOutput(),
+              ElementsAreArray(DequantizedArrayNear(
+                  {-0.5, -0.5, -0.5, -0.5, -0.5, -0.8, 0.2, -0.5, -0.5, 0.9,
+                   0.7, -0.5, -0.5, -0.5, -0.5, -0.5},
+                  -1.0, 1.0)));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1}));
+}
+
+TEST_F(QuantizedPadV2OpTest, AdvancedConstValuedTest) {
+  PadV2OpConstModel<uint8_t> m({TensorType_UINT8, {1, 2, 3, 1}, -1.0, 1.0},
+                               {4, 2}, {0, 0, 0, 2, 1, 3, 0, 0},
+                               {TensorType_UINT8, {1}, -1.0, 1.0},
+                               {TensorType_UINT8, {}, -1.0, 1.0});
+  m.SetQuantizedInput({-0.8, 0.2, 0.9, 0.7, 0.1, -0.3});
+  m.SetQuantizedPadValue(-0.5);
+  m.Invoke();
+  EXPECT_THAT(m.GetDequantizedOutput(),
+              ElementsAreArray(DequantizedArrayNear(
+                  {-0.5, -0.8, 0.2,  0.9,  -0.5, -0.5, -0.5, -0.5, 0.7,  0.1,
+                   -0.3, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5,
+                   -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5},
+                  -1.0, 1.0)));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 7, 1}));
+}
+
+TEST_F(QuantizedPadV2OpTest, AdvancedDynamicValuedTest) {
+  PadV2OpDynamicModel<uint8_t> m({TensorType_UINT8, {1, 2, 3, 1}, -1.0, 1.0},
+                                 {4, 2}, {TensorType_UINT8, {1}, -1.0, 1.0},
+                                 {TensorType_UINT8, {}, -1.0, 1.0});
+  m.SetQuantizedInput({-0.8, 0.2, 0.9, 0.7, 0.1, -0.3});
+  m.SetQuantizedPadValue(-0.5);
+  m.SetPaddings({0, 0, 0, 2, 1, 3, 0, 0});
+  m.Invoke();
+  EXPECT_THAT(m.GetDequantizedOutput(),
+              ElementsAreArray(DequantizedArrayNear(
+                  {-0.5, -0.8, 0.2,  0.9,  -0.5, -0.5, -0.5, -0.5, 0.7,  0.1,
+                   -0.3, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5,
+                   -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5, -0.5},
+                  -1.0, 1.0)));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 7, 1}));
+}
+
 }  // namespace
 }  // namespace tflite
 
diff --git a/tensorflow/contrib/lite/kernels/register.cc b/tensorflow/contrib/lite/kernels/register.cc
index 29ea718a96..a6ea874546 100644
--- a/tensorflow/contrib/lite/kernels/register.cc
+++ b/tensorflow/contrib/lite/kernels/register.cc
@@ -60,6 +60,7 @@ TfLiteRegistration* Register_LSTM();
 TfLiteRegistration* Register_BIDIRECTIONAL_SEQUENCE_LSTM();
 TfLiteRegistration* Register_UNIDIRECTIONAL_SEQUENCE_LSTM();
 TfLiteRegistration* Register_PAD();
+TfLiteRegistration* Register_PADV2();
 TfLiteRegistration* Register_RESHAPE();
 TfLiteRegistration* Register_RESIZE_BILINEAR();
 TfLiteRegistration* Register_SKIP_GRAM();
@@ -121,6 +122,7 @@ BuiltinOpResolver::BuiltinOpResolver() {
   AddBuiltin(BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM,
              Register_UNIDIRECTIONAL_SEQUENCE_LSTM());
   AddBuiltin(BuiltinOperator_PAD, Register_PAD());
+  AddBuiltin(BuiltinOperator_PADV2, Register_PADV2());
   AddBuiltin(BuiltinOperator_RESHAPE, Register_RESHAPE());
   AddBuiltin(BuiltinOperator_RESIZE_BILINEAR, Register_RESIZE_BILINEAR());
   AddBuiltin(BuiltinOperator_SKIP_GRAM, Register_SKIP_GRAM());
diff --git a/tensorflow/contrib/lite/kernels/test_util.cc b/tensorflow/contrib/lite/kernels/test_util.cc
index 0bb28b50b2..5a6c85e97e 100644
--- a/tensorflow/contrib/lite/kernels/test_util.cc
+++ b/tensorflow/contrib/lite/kernels/test_util.cc
@@ -22,23 +22,6 @@ namespace tflite {
 using ::testing::FloatNear;
 using ::testing::Matcher;
 
-namespace {
-template <typename T>
-std::pair<float, int32_t> QuantizationParams(float f_min, float f_max) {
-  // These are required by many quantized operations.
-  CHECK_LE(f_min, 0);
-  CHECK_GE(f_max, 0);
-  T q_min = std::numeric_limits<T>::min();
-  T q_max = std::numeric_limits<T>::max();
-  float range = q_max - q_min;
-  float scale = (f_max - f_min) / range;
-  int32_t zero_point = std::min(
-      q_max,
-      std::max(q_min, static_cast<T>(std::round(q_min - f_min / scale))));
-  return {scale, zero_point};
-}
-}  // namespace
-
 std::vector<Matcher<float>> ArrayFloatNear(const std::vector<float>& values,
                                            float max_abs_error) {
   std::vector<Matcher<float>> matchers;
@@ -49,69 +32,8 @@ std::vector<Matcher<float>> ArrayFloatNear(const std::vector<float>& values,
   return matchers;
 }
 
-int SingleOpModel::AddTensor(TensorData t, std::initializer_list<int> data) {
-  int id = tensors_.size();
-
-  // This is slightly different depending on whether we are adding a
-  // quantized or a regular tensor.
-  bool is_quantized = (t.min != 0 || t.max != 0 || t.scale != 0);
-
-  flatbuffers::Offset<QuantizationParameters> q_params = 0;
-
-  if (is_quantized) {
-    if (t.min != 0 || t.max != 0) {
-      if (t.type == TensorType_UINT8) {
-        std::tie(t.scale, t.zero_point) =
-            QuantizationParams<uint8_t>(t.min, t.max);
-      } else if (t.type == TensorType_INT32) {
-        std::tie(t.scale, t.zero_point) =
-            QuantizationParams<int32_t>(t.min, t.max);
-      } else {
-        LOG(FATAL) << "No support for the requested quantized type";
-      }
-      t.min = 0;
-      t.max = 0;
-    }
-
-    q_params = CreateQuantizationParameters(
-        builder_, /*min=*/0, /*max=*/0, builder_.CreateVector<float>({t.scale}),
-        builder_.CreateVector<int64_t>({t.zero_point}));
-  }
-
-  int buffer_id = 0;
-  if (data.size()) {
-    // Initialize buffers list with empty buffer to allow for non-const tensors.
-    if (buffers_.empty()) {
-      buffers_.push_back(CreateBuffer(builder_, builder_.CreateVector({})));
-    }
-
-    // Add data as a Buffer to buffers list.
-    buffer_id = buffers_.size();
-    auto data_buffer =
-        builder_.CreateVector(reinterpret_cast<const uint8_t*>(data.begin()),
-                              sizeof(int) * data.size());
-    buffers_.push_back(CreateBuffer(builder_, data_buffer));
-  }
-
-  tensors_.push_back(CreateTensor(builder_, builder_.CreateVector<int>(t.shape),
-                                  t.type, /*buffer=*/buffer_id,
-                                  /*name=*/0, q_params));
-
-  tensor_data_[id] = t;
-
-  return id;
-}
-
 int SingleOpModel::AddInput(const TensorData& t) {
-  int id = AddTensor(t, {});
-  inputs_.push_back(id);
-  return id;
-}
-
-int SingleOpModel::AddConstInput(TensorType type,
-                                 std::initializer_list<int> data,
-                                 std::initializer_list<int> shape) {
-  int id = AddTensor(TensorData{type, shape}, data);
+  int id = AddTensor<float>(t, {});
   inputs_.push_back(id);
   return id;
 }
@@ -123,7 +45,7 @@ int SingleOpModel::AddNullInput() {
 }
 
 int SingleOpModel::AddOutput(const TensorData& t) {
-  int id = AddTensor(t, {});
+  int id = AddTensor<float>(t, {});
   outputs_.push_back(id);
   return id;
 }
diff --git a/tensorflow/contrib/lite/kernels/test_util.h b/tensorflow/contrib/lite/kernels/test_util.h
index 6fb6fe27eb..6a9fdf1112 100644
--- a/tensorflow/contrib/lite/kernels/test_util.h
+++ b/tensorflow/contrib/lite/kernels/test_util.h
@@ -116,9 +116,14 @@ class SingleOpModel {
   int AddInput(TensorType type) { return AddInput(TensorData{type}); }
   int AddInput(const TensorData& t);
 
-  // Add a Tensor containing const data and return the tensor id.
-  int AddConstInput(TensorType type, std::initializer_list<int> data,
-                    std::initializer_list<int> shape);
+  // Templated version of AddConstInput().
+  template <typename T>
+  int AddConstInput(TensorType type, std::initializer_list<T> data,
+                    std::initializer_list<int> shape) {
+    int id = AddTensor(TensorData{type, shape}, data);
+    inputs_.push_back(id);
+    return id;
+  }
 
   // Add a null input tensor (optional input) and return kOptionalTensor.
   int AddNullInput();
@@ -224,7 +229,79 @@ class SingleOpModel {
   std::unique_ptr<OpResolver> resolver_;
 
  private:
-  int AddTensor(TensorData t, std::initializer_list<int> data);
+  // TODO(gavinbelson): sync this method with
+  // //tensorflow/contrib/lite/kernels/internal/quantization_util.h?l=31
+  template <typename T>
+  std::pair<float, int32_t> QuantizationParams(float f_min, float f_max) {
+    // These are required by many quantized operations.
+    CHECK_LE(f_min, 0);
+    CHECK_GE(f_max, 0);
+    T q_min = std::numeric_limits<T>::min();
+    T q_max = std::numeric_limits<T>::max();
+    float range = q_max - q_min;
+    float scale = (f_max - f_min) / range;
+    int32_t zero_point = std::min(
+        q_max,
+        std::max(q_min, static_cast<T>(std::round(q_min - f_min / scale))));
+    return {scale, zero_point};
+  }
+
+  template <typename T>
+  int AddTensor(TensorData t, std::initializer_list<T> data) {
+    int id = tensors_.size();
+
+    // This is slightly different depending on whether we are adding a
+    // quantized or a regular tensor.
+    bool is_quantized = (t.min != 0 || t.max != 0 || t.scale != 0);
+
+    flatbuffers::Offset<QuantizationParameters> q_params = 0;
+
+    if (is_quantized) {
+      if (t.min != 0 || t.max != 0) {
+        if (t.type == TensorType_UINT8) {
+          std::tie(t.scale, t.zero_point) =
+              QuantizationParams<uint8_t>(t.min, t.max);
+        } else if (t.type == TensorType_INT32) {
+          std::tie(t.scale, t.zero_point) =
+              QuantizationParams<int32_t>(t.min, t.max);
+        } else {
+          LOG(FATAL) << "No support for the requested quantized type";
+        }
+        t.min = 0;
+        t.max = 0;
+      }
+
+      q_params = CreateQuantizationParameters(
+          builder_, /*min=*/0, /*max=*/0,
+          builder_.CreateVector<float>({t.scale}),
+          builder_.CreateVector<int64_t>({t.zero_point}));
+    }
+
+    int buffer_id = 0;
+    if (data.size()) {
+      // Initialize buffers list with empty buffer to allow for non-const
+      // tensors.
+      if (buffers_.empty()) {
+        buffers_.push_back(CreateBuffer(builder_, builder_.CreateVector({})));
+      }
+
+      // Add data as a Buffer to buffers list.
+      buffer_id = buffers_.size();
+      auto data_buffer =
+          builder_.CreateVector(reinterpret_cast<const uint8_t*>(data.begin()),
+                                sizeof(T) * data.size());
+      buffers_.push_back(CreateBuffer(builder_, data_buffer));
+    }
+
+    tensors_.push_back(CreateTensor(builder_,
+                                    builder_.CreateVector<int>(t.shape), t.type,
+                                    /*buffer=*/buffer_id,
+                                    /*name=*/0, q_params));
+
+    tensor_data_[id] = t;
+
+    return id;
+  }
 
   std::map<int, TensorData> tensor_data_;
   std::vector<int32_t> inputs_;
diff --git a/tensorflow/contrib/lite/model.cc b/tensorflow/contrib/lite/model.cc
index 590f042e21..6253570fa2 100644
--- a/tensorflow/contrib/lite/model.cc
+++ b/tensorflow/contrib/lite/model.cc
@@ -569,6 +569,9 @@ TfLiteStatus ParseOpData(const Operator* op, BuiltinOperator op_type,
     case BuiltinOperator_PAD: {
       break;
     }
+    case BuiltinOperator_PADV2: {
+      break;
+    }
     case BuiltinOperator_RESHAPE: {
       auto* params = MallocPOD<TfLiteReshapeParams>();
       if (auto* schema_params = op->builtin_options_as_ReshapeOptions()) {
diff --git a/tensorflow/contrib/lite/nnapi_delegate.cc b/tensorflow/contrib/lite/nnapi_delegate.cc
index 6eac18c4f5..b4c46917bf 100644
--- a/tensorflow/contrib/lite/nnapi_delegate.cc
+++ b/tensorflow/contrib/lite/nnapi_delegate.cc
@@ -347,6 +347,7 @@ void AddOpsAndParams(tflite::Interpreter* interpreter,
       case tflite::BuiltinOperator_L2_NORMALIZATION:
       case tflite::BuiltinOperator_LOCAL_RESPONSE_NORMALIZATION:
       case tflite::BuiltinOperator_PAD:
+      case tflite::BuiltinOperator_PADV2:
       case tflite::BuiltinOperator_RESIZE_BILINEAR:
       case tflite::BuiltinOperator_CALL:
       case tflite::BuiltinOperator_SKIP_GRAM:
diff --git a/tensorflow/contrib/lite/schema/schema.fbs b/tensorflow/contrib/lite/schema/schema.fbs
index 265b1dd3fe..84ff3b16bd 100644
--- a/tensorflow/contrib/lite/schema/schema.fbs
+++ b/tensorflow/contrib/lite/schema/schema.fbs
@@ -137,6 +137,7 @@ enum BuiltinOperator : byte {
   MINIMUM = 57,
   LESS = 58,
   NEG = 59,
+  PADV2 = 60,
 }
 
 // Options for the builtin operators.
@@ -163,6 +164,7 @@ union BuiltinOptions {
   EmbeddingLookupSparseOptions,
   MulOptions,
   PadOptions,
+  PadV2Options,
   GatherOptions,
   BatchToSpaceNDOptions,
   SpaceToBatchNDOptions,
@@ -316,6 +318,9 @@ table CallOptions {
 table PadOptions {
 }
 
+table PadV2Options {
+}
+
 table ReshapeOptions {
   new_shape:[int];
 }
diff --git a/tensorflow/contrib/lite/schema/schema_generated.h b/tensorflow/contrib/lite/schema/schema_generated.h
index c172f77aa9..8855e4ad58 100755
--- a/tensorflow/contrib/lite/schema/schema_generated.h
+++ b/tensorflow/contrib/lite/schema/schema_generated.h
@@ -88,6 +88,9 @@ struct CallOptionsT;
 struct PadOptions;
 struct PadOptionsT;
 
+struct PadV2Options;
+struct PadV2OptionsT;
+
 struct ReshapeOptions;
 struct ReshapeOptionsT;
 
@@ -276,11 +279,12 @@ enum BuiltinOperator {
   BuiltinOperator_MINIMUM = 57,
   BuiltinOperator_LESS = 58,
   BuiltinOperator_NEG = 59,
+  BuiltinOperator_PADV2 = 60,
   BuiltinOperator_MIN = BuiltinOperator_ADD,
-  BuiltinOperator_MAX = BuiltinOperator_NEG
+  BuiltinOperator_MAX = BuiltinOperator_PADV2
 };
 
-inline BuiltinOperator (&EnumValuesBuiltinOperator())[59] {
+inline BuiltinOperator (&EnumValuesBuiltinOperator())[60] {
   static BuiltinOperator values[] = {
     BuiltinOperator_ADD,
     BuiltinOperator_AVERAGE_POOL_2D,
@@ -340,7 +344,8 @@ inline BuiltinOperator (&EnumValuesBuiltinOperator())[59] {
     BuiltinOperator_ARG_MAX,
     BuiltinOperator_MINIMUM,
     BuiltinOperator_LESS,
-    BuiltinOperator_NEG
+    BuiltinOperator_NEG,
+    BuiltinOperator_PADV2
   };
   return values;
 }
@@ -407,6 +412,7 @@ inline const char **EnumNamesBuiltinOperator() {
     "MINIMUM",
     "LESS",
     "NEG",
+    "PADV2",
     nullptr
   };
   return names;
@@ -441,31 +447,32 @@ enum BuiltinOptions {
   BuiltinOptions_EmbeddingLookupSparseOptions = 20,
   BuiltinOptions_MulOptions = 21,
   BuiltinOptions_PadOptions = 22,
-  BuiltinOptions_GatherOptions = 23,
-  BuiltinOptions_BatchToSpaceNDOptions = 24,
-  BuiltinOptions_SpaceToBatchNDOptions = 25,
-  BuiltinOptions_TransposeOptions = 26,
-  BuiltinOptions_MeanOptions = 27,
-  BuiltinOptions_SubOptions = 28,
-  BuiltinOptions_DivOptions = 29,
-  BuiltinOptions_SqueezeOptions = 30,
-  BuiltinOptions_SequenceRNNOptions = 31,
-  BuiltinOptions_StridedSliceOptions = 32,
-  BuiltinOptions_ExpOptions = 33,
-  BuiltinOptions_TopKV2Options = 34,
-  BuiltinOptions_SplitOptions = 35,
-  BuiltinOptions_LogSoftmaxOptions = 36,
-  BuiltinOptions_CastOptions = 37,
-  BuiltinOptions_DequantizeOptions = 38,
-  BuiltinOptions_MaximumMinimumOptions = 39,
-  BuiltinOptions_ArgMaxOptions = 40,
-  BuiltinOptions_LessOptions = 41,
-  BuiltinOptions_NegOptions = 42,
+  BuiltinOptions_PadV2Options = 23,
+  BuiltinOptions_GatherOptions = 24,
+  BuiltinOptions_BatchToSpaceNDOptions = 25,
+  BuiltinOptions_SpaceToBatchNDOptions = 26,
+  BuiltinOptions_TransposeOptions = 27,
+  BuiltinOptions_MeanOptions = 28,
+  BuiltinOptions_SubOptions = 29,
+  BuiltinOptions_DivOptions = 30,
+  BuiltinOptions_SqueezeOptions = 31,
+  BuiltinOptions_SequenceRNNOptions = 32,
+  BuiltinOptions_StridedSliceOptions = 33,
+  BuiltinOptions_ExpOptions = 34,
+  BuiltinOptions_TopKV2Options = 35,
+  BuiltinOptions_SplitOptions = 36,
+  BuiltinOptions_LogSoftmaxOptions = 37,
+  BuiltinOptions_CastOptions = 38,
+  BuiltinOptions_DequantizeOptions = 39,
+  BuiltinOptions_MaximumMinimumOptions = 40,
+  BuiltinOptions_ArgMaxOptions = 41,
+  BuiltinOptions_LessOptions = 42,
+  BuiltinOptions_NegOptions = 43,
   BuiltinOptions_MIN = BuiltinOptions_NONE,
   BuiltinOptions_MAX = BuiltinOptions_NegOptions
 };
 
-inline BuiltinOptions (&EnumValuesBuiltinOptions())[43] {
+inline BuiltinOptions (&EnumValuesBuiltinOptions())[44] {
   static BuiltinOptions values[] = {
     BuiltinOptions_NONE,
     BuiltinOptions_Conv2DOptions,
@@ -490,6 +497,7 @@ inline BuiltinOptions (&EnumValuesBuiltinOptions())[43] {
     BuiltinOptions_EmbeddingLookupSparseOptions,
     BuiltinOptions_MulOptions,
     BuiltinOptions_PadOptions,
+    BuiltinOptions_PadV2Options,
     BuiltinOptions_GatherOptions,
     BuiltinOptions_BatchToSpaceNDOptions,
     BuiltinOptions_SpaceToBatchNDOptions,
@@ -539,6 +547,7 @@ inline const char **EnumNamesBuiltinOptions() {
     "EmbeddingLookupSparseOptions",
     "MulOptions",
     "PadOptions",
+    "PadV2Options",
     "GatherOptions",
     "BatchToSpaceNDOptions",
     "SpaceToBatchNDOptions",
@@ -661,6 +670,10 @@ template<> struct BuiltinOptionsTraits<PadOptions> {
   static const BuiltinOptions enum_value = BuiltinOptions_PadOptions;
 };
 
+template<> struct BuiltinOptionsTraits<PadV2Options> {
+  static const BuiltinOptions enum_value = BuiltinOptions_PadV2Options;
+};
+
 template<> struct BuiltinOptionsTraits<GatherOptions> {
   static const BuiltinOptions enum_value = BuiltinOptions_GatherOptions;
 };
@@ -948,6 +961,14 @@ struct BuiltinOptionsUnion {
     return type == BuiltinOptions_PadOptions ?
       reinterpret_cast<const PadOptionsT *>(value) : nullptr;
   }
+  PadV2OptionsT *AsPadV2Options() {
+    return type == BuiltinOptions_PadV2Options ?
+      reinterpret_cast<PadV2OptionsT *>(value) : nullptr;
+  }
+  const PadV2OptionsT *AsPadV2Options() const {
+    return type == BuiltinOptions_PadV2Options ?
+      reinterpret_cast<const PadV2OptionsT *>(value) : nullptr;
+  }
   GatherOptionsT *AsGatherOptions() {
     return type == BuiltinOptions_GatherOptions ?
       reinterpret_cast<GatherOptionsT *>(value) : nullptr;
@@ -2873,6 +2894,46 @@ inline flatbuffers::Offset<PadOptions> CreatePadOptions(
 
 flatbuffers::Offset<PadOptions> CreatePadOptions(flatbuffers::FlatBufferBuilder &_fbb, const PadOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
 
+struct PadV2OptionsT : public flatbuffers::NativeTable {
+  typedef PadV2Options TableType;
+  PadV2OptionsT() {
+  }
+};
+
+struct PadV2Options FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
+  typedef PadV2OptionsT NativeTableType;
+  bool Verify(flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           verifier.EndTable();
+  }
+  PadV2OptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  void UnPackTo(PadV2OptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  static flatbuffers::Offset<PadV2Options> Pack(flatbuffers::FlatBufferBuilder &_fbb, const PadV2OptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
+};
+
+struct PadV2OptionsBuilder {
+  flatbuffers::FlatBufferBuilder &fbb_;
+  flatbuffers::uoffset_t start_;
+  explicit PadV2OptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  PadV2OptionsBuilder &operator=(const PadV2OptionsBuilder &);
+  flatbuffers::Offset<PadV2Options> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = flatbuffers::Offset<PadV2Options>(end);
+    return o;
+  }
+};
+
+inline flatbuffers::Offset<PadV2Options> CreatePadV2Options(
+    flatbuffers::FlatBufferBuilder &_fbb) {
+  PadV2OptionsBuilder builder_(_fbb);
+  return builder_.Finish();
+}
+
+flatbuffers::Offset<PadV2Options> CreatePadV2Options(flatbuffers::FlatBufferBuilder &_fbb, const PadV2OptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
+
 struct ReshapeOptionsT : public flatbuffers::NativeTable {
   typedef ReshapeOptions TableType;
   std::vector<int32_t> new_shape;
@@ -4258,6 +4319,9 @@ struct Operator FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
   const PadOptions *builtin_options_as_PadOptions() const {
     return builtin_options_type() == BuiltinOptions_PadOptions ? static_cast<const PadOptions *>(builtin_options()) : nullptr;
   }
+  const PadV2Options *builtin_options_as_PadV2Options() const {
+    return builtin_options_type() == BuiltinOptions_PadV2Options ? static_cast<const PadV2Options *>(builtin_options()) : nullptr;
+  }
   const GatherOptions *builtin_options_as_GatherOptions() const {
     return builtin_options_type() == BuiltinOptions_GatherOptions ? static_cast<const GatherOptions *>(builtin_options()) : nullptr;
   }
@@ -4432,6 +4496,10 @@ template<> inline const PadOptions *Operator::builtin_options_as<PadOptions>() c
   return builtin_options_as_PadOptions();
 }
 
+template<> inline const PadV2Options *Operator::builtin_options_as<PadV2Options>() const {
+  return builtin_options_as_PadV2Options();
+}
+
 template<> inline const GatherOptions *Operator::builtin_options_as<GatherOptions>() const {
   return builtin_options_as_GatherOptions();
 }
@@ -5572,6 +5640,29 @@ inline flatbuffers::Offset<PadOptions> CreatePadOptions(flatbuffers::FlatBufferB
       _fbb);
 }
 
+inline PadV2OptionsT *PadV2Options::UnPack(const flatbuffers::resolver_function_t *_resolver) const {
+  auto _o = new PadV2OptionsT();
+  UnPackTo(_o, _resolver);
+  return _o;
+}
+
+inline void PadV2Options::UnPackTo(PadV2OptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const {
+  (void)_o;
+  (void)_resolver;
+}
+
+inline flatbuffers::Offset<PadV2Options> PadV2Options::Pack(flatbuffers::FlatBufferBuilder &_fbb, const PadV2OptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) {
+  return CreatePadV2Options(_fbb, _o, _rehasher);
+}
+
+inline flatbuffers::Offset<PadV2Options> CreatePadV2Options(flatbuffers::FlatBufferBuilder &_fbb, const PadV2OptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) {
+  (void)_rehasher;
+  (void)_o;
+  struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const PadV2OptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va;
+  return tflite::CreatePadV2Options(
+      _fbb);
+}
+
 inline ReshapeOptionsT *ReshapeOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const {
   auto _o = new ReshapeOptionsT();
   UnPackTo(_o, _resolver);
@@ -6432,6 +6523,10 @@ inline bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *ob
       auto ptr = reinterpret_cast<const PadOptions *>(obj);
       return verifier.VerifyTable(ptr);
     }
+    case BuiltinOptions_PadV2Options: {
+      auto ptr = reinterpret_cast<const PadV2Options *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
     case BuiltinOptions_GatherOptions: {
       auto ptr = reinterpret_cast<const GatherOptions *>(obj);
       return verifier.VerifyTable(ptr);
@@ -6618,6 +6713,10 @@ inline void *BuiltinOptionsUnion::UnPack(const void *obj, BuiltinOptions type, c
       auto ptr = reinterpret_cast<const PadOptions *>(obj);
       return ptr->UnPack(resolver);
     }
+    case BuiltinOptions_PadV2Options: {
+      auto ptr = reinterpret_cast<const PadV2Options *>(obj);
+      return ptr->UnPack(resolver);
+    }
     case BuiltinOptions_GatherOptions: {
       auto ptr = reinterpret_cast<const GatherOptions *>(obj);
       return ptr->UnPack(resolver);
@@ -6792,6 +6891,10 @@ inline flatbuffers::Offset<void> BuiltinOptionsUnion::Pack(flatbuffers::FlatBuff
       auto ptr = reinterpret_cast<const PadOptionsT *>(value);
       return CreatePadOptions(_fbb, ptr, _rehasher).Union();
     }
+    case BuiltinOptions_PadV2Options: {
+      auto ptr = reinterpret_cast<const PadV2OptionsT *>(value);
+      return CreatePadV2Options(_fbb, ptr, _rehasher).Union();
+    }
     case BuiltinOptions_GatherOptions: {
       auto ptr = reinterpret_cast<const GatherOptionsT *>(value);
       return CreateGatherOptions(_fbb, ptr, _rehasher).Union();
@@ -6966,6 +7069,10 @@ inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) FL
       value = new PadOptionsT(*reinterpret_cast<PadOptionsT *>(u.value));
       break;
     }
+    case BuiltinOptions_PadV2Options: {
+      value = new PadV2OptionsT(*reinterpret_cast<PadV2OptionsT *>(u.value));
+      break;
+    }
     case BuiltinOptions_GatherOptions: {
       value = new GatherOptionsT(*reinterpret_cast<GatherOptionsT *>(u.value));
       break;
@@ -7163,6 +7270,11 @@ inline void BuiltinOptionsUnion::Reset() {
       delete ptr;
       break;
     }
+    case BuiltinOptions_PadV2Options: {
+      auto ptr = reinterpret_cast<PadV2OptionsT *>(value);
+      delete ptr;
+      break;
+    }
     case BuiltinOptions_GatherOptions: {
       auto ptr = reinterpret_cast<GatherOptionsT *>(value);
       delete ptr;
diff --git a/tensorflow/contrib/lite/testing/BUILD b/tensorflow/contrib/lite/testing/BUILD
index 211de63d58..ca1390fdeb 100644
--- a/tensorflow/contrib/lite/testing/BUILD
+++ b/tensorflow/contrib/lite/testing/BUILD
@@ -45,6 +45,7 @@ gen_zipped_test_files(
         "mul.zip",
         "neg.zip",
         "pad.zip",
+        "padv2.zip",
         "relu.zip",
         "relu1.zip",
         "relu6.zip",
diff --git a/tensorflow/contrib/lite/testing/generate_examples.py b/tensorflow/contrib/lite/testing/generate_examples.py
index 7e892769bf..6fe0f491d0 100644
--- a/tensorflow/contrib/lite/testing/generate_examples.py
+++ b/tensorflow/contrib/lite/testing/generate_examples.py
@@ -1391,6 +1391,60 @@ def make_pad_tests(zip_path):
   make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs)
 
 
+def make_padv2_tests(zip_path):
+  """Make a set of tests to do padv2."""
+
+  # TODO(nupurgarg): Add test for tf.uint8.
+  test_parameters = [
+      {
+          "dtype": [tf.int32, tf.int64, tf.float32],
+          "input_shape": [[1, 1, 2, 1], [2, 1, 1, 1]],
+          "paddings": [[[0, 0], [0, 1], [2, 3], [0, 0]], [[0, 1], [0, 0],
+                                                          [0, 0], [2, 3]]],
+          "constant_paddings": [True, False],
+          "constant_values": [0, 2],
+      },
+      # Non-4D use case.
+      {
+          "dtype": [tf.int32, tf.int64, tf.float32],
+          "input_shape": [[1, 2], [0, 1, 2]],
+          "paddings": [[[0, 1], [2, 3]]],
+          "constant_paddings": [True, False],
+          "constant_values": [0, 2],
+      },
+  ]
+
+  def build_graph(parameters):
+    """Build a pad graph given `parameters`."""
+    input_tensor = tf.placeholder(
+        dtype=parameters["dtype"],
+        name="input",
+        shape=parameters["input_shape"])
+
+    # Get paddings as either a placeholder or constants.
+    if parameters["constant_paddings"]:
+      paddings = parameters["paddings"]
+      input_tensors = [input_tensor]
+    else:
+      shape = [len(parameters["paddings"]), 2]
+      paddings = tf.placeholder(dtype=tf.int32, name="padding", shape=shape)
+      input_tensors = [input_tensor, paddings]
+
+    out = tf.pad(input_tensor, paddings=paddings,
+                 constant_values=parameters["constant_values"])
+    return input_tensors, [out]
+
+  def build_inputs(parameters, sess, inputs, outputs):
+    values = [
+        create_tensor_data(parameters["dtype"], parameters["input_shape"])
+    ]
+    if not parameters["constant_paddings"]:
+      values.append(np.array(parameters["paddings"]))
+    return values, sess.run(outputs, feed_dict=dict(zip(inputs, values)))
+
+  make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs)
+
+
 def make_reshape_tests(zip_path):
   """Make a set of tests to do reshape."""
 
diff --git a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
index 0673a3bb46..96681952c9 100644
--- a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
+++ b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
@@ -54,9 +54,11 @@ std::map<string, string> kBrokenTests = {
     {R"(^\/div.*int32)", "68808744"},
     {R"(^\/sub.*int32)", "68808744"},
 
-    // Pad only supports 4D tensors.
+    // Pad and PadV2 only supports 4D tensors.
     {R"(^\/pad.*,input_shape=\[.,.\],paddings=\[\[.,.\],\[.,.\]\])",
      "70527055"},
+    {R"(^\/padv2.*,input_shape=\[.,.\],paddings=\[\[.,.\],\[.,.\]\])",
+     "70527055"},
 
     // L2Norm only supports tensors with 4D or fewer.
     {R"(^\/l2normdim=.*,epsilon=.*,input_shape=\[.,.,.,.,.*\])", "67963684"},
@@ -268,6 +270,7 @@ INSTANTIATE_TESTS(minimum)
 INSTANTIATE_TESTS(mul)
 INSTANTIATE_TESTS(neg)
 INSTANTIATE_TESTS(pad)
+INSTANTIATE_TESTS(padv2)
 // INSTANTIATE_TESTS(prelu)
 INSTANTIATE_TESTS(relu)
 INSTANTIATE_TESTS(relu1)
diff --git a/tensorflow/contrib/lite/toco/BUILD b/tensorflow/contrib/lite/toco/BUILD
index ce0a74724a..01ce0d9db2 100644
--- a/tensorflow/contrib/lite/toco/BUILD
+++ b/tensorflow/contrib/lite/toco/BUILD
@@ -280,6 +280,7 @@ cc_library(
         "graph_transformations/resolve_mean_attributes.cc",
         "graph_transformations/resolve_multiply_by_zero.cc",
         "graph_transformations/resolve_pad_attributes.cc",
+        "graph_transformations/resolve_padv2_attributes.cc",
         "graph_transformations/resolve_reorder_axes.cc",
         "graph_transformations/resolve_reshape_attributes.cc",
         "graph_transformations/resolve_slice_attributes.cc",
diff --git a/tensorflow/contrib/lite/toco/export_tensorflow.cc b/tensorflow/contrib/lite/toco/export_tensorflow.cc
index 99ccfaea64..9e899cf977 100644
--- a/tensorflow/contrib/lite/toco/export_tensorflow.cc
+++ b/tensorflow/contrib/lite/toco/export_tensorflow.cc
@@ -1492,6 +1492,37 @@ void ConvertPadOperator(const Model& model, const PadOperator& src_op,
   shape->add_dim()->set_size(2);
 }
 
+void ConvertPadV2Operator(const Model& model, const PadV2Operator& src_op,
+                          GraphDef* tensorflow_graph) {
+  auto* new_op = tensorflow_graph->add_node();
+  new_op->set_op("PadV2");
+  new_op->set_name(src_op.outputs[0]);
+  CHECK_EQ(src_op.inputs.size(), 2);
+  *new_op->add_input() = src_op.inputs[0];
+  *new_op->add_input() = src_op.inputs[1];
+  *new_op->add_input() = src_op.inputs[2];
+
+  const auto params_type = GetTensorFlowDataType(model, src_op.inputs[0]);
+  (*new_op->mutable_attr())["T"].set_type(params_type);
+
+  // Create the params tensor.
+  auto* params_op = tensorflow_graph->add_node();
+  params_op->set_op("Const");
+  params_op->set_name(src_op.inputs[1]);
+  (*params_op->mutable_attr())["dtype"].set_type(DT_INT32);
+  auto* tensor = (*params_op->mutable_attr())["value"].mutable_tensor();
+  tensor->set_dtype(DT_INT32);
+
+  CHECK_EQ(src_op.left_padding.size(), src_op.right_padding.size());
+  for (int i = 0; i < src_op.left_padding.size(); ++i) {
+    tensor->add_int_val(src_op.left_padding[i]);
+    tensor->add_int_val(src_op.right_padding[i]);
+  }
+  auto* shape = tensor->mutable_tensor_shape();
+  shape->add_dim()->set_size(src_op.left_padding.size());
+  shape->add_dim()->set_size(2);
+}
+
 void CreateSliceInput(const string& input_name, const std::vector<int>& values,
                       GraphDef* tensorflow_graph) {
   auto* params_op = tensorflow_graph->add_node();
@@ -1795,6 +1826,9 @@ void ConvertOperator(const Model& model, const Operator& src_op,
   } else if (src_op.type == OperatorType::kPad) {
     ConvertPadOperator(model, static_cast<const PadOperator&>(src_op),
                        tensorflow_graph);
+  } else if (src_op.type == OperatorType::kPadV2) {
+    ConvertPadV2Operator(model, static_cast<const PadV2Operator&>(src_op),
+                         tensorflow_graph);
   } else if (src_op.type == OperatorType::kStridedSlice) {
     ConvertStridedSliceOperator(
         model, static_cast<const StridedSliceOperator&>(src_op),
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h b/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h
index 72ffd51db4..4e3ea72182 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h
+++ b/tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h
@@ -174,6 +174,7 @@ DECLARE_GRAPH_TRANSFORMATION(UnrollBatchMatMul)
 DECLARE_GRAPH_TRANSFORMATION(ResolveSpaceToBatchNDAttributes)
 DECLARE_GRAPH_TRANSFORMATION(ResolveBatchToSpaceNDAttributes)
 DECLARE_GRAPH_TRANSFORMATION(ResolvePadAttributes)
+DECLARE_GRAPH_TRANSFORMATION(ResolvePadV2Attributes)
 DECLARE_GRAPH_TRANSFORMATION(ResolveStridedSliceAttributes)
 DECLARE_GRAPH_TRANSFORMATION(ResolveSliceAttributes)
 DECLARE_GRAPH_TRANSFORMATION(ResolveMeanAttributes)
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
index b02b02c5be..9b0e232132 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
@@ -1146,6 +1146,32 @@ void ProcessPadOperator(Model* model, PadOperator* op) {
   output_array.copy_shape(output_shape);
 }
 
+void ProcessPadV2Operator(Model* model, PadV2Operator* op) {
+  CHECK_EQ(op->inputs.size(), 3);
+  CHECK_EQ(op->outputs.size(), 1);
+
+  const auto& input_array = model->GetArray(op->inputs[0]);
+
+  // Yield until input dims have been resolved.
+  if (!input_array.has_shape()) return;
+
+  if (op->left_padding.empty()) return;
+  CHECK_EQ(op->left_padding.size(), op->right_padding.size());
+
+  auto& output_array = model->GetArray(op->outputs[0]);
+  if (output_array.has_shape()) return;
+
+  Shape output_shape = input_array.shape();
+  std::vector<int>& dims = *output_shape.mutable_dims();
+  CHECK_EQ(op->left_padding.size(), dims.size());
+
+  for (int i = 0; i < op->left_padding.size(); ++i) {
+    dims[i] += op->left_padding[i] + op->right_padding[i];
+  }
+
+  output_array.copy_shape(output_shape);
+}
+
 void ProcessRankOperator(Model* model, RankOperator* op) {
   CHECK_GE(op->inputs.size(), 1);
   CHECK_EQ(op->outputs.size(), 1);
@@ -1628,6 +1654,9 @@ bool PropagateFixedSizes::Run(Model* model, std::size_t op_index) {
     case OperatorType::kPad:
       ProcessPadOperator(model, static_cast<PadOperator*>(op));
       break;
+    case OperatorType::kPadV2:
+      ProcessPadV2Operator(model, static_cast<PadV2Operator*>(op));
+      break;
     case OperatorType::kStridedSlice:
       ProcessStridedSliceOperator(model,
                                   static_cast<StridedSliceOperator*>(op));
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc b/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc
index 347302c7a5..58e214b76b 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc
@@ -48,6 +48,7 @@ bool SupportsQuantization(const Operator& op) {
          type == OperatorType::kLogSoftmax ||
          type == OperatorType::kTensorFlowSplit || type == OperatorType::kSub ||
          type == OperatorType::kSqueeze || type == OperatorType::kPad ||
+         type == OperatorType::kPadV2 ||
          type == OperatorType::kTensorFlowReshape ||
          type == OperatorType::kTanh || type == OperatorType::kMul ||
          type == OperatorType::kSpaceToDepth ||
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_padv2_attributes.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_padv2_attributes.cc
new file mode 100644
index 0000000000..ebb023e342
--- /dev/null
+++ b/tensorflow/contrib/lite/toco/graph_transformations/resolve_padv2_attributes.cc
@@ -0,0 +1,55 @@
+/* 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.
+==============================================================================*/
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h"
+#include "tensorflow/contrib/lite/toco/model.h"
+#include "tensorflow/contrib/lite/toco/tooling_util.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace toco {
+
+bool ResolvePadV2Attributes::Run(Model* model, std::size_t op_index) {
+  const auto pad_it = model->operators.begin() + op_index;
+  auto* pad_op = pad_it->get();
+  if (pad_op->type != OperatorType::kPadV2) return false;
+
+  auto* op = static_cast<PadV2Operator*>(pad_op);
+  if (!op->left_padding.empty()) return false;
+
+  CHECK_EQ(op->inputs.size(), 3);
+  if (!IsConstantParameterArray(*model, op->inputs[1])) return false;
+
+  const auto& array = model->GetArray(op->inputs[1]);
+  if (!array.has_shape()) return false;
+
+  const std::vector<int>& dims = array.shape().dims();
+  CHECK_EQ(dims.size(), 2);
+
+  std::vector<int> buffer = array.GetBuffer<ArrayDataType::kInt32>().data;
+
+  for (int i = 0; i < dims[0]; ++i) {
+    op->left_padding.push_back(buffer[i * 2]);
+    op->right_padding.push_back(buffer[i * 2 + 1]);
+  }
+
+  // TODO(dkalenichenko): Delete the extra input?
+
+  return true;
+}
+}  // namespace toco
diff --git a/tensorflow/contrib/lite/toco/import_tensorflow.cc b/tensorflow/contrib/lite/toco/import_tensorflow.cc
index 453ff29b0d..532fcdd808 100644
--- a/tensorflow/contrib/lite/toco/import_tensorflow.cc
+++ b/tensorflow/contrib/lite/toco/import_tensorflow.cc
@@ -925,6 +925,19 @@ void ConvertPadOperator(const NodeDef& node,
   model->operators.emplace_back(op);
 }
 
+void ConvertPadV2Operator(const NodeDef& node,
+                          const TensorFlowImportFlags& tf_import_flags,
+                          Model* model) {
+  CHECK_EQ(node.op(), "PadV2");
+  CheckInputsCount(node, tf_import_flags, 3);
+  auto* op = new PadV2Operator;
+  op->inputs.push_back(node.input(0));
+  op->inputs.push_back(node.input(1));
+  op->inputs.push_back(node.input(2));
+  op->outputs.push_back(node.name());
+  model->operators.emplace_back(op);
+}
+
 void ConvertShapeOperator(const NodeDef& node,
                           const TensorFlowImportFlags& tf_import_flags,
                           Model* model) {
@@ -2169,6 +2182,8 @@ Status ImportTensorFlowNode(const tensorflow::NodeDef& node,
     ConvertMergeOperator(node, tf_import_flags, model);
   } else if (node.op() == "Pad") {
     ConvertPadOperator(node, tf_import_flags, model);
+  } else if (node.op() == "PadV2") {
+    ConvertPadV2Operator(node, tf_import_flags, model);
   } else if (node.op() == "StridedSlice") {
     ConvertStridedSliceOperator(node, tf_import_flags, model);
   } else if (node.op() == "Shape") {
diff --git a/tensorflow/contrib/lite/toco/model.h b/tensorflow/contrib/lite/toco/model.h
index 482cc71d8b..7ee7841511 100644
--- a/tensorflow/contrib/lite/toco/model.h
+++ b/tensorflow/contrib/lite/toco/model.h
@@ -82,6 +82,7 @@ enum class OperatorType {
   kStack,
   kBatchToSpaceND,
   kPad,
+  kPadV2,
   kStridedSlice,
   kSlice,
   kSqueeze,
@@ -825,6 +826,29 @@ struct PadOperator : Operator {
   std::vector<int> right_padding;
 };
 
+// PaddingV2 operator. Pads a tensor with the given constant value.
+//
+// Inputs:
+//   inputs[0]: required: the input array
+//   inputs[1]: required: the padding array
+//   inputs[2]: required: the scalar constant_values
+//
+// This operation pads input according to the paddings and constant_values you
+// specify. paddings is an integer tensor with shape [Dn, 2], where n is the
+// rank of input. For each dimension D of input, paddings[D, 0] indicates how
+// many padding values to add before the contents of input in that dimension,
+// and paddings[D, 1] indicates how many padding values to add after the
+// contents of input in that dimension. constant_values is a scalar tensor of
+// the same type as input that indicates the value to use for padding input.
+//
+// TensorFlow equivalent: PadV2
+struct PadV2Operator : Operator {
+  PadV2Operator() : Operator(OperatorType::kPadV2) {}
+
+  std::vector<int> left_padding;
+  std::vector<int> right_padding;
+};
+
 // Strided slice operator.
 //
 // Inputs:
diff --git a/tensorflow/contrib/lite/toco/tflite/operator.cc b/tensorflow/contrib/lite/toco/tflite/operator.cc
index e18ae805c0..df784a2a76 100644
--- a/tensorflow/contrib/lite/toco/tflite/operator.cc
+++ b/tensorflow/contrib/lite/toco/tflite/operator.cc
@@ -465,6 +465,21 @@ class Pad : public BuiltinOperator<PadOperator, ::tflite::PadOptions,
                    TocoOperator* op) const override {}
 };
 
+class PadV2 : public BuiltinOperator<PadV2Operator, ::tflite::PadV2Options,
+                                     ::tflite::BuiltinOptions_PadV2Options> {
+ public:
+  using BuiltinOperator::BuiltinOperator;
+
+  flatbuffers::Offset<TfLiteOptions> WriteOptions(
+      const TocoOperator& op,
+      flatbuffers::FlatBufferBuilder* builder) const override {
+    return ::tflite::CreatePadV2Options(*builder);
+  }
+
+  void ReadOptions(const TfLiteOptions& options,
+                   TocoOperator* op) const override {}
+};
+
 class Reshape
     : public BuiltinOperator<TensorFlowReshapeOperator,
                              ::tflite::ReshapeOptions,
@@ -832,6 +847,8 @@ std::vector<std::unique_ptr<BaseOperator>> BuildOperatorList() {
                                OperatorType::kMaxPool));
   ops.emplace_back(new Mul(::tflite::BuiltinOperator_MUL, OperatorType::kMul));
   ops.emplace_back(new Pad(::tflite::BuiltinOperator_PAD, OperatorType::kPad));
+  ops.emplace_back(
+      new PadV2(::tflite::BuiltinOperator_PADV2, OperatorType::kPadV2));
   ops.emplace_back(new Reshape(::tflite::BuiltinOperator_RESHAPE,
                                OperatorType::kTensorFlowReshape));
   ops.emplace_back(
diff --git a/tensorflow/contrib/lite/toco/toco_tooling.cc b/tensorflow/contrib/lite/toco/toco_tooling.cc
index 6973b22c5a..58c99051bd 100644
--- a/tensorflow/contrib/lite/toco/toco_tooling.cc
+++ b/tensorflow/contrib/lite/toco/toco_tooling.cc
@@ -106,6 +106,7 @@ void MakeGeneralGraphTransformationsSet(
   transformations->Add(new ResolveSpaceToBatchNDAttributes);
   transformations->Add(new ResolveBatchToSpaceNDAttributes);
   transformations->Add(new ResolvePadAttributes);
+  transformations->Add(new ResolvePadV2Attributes);
   transformations->Add(new ResolveStridedSliceAttributes);
   transformations->Add(new ResolveSliceAttributes);
   transformations->Add(new ResolveMeanAttributes);
diff --git a/tensorflow/contrib/lite/toco/tooling_util.cc b/tensorflow/contrib/lite/toco/tooling_util.cc
index 341d45e753..f82bb33535 100644
--- a/tensorflow/contrib/lite/toco/tooling_util.cc
+++ b/tensorflow/contrib/lite/toco/tooling_util.cc
@@ -356,6 +356,7 @@ const char* OperatorTypeName(OperatorType type) {
     HANDLE_OPERATORTYPENAME_CASE(TensorFlowMinimum)
     HANDLE_OPERATORTYPENAME_CASE(Neg)
     HANDLE_OPERATORTYPENAME_CASE(Pad)
+    HANDLE_OPERATORTYPENAME_CASE(PadV2)
     HANDLE_OPERATORTYPENAME_CASE(StridedSlice)
     HANDLE_OPERATORTYPENAME_CASE(Stack)
     HANDLE_OPERATORTYPENAME_CASE(Range)
-- 
GitLab


From f21beadedc81460f211d66c20739d38d65486002 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 4 May 2018 18:49:08 -0700
Subject: [PATCH 281/395] OVIC Benchmarker App (current without the
 functionality to bind to a CPU).

PiperOrigin-RevId: 195503895
---
 tensorflow/contrib/lite/java/BUILD            |  27 ++-
 tensorflow/contrib/lite/java/ovic/README.md   |  35 ++--
 .../java/ovic/demo/app/AndroidManifest.xml    |  48 ++++++
 .../contrib/lite/java/ovic/demo/app/BUILD     |  29 ++++
 .../ovic => demo/app}/OvicBenchmarker.java    |   6 +-
 .../demo/app/OvicBenchmarkerActivity.java     | 163 ++++++++++++++++++
 .../lite/java/ovic/demo/app/build.gradle      |  58 +++++++
 .../app/res/drawable-mdpi/ic_launcher.png     | Bin 0 -> 2381 bytes
 .../app/res/drawable-xhdpi/ic_launcher.png    | Bin 0 -> 5201 bytes
 .../app/res/drawable/start_button_color.xml   |  39 +++++
 .../demo/app/res/layout/activity_main.xml     |  54 ++++++
 .../java/ovic/demo/app/res/values/dimens.xml  |  20 +++
 .../java/ovic/demo/app/res/values/strings.xml |  22 +++
 .../contrib/lite/java/ovic/demo/build.gradle  |  23 +++
 .../lite/java/ovic/demo/gradle.properties     |  17 ++
 .../demo/gradle/wrapper/gradle-wrapper.jar    | Bin 0 -> 53636 bytes
 .../gradle/wrapper/gradle-wrapper.properties  |   6 +
 .../contrib/lite/java/ovic/demo/gradlew       | 160 +++++++++++++++++
 .../contrib/lite/java/ovic/demo/gradlew.bat   |  90 ++++++++++
 .../lite/java/ovic/demo/settings.gradle       |   1 +
 .../org/tensorflow/ovic/OvicClassifier.java   |   4 +-
 .../contrib/lite/java/ovic/src/testdata/BUILD |  19 ++
 .../tflite_ovic_testdata.BUILD                |  12 ++
 tensorflow/workspace.bzl                      |  11 ++
 24 files changed, 814 insertions(+), 30 deletions(-)
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/app/AndroidManifest.xml
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/app/BUILD
 rename tensorflow/contrib/lite/java/ovic/{src/main/java/org/tensorflow/ovic => demo/app}/OvicBenchmarker.java (97%)
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarkerActivity.java
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/app/build.gradle
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/app/res/drawable-mdpi/ic_launcher.png
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/app/res/drawable-xhdpi/ic_launcher.png
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/app/res/drawable/start_button_color.xml
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/app/res/layout/activity_main.xml
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/app/res/values/dimens.xml
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/app/res/values/strings.xml
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/build.gradle
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/gradle.properties
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/gradle/wrapper/gradle-wrapper.jar
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/gradle/wrapper/gradle-wrapper.properties
 create mode 100755 tensorflow/contrib/lite/java/ovic/demo/gradlew
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/gradlew.bat
 create mode 100644 tensorflow/contrib/lite/java/ovic/demo/settings.gradle
 create mode 100644 tensorflow/contrib/lite/java/ovic/src/testdata/BUILD
 create mode 100644 tensorflow/opensource_only/tflite_ovic_testdata.BUILD

diff --git a/tensorflow/contrib/lite/java/BUILD b/tensorflow/contrib/lite/java/BUILD
index 1dda55b8ed..1e57922603 100644
--- a/tensorflow/contrib/lite/java/BUILD
+++ b/tensorflow/contrib/lite/java/BUILD
@@ -46,12 +46,27 @@ android_library(
     ],
 )
 
-java_library(
+android_library(
     name = "ovicbenchmarkerlib",
     srcs = [
         "ovic/src/main/java/org/tensorflow/ovic/OvicClassifier.java",
         "ovic/src/main/java/org/tensorflow/ovic/OvicSingleImageResult.java",
     ],
+    manifest = "AndroidManifest.xml",
+    visibility = ["//visibility:public"],
+    deps = [
+        ":tensorflowlite",
+        "//tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite:testhelper",
+        "@org_checkerframework_qual",
+    ],
+)
+
+java_library(
+    name = "ovicbenchmarkerlib_java",
+    srcs = [
+        "ovic/src/main/java/org/tensorflow/ovic/OvicClassifier.java",
+        "ovic/src/main/java/org/tensorflow/ovic/OvicSingleImageResult.java",
+    ],
     javacopts = JAVACOPTS,
     visibility = ["//visibility:public"],
     deps = [
@@ -170,18 +185,14 @@ java_test(
     size = "medium",
     srcs = ["ovic/src/test/java/org/tensorflow/ovic/OvicClassifierTest.java"],
     data = [
-        "ovic/src/testdata/float_model.lite",
-        "ovic/src/testdata/labels.txt",
-        "ovic/src/testdata/low_res_model.lite",
-        "ovic/src/testdata/quantized_model.lite",
-        "ovic/src/testdata/test_image_128.jpg",
-        "ovic/src/testdata/test_image_224.jpg",
+        "//tensorflow/contrib/lite/java/ovic/src/testdata:labels.txt",
+        "//tensorflow/contrib/lite/java/ovic/src/testdata:ovic_testdata",
     ],
     javacopts = JAVACOPTS,
     test_class = "org.tensorflow.ovic.OvicClassifierTest",
     visibility = ["//visibility:public"],
     deps = [
-        ":ovicbenchmarkerlib",
+        ":ovicbenchmarkerlib_java",
         "@com_google_truth",
         "@junit",
     ],
diff --git a/tensorflow/contrib/lite/java/ovic/README.md b/tensorflow/contrib/lite/java/ovic/README.md
index 76c33838bf..373a50854c 100644
--- a/tensorflow/contrib/lite/java/ovic/README.md
+++ b/tensorflow/contrib/lite/java/ovic/README.md
@@ -37,7 +37,7 @@ unzip -j /tmp/ovic.zip -d tensorflow/contrib/lite/java/ovic/src/testdata/
 You can run test with Bazel as below. This helps to ensure that the installation is correct.
 
 ```sh
-bazel test --cxxopt=--std=c++11 //tensorflow/contrib/lite/java:OvicClassifierTest --test_output=all
+bazel test --cxxopt=--std=c++11 //tensorflow/contrib/lite/java:OvicClassifierTest --cxxopt=-Wno-all --test_output=all
 ```
 
 ### Test your submissions
@@ -56,28 +56,27 @@ cp /tmp/my_model.lite tensorflow/contrib/lite/java/ovic/src/testdata/
 
 The test images can be found at `tensorflow/contrib/lite/java/ovic/src/testdata/test_image_*.jpg`. You may reuse these images if your image resolutions are 128x128 or 224x224.
 
-* Add your model and test image to the BUILD rule:
+* Add your model and test image to the BUILD rule at `tensorflow/contrib/lite/java/ovic/src/testdata/BUILD`:
 
 ```JSON
-java_test(
-  name = "OvicClassifierTest",
-  size = "medium",
-  srcs = ["ovic/src/test/java/org/tensorflow/ovic/OvicClassifierTest.java"],
-  data = [
-      "ovic/src/testdata/float_model.lite",
-      "ovic/src/testdata/labels.txt",
-      "ovic/src/testdata/low_res_model.lite",
-      "ovic/src/testdata/quantized_model.lite",
-      "ovic/src/testdata/test_image_128.jpg",
-      "ovic/src/testdata/test_image_224.jpg",
-      "ovic/src/testdata/my_model.lite",        # <--- Your submission.
-      "ovic/src/testdata/my_test_image.jpg",    # <--- Your test image.
-  ],
-      ...
+filegroup(
+    name = "ovic_testdata",
+    srcs = [
+        "@tflite_ovic_testdata//:float_model.lite",
+        "@tflite_ovic_testdata//:low_res_model.lite",
+        "@tflite_ovic_testdata//:quantized_model.lite",
+        "@tflite_ovic_testdata//:test_image_128.jpg",
+        "@tflite_ovic_testdata//:test_image_224.jpg"
+        "my_model.lite",        # <--- Your submission.
+        "my_test_image.jpg",    # <--- Your test image.
+    ],
+    ...
 ```
 
 * Modify `OvicClassifierTest.java` to test your model.
 
-Change `TEST_IMAGE_PATH` to `testdata/my_test_image.jpg`. If your model runs inference in floating point, change `FLOAT_MODEL_PATH` to `testdata/my_model.lite`. If your model runs [quantized inference](https://www.tensorflow.org/performance/quantization), change `QUANTIZED_MODEL_PATH` to `testdata/my_model.lite`.
+Change `TEST_IMAGE_PATH` to `my_test_image.jpg`. Change either `FLOAT_MODEL_PATH` or `QUANTIZED_MODEL_PATH` to `my_model.lite` depending on whether your model runs inference in float or [8-bit](https://www.tensorflow.org/performance/quantization).
 
 Now you can run the bazel tests to catch any runtime issues with the submission.
+
+Note: Please make sure that your submission passes the test. If a submission fails to pass the test it will not be processed by the submission server.
diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/AndroidManifest.xml b/tensorflow/contrib/lite/java/ovic/demo/app/AndroidManifest.xml
new file mode 100644
index 0000000000..55f2961fd7
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/AndroidManifest.xml
@@ -0,0 +1,48 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+ Copyright 2018 The Android Open Source Project
+
+ 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.
+-->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+    package="ovic.demo.app"
+    android:versionCode="1"
+    android:versionName="1.0" >
+
+    <uses-sdk
+        android:minSdkVersion="19"
+        android:targetSdkVersion="21" />
+
+    <uses-permission android:name="android.permission.READ_EXTERNAL_STORAGE" />
+    <uses-permission android:name="android.permission.WRITE_EXTERNAL_STORAGE" />
+    <uses-permission android:name="android.permission.READ_PHONE_STATE" />
+
+    <application
+        android:allowBackup="true"
+        android:icon="@drawable/ic_launcher"
+        android:largeHeap="true"
+        android:label="@string/app_name">
+        <activity
+            android:name="ovic.demo.app.OvicBenchmarkerActivity"
+            android:label="@string/app_name"
+            android:screenOrientation="portrait">
+
+          <intent-filter>
+              <action android:name="android.intent.action.MAIN" />
+              <category android:name="android.intent.category.LAUNCHER" />
+          </intent-filter>
+      </activity>
+  </application>
+
+</manifest>
diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/BUILD b/tensorflow/contrib/lite/java/ovic/demo/app/BUILD
new file mode 100644
index 0000000000..be5aaa5ece
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/BUILD
@@ -0,0 +1,29 @@
+# Sample app for OVIC benchmarking.
+licenses(["notice"])  # Apache 2.0
+
+android_binary(
+    name = "ovic_benchmarker_binary",
+    srcs = [
+        "OvicBenchmarker.java",
+        "OvicBenchmarkerActivity.java",
+    ],
+    assets = [
+        "//tensorflow/contrib/lite/java/ovic/src/testdata:ovic_testdata",
+        "//tensorflow/contrib/lite/java/ovic/src/testdata:labels.txt",
+    ],
+    assets_dir = "",
+    custom_package = "ovic.demo.app",
+    manifest = "AndroidManifest.xml",
+    nocompress_extensions = [
+        ".lite",
+        ".tflite",
+    ],
+    resource_files = glob(["res/**"]),
+    tags = ["manual"],
+    deps = [
+        "//tensorflow/contrib/lite/java:ovicbenchmarkerlib",
+        "//tensorflow/contrib/lite/java:tensorflowlite",
+        "//third_party/java/android/android_sdk_linux/extras/android/compatibility/v13",
+        "//third_party/java/android/android_sdk_linux/extras/android/compatibility/v4",
+    ],
+)
diff --git a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicBenchmarker.java b/tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarker.java
similarity index 97%
rename from tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicBenchmarker.java
rename to tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarker.java
index d0102883e6..113ab74a20 100644
--- a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicBenchmarker.java
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarker.java
@@ -1,4 +1,4 @@
-/*Copyright 2018 Google LLC
+/* Copyright 2018 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.
@@ -12,7 +12,7 @@ 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.
 ==============================================================================*/
-package org.tensorflow.ovic;
+package ovic.demo.app;
 
 import android.graphics.Bitmap;
 import android.os.SystemClock;
@@ -22,6 +22,8 @@ import java.io.InputStream;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.nio.MappedByteBuffer;
+import org.tensorflow.ovic.OvicClassifier;
+import org.tensorflow.ovic.OvicSingleImageResult;
 
 /**
  * Class that benchmarks image classifier models.
diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarkerActivity.java b/tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarkerActivity.java
new file mode 100644
index 0000000000..a871b869b0
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarkerActivity.java
@@ -0,0 +1,163 @@
+/* Copyright 2018 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.
+==============================================================================*/
+package ovic.demo.app;
+
+import android.app.Activity;
+import android.content.res.AssetFileDescriptor;
+import android.content.res.AssetManager;
+import android.graphics.Bitmap;
+import android.graphics.BitmapFactory;
+import android.os.Bundle;
+import android.util.Log;
+import android.view.View;
+import android.widget.TextView;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.MappedByteBuffer;
+import java.nio.channels.FileChannel;
+import java.text.DecimalFormat;
+import org.tensorflow.ovic.OvicSingleImageResult;
+
+/** Class that benchmark image classifier models. */
+public class OvicBenchmarkerActivity extends Activity {
+  /** Tag for the {@link Log}. */
+  private static final String TAG = "OvicBenchmarkerActivity";
+
+  /** Name of the label file stored in Assets. */
+  private static final String LABEL_PATH = "labels.txt";
+
+  private static final String TEST_IMAGE_PATH = "test_image_224.jpg";
+  private static final String MODEL_PATH = "float_model.lite";
+  /**
+   * Each bottom press will launch a benchmarking experiment. The experiment stops when either the
+   * total native latency reaches WALL_TIME or the number of iterations reaches MAX_ITERATIONS,
+   * whichever comes first.
+   */
+  /** Wall time for each benchmarking experiment. */
+  private static final double WALL_TIME = 3000;
+  /** Maximum number of iterations in each benchmarking experiment. */
+  private static final int MAX_ITERATIONS = 100;
+
+  /* The model to be benchmarked. */
+  private MappedByteBuffer model = null;
+  private InputStream labelInputStream = null;
+  private OvicBenchmarker benchmarker;
+  /** Inference result of each iteration. */
+  OvicSingleImageResult iterResult = null;
+
+  private TextView textView = null;
+  // private Button startButton = null;
+  private static final DecimalFormat df2 = new DecimalFormat(".##");
+
+  @Override
+  protected void onCreate(Bundle savedInstanceState) {
+    super.onCreate(savedInstanceState);
+    setContentView(R.layout.activity_main);
+
+    // TextView used to display the progress, for information purposes only.
+    textView = (TextView) findViewById(R.id.textView);
+  }
+
+  private Bitmap loadTestBitmap() throws IOException {
+    InputStream imageStream = getAssets().open(TEST_IMAGE_PATH);
+    return BitmapFactory.decodeStream(imageStream);
+  }
+
+  public void initializeTest() throws IOException {
+    Log.i(TAG, "Initializing benchmarker.");
+    benchmarker = new OvicBenchmarker(WALL_TIME);
+    AssetManager am = getAssets();
+    AssetFileDescriptor fileDescriptor = am.openFd(MODEL_PATH);
+    FileInputStream modelInputStream = new FileInputStream(fileDescriptor.getFileDescriptor());
+    FileChannel fileChannel = modelInputStream.getChannel();
+    long startOffset = fileDescriptor.getStartOffset();
+    long declaredLength = fileDescriptor.getDeclaredLength();
+    model = fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength);
+    labelInputStream = am.open(LABEL_PATH);
+  }
+
+  public Boolean doTestIteration() throws IOException, InterruptedException {
+    if (benchmarker == null) {
+      throw new RuntimeException("Benchmarker has not been initialized.");
+    }
+    if (benchmarker.shouldStop()) {
+      return false;
+    }
+    if (!benchmarker.readyToTest()) {
+      Log.i(TAG, "getting ready to test.");
+      benchmarker.getReadyToTest(labelInputStream, model);
+      if (!benchmarker.readyToTest()) {
+        throw new RuntimeException("Failed to get the benchmarker ready.");
+      }
+    }
+    Log.i(TAG, "Going to do test iter.");
+    // Start testing.
+    Bitmap testImageBitmap = loadTestBitmap();
+    iterResult = benchmarker.doTestIteration(testImageBitmap);
+    testImageBitmap.recycle();
+    if (iterResult == null) {
+      throw new RuntimeException("Inference failed to produce a result.");
+    }
+    Log.i(TAG, iterResult.toString());
+    return true;
+  }
+
+  public void startPressed(View view) throws IOException {
+    Log.i(TAG, "Start pressed");
+    try {
+      initializeTest();
+    } catch (IOException e) {
+      Log.e(TAG, "Can't initialize benchmarker.", e);
+      throw e;
+    }
+    Log.i(TAG, "Successfully initialized benchmarker.");
+    int testIter = 0;
+    Boolean iterSuccess = false;
+    double totalLatency = 0.0f;
+    while (testIter < MAX_ITERATIONS) {
+      try {
+        iterSuccess = doTestIteration();
+      } catch (IOException e) {
+        Log.e(TAG, "Error during iteration " + testIter);
+        throw e;
+      } catch (InterruptedException e) {
+        Log.e(TAG, "Interrupted at iteration " + testIter);
+      }
+      if (!iterSuccess) {
+        break;
+      }
+      testIter++;
+      totalLatency += (double) iterResult.latency;
+    }
+    ;
+    Log.i(TAG, "Benchmarking finished");
+
+    if (textView != null) {
+      if (testIter > 0) {
+        textView
+            .setText(
+                MODEL_PATH
+                    + ": Average latency="
+                    + df2.format(totalLatency / testIter)
+                    + "ms after "
+                    + testIter
+                    + " runs.");
+      } else {
+        textView.setText("Benchmarker failed to run on more than one images.");
+      }
+    }
+  }
+}
diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/build.gradle b/tensorflow/contrib/lite/java/ovic/demo/app/build.gradle
new file mode 100644
index 0000000000..c5d19bad89
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/build.gradle
@@ -0,0 +1,58 @@
+apply plugin: 'com.android.application'
+
+android {
+    compileSdkVersion 26
+    buildToolsVersion "26.0.1"
+    defaultConfig {
+        applicationId "android.example.com.ovicbenchmarker"
+        minSdkVersion 15
+        targetSdkVersion 26
+        versionCode 1
+        versionName "1.0"
+        testInstrumentationRunner "android.support.test.runner.AndroidJUnitRunner"
+
+        // Remove this block.
+        jackOptions {
+            enabled true
+        }
+    }
+    lintOptions {
+        abortOnError false
+    }
+    buildTypes {
+        release {
+            minifyEnabled false
+            proguardFiles getDefaultProguardFile('proguard-android.txt'), 'proguard-rules.pro'
+        }
+    }
+    aaptOptions {
+        noCompress "lite", "tflite"
+    }
+
+    compileOptions {
+        sourceCompatibility JavaVersion.VERSION_1_8
+        targetCompatibility JavaVersion.VERSION_1_8
+    }
+}
+
+repositories {
+    maven {
+        url 'https://google.bintray.com/tensorflow'
+    }
+}
+
+dependencies {
+    compile fileTree(dir: 'libs', include: ['*.jar'])
+    androidTestCompile('com.android.support.test.espresso:espresso-core:2.2.2', {
+        exclude group: 'com.android.support', module: 'support-annotations'
+    })
+    compile 'com.android.support:appcompat-v7:25.2.0'
+    compile 'com.android.support.constraint:constraint-layout:1.0.2'
+    compile 'com.android.support:design:25.2.0'
+    compile 'com.android.support:support-annotations:25.3.1'
+    compile 'com.android.support:support-v13:25.2.0'
+
+    compile 'org.tensorflow:tensorflow-lite:+'
+
+    testCompile 'junit:junit:4.12'
+}
diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/res/drawable-mdpi/ic_launcher.png b/tensorflow/contrib/lite/java/ovic/demo/app/res/drawable-mdpi/ic_launcher.png
new file mode 100644
index 0000000000000000000000000000000000000000..715d1b6d69c0f4dc4d1ae58c8262c22856b20f43
GIT binary patch
literal 2381
zcmeAS@N?(olHy`uVBq!ia0y~yU@!n-4kiW$h8qca7Z?~gMLb;`Ln;{eMkbf%c#6Ee
zZ@%qp`t5DG${DWunXbJmxwp58o{w8^{54NI%C<PXcE`#)t8SI<IG=gz&#HN!N(|?G
zoqc$g;c|zBcdssQ|0-?OwstSq`V*qtmvnucdGqDfw|k5nb{4sbPuO_He;uE!&5h5`
z%Z_cIZDkjK|KH19+xiX7w$$(1`T2D08u|Y3`Fr&3PMaFWaK65BCge5a*|X<eiw`n7
zKbJhW?d;t6HI;uAkKMX1zc%~uj~!`ia(};g%P#+*dF~qdcirbQ{s=F9nRJ>V^>k+6
z#JU~jynV+#EsB0JKWN}d+HF~MMj}YicAe0LJ9p;jX6-PlNT|K7p3S)Q=gO%U0}>>r
z&dy1%xvsWf_<VafXGYO3p6t7a|1Dh@Vh}m={*{lcI+^O9f6JbEzuj=(%^UwDCxnVb
z^-FKJI(;{<Gb52%n>l38smoLO&C6=qYG1!i_I@6&{>$8U%AEJt_R4JBD_PEe(q?hc
z`4cm|6*iPy4(-qDm1{nE_}FXx=9znA_xtT=OM9L+oqL7B%xcBGjHYYXhB3ZlS8n~y
z@s@enhFfyRw;sDipGoLVd|~(@FM5VaR`aI1<mayc=K0QzPe?p#7a(solX;mcqwl#X
z)>FPWwq*%SxO}Yde14BitAW<?s71SuwwtbEPbzjfY-;HLC4bF+=@6csdW;;aTJJu1
zv-N6*K$JlB8A06(4>pMMF1XI3_OA2XmE7MC_pQ616Zki&VL`{;HS-%9x}G_lpTD6e
z)IpK?R?CW8qTjb&*uN;L$u?H#jjy<HFXPiwFTI;;e{@uJHrv*B%EUAslL*UXx)iFM
zylVad#Z_0o2X6nyYI2!D;QW!O`~N(9m$9u4-PCxEWtv`T0+ZS%E1&gG%gwgbyV<f%
zJSw^U3X|6p<>U9(Ebn4xY+G8%m3f^*o+){SJmbw(3@?2JzplOb*tDzKs!q~z50^mB
z&Rg7ehm3+<g7<5x$$RPX%y@Ba`y^xc)e|0`SAFaGP+0cW-`Npb2OLG;wS1LTV427L
z=dOd-g*R#cTQt&hQ%+n^xPQj(@Ljd<T{fQ99vmmwKE*DZ;w7@EvFEs;x&G5~wi9xR
zleJW~J$q^G@N7es_TTIq;*2bk5woH@71!+1xbeB-O;Rmid*!s{-V<X4mg@;E+T!?f
zqq+106P62^=fWzkGAV3(Cb_JGVdCr-zP_zL`veZRR9C2fZZEu^V7!ZIrmm9Kf!y$W
zi={aC@H?_+XFjM)h*scad1HCcZ%6s&y?ZpDe|aM^N!IDa!g>4uXkJ&!;mceyQ}v*V
zz?#ld=cotK5w6^+d)RmVS5RvcSJ{-twEej6(Une5Z8kP{zklE5H+kkS^|%g`>$|zK
zlx8)$I6QrjZQ>lj$gH%;dBTGk5@#6>TAWCn8T9GO0WlrNZCP&W3f|$DZuxldA9>E9
zu&1HVX@$uL_J)ddhcoYw_}T5tTk+#i$h<1K1$#<Q&UY-F#i==qJ95#jw%v-|!M2=T
zT&Xo%;_hl^{a}m?{H4;%v{9_5cGs-tri|uW|Ng%E+pVH%nslJxMq5L`iUU78bq?gt
zDKzd~S+!>xldxojtb}{lF{W1yXZ~b{+pc!{v!X&j_Ui`4R6SqOgvx8L?iv`J+`i!;
zqi4|VV?V?uElX6^Qn27ll6@m`Yo#QI!Pn<MgZxyKIXLx#&#QiTz~@`0eadz*i@>q4
zU*;NLGCa09xT{KY?D%}Tw&30w@0(|pYqQ(e?CtpRS-bLyoMLF|4j~KK-t?qND*cfk
z6osOu{XFt{qQ;*5r~H3b-m0nps48pH8s2fM{}1!~V-t%0)mJ+7dpNWyHsz^s8%lW<
zszfKbRXDy|vz_@$##+~AA?d5ij4k3SJTo+PqrK)n%USH^bH&_vvvEqQ*4s4ixB}BV
zB`dEKymq;i+`mK2y4q@a>FV@V`68!V1GOB3pWN|#%6#T>!Pj)*>mSUou2a6r5Gl#A
zy83B83)ik&tL4p$j^_mJ&7QOP&BdQ*%jP>Rc<3WhalY*x*V&}S(dUA{^1QTQ(!DPI
z+2f1M)ad0~__(*!#u{zAx3=S%QOa%3?*db;4((d{ZenYWe&MH;Pee^GOxW7KhNmsG
z|7ZCvH~EdK=jA@E+b^^8f6|Y4J8I|XRUMK_XPe}>wDbP*H`0Nd<9O@@X0_dJ^;dPf
z*RW}M{-dYP_ts8#=I>$rJ9l1le_Q<2w6y<MuRn=>sJWt5>i;3hcvgWW;!HD-zuCRz
z(A{}@YTDh-g1O6g<U07hpK563a{h-Dpa0FLUKQ*qU*>iCZLADg^H?J$Yi;j+y_#I^
z==M^k-Tq&zUij_)@UiwC?>83fS-eKB{kzp)H2Un?v|MqC%U#~;r8iEp9=XcuaN$<=
z&D;kIxHx6|TQ3zFrGLzPllc6}mnHUT-&LDm)+{m!^?Ew_Wy}$=O)ABmF`D!Bmo0kf
zQMhTwspVZyW0a&z+t=>fp+B$JS#HJ|F`=753xk#~ezxY>`}vD6haO!iIMHFp3dQRB
z@A)Tm{<m&ld`kJ6^D1MOt6B~M*NbPEv6|k~=YQv(kfAHZveWj>{KCXl6)Vn6c023!
zdw2bEGd6_{zZY#zc)!-^amUQQZ);cY`^H}0qJHqK^z=&{ohmygG?ux0KS}9*;qZo2
z>QI@CVe;i$av_X?6W1sfISHSgykV^r*Bi%Ia&{L!AD@-=aq-`VV{e7tN-R0~c~7kA
zj1TG=Cs(YUf8qI4`;HR{Z3!n!eGmFysB`$9{<2LpNbf`uuUZcKJEn5h`7EDZ*F<zp
zJEOE`p@Lht#45#q-p4L~)c&yM>9#L>JlQ7(_pS)|tem41xVmOi!ZGf)E(7K;mc0t5
zi@IC`b;>?W=;QnPFLbA8FQbC?o+<YZtPQ*N#l_0mY+X(uKQ}WI^Im()4V+?Z{*L+#
zyfy|~WzW192-mKC^{X}dMSSyDp{x4h0xu;>Gv_Uwpt>mQn(}&{J?oEV-1B9Ri7?k{
zy=s2$pz{_+O{vI^wKpbCxSPty#QNpS)Hh<Q=3M);WrBYCM`dm%^Qt1z+O@G<l3zkE
zN?91*vR|<}Aw~D!M^?{^B0-IAbDzF*ST3`z^RvMMwO8C54()MJU&kx&*4B1jNOr4)
zz|Xymd`;5DJw4)=UYoo39g@H3?!LrfC&#qQ>vm{wDGtz_uqo5^y#4e$nV&!3FsWf+
PU|{fc^>bP0l+XkKU=w}s

literal 0
HcmV?d00001

diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/res/drawable-xhdpi/ic_launcher.png b/tensorflow/contrib/lite/java/ovic/demo/app/res/drawable-xhdpi/ic_launcher.png
new file mode 100644
index 0000000000000000000000000000000000000000..9beff0885fd4c8c65ea30c99c838370dcd745f3c
GIT binary patch
literal 5201
zcmeAS@N?(olHy`uVBq!ia0y~yU`POA4kiW$hTLBd@);OJBs^UlLn;{W&Ml3OJ}q|a
z_sw_vkKc`5&A#%Y)7@DAy|x(&NAKR9tH7zp_^9jk-u)%l_1b;P%vv3OpE*2PXm(}#
zOY3a0+yA65)So`>&F%g9b9eD(Ze!!m!IwWzzWh1ZW?uB1r{BF#cmL=8vcH!tqcHQJ
zVXEw=uNlW)ovtnpx~;VMb)L-%tNPVnZY({-yu7&i-1GN)<Ae&=KRRQ-HGICAz0Nzo
z{EOk=qqki1RJLGx8O?Fw(50r-LpM@0PQLcq6J7W#<=ndsTjc*&{kmCp^>_WMd!H@K
zR@>W6`1JSV8d-C%{$<--x4y_eHpjW+_B8g4%*vBinX6=V>H<n1x_mpQAkbOg$8q6J
z$A5=^6`ZR4WvVqb7JsK_AAM*!_f_C?g;_hkG92o7(UbdX*>8^1d0*8p-8v%hqyPGD
z8wO32Eya5{8*fQ57zXxmo2BM|GS^Ao9TPJ}`6!3yuUCFYUp||-YUR$(v)lq=TWZ%l
z^!ur~tb5M_U!8=-ZSIWC!uxkbR&B{#BJ}fa34>&e@~UIj9=#H?zD`{xZeY7^pGm^i
zL`HK~2NoSu=GGetF15-F1rM~u$$R~_j1|tFZNe&{@VP~C-NW|_j+}2bvtV%kH0k*#
z=i6@CoqL!V1QzzDt?66fcgv(M@6f?*Z~cGjZU6l8%AF@|OfF6nt_tqge`;-(#XGy_
zef@OPH(ozptY1E#&GxU%XTdo?RP8H$Bs^;vzSJMtyIXAcF`n-$>??zck3NyRt<a#B
z7Q5=$?wZw1B@wGyJTg9APygIhSec;_G5uRg*th3`L1sMDuk%g~;T4(k_1Tl`$-=Yc
zZP)Cp*r~xYn<=hB+clxGN4b(`TJ>E<rQWKb;>z?x?lnvI7aGi%cj~5aWcAvB=}v2|
zB)^l`^0mj}SNog2%-h~E&R%za)$7e!zipk5=U0^7xM-jA{`KaJt;eptYGex9yxO(>
z!_-NKKZPgwr@Vh<WU^k~VaAh>`%nA-(tfDZR`GrQEbXhzt+mti_7)dhxOcc&&f9(a
z&b_zSJY2qA(UyD4-f4acm#x_jtS;SRw)3kzOF-8}L4lv2r00FqUiscg^1U?s3A4<7
zv1M~=KOFC`sWAD>zisd33p{!rF8xb`dpas!*t~!1aAbnDg5HIt3oIhEe{N=4ykX1y
zTW+uK891zyJbri;ciqX)2d~HrN6Sr|$ZP%U*P7T<0)if5Nv%GNVV?fSe*U>I_oKYh
zDx>FRCZEsC1^<?2=3wIv?z%fsa`{2OJ6hI$C;AItZuPjZ=6XfiiCvWqbKXsV8(JwT
zWtH=~mYLz;%co_lHm+Y2bMBdew}pY(w1a!h(-!P}k$*Pf;><k?H;cafyZY$xhLCL5
z4v&)fQ`?Km+Gfr@cQ^Rr{C57N<eA^u__r;o-Se^TsikJaB;$uOenl5^8=qp7D(Bmr
zCdqfzSXyGuwvD&!JPVji3d(qL<QbRu_}0oZ^S@m!ps+)fXS2*r#Sgo~E8UFm9e6S$
zb%8s(o#JBKy!cJ|2cNSFO70hQ&|+M%##}mbYO`O+U0x<j&2^E6R=v7Cz3rTiY8)0b
zel`2|T<HHjHT~hwKg=gy$l1=^e0q7&!x;I{_$xDSg}Jt0Ow7=FeflMrfZH6w7PVU4
zZSFSRi>~sB{A3jPbKvqb_M)1Dy_1w37_68Y&tF#CG;_w=hJ-20I@P|ivThW3_|-Cg
z0=K}gN5;zzES`UoH|tvX<D{d7T<dp+S56ZVkzsBMj8S&H({5*Oy;fIB=#sgEK}*22
zq^C^QOl3uRGYb=HF7Mqfn$DEeb3;*wA!chtM}%O@!3A4yn>*ZHR`$kTUOntqi7(G4
z<>Icdhf`B7rOasj^w-1idEWvBGq+CNe$F0OhJ(V#uHFCi({V@r-1(O2DR)h`br={0
zbcJ1+D>mte>SS>aLCaYOmx?o*q`Y=5KQdov-pA<29!uk2KRNp1apRTu3-_5k{(5({
zL=JC<ahP3Qar2w`m#qr#^{+hj)=Px3<<XCxs!7rFH8?m*)p*l7B#)TgdG@jEt?JrV
ze#!D)fw&Fri_a;pFq%|)CAjHLuZGmm&kMiIXZ_Xqv{j<O?{+vlhu>qTgLZK$75``6
z*?)_-_e`{-UF4<7Cub_vB<#~;otgBPZTmgb<oS^^p4HS=R{s3&%=;=nZDIM@$!s1*
z`wHJ0Y!v?3`~8SeL#6897}*7VAKjeTHvU+Wxk$uVqNV%FR);I|&;8|Ye6A9vw%I-A
z*_@5u25Svp9*=U)=W5}Q?aAJ8Hg7FUM|`2lsmn=z+jH0Wz0{d0@0c7E@JRf~{S6<v
zyZ?ozDx_}H<PJF~)!K1{Ls>iKnP0|)sSG!loiTXD^fxKA#kk{q)C-O^JTX#EBD21P
z{tcMPz~{)-tFiW|kYKR@8@GWHbFR#~2+j%DE*xT9*&x{<`e?DJbta##fkus>qUrL4
zH!M>Plg}llZZ(bgvbCk}NUx8>gO$Rw-mPWPh@Y)u_I?Kc+LfUz-@H(Je&>?-&U6`Z
zZa$Bt0sX7>4?ow^;+S#Snq4G2(ydx2G~%B4Z@q^c!quz-VH0LKFs`VPDcoT6_fhPR
z<JYIRTR&zNEShmgDOT&5|4bRj&05B-GqTQ!%xL(&<?(XC>i^meZ-UcwCYv|>P^c-s
z8?e&o@1g^VtSnE?-mpHiFr@8b_Pc;GjW_d3wHCkFvdm1xNkc&1G(jnzsU?c*hj|@;
zjmrV9KjMN4?`n=7I-tEe_UonW12Zj_=bp+^=v*>oXSYP6+tJ_;D(e{zWi-z`z)<M0
z?u?ApvbC$%^%{ll@VhW^8}};5$cbFjnyzYnNM(^Z($b{r%ra?*9s3H0**aTZy?^*X
zu}ZXq|6up_v*G5p6BwLYzpK_C?%c4*>Pu>K^_BAq(Mc;;eK=jS)x1GZ?9FsOn{69Q
zL$Wy-ZY}<2ve|j5!)@Kki-Ofxe40LI<B^w=v5UD3l&U_I=P|S#6==C~=6S@mGVRao
z3U>~^-6VXE;gK@O>Q(IAf45%nsVd5IifQ`lfA;G2{nHQee=@3Pbn^IHo-Qlf<|cIh
z?Y;0?g|F-{w<Jqn4QP@HU35!h#fxv}ow8?aIWLs|YPM&5{3%OD6?aXpEjedQ7d6CI
z1XTJx*x@{B-;!sCm<100emTdMsX@T?+@3$Whc8&O&D`ZLT)q2^O8nnjnpfX{nD;Nq
z<<5ikl1_?Buh)IP+GOmyk|RYr@Xp6KvPliM{<3s1ZdQK&h*MHBxbm)IfKXrLM{b3p
zTGof!^WtQr56f3F91?Lf<=SFs^EbTrKSN-y>iHv9ov-8PYsAf8WV3(O9L6j6C3u(n
zob39MvnX%F5oUp_z3xW;Kk!|sKEAm;Nd4WMjh4&)N_?8U=F&Z#ZP~5S3*64@TlWh5
zG`0Prw(_fC=5z&-MgMLyPpCcB<o)<RQ@eTRnf15rXNYNRdAmMR$%kPLFMGi2s^0cD
zK@0+m>sDN07Wl*|{B3zTr`4pN){Gy2`HI8|&Yk3SL}jx3#678pHmrNG;81}SQ$s@Q
zk{hD>W|eySE(`)@C-$7XEL+|ubMnVruFNdkk4XZTqzjJg?|Qe)yy1q{jPu_%y74=R
zn8j#v30D}-XiThrU**4jk-qqPMUC{T$L6;*UI_?E#0Mw5WX#A*6f^RWG4fRVocpUi
zap#2n;SYB2G&yxNir?YMnkz*dGICF6Pf6UH${-lsRcu&e{_b_9S^S;3Y788%Pqp4g
zvjwc*DXc1#@;M~ovD<v+BdKm#<*b5?tlws=+RnM)l3bJi#YsC0KX&my3bSH=eq#25
z3H=Sf*E+W6vCV1VuCeiMdsul?V^{i|5NFOeCa#O;?C!dBZso;W=k6r=mPLmD?6|x#
z;nL)kCv69}?)=^==zM&}?eBZzu646*jB^xVY5d5gkn&R1^l@&$b%h_NCNh4VZ@zEe
zl&V*cF7Iu>`kX!Ha-DSRw-=vh@?T$e?A?qt%qQpWdv7C{#hJCBAeSX|PT@Wlk5!px
zg~C5F?ySDHI)mr?YX7CSGaVSj<{F-Ke5JQm{o0I;{F1xdbK7<<efsohX7}kgGYoj8
zs&yv|2hFlK&XqKI^`-3cBK8F{K1+UglljGX`N5GpHoCSsdW(xM*`&PL>gTI_W~ZCa
z-Nk;lmEJ$Gj#>Qo@m;s=**8Dy@^#-%+Qch4=dfWU`+UYNIpsYw)b@4>HSR9h_`ayG
zq4l4A!-iLnngX`?6x`%$=u^LI>gTsJ_V0IL=S}Qv>1w~{rL?#nX>D_7KW=v2_lD{6
zMRoRbYu;3R>D$n<nOXR<UiRa+M;0vJxG-^IwUFuAU41w1mhZQnTk+z;<am8iwcB@-
z>sSA?SlDMFy8Cm1(~rM%J}(ccc3S)`tbex8;-%uDj_Qw{4n0<S#XIev-qn4__q(~g
z-xWSnY^vn6ODi|`Oq~<<R$um7b$5fxKZgLtX_v%yeoC2YG^O~b;-l@VCij$EKQ2o0
zuiyFU-@+RU{_@-17c6zMw~dZ2Ra5%9?3%@eb)oNj@7%u=zU%)(bIo1zUTXerUUK)N
z?5s6=Hb|#NuDhA}H={JpJ=gvAWZCcKTV|~-OTK#Y#RIF__fwD0x6N8Le~#2%!Jj`r
z)@&`lHZ`>1OPK4#%NOsPTrkU%Kc{b1wozsE-nS>uvG0!ZU3RlFJE!c%;!875f7~r(
za`qmdr@^u(CTlzn9CV(#Ht64^E1x7@>*{$e-*R#`OY3UAE&X4=y?XtAsf<s?>Q4WA
z`(A%Mmo0gG{j%2i9yYhHOs<cTUH4VvVM)}%i|VF7%Y_bQeOVT6<{?~b_V}Jke;Qk!
zTC3Q%%>30&S-L7GKfn4?%Q;tSUgoOZwR=@k1LBKpe}0(yMR@t5I-c`>8Cji{(|=Ad
zPkIrWxxc>dpZ)#2_wFT?R==)m_<!%+YPK$cLeYj@zEYdG3QLvPBktt}sGL6TFeQBM
zu7C;C=GH$c3hXK_JN5F?cHYvc-qPh)-FK_oGAtFB>0PrvVCmb?ty?R6u5Q_Q+~CU1
zZ+ASHFRXpe&}tHqBiCTWXLq&CmVHHl=v<+a_4`VSt_23j2+mu|w4nQ@|FZd+{ywkf
zr*B$3iMwyUzR#Q}-)lS7)aHKU@V1m+8|T1kS-bYfh8=GhUnO5L+|kwbRXu0_&%<lq
ztp0A~FW>WLb%95|aoD`hyQR9XfAvn;;LXbEmV9_+3x_O&qT^4&iN0%p9qf60u{!Vh
zy&Y2oYL*-MUzV!Xd)&x!;p~wIhvMF@-Lc?Ty3mvfbq;r*GOpP8?%7tQ(!i^&&a0FI
zk8ba`{b!^TeD#V<t=;R3#tnBa&ip;uy879jSn=fY=U;gLMAy_<WxpyCm~xeIQkKnw
z)FWSy7)<?dw{FH937xEv<=53y%$96>(JnU4_F(hXTakC>1l<VCoHZ|Y*4Ax&Cf2t8
zI(+@6-=p5zNSK?5ny63DdOf$8JFj|v+LG#n%1y$t&85kC%eU=6XJmZ(IZuxXkCm)V
zvT1p#R9pIUnci5-pR<zQdU)S{vEb^p<rlwya1(fO*__96&K_6$sQXsS;vG8=zxQa@
zIHWYKfGcnhQ%cFp{<OcEkxGt=%)UO`V!oXD_<!mm=cJE{c2DkjhCDsx$XmJcOIG1p
zfu#$gIb`SP-<@H>ClJ4`!J@_bYhdR2tJxYZEgmucJPW%Q<%%{~ENuBEtk3m>t90eQ
z^-2!qL5VLva|xdES=^%A$XCXE(ZS(J*1>~}Dihb-wwuCNwAt}br_JfaiRV__eDb1T
zk8|VR_0O2!uis*q)L@aOacQY%NvFi(V9VW}f6R2PSoCD(Kbl?}6?)3NAg+jUQp58}
zdhfqfKA+5IIb*GF2OrP;T|XOd7<fNxIy29Tb^rS{k4}GlQS|QUp}J!3wBUu!77ctR
zny<fe1~^@=?OO7OjU$_vUFG?@c_|M&9+~@}IDBx;95GX&Ugeu<f(}imn)aXeOUQnb
z;36-)>;DGPM%VUN?ze<&H(T#Zwb%LTzIJ1(?_6;U>8QZpiW_bkcwNbG-oe5m<+3j{
z-nnt(89TuXQs-A~nSSlooffXlGc(TQAN^jK#JPrLz4YsIZ$Gq5X+4oMwYkIf`hAh<
z3r?_d9GfL(a&Xb_p4E4}RNfyI&0{dRBlYNkxLne7``V3iA@in3&e`jmHr=;K((B-p
z<Q9#Io3{2{+u;3rd9d=geLs%v4r=!PxwkiOn%w02uC2cRGVLqZ#P7X3!=&i=RjFT6
z7Y=v@Y+A8^SAHt%9s6fLX4od5m@gE2!Ncy_(soIS)bb<FFPbLCzjf|(@7S5i$Me>|
z+-Hi@%}n{{kt{J&_6i?qJM`wq&9a#%6|C=c#va$JNcnK*s^E>BgV#Qat*a^W?fG^|
zBWL--PAC7K*KZD)i2c3)lCP@Q*+tyq$5IQnO2xEBgQNR?-Pie0d$0BT1gCkWEFPO~
zKKU3Uk;E*zy(w*%htkw<>7Si6_BAeTv1i=+=pD~3b)f}9ckR~g^m!=sP43x)(BLdZ
zy{2`EAuF|a8*TQTxuw{^>Ah;$?r+Do=t!S<vE$@&&&J(f;_LRkh~2lWZvKhtwB1=n
z*Y+Mfn>YQXzw31VZ(#~vjIwfH{GYyS>{8pu?_-`2V{wGztdv=_dBQX2fKMxa3GJ46
zP<P29|5C^vX2An#4V+Vt9GH8Od)de5Xa0Oid+EWq*TGY0-y9jM04epYu1kNiOlM^^
z-^j!~;mM4hbrY^nby+DGXn$$T#S2UR_Pkf-oucEXz$lb0Gxc*tcHSG$*f7c0w~ov^
z`=C!?u0^H=gGux={{8~{IO8aGRi6Lnr5Zdx-OPW&ui)|M<oxV~6+)ZlUP&yhEcE{O
zQpzdi%eurJK35Cc&lz73)DPEC<+3|J-NyOrE8QcHV(u-vtHwLoLSNx`rT4!df~(wD
zTuXnq=%>X+;VoiEO=7B-%wzffeCLXfOhzvoOOm_uB%g?-EQwjMGqA{Hvn^j*^AqQS
z@I-A!6$SqQhaG`)H%{d)-C-{!H!03x^YJV9mvH)eWv7(6GZ{6lW)pm)SJpeFm!s}+
c`fvNNxE*3^9;}ICU|?YIboFyt=akR{03$&E>i_@%

literal 0
HcmV?d00001

diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/res/drawable/start_button_color.xml b/tensorflow/contrib/lite/java/ovic/demo/app/res/drawable/start_button_color.xml
new file mode 100644
index 0000000000..93f5c6a016
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/res/drawable/start_button_color.xml
@@ -0,0 +1,39 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+ Copyright 2018 The Android Open Source Project
+
+ 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.
+-->
+<selector
+  xmlns:android="http://schemas.android.com/apk/res/android">
+  <item
+    android:state_enabled="false">
+    <shape android:shape="rectangle">
+      <solid android:color="#808080"/>
+    </shape>
+  </item>
+  <item
+    android:state_enabled="true"
+    android:state_pressed="true">
+    <shape android:shape="rectangle">
+      <solid android:color="#44ff44"/>
+    </shape>
+  </item>
+  <item
+    android:state_enabled="true"
+    android:state_pressed="false">
+    <shape android:shape="rectangle"  >
+      <solid android:color="#227f22"/>
+    </shape>
+  </item>
+</selector>
diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/res/layout/activity_main.xml b/tensorflow/contrib/lite/java/ovic/demo/app/res/layout/activity_main.xml
new file mode 100644
index 0000000000..e9d83bae54
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/res/layout/activity_main.xml
@@ -0,0 +1,54 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+ Copyright 2018 The Android Open Source Project
+
+ 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.
+-->
+<RelativeLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:tools="http://schemas.android.com/tools"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent"
+    android:paddingBottom="@dimen/activity_vertical_margin"
+    android:paddingLeft="@dimen/activity_horizontal_margin"
+    android:paddingRight="@dimen/activity_horizontal_margin"
+    android:paddingTop="@dimen/activity_vertical_margin"
+    tools:context="ovic.demo.app.OvicBenchmarkerActivity">
+
+  <TextView
+    android:layout_width="wrap_content"
+    android:layout_height="wrap_content"
+    android:text="@string/initial_status_msg"
+    android:id="@+id/textView"
+    android:layout_above="@+id/button_start"
+    android:layout_alignParentTop="true"/>
+
+  <Button
+    android:layout_width="wrap_content"
+    android:layout_height="wrap_content"
+    android:text="@string/start_label"
+    android:id="@id/button_start"
+    android:layout_alignParentBottom="true"
+    android:layout_alignParentLeft="true"
+    android:background="@drawable/start_button_color"
+    android:padding="10dp"
+    android:layout_marginRight="30dp"
+    android:layout_marginLeft="100dp"
+    android:layout_marginTop="10dp"
+    android:foreground="#000000"
+    android:textColor="#ffffff"
+    android:enabled="true"
+    style="?android:attr/buttonBarButtonStyle"
+    android:onClick="startPressed"/>
+
+</RelativeLayout>
diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/res/values/dimens.xml b/tensorflow/contrib/lite/java/ovic/demo/app/res/values/dimens.xml
new file mode 100644
index 0000000000..250b581430
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/res/values/dimens.xml
@@ -0,0 +1,20 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+ Copyright 2018 The Android Open Source Project
+
+ 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.
+-->
+<resources>
+    <dimen name="activity_vertical_margin">20dp</dimen>
+    <dimen name="activity_horizontal_margin">16dp</dimen>
+</resources>
diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/res/values/strings.xml b/tensorflow/contrib/lite/java/ovic/demo/app/res/values/strings.xml
new file mode 100644
index 0000000000..d26beb1d27
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/res/values/strings.xml
@@ -0,0 +1,22 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+ Copyright 2018 The Android Open Source Project
+
+ 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.
+-->
+<resources>
+    <string name="app_name" translatable="false">Benchmarker</string>
+
+    <string name="start_label" translatable="false">Start</string>
+    <string name="initial_status_msg" translatable="false"> Press start to run the benchmarks.</string>
+</resources>
diff --git a/tensorflow/contrib/lite/java/ovic/demo/build.gradle b/tensorflow/contrib/lite/java/ovic/demo/build.gradle
new file mode 100644
index 0000000000..b78a0b86c9
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/demo/build.gradle
@@ -0,0 +1,23 @@
+// Top-level build file where you can add configuration options common to all sub-projects/modules.
+
+buildscript {
+    repositories {
+        jcenter()
+    }
+    dependencies {
+        classpath 'com.android.tools.build:gradle:2.3.1'
+
+        // NOTE: Do not place your application dependencies here; they belong
+        // in the individual module build.gradle files
+    }
+}
+
+allprojects {
+    repositories {
+        jcenter()
+    }
+}
+
+task clean(type: Delete) {
+    delete rootProject.buildDir
+}
diff --git a/tensorflow/contrib/lite/java/ovic/demo/gradle.properties b/tensorflow/contrib/lite/java/ovic/demo/gradle.properties
new file mode 100644
index 0000000000..aac7c9b461
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/demo/gradle.properties
@@ -0,0 +1,17 @@
+# Project-wide Gradle settings.
+
+# IDE (e.g. Android Studio) users:
+# Gradle settings configured through the IDE *will override*
+# any settings specified in this file.
+
+# For more details on how to configure your build environment visit
+# http://www.gradle.org/docs/current/userguide/build_environment.html
+
+# Specifies the JVM arguments used for the daemon process.
+# The setting is particularly useful for tweaking memory settings.
+org.gradle.jvmargs=-Xmx1536m
+
+# When configured, Gradle will run in incubating parallel mode.
+# This option should only be used with decoupled projects. More details, visit
+# http://www.gradle.org/docs/current/userguide/multi_project_builds.html#sec:decoupled_projects
+# org.gradle.parallel=true
diff --git a/tensorflow/contrib/lite/java/ovic/demo/gradle/wrapper/gradle-wrapper.jar b/tensorflow/contrib/lite/java/ovic/demo/gradle/wrapper/gradle-wrapper.jar
new file mode 100644
index 0000000000000000000000000000000000000000..13372aef5e24af05341d49695ee84e5f9b594659
GIT binary patch
literal 53636
zcmWIWW@h1HVBp|j(AqTBoq>UYfeAz~Ffed3FfjPKhB)ea`nl;dGoUKGK6OTrJp%(n
zC<6n72(m(7M?X(D*WeI6U$@V`XHNTg>*`(P_14uocjo-&AcHH$51xKHqkF>htnXQE
zPaQ_CS8XZNo-B#d+##;I?8%a(6Nk1+y_977`l*N!$wDzSm$5J~Fyt4dqc{p(4L4Lx
zdQoCZPAXod!l+8iixLY8Qj0LOWHhCuEoNX~xXQr5pp0FMOMZD?PJUvFilJU|PGWI!
zZI3V4Ap?Qd`x&ND+GYAp+}GRY9h5In)U$ESan9lN^jx)fHaGu+g-1jRU)wyhl{-_j
z{`+H21?NKtB$AwJwwX^qUAs~>ao5(h7sEted);A8+-AIU+dh+58najHN~pE8mUd~Y
zbLm#Tc8T>qUSGeuhry}Hz-?Er*gbE54{NFGhcxPTg&5^?e72uZA}L^7vs}LAf)bCD
zn*^JDd+%Z1QeD|vv`{{Jul1S;69Yp53j>1-0Y3&;7MG;v1{CENq!yKArWOYj<yTZX
zmX>6s=9Oe7Czj+FK>}fEaBjYkqd=XQM*H!Lk5(qEoqFq9Rmjt>{EG=voV}f#hQ6zO
zee81+nbX&mr{qukwEjzXuE|ICAB%f7J4{tz6n=mB+#8$EXKbu@e}DY^`g?{`6<JEX
zO`P1U%-NcE%r^G*ikJmQTze-g9K9o3iSPJ<ZH?7DC&l9)tc-2j5p0v<aplmmgostu
zjTM(|EyX_R+JDnbx_+4LJO9ovU+P4BR%{bDTqFNQS=H}b>J7o5N+qTDTjJlY+Qu>E
z`U9;gJl9?;2yA-xSwrYZWh8&<>0HlyUs`WZ-s2onuK$L$?!@fbpSi;%rbRow4c-~2
zF*|vNI2(_&-^PE9^-uB?rdZpqOnP9)?6J={*n8(9c46zgi5h1-Qw;bIt;^{MpC{-4
zj7cj~Q_Lu0iBxflq0Y9cuclc2I`Ljfk=JHJK$Ca*2lgEgvQmya+&tuLb4gZCy5i7=
z-}Q%6Py7wDWL$G@4uf-<!<WhnMeVf`mwwG_yZehL+OJl})l)wsw5z^dVxF((ANQ;>
zBR8dxi0tCXsgovmv;^C%GJLXpi^Zy!uj6JG`mDJapXOL!c;w64OUhA;=6X+y6A{~X
zUf{QnA5Ub|<YtBAsZnbanQ~&3XFKry6x`_**!y(9;=g--bCM7G-8yC&6kqU?=h^bk
zqNFu3brF3Bc=CEq*{VxRXC_|W`lO9hS@%}-!Ly9Hr4#D8!%yCNvNZ6B=!_kHVmGho
z2f#}mP`-WsCZc#10|Uct1_lNt?5WNZmbD;h%C+|(*C7L*w*O~dJ`)k#!xi^6>EoW&
zT+EIEU*fhKnR!{iDzTf*E`Ok1$a<<I$I0n`etxdGE6*VIqWVP@2S?o^ZB<i$j>#|8
zo`;{CqH!}(yfyxmNpyNp%VuT8h`ULFQ)Y&x^V^=Tt~K1a;(LqZmOqhaUppU)m_21_
zn|jcOQy&%yw3$jB`jD)@)9dE9Dcw6%dj2hb7ToSo5Wv2(MJ(oYu6OC%8((tXAE;+W
z^(o8Bs=F~v3=DxR3=F#1ed-P_utJMdi#+mkQ+@LDvr7vgp6(5f_Lp)L`R6t<_S)B=
zH5yyDWks_@-;8=2diludRh=Om3Y`8at(@!hcBe$<{0ZNyl9hGoxZM|bQ`^Hq+9#TK
zpWJ!2d|&aiGiTo0+t)Eg99q>Bc;;lqq6eDuJUcWw65`xqesl`BZBTf)vuOguUV*AW
zx1bE)L*AU;Jd>k;ef#C86@TmGy8NYIE;;(0pH_Bxeo0%ws~Z-kCZ{*7jyhNDyL;Cu
zgS)-jsheK*%sVUYU0pRXG2`+^>3NCOtcSCT9_??v{+-*-WJlVqW#1<*w4bM8z+yI+
zOSNZ-&+PZ++3k1EO1F!N*IJ|=){<DK`KQd_f$U7)HkVU{tD~-(YKdL9+wHdZTx{<0
zr*bAccE7$-ajMX7R%ur4j$4B1620nwbZ7N#%Z-|2@G0tMR{7E=bCyLvOKMcS)orXf
zDb&TlF~+WPlV`^hK}VH6K^uES*}kUnGS>Tkc(-vyMYE{rcY(ASFXsJrySzo~PAW(2
z<<%A4<yY2E`@+4tqMx7VbGZzA>4*Iuzdp>EzPB&n-I|9Vr4A=dt(Ls5DD}QaQEY#$
zs?grNlNI_3a|<WD*iq}%mT`8SfzG|oO|}e$4)Y$jOT0T?edU<W-Cu!KvVS|eGAD^1
zmB>6a`L27t$Yb+A3W_r=oF`2dx#(ehiJ4nGAwGJ7#>oV2i#~RhQ27Q;8x`3@CURRR
z`$wgGF*#g2{e}L!zfD!G4ii_f?Q3bCKEcZQ&=0oA8P{8tTzHhiI?8zZJ<ln$?R<Sm
zNK8ZRo9g9P``1m^jIId|=MdFj(89j*pXHSg0-eP?_3b*}4&}X3bt+mGTJ{YsL&`jg
zy?vCGf#DYy1A{U4Wa?UxoKsqyS(fUSnUm_6oSa%*?3<XEn4StP^{0k;7YL`y{Lh`G
zmboie(LpgFhoN;#K$C{p<)&#jk}{$e@GfN)-G0(6cGunB*}mSj3U;r4mwtFLcd?PO
zZUp-=Nq6_-%~p>aKOQ<NsXtFbU+&2L?X#!tljZv<_y76(|G&)a|NXwo%fSD5vck57
zo0$Zc9J%z=y<8wp=KEt80lSi})1HN|I2VM!$#k?-owiTGZ}I*9Muu}Nfyy1G7SczI
zE6hBC&$AZto&UX%#Xv*Em@%bwmXwQbo7Tj`1p-p%!jo8b$uF*xsGA{lZAqWH)%jAD
z;J44K{7Xwzeyg?wJ)df*x=h<m?^<^BtclqZ1Ew1EHm#K?o4-grlv8w@*{WHtyH_hX
zdU>rlH+$DSlbowk!6lAgwq~`Z$xjQA^Hr1nrBWfBHm%E4^K8KG?#0ckriQVVhW;=%
z?z*=;V%Co79dXqWi>F4M+ETP}&eD`pQ?oqHpL?>7-qKI`lH@2IY;6*odphNc(?V~<
zDRZ&}SNIol2;EwgCo@Y@xJdWi+KlWfon)z*i&9s%$$D$-(zDrSa;E#&j|Y{e$(Lu$
zSBx&3eO97haB^Ono641)lh}CE+e8lsd~JH%nAp?dv-0ypLu2m5jtf(m&R+S?ufWOp
z@Vx1!V?jM@EMKv`o@93Z=G(|4!Uw!hJiDXgIJ-c9%G%f&|2zNW@XE!=CT3gREt>Jm
z@tX75+fvdkwHH@E<vFSsb>)4Q@A2trUq94dJz9F{K<J@cfvHIWzAw@oZ6mLQM9)|F
zQGHFOE!unA?0yTM#-zhG9K{})Yn6<(^%L(Jxk*H<kKAIk)acs!JEhgFn|d@p8Oa^;
zlia-_xXH%eOH%*9qz#21Q)D%>s=~IWh$NiWNa}gJG5mbb2?_BQ4ZRaH4s_S;QF)xY
zM=9~Lm+SNg8;_;N99dl#GwGvL-1LbDg6sAuKKAw56#wq$J58yin&+Do7aG0oD?TFA
z#UYbl_(f^js+6{?IVmw0bCP;_W`?D#+o}9+j%$vwk<0hOReRq$)gA9IPu<i0{qGNl
zW9=t?f1KFUKQVpYbj{Y5Wxr<?@*RFOZ&HAmsOkDsE5C8Ne(}%DTEyGke%^m#Uh0%H
zUoW;jm9v`vAwQ*4YuA*$A&WyC4A*XZIZ5SB#I~d25<;?8R+Fx-X=46&rfEf$I@dO<
zQ*HXb%^H5wm%W~)eP+4F2?L?c%KLK7CWW@uyk5n0raRYn$(-kJ8%5VnmD^G2aQW~X
z!Dp-XELg4e`gO0I@WiXGlUJv%YI|^-jeA+#maD>Dwr3<`^{&@$UJ+g;{i2_L`o0}g
zir0UCHuZ4AfoZE&do{XU)0<h#_QZzoPlWpq=2HU4%HA;Cce!W7p`@ZV_xS{gGIyt>
z*X9MePjVZKQtxC-Emt(n+sWmhd^dDTY~hr0>#VqSw|C^Zv+{;t{va2mv`%W)2Jr&F
zL%mPtv?%kRH+-{L<Zi{Q9}A~5o-}5Df8d?wl=BD74S4PQ8#O~(nqrm(Sn}@N-FZhd
zWbU*_ahh>;yjRcKY07c1?{jC@-@5OZ{==Cec7J)!8n2Md;95KL*sUGrZ{|8_#K|pt
zW##_s-T6eV_59U41D57|k}`|gXRZ9k<m`g71uP0%CeCX;shPvoTB)|^?B;}RckUF+
zmI^%*O<Z*6m4u4gakYnciv`yGP?C#@Q>!qJv-IEaNzM4D<+8%=nY&u=&#jfOJ0$!3
z-}cSZEobiUs9Lvs#&?b48wY&MgEn!p&-mDEv-F4gTKV5U^Ir+_Otfk`KEd7i{*mS#
zEnz;71Z=ln3}=d8EPH<e^ZiE}bq6kgx!xI{@ldRC`^n$Ud+bmC?%rel<gfF*6-E3^
zTWXr;ecRl`Zs;7}Kk1m)xv+W5nZM{*E>AmE(ecAfPJL~U(SF4~v15Dq_X}NCobuRU
z>J!VK>Kgx5G~<F^b%uSpE4WQ}KgUur+tYU+Z+TI^B5O|{yGl%(d1-5kKI`fg&$eah
z_A;w`zE$1ZKgZ<u*87}GrWswmb6ax7?IrG#T80WQci4QBM{T^E;f-^f$i%?#lMQF%
zB^=Usa?a07%S<mVN-W9D&(i~Ug?cCF`U^V>9RGj&`L$-driMj}l-?@NS@6J^lRwlj
zlt<K~#(7=AH07s#cVbOH24^Z6er%OLz@Iv2Z}fsSw~v?`Z#Mos|L@zF-@nhFk7s!H
zL8IM!#>oV;ZkZmQcZK%H&)S%O-&uL!tdmrCrcKI)Z@Km#-h`QNh;2{YP?%9?#Cx2p
zXYJXq%XwPw{pxxd&9ijH@?@*r#h#|ib-(X5Zf0&<crWDhfjRLJIc5D4m*U?%&JgNf
zbnlAs?TyW<Z53O$ot+R-`hDfR{2P5=w^d)+wC8);YT<it-_4wxdwKKENEP9Ojc<*w
zNUghWc5B|w`x|@C&#SD-sXXTQBs}){;=Avv)Y292)ZFH}rnhHLz}l1f=YJOkZ>-wb
zzIocxz>Al)mOs{8xAM*ES?#O-e_Q3|zI#%-%T2+R+tVK^*WC~53z7G;cR%;u)@5qO
zjujowm)}-ZU*n$rzAj|zgUEkUyrMi8<Kpx5)|<ZC9ewuKvv;d|%}aTkPZ;yesZ;!?
z{n1op!u5bh%g*$+G4i_Jv-mnce(E}__(yK9FWllbVEot45Y+LpQ{?E6w+r<-;zjnU
z)L4FDPM9<QjltY9jq@k^U)vw}67+_paPJfgN28wL#5w=OPU%mUb(@sK%Gp&kMKfW_
zBEz15ux_@%v>uzoqA{Y=C7x93uB>-|_WI_ADFSo4C#OFuj1hjM`l|0^UhOZInZ;9L
zX3SMMd$!<E!4*}t4<6k2PZh|TizpOIrUY>sIX)IuD3n~{CbgHJ<G;b5N6}srmhc!Y
zX4$!x&GVE(@sgSEnS%rrCVNd-!?S4t%g)7YCr>H#dwyQR?-AJGqv;rAbfQ?*Bb0a2
zB*$r<n=i31QE|xfny`f@D7@)K)`TrQoBU>7uXQ-Zlwx(|&R=fSv>jM0oB5Q5fkBds
zfx!TKiwC{G>yen3k^^h*^oG4I77i8ppF2C{aM|g_0bX9M988(NIvPZCxU?>KE#+4U
z&}rJ`n?JcV#XQGsi>T=QA8vmS@$?-Ix3g?MEOFr6iSqv%U(3xGp6$MsR`qW0^V0J_
z=N6xTbN~OJpW+RFe?%VE>se9Z>@4hNGUv#OE_2f#s-LH@i6(zC&?^_#QVX3E_HkNU
zSPx5~+d59qb6!P$&3!8s&Sh6fU%vFliTmEcqhc|Kc2C-)u5Ycr^nJF0fwE`NVV#zc
zdD)I!nWD?OPIECiFH7C!E^6YwYq`jezT2Cf4z?Ly77}iK(^GIx%2=}K*0DEBCeAdH
z37&V_XZ~!n&s<f<bJRt>Esy269eF0JJYCpNXXA`=*Kb>rKR=kIRDWrMM%g1Sv*j((
z+MyeajmmEBNIT+VZI+bgvS)XR&W<!Sv%Du6;wEcbUUMj~RtQ;Yo%NI}&(p=UwQ;Nd
zX?OK@@kvV0XKizh{4Barqxj<z$-`bpGt#0oyRNKxsCM~KWZApBZ#t$#Jor1$<5+^l
zN~f@hysgtu&OEL3GU#5)z6A;1t<8bq7H!$yx4-UuT9SGtB_w^#f~9LVo@hxG%o4m%
zB-)U>R$x7YrtwKhz5K0GsYSa~y1YNHYSMc)PpovouUiYNHoUTQe_0TF$S1Nby5L1}
z#n+^DH<wS_!0i__<-w~eH}9|B!Lqv@Uzo4GQ+(;f^2+I5+Ezvhmvb^(rZ@GaZ<+KZ
z<K>2lmkUFG-Mh13xk$-6mbzsLJJZ5KIIOcaTZI`IZ{BpWjNiKNZOik^hq%&nlBaF@
z;8l9KxaVmHAA6-|j4W%+WX7}FH)h1QJxr=!Vt;FGlH+kDy9<S~@~eJ$96!Ee;Ui($
z`Av~EZRbm64v1^*yHXfjwC~V$t$i0i>h8Pz(Y7w2(EE7*4^PLExBN33uYB9jv{%TR
zPki&8moj?d5p!;>o51X@wm#u@*NRW3-n{ij$<LP>OZK_m-RDuex#arBQy#~&zns=D
zHT}CFdwN99Opm=vri*x$Ry>&FYyE9cjZfSk+4B|6H7h)=XMg?>sk*eS-~52(G=X(3
zTY`FZxT04}6y;4z7m%-=T6HuhYxQ%ZeOo5Z7vL^=xOvY~slKNR%QqPPdYpOKvtu=D
zqN;1NOqg4ytlndP_6_;xrj&kLF0=c9<qbK>+=$lRtu~J)35K2Wt>!iNnl<U?`@i#l
z=|2x-oMYZF`(_6hpSeN&WM`Z6_w`QfSJNo6oA~ab=CRO+ulns;KJ8Y&=PmVN#_f&$
zA^MsMO?m5H3%K%c%~)g8!0-6!Uxd@;FU}_YuH7I0GS#2|tNP<4M~}Ya+-1)nxa7wR
zzW%ZPZnNEs$+IWCKWW>dK5c=&{=%uZUf(to**H<pO!VTl?*g;uZ^^A}`PO+)X`7u}
zdy`}QiSH7(MeZEr+{r!h>3i;<-7j>KekUl1Xe@fo_pOC#pVB<(5a%pG>vxOxsxSHG
zwRUbuO;3UGx%<Dr2cEAr5uRB4BAw}P#l>i6ix$T%0ZNWDoUC8+E!!b1EpSZjReQ%B
zJ+JSAeR52f&VLoC$`o~-Gw1V@7^PL$Uw%}x-<DT+=k77bJ4d~D_@8CBc++Ub{xk84
z-)XyhV(#K^lkVJgQ`i0e`SNQuTe$_P4g1(KY!3)6Ddf_>-0k1w=fA#y_mTNMx%7!n
zKSZ4^4n2x_^!-}k9<TcKQ@JB9Rn6U<95=DQxcty<-v8_YD5Ew@D-Cvqfd-3N85rcT
z7nwf!>F^F=>*U}0mjXnN@1K5LYrWyI0-qVebGcl@+wu+7+EO?L<vG*Tey^P}d!2XY
zLXqD8UG=R}M~}$=Vg9qY=lM2Xp@_{B=FKa9Z#n<7R{g(!zjz-A_)qqDkj6B<K)dnw
z3bnvy9hcY+Yo|L?3ltnDUP$QZU(Qx8d4u!tE#t#Rcb@ZXeC7FeMqF<70m)6Vhaat9
ze%P-3-pbm^erCJ3Et_?GRWgUQMbFOV$66L%*(w`*QMNRp=;`9ceQtLGcWql&y!@zu
zY3|&G0dME6ln*%?q8s~iXJ6t=-H8EGpH8}5-OHoR7c##glDX|zQqAv>(yAJpr_+u_
zZ*A=spSNAZPvx3ga$8b_jI;5Sn=e~}9t($;9{FN5<)C)%@>_gg)1+*3+fvTYvy!qm
zI&vfEUk;y(kCdd$vWZSnPR6~}JM<qqO7RDJEbef3vJrd4`IK|6(cxqN|K$Hlk6FyL
z{hW4Z*1E2Z`aizg`Nr5S;jZ=yWZf@(IPBgfu9zs6@b)Vn|K=p$y5|wG?C64+E%QYm
zJv@3jzEks@``(gPuG+tx7Z$v{De`zpz1y5GJl8zaPFgv&)HwwG<5Zhor|@a&8V)m$
zb;@3!jFN7sc?nvc|Foxb8oP~eK<3;VPKpWpcxFWYv{srw@lRVuZQBi_9W0!wm)7ew
zY%VUC=C{qZK+u};M<~O&pAV++_wAp)UuU6q%RWEF-`p?m$jm(~>h*5%X~w9<ddXkT
z{eIEoFugoIBDyu4Gfvs#Zu9B78#43wA2{7!dQWNIHutD2-W9rgreET@lx5X=$y8&B
z<R$M`E784oK0NyNpAj`lo-2CYCCbjgkjKZspoTq31|*hbI2ISD<|e`G?x|rvlcmE2
z|6RK#S~=OGtN6y&<|OrAott8w!5W^zAFjHrUH?nV@79j&Y~_5d_q%qlKGIestns~;
zg=xxfHt{IUBU>2jp7b;AyBQn3P3!VE_PyrsGw+=-H@-i6`~QFE|C%$-Ikx|iMb{y@
zIY)zJ*z%5O?p1IOY<sHVb8}nDMxmqKg+~G>$VRKEvxPqkk_o&s<!!Xl8n=yN#~o`_
z`ub1Rr+*CMk?(%=xUhz4e$T_;!Wyah6OX35?vg$FJmp`F%<&JKc;Y)AbqCgHw&@?g
z*t-76htR-(vv}ggXWtiiIoq(t`&s_+4|65<Ngx0Iq4r>*dYk-_56+gA>Rr3J_<AC`
zqY6(H#N3Y1*x8n<6;Qlv<?UwIvxOz@2Pb+@{_CuK__J~F;mMz;A9rtmzI>T_aq@H#
z6}P+$^Q}37#@shGr3|>Ai@uTF;&c3}PkEU3v=Y|eYHG*l%gD{SXSH>GLiW!^c8_0t
zOU;fnsWEPwT6;+2m5E^G9RuF)i<*vwd|%UaVZqepxBdwpzH;2_H1{sswYnmuAN3!k
zo8F#Kx<RsNI?tkaG8^AGT`QluH#J*Zd)M5#7w#;rbKC0UJG+l>N!c#3{NLQWb{oro
zURkwm-p6#e>bn|Ka(qk9m+xPDT{z!!$CI4G-J-0eXRnm5P}sJ>^4AQt&Yq{4*TiS%
zoKBM1w!iC?d<^R)(PY<y6|dG_trjo5wz6hRinxCEIU&_&7pBXtd%X1B!@Z|>J<edQ
zJ)!ZUZ}o!1;YDALPRiYpee>em`e6G#8=i4lORwY82#|QZFgL6{T>R@)1$)8c#Rm`F
z`n9PvSU#?_^zz3`GDdsfXq;)Alzgt~)NjkzJQHT`nmu>snYTilmJ}b;%bQpcJo{L9
zL>X(Llx#xlo;T(1=Jjg*b`A=Vo3qL5+akT$DZlvmzMkY16)y-pT4^J9v_FLB@o}wr
zhfPI`E254DOFVYHmlt`||3g<7tIj8{>x(2bB#-XNTm56@i$6T}>28~}>J(RV|5(};
z-}xnb&!PuscBk*S@xxbSf2d*O^2lA2{l2MqUypG$+*jZnU$Xee+^+kN`S<)-?A-76
z(OtCu(EQLfk3&W455|jbsc@h6QqMX;q-fb{ah9*yiwYRmb3fYD)Lc4a!Gg8BtRszL
zlOykcOaHwfRO{K~Wlz<nd)!%)xNgdwCH`~mHpV8pY1}z{KjhBwwyyo&ZtuExKfA@Y
z@)7SI&mxw&{o<RCq+OqG{zgM?#gBH87fk;=^SF%9se0~R@=(S58TYc%?u1GC)00@2
z-+Sn*x+|ny*OdR-mK}`!1-dJDiL4Ku_aQ@8Fgz{J@xAq~2aNld#O}NnQlPjyYgfLk
z^X7l7NryifvT@5h*WcZ@XY=vvpEvBiz4~_W=H^MWHI%RG&Ta|*&}cO=PIl$fg1>ej
zFZ|C;t9yL1vj3S<vd!LVxAM1bhA)lkoOzb+TB_>irCoV;>CS>zuNqvWPO{%qS~+Kb
zvBuQZ+-9-wQ&+t6(Z6|3w=#P#i`hrx#s{B&z4}pV7|Cw^dDn4G-JEX^53h`=-Tx_z
z<;|sdolMu|$L^&yt$(xT@~p?NPS**k=X8q{TrXoZ>VEv^p4F-R-ELbR{jK0xqZQ73
zqM$bK_O9vnWnrfb-)+2D^hY7QzaY$_{j#XW`l?m(k0hS+lqei|zq!t9)%q4jo+2@3
z=f+0$6`H#W4JVx2)7Q3YYHbUvO6F6GZwoGOF1c}?Pp4C51IrPE+@EIuz5N5vZNEIz
zNJ^WdcUiR-NB`1lV^NoLt!4sqC+|-ZeRI>HW?hrox+_NC)V>vF@D^$$*R!5|erWbC
z_OcghWiP6N{J#dIFtb0uwD#dO4_4jR(GRyM9x^-{uxq1yrcB~A_KVG-Y2i-0-!$Lj
zdTXXv*dKGy`@}{uUYk;3|DdcVk(2#zCHtIJ`075R$>8mqqZ&&&On=^CJ-Opxz0;ij
zrJWl#_Be~m$a%|f?r&UN+&gpS&Vct<Y`3htKINT>ciA!Lso(tb7g{}246yP1U>acD
zSjZpRyeGZ!{%#v}_1{9d-#=W?n0(Fd+U01s<7|CrTRWcGimJ@^pHf&;Tf!V$+$p@~
z;z{;-Pao@)J?~rdwOGbM*wAFM)m)u2@zc6L&Kd3FxMkR>IWHqAZJD97PVAKQk3HhP
z_r-qM)!cBtuMp45cq#je*rKf|Lj2PUUpf7bZ83g5^Vm0|psRaS_GHXU>UGPgoU&8y
z`ychco1^8T3XXmjSb8kf=Jmzd^A#KQXWvR%-uv@~+*`qCb9UWuY`1-FE4%RdX^(JQ
zxwMpokVLm7M$6KUO;-4q6?$vio^6S8=P$cUv@fWP39PI=vOlEtzD48z>6Ly)MpHZm
z9M8)hxwL*^qV41tU&GEieRsKUo1$Iue%BP$2Q`mOUd|O<H~;37pUV7M6YRW~*}1R3
z9?@JmN1^XktL!Z&tCpGo#=F+_vhP1f9B-exwMOLBf3$gFi+kJdN((bEywPM}P{3ZX
zdgc|EB<AEmMtkN)a@UAF-D<DC`{qrZ!`cFePG!3&RkoW;3GjMOxFIm%gMq@Ul=4YE
zZ<ODio8+=Ks@v=8zJTb9VQc!6bVCC@FKFDzs@?W&;nuBNuZCskzka)S@Aq%_ZvWQ$
zUY}NGK08I=k^cU_Y0s+XJpVuEeNBAa-`|(rJ2MpP1s2$yn|@5aSE)CCRuhv6W9G4w
z%cT|=$*upu!vAshc~*@B!l^lvOeQnOoCsbrQ|f`zJh27VJ0fPi<YH$wnVNp$&rGiK
zpFY1~IG^bI@G^6W)Op?yk~aQTUP~;4AO2*Vb@(xJ%&AnyRga5)u9>hbXx6I2n6su_
z`X4?!wThqYIm$Tey!Vp&mGh^cEZ*5sJ*{W1?1H3qUq5qIuG4u?ToKxE{^@6?iq*jt
zr47@X#gDyOnmJ)vpvD?OuIIli*M2e4tZ6(n|5Nei`PV(&y&KLw{(32g<3sD7z=ziu
z);|4McTeT~*-Ig;72NruKa^T;r(D{`u=e>|#;k|kqW1(Bggv}{X!?(zK|go+%+hN8
z%CNO?_o0OhxzEe&+yCbL*y;S_t<D2IGuPFp|LzH%Cn%`4szAp<c7w{S9~SSPX)p2A
zn8S0yFK^1RNu8Rl<(hYd7v#;9ejXQjLA@(Z^1<pKSwHtg-LuoR+ZpvRdqxv)Z}uxQ
z^ZZ>~7OpJ4d~Dg(+2*Ee^LLgw-(9xVr}gag*So*W>X3WACUdJ$+L~h4%Rhcq=w!Y$
z%<h)+bo`V(vu0vN`tpehvo6Z`ms~4R`%-XZ#Y@i@ckj9_T6R<Jh*!PEJ>l*nUkxtj
zZ142teOdVM(%R_Q(vwRBbNe}-?#$yhdvfmJqRTz^0#v4&luKrx=?E0(l05mu%68f9
zO;>jWX4pJrn(%q8N%FF-T`wnH>c9B<v8m<G8>!JV_I&-g;h@RJn>~`cD>HUays;!X
zz(#=o<wV`8LsQ<DEWiHBXtIxStYvq;lAwn?hbqrqBdIe>j+Sg&%@#dZUFKF|*)pMB
zsmGGFPsuEw{ZVRT=MJSQ*5^XXPAypwDDg%(t=em!xw?{%u*@gJ!>exY*~l5(RkVCd
z*~{)}dfC3NGu7N(zP8PAY+dhM{`o|n+Ez_5=Jy?Ziw_xi`(=HpT-@Ox;L4I|ARH;#
z?z7<5rX1tlPo<Z!=Je!6P7C>A{6UD#c_qi1xyt|VTzVLC>Q2+i%c)<CHg?7f%iK<!
ze7;-BDb4coWH*he25a}tDs50JxS3EU<>?((#=+LL$!1p4T}|#1UfZs~w$PflEBp8y
z>+XHK%CYa8U4!HHl4VQ=r78)t?7UW8{ibzc#?8!#(v2U)loQLoJiV2bY{TcFE3xo?
z##cARl@`~p^12yKG1x83Uv)T*>6K4WSb|N^<(ykpe$sM`wRdM~TyQ%RYdPn(R^;x}
z6K!RCzQ(#`qzBb(-cqt<i~3cz!c{J^`Ul#oPI$c(ynUJVWQCTXQBFm_an<hL38^h_
z1?!JAiiBLg^KioIRpI_B&lJ<P^>3MNo^-@uLi!@j3{&gmT~oaz7RMS)TGs8kOr3Z2
zM)3ukj&w=)xLom3V?EKaX5y7?9H&#nlwXu@NNh2es=DDh(NkN-KQdC`)|%@_gahwx
zdc%2-BS`8+$#!Mi=MtsPN_r>DN>BAzOxdvT;*^WRDMuvQK7YOWAbKKe<K9J@Y|}bJ
z%zNfrt3FU^+WJxBMNhfitH&2t$T?4086oJu=5*CnpFGp#8_O1KUA`qpTBJukx+F4E
zz)Y!a$>o;rr3QhX%ig?Y{`Bay{HvxUPUZ{qHr+M-aM(%1UUc`RmPf@g$(q*%f4ZIV
z`jT3kcw5ptd%|R?RF_F2G84}-3Ry@hy)ORMy7lZOX|t5%fR=Tdk(({ri;{IjO*eVQ
z>@=GEqszV6-?=17<KErrlbdI(@{TQ@ygP*NT7qtWubRJ$j`)S#?sL%_5*ULwzvi2E
z^WCvIU7xDNJvIl-a&=xK#-;uJtkL09rN2*KvygP^Zn<!3qwuM<jYeC<=Fhq2S?t|*
zO=w}LsaVwQHZlFFN5$tIdbIJ`)jju)I>$`ooGzMNR-U&d#Y0zq^R`VpuKGRadvx{7
zo~)V6E`9#e<h6fcd8n1oU#-7$>K0y7uw8tq^o!)wbqljYt6JnjU%Ks+>p31$aB9){
zRa<7~1-z{KwdT<7(0z^hAwSjbE&lNQ$>X%C`46K*|JmH{di(Q;*IDK@o-b-0pXz^_
zzQA?eB#sx?ww8vj2p0Lb=~VSkGnTS<EprN31*;q%-VWW@ynoTB@SlEDzkiH9q4zGe
z`B3@lKc@e7o%;TvZtDA|`5wjdPkrwe;$QJc;V+N$!3&RiSI$Uv>&-cow(Ctzw7Ty0
zunW(ot31zsxMSj#q(X0J$E=rWK9`egXG}X#Cp}#zY^~`s^Z2t}*|Hu_W?k&AwYs!i
z_sqH1tKDWA9;vVW{q#lH&)GE&Jl`gqy8q~}kF}qM%dB!Cokvk$CVMXUbofv1)OU(A
zUMSSNKXp$KeYZ(l*EnT)tnR&qJ5OF%c-ZmbW$z@_jb;a~l=_9dzr5>9Y}>Uw3B$9O
z7T?{Jo^8AAYG|y`X74X9A)9xVaQ7}eTdsTHdG_5-3$dtLDV@9Odt^3*E<G4FSFZim
ztSf70&1@+YTBIx^JWX}hyoQXIEo?nE-dsApwa>+7)yr8^Jyegz`SFJ8c4h87x~y%}
zsxWQOX-AA^22Sn@Xq)qH%~9u<pVnO~6FHGlTGnfp<?z&<wYPs|O^1B=56-FcMUKlq
zt9x_JdFE!li5GsW)UL8|un+#pVZZd#<xgfD^DoYBdg8EK(wlwZf#;Gr7t~i3wjI<9
zu6y3TYf<49#d%M?r?2{>^IlNPJ=E^ldCfnuQ}?SU=*Y=tXLd!cPGn4+S>*RxZq@TG
zeHDwsCcmhTdng|A&(3~I=dHk2ncxeltILjs)J=>RS{XO(+0;a_Q({$G`%hH|1b5i0
zOp811BOSM4<sasmX1^b2n=SkFt3pd?o3m}D5bMV1{0@n#pn_#5cSa-$`MDpNDQm`7
z{bh#g#TkcW4#rz0nVBmJYJ5y7<#{^yl>L+IqW>nIvVS&z#h>6h=44wx)fJy+^UU9W
z)ZSh4Rn)%~Uljj4J=`ApZ|Q=^;eY&&U0Qs~WNJOH{O5cYYh!n-<i8t!<yPf+os$)^
z-=m#Z>A%$eqC40Bu&4Jwe%0E4zFcYhyozXtpR-?1oZx<Td&V2~r#<Pp-=F^AdAe!J
z%B8-(pXSteg|leL6rPgts}IV&dBVEsrLx*nna#_m?UGxwb#mshl9v(hT<k9PIz8mr
z`1Y`K*xrRnCJmK0Pq*HkzvGi+-qzr-b9+@jzkK=FXYT8pN8HT59SY!*v`@Qw#o40s
zLgICqbI-U~(`(MevYk5<aQ~XVSkpi5*mKWP)pzpDd^Bm_guV|P+-lp`>t#2&t$ep3
z&+}(9AN#U>1yL;dnKNT1`N++i`KJ1CVOY)E3$>nsnUl5e)mUsSiHa?g;n3f>S;ioJ
z%C-J^^Y<pLwAr{W@ABU;8|C9I4@K>!7)|cIn8)qz65l>|+S}qY8n14<?>-(nU&;LB
zlVwJkah$?c4>koJJF)9_dQ79p92rN|!+w5O&n4dK{I}~{ip``=Je|uQ&Q6inV*YWl
zSy$;hR|4}ANvm^yg)iGIcRkfs3iozfpILq5qs6{ASHcy-jhQljKU<u0W$tMU<C_yS
zZzwd$wHcqc_&HTP!l%jaYn6_iv!;0Alv!b~iZ~OVdoQp{j5fMH!MR+1vU};NV@qa>
zZHxPzQtb8N<eiNy{uM=2mCIuiQrnaKctoG4tbQ2YbNTN1?~ATZS-9|v@5C>^RApwr
zeZ636veUB_-l`|3-?WxKDJOktsn+G0%Tp$aO3txlR9P9Vx6*8T+W+ShjW--R$9VRm
zoQj`_l(S9I^Btc80vsjwzP9!g@bmrlHR4j?{@1#J8%(aV9L`VeieVQBimfV&wMjH_
z>(MYWW-!)E`x5o?{6y(Avoz5e7k9q=<XUdEed9qz^S)hcR-W8?a6|VB;fJZKq;*#|
zC>I)L)$tqcc3d7TKUFZ$j`gM2uC5vDj)lwB&f<Dn(fi%WGW)jwoJnuh`_5GO9lRs`
zQf*81hf-OMijQ-`?HpL2Hj89v%b3rMlT-L=(5t`R!pE!fY2v4;;j+s(+jIElJ6B8n
zym0QQSMRb*k`of;&Mp($GuP8RdC8rHp4%_1b@?{?#F1~aJ*KAb?~IZAQa5j3gWSYt
z9IPG(i$h)QePi0wg1moDHeuACoW0jE-(K|Ev(w35Yx(~^Zr<lO>sFZD7uL_(+N;=B
zteC&d*8OOI?}W`fR@0h`R8!V2k__QJDSo49UFnY0r!lVaOY|KV)mY8m=Tq&d>frst
z?8f!xtxrp=?{c0^b2wjHBA#}3>C6)Ido^E7*L+KQdVM3?^H6vBJg0aDWpfcfd#8h?
z%_ULMM;Xp$eAM~uH_zq#kL=Igdn$itriB{ppP#1pS*-HC<-h5t680_U@H%Hz@Td7y
z-K<z{-RE0)Cxor8nD(>#kK_G|zcUV|Sjl}6{V5(_+~T)hmqYB!<(T(-P4rgyL|yPt
z&*9a&WSajv&tBf*Usb@%<&zXQuUU99W4qIv?hK`^j8eBB9C~ZguD9&x7Us5HE$5VE
zpK)CLDYC3L<wEKeN43&Rq8Yz0l<T_Y2$wk=JiXypoW#loy#-6Z@b2pWv#8|S8uJT(
z4TO}~R?975S3K%+%k};eW`k@+)t8EEJ~X+t>~rNmyk_-*ueU2cKit{Hcd7irvKVH2
zr~YM(pLyI^vUHzr$gcCg8eF1w*!(hkW?gjR%I}d{*`GtU$tf<_{Yb)1d!pC3<?XB8
z_^RhEdS&Cqd}1>50__WG3dKSTd>QR_ub5X=+;DHw)h$Qa*ly?dA5^N0*g1{w;^PJ1
zSx%-3FRz~G;=j(XeWm@{%bPbe6y&a{n)vFpHv5WY*P>Nd#21L`vMGq=>BcO%u_%1Y
zw#thecWl+&ziR7Q@h@FkTYvUn7CLx2=)GW+naJki+~qO6ZPzbYSR`pF8Fob2a4(wQ
zE|J6VDSun?fp)v8YzlRDY4dtl>^T*-*>H85vTf>fnM&b5&u`}Y-!7cA`ig?g6OFIJ
zHmcWNv1M9rsF}v{{pRlsL%FTGq3+WwZBKvd7U2CpF)Vb-{uS56zkI#iRl6jNt=9kI
zDVtlntmQrGwilf+-<;dD-Er;hhqJ9$Oxf48;hz@EzUFhS5i=R|S#Q1H-dv&;dSI9P
z#yYElx-)Kfd47Dc-1D@%Ad>a<rMfT8vnurUU)wH!?)jzWy4l<BAFAiY6<pHX&#cG#
z;=e$LRl0m7KW|WBwClrFj8^5c+j%Zk=I%>hx$bQ64ykq8YuX<_*xz=IF{|X$ZK)^b
zJKp`-{enluB6Y&GLlM5M;$JQ=KXS3+OK`FMT5E@Fwv|hBCtTZp@((X_$wP+RDLxmw
zSZ8O68s2}lE$V-Fh)f9Ac8dW1U(vraUxz&r_2dxWDzhMy?{;p5?DIWwcbINWeQI>y
zaMDg5vACw`&KLG_eEcI&)L}E<fv@mD8K>ytyZo>3b9diu_kCVheJshPR`*ez>9sA&
zhpGjhUsv3h@96X1{fJG6Mg_;WI*BcEhkpxi_>%aCm#ObK!;x}^J@+4b*WBhU_@*GW
z-(UaGtLursD?8)Zj@z1a$1#fheK2i7_)F7@ozu!L@xN&QuqkEg^_T81taJANz0vlq
zJ@R}#$BM4!%P&ZO36{R)`l9l1TiLvrw58_fa(@Wjm>*X0QgGLs!b`tb?%lb_;EMml
zJ?t9yt1lEANLRI*^)Ehf-f$M%7GL%(m*z~qR+IdxY8Rr*BA%XXt(<kfc|y>JoqG=T
zwP)G%l+Fv?thJGIPF&vM)kW8o9%MJ{+0>X2{YmD|+ldheue3TlS$zDj#PdZcVsB4E
z1*=-MsOB3}l|R!n1+R!KoYoM_nYPz=f6EPF0l7&RY$J-*{&MjNOq`y6cI%wF+}Yax
zdA~33zUTTt#c@OV1gn(&uS+IeRegDL%9or+Pm9_V%d86e-z-kLxFpDK(S@b+lN?@3
zH;b?z*?qgMm-m7w^Arb0fiGU+MlYllydTW7a(FEJV0qK2jm*dM#h2|~ayfrL>)+#i
zj*DeWEE61mM1C`Fh^Y{sXMWJjy!VFE{X;Kqn{ZaKFZ^lqB7#%*{KD@O4DYiqz2AII
zUc>J{gIB#n-e2X6+7_$V@{TW+Lj>y2sV%I$k#&FZ{X^TVH=O;)w_5e+`?I$S^=l{G
zv-sAu>Oh{*tHu0|SD7WhH5`^#tXA8;E%A^4g9lHpMPE=C(46_k)%RjhU~AN(+?n6=
zW(7p1h6~xU_x+blc;PCt+)b>o-tl*Sz_;}pssDtZ-np;vsy?-mVNs-@_kwGT(MN<d
zzU|w){p`l_9R7qLpZ_V(FR$sC{O8O01@Ag;bXdySO#Q@_lksn1(2;cEIf=$s=4gJ=
z&uruRn|(`-%Vo-0lik}Je#F1Le_1i8dHI=B(~msyiz-#KShDel{essoj_wPN{$nn2
zA~I&SACukgKLzn`S$68Lo_hWwSINKG|Bm<eeSaYF@{j)C_*-YB6J$M2>t;xP-SB*Z
zw(9bXkLj8}9Z%jo`**|U&9i^Me1EvS_=DxSC$G-gl$*)FKexN+4(A$!;(4<TZ_NG1
z+2dHLWAph&x0S=EJh{6ET~}}ytl9MZ;(un;8LPVod^(S?FfdHz#JQ%*DL=oYxTGkt
zz&9~7FSsPJs3aA-#<n*sGWvFyfbF48tgD0$<QTFPFp11Q>KE)Ky}*RCb%VsDglSew
zx;N;%xlCK3Bm82TUgmG<U$>_3Gt+NdU*miKxBjmy^%u+5r8Jv=y2@WZbEb9ixifo{
z|NcCmug|!rQ9kuYVYr)qn#vyi#4tDQXDN3M@9AOH`IP9{#%?HiI^yBszg0F5CbfT!
zm?|*cKxMIDw$QcIoP^**>|Z?!%=>?AINE;v>yLnA?mr?9secXlSrqLiK1nEusU~>V
z!kd?wY!-fKV<<f4W60QRcs4<M=NYrj*RJHwRrxn#X2@6GtCxQ4QxJI@tE;UbW4yM+
zBzL9eT-W(gum0I>UR3UT`P{87kzuCFd8?<LUE|4kPbZ%11%K|*Z&#R2SFLullG?6&
z*2907(Al1InO3t)KCB3vrL^UWftbQ>#ph>v5_iVbML0CCjt$Xuns358TlSXmu32Y~
z`IsK{kSOq%IMcUo{V}hQKLu?8)=G?LRX55C@}K>B`pi93&GxyfY+|=BMqcXMWwj|s
z-~Z0VI>XAfQ}dR$E!+A0=3$lF+f6T@@OrTH#>JixpR<xDW78~Cs+>eCSEcB^%kV1=
zIsHnVJvy{%uH+H{F_tTtTP8eCJLk2@@_ERDu!l-_nwu2XX;!drXDNx9V8!UQ((NGM
zC809Wj2n|bxHR_lFg5c#&Jg-0TEpe-C|06g^+L+^&MMb0Opj04rk~%aapGaah0j-n
z*LJ<ptmtxP^nL#5_wDKsPEL0&yTi|2Q(8?fv|itu+xLi@HU5z{>-i6+ABzt6iyiKF
zYm*bM^gd+HweM(o$eYg$t1S;-GkUWofhjHh%$#{#Qp;qotbAj}eQUSQiy#iMS~l_b
zYuHQ|YxreO$<4JsZp*^XzNuM|_54TM_HWG%2mk-v@}hSRv$o@tt!KZ6YLx_Q#8llC
zS39k}HDE(eXymaIo7O5_KN_`t*-TmcQ@3=2|83}!)AF}iKj|iaNVjCj`+|+^BFeD^
zUFJ=bcCC9G<+C)sTyy5;9g83B+p4oh`{{$tXXdT^_{Z9LJ;(J0mv?>m$L@6X?DDRJ
z#;A>|ch(iAt$yY6BeeJNDOankTeAui>=s(rtTySXzJL0Vh|9yQi4FU@%z~PiOy<}%
zFXp^-qDIVBZxyAzi$1Z%%=@=l?QV$XVq?jx%i0~}Zf7n1CMbBvP}gXcLWFX^DaSn5
zbuVj}ufO<zjsM?ghX=a7$6jygzNB;EjjUDMhbFPo7mF`G=dixuzxa_<0>k>tKK;Uq
zn;y0QSF`SXc6Q;!<ST3E9GWap-#+6*&$EYT4~yG3PI3QhD5Z8p_ynV2|8=HC@h2Z&
zYB7KCHD2xYv6khmz8`I7sn;;hcMz&SWbE^M+F5bFD8A*#IfWM#%-!<-kiohnp1WVp
zUT~IY*vogOQrwGCExy(8iNk%B--1nV4m<Z;>iR9?+_IzF<oiRm+Nl*M!-ex6`EHWZ
zo={erKQV8PijwC&HT^IC0VnPnx_-I1r=;uJkq!IbO1!Ns(dpka=UY#AZqoPjFJv=?
z^NUth#(ZS|^x_ZazJ~ADpD+4wdVl@gMcNPdN2$~_y-WV`r!cIsjEkE~c*$K2!99+z
z3s2cDxAIMFW+`tvQBgSam0_o*)pegI_f^-n+cm6vY#Y@nb1%|s+8mb;P4~+l-TBFn
zS{!;@+pcn%i-AE|f`LH;`#eK9WQnP3MQU;>e2>}G@X7?4P|5$kn@;ALU2e&nv`te}
zM>Y4-S=ZBdrzv{|OP#oyb7|J1YbW2CMBaGQa`UDf2bYG1)*=}t5f|6B+Z|uBbu=sz
zJ1YJDx827>%=sU_pY<`@d^?x7{>|*~d%y2}|D*PJ_3tOe{PtW4dxbtd)~&lEcp$o2
z<KW!7JEAU<d~6FJ<VEg0-oKr-V7+`k|0CIobOF&F`W;>}UmnLghCEJSwY;x;xH`o`
z<#6@WFTYbN_HU?BXy8Bf-Fefe#F@1!4U#|a=^u7y{xPFQe!|<sMyJ^fwC}U0R%kS|
zKebjreEe^PN`w2+uH!9#g&pVg`FG9jKm8`y=ywUjd_noXA0Jyj$V{C3JLOxXy`A`x
zn-&rWCV%={!w}E$pqsHqvq7Kz!&HWSTBrSge%z|CtL5$MPk-BZ{6inj{xQwt=QjEK
zoQuToi>lt|e)RiC<BV^I73x~#**^;ZF=?2ux#IpMX<5m4Z>3~k%F5iZwYsytI-{as
zMS4YwkK>B>M!D;+O}r8$EE{;*XNj7bX2hl3C6}g6+f<fuvS(dF>b!<^=Q8@;o&G%w
zcz(j{{Ijk-XJ2ePcUStHO?G+7miU98V)r`lUaZ`F*tS*3)|&NE)8%D`YiEZ!`D;f!
zn=x;r(8<1qcdX8?k9u%xmdB-kTJ6S(B^<hZ?=GJasg*qwaamGZ{gm6n;;JbBqgxAS
zZ?s8U@=ZleWX)^iGo5{$!QJX#`)ZfDsclYP{`u{;N~2|Isy!n7SKk==NR`Z#+q_N1
zr?0Kj;o#A;Kf^@AD{X9Te}9c=JGA1^_1q&bBTD0ymhz`8J62lk(&uf{J7e0Ng)=!9
z-gRv}yR}VmVd=U93z_DeGFK1cd|ku(Ty*PJftSe@8=M=R5=-a3o%h$U_Ll4MbYC@z
zy~mDP?#j7!Ww~ka{Yf6i*2j&NQm*)DSS+lZyI5gqwCv);dxh3)nP`~I{W0za=QAFU
z-ye&Q-M_MRhwdc_)5&KZ+EjJ>Ej2uQ>9Msb+q$L(-^^)CXIn1MQQ3A*KiJW=#xWxF
z(uGIP^IUSHL}oQbD@;rhe03|tPx$<oWiu9^`=n-HcW{OOi)rHgO$Vp$(FwV5>5o&;
zwe`0r2=Dfut#7+aRb=uNZC$0AI=il?N{GF^yzIB>M-O9G)xrzQX9~M4S=PZPHL<`&
zE5l-*6YJK;u219_yLHbx)KhR{n`H2pos&N7oS3lZx|w|Lon1-aR^`pNG+59%rMPQO
z^qM6*_PL){dH(zD&hANDR<BsmurhvP;`PqWT;-Y?w_kUxNtTpcc6GkT@{`*i<QQaK
z?+O0DmG$d&y-%WR{{5foCB8><-GZGB_Z4r)Y^q)8yn4nF{gt~Sd{5k4pW?CNhiTnm
ztJs1Hi^-dWf~(eT5Z$uNtW~_F{riO(1(#=)1Y6`y@R1CQxPSgqX8Kv36H+qM-@Q1V
zA8ONkzw^-j%~|4+wpwM6Q?-t5@IJb<Vq@3z2TQg19(ajrN@)FSun&s76TCh)Dk}V=
zy3OK8r8_*2&UcNKJ3b-Jch!&Xf<K{VulLNp=-E3x!^C*$qx`;}*`YStIj^R6zSLN}
z&|2!B-tPY=ZWplp4>tR!DsAdoWis)NnbFA|z0<yI`mj&wz^r3acyo3i_?j;|?Z+>9
z!9Vdf%O6#q6FHEuz3chbcS|Jx{Q6l_7`>&dXM>*pndj?#lq9EL`j&p?iGg-xsL0Zm
znRf#jW&}5SN!xzXdw=2VmbVv;-e;8uC5yiC>_~{WUKltv$i|r8tu8<O!;iwtUGJH@
zdsSuop6m5)n72r)$aJTh`PR1CT)nz6#k#p%*()ploom|a`}W~PPusQa-&H4_{IEmL
z?5k$7zUVSu<*RaXUaxqnEhd-wPR=`(y|$J&_)F%|=i#SV_7;h5=lXYg@7Gs5lO+6(
zuGeB*dT2dM=+Ypr526Ntx>jj<Ptn>Rv{p-9!&|h3Pa<*m&I>ZxMm0O@4g}7gcP%+7
zZw{Nt**TSOBAV5!CM#4m1ibscB+>UBo6MIr^QzvIU#aPS>Kc4YCwXV=o6OmFriOpw
z)3%&*dT+(~ulIud`ZezDd);+cA|!sU-kufH&7Rx)wya+Ez-s!VROcMMRjbrudv*r2
zO!;|n&$(A!O!s1}ABO#!5$!nd*3=ngS6DB{9lBZAvfG+Z(DI+DnNg<H?;`(WTwR^V
z)Jj)ARxOhMd~y1wIj2>=UyH6dc>U>(#eVyaRx~8)Z031%YDd<$?hwb@i}H%MUp(<Q
zQ}Xl9jXznJ9tq4pbSvq6GVAinrDgtEX;Z|v&FgwunfCNAcVyIrPv5+~tc;zn#eA8Y
z@jP@QFO!J3l+d=R$LkUedkl@Dzb(GKuIH$oN<UZdM(-^*x_io&>RLDSy5679`tGsF
zzu#+F1YB&lCC><%62f_HvfRfD+8+H?FW4?w+k82>q`c}yu;>00&g)%0^2*LR7wx<q
zfA9Ub>DTS{hpuR7E1cIcQOd#6pHchyf}XPrQ*?y1=C<bZr6zCI$e-C-e{7oc^N4fC
z4#p*Yr{@IUNPazGW=3LYf|9y5<NmW(#ZD>}h=-k7n4}}}#kT3x#!DVw99L|P_%!21
zuSjTkYf#@>kxOn8RyJ9;W3Hz>ThY2_r_9z<W%D#1pUO!Q3l)j=y{Im8?2^ap#1<vp
z_CVoG>z-R;Ql;v7*_pKhr4!kA*mkEz>TF-WQ1sKuMRj7@n=LLzW*1xwyP2(SnXPV}
zotRYSCANC*(crm9-DVf|Zr*k+GJDhYS^NtV7oTraestiXQo)tqX%|1;_*;_9`|gl!
z#B-*NqTZHy4gIgLTF7p-Xb!X&R)46oZca1DIn}L#dvCZ{m>oV+xr6a}u$%hQn1<>?
z#Z51^G1@d=xYPY~2ha3iIY%d<V6!Qk%~o{TPUg|tn9h{1`MgH5Xv&IXm44r(I3u#<
z19PgGzLq_+X_%#RYv-kdJf|LQnK`XBp)DsjdgWiG6<bP7^**v}xwv&v&rjQ4?Tl8>
zQ^)gX{SNc{URKR3aKm}w*T@`;rlrLeb0c@xiKU-;bVf3&G%S&iH}GlMvKJj|gGya?
zXzlGXvUELfbo1#GUa$Ky=hB$w3Uki&=8eqZ$$ccinOnB<bIz+jHm{hb@0#E&Hdo}L
ztKi+gYZvOK`h62#^<swp7Q0!l=lz1Vu8qv8<q)5m@LVzE+wFHbw>9T&?FinozGYfk
zk><u}t!!d-oXl4b3w+(Y<FQWu<m9%ed#A_6_~swI&++$l`l2^S?LNIvE%o@o`b|29
z$LmAZo9(ktFiAQ7bbqkuqgRsR!)Vj?dxlIkh4tKT_qOcc=gh%(Vta38wA12M0q)1F
zj&0QwZ<)#xyvvdGYFJ=m$IX!A`OFW+m=Bv>SzYA1_tZK2m#SMUEZ1#{EkCnsmi@lB
zcm8c^jOR=}uyXCcLw-{~wpc&dwd%sm>^q-j-~BYQjIa2ib^hU|I}fu9=XQHea`T?f
zxz*|xyV<W>QLkoxYyK{ta^LjGr>JQkN_R3%DVm|bTOiK$i)eM3t#VM(`MBompK>Oo
z-S-Hnv1CyB_s9H^)rOn8M;(^9Ccd{=dZGQbt*Y3GOr8h4Wouk7`L0{Cc1p-Ak1Soa
zmu~(G`lhLDHZKZGI=$!S)JpD^M!N$yMBEfL%lUHkci)0N_slO>v81jwUj5F+wdLpD
zT(`&z+h=WZ*xfWG!Z%WJ&rG>Hhq^zSshxRyXaC{m<m7YHj?K~SU8ieQE$q+Tr>&WI
z!|nAonJVXZU#}#XKU!$CU9kV>n^z_4mov${`=)Sf!sWXUUWH{W{&4@FJ!-?aMQX|F
zO<W8NP7(|ZTG$)Lkd5Qe<}qlj0J;<>g1tiIxadE2zgw5HLM=F#9G=3=ePzPAhTe?_
z1R{kLCI~D@Ys}p<N9^T?mtl9_20ogWwERZrWsi-MEo}WBiOjK3zuD#cMnv7Oyz2eg
zmwT;$rk9sm-PLu9-^YLU?VifNdn@1b`rCc}u$bXmcdypAhkByYAB=YHiP2jh@WD%z
zRVbz*T+`z5@zsATdhhHFJNmj{yZrYbAFEjN0|M1Hp8DWcS+Ty${NY)ydy5ubV%#5k
zly8skr<&a{_Y|j>uiC-?KIlUp>-&$7>o}{{cRj6`-}HTTL@WP_h_=mKash$LJMXmf
zYyFrY&%bKL(e}_KhKl>uo_?ATBlj)&RDb`9503?R)`uLl6qP@ktR;8+cu0-R{D6<1
zHVYpFbJ?v;GM*RmQL^S*|9Q_@=YM}>WsRRQKYmiX|H;Gp-fsRsKZ>4wvAV9t_u%!V
zy5%RIO4r@dtKa|f`1X(*>HC2nV(s^-ckW;Ri2cXHn7Sp8j{oTFn*S)Bt8U#Q{`~?g
zZdESWx_9;Cg8Ka9jfZzv?Xu?k{^s@TH#fBfURW)gf8ns9t)0-jtwI-noDgfAWM({X
z%a<Ld3yuf1IcW6n*}HiC*3GLI+a6wi$o1y^`!_F6Z`7O3dRAHMj#RJv+FT~?*RR-F
zYI2_$ES*~F`uNzrD>np{FD{qwlnnATUUZ!8i%ZD6*ahdMdevC3mCfv&_`$<TzEASu
zWX{@}J030Nw{AP8UtjC~*yNP?vG_*S()Lfj+m^oW6#TUMgN0Z??77`vTO3w43z{Cb
zUS!z3eCcD($?aiRJiIy@CK>pyw)MJmsL*zz$&Zaz(dRvMr1Ff9AN%rSgGt%^US~H~
zLF2ln9+iWsJ6ZNFpFGWQo4&GEx8rRS$x4l9J493e*tpy(YTNvxBW~WAL~n7+Z%fx?
z=-4n%FSFY7Yvmig;+`_GV`nA`u6*Mf7x-cmzi0R}Co4b6iBnRA_8;S}FEE={Qhj`7
zhheT?tDVi9*UXJmOE)hztX5}UZIvX`8~$X$_2<WmCa(6~*7IqXjfrF2Y1xE38xH+t
zx2pC(uKPE{E%DKhWS{7tb3G-+Lw8*gHBL~OWajpCOV6Cu6^9%P-Z6)IN_ZMi?T}dO
z`AFi|j02B)^5%a0IOS=B=PFrG*W|<OJspCw+l?~Je&^_#Ux~Hw{dDHOo&KVoMNV(G
zo-KHHdeR>5+d8aSJJ_ygNgULb$~)#O)%Q>>Y%zDN3h(|SmrFv|SxcJ!n`*tpS2M}1
z*GSXk57+*%57NB*kG}6Hj#==aLM!V+y7d2k?G+Ed|7dQrf8@{gZ_PvfAKH_%o;kg}
zG<oseUpw0NKYTCsudzPM;}IY4@{c*ZV(&se==1(RR?jl6YRCHEYfmG3#Lqq6aMwn%
zf95ZP?1ti^pu?XFXMcHQy!oz~??(UaZqbuHlY2x>6+iwv-KyuXmi<%cnoQ4bY03F<
z-rIXcopcgbr7mYbTdFZT{nVnA-2sAjz9Jn?z8XC%J_MTvw%m?zTbH;)*p%t$O%*56
zqm%rSch2zd4R`5O^KzZ3bjoz22a8DH!cMP^5rW<(fvT%sZn1o){nvJ~K6_B1wpCgR
z-%=?@SDmXzT!Z582hLjAnVPioz&0~mIrFs_HaM2|7z#Q~oH5U*@ABq_i_?^ztddT8
z&LQhNBY7dm>^p+CnJrn@RkCNNFU>qZG1=|Nr6XJ$gS>b$lctJ@^V~MysrfMGX{240
zu3lQk_gRY4s!tZREL+VI>T$GPbLpkkOR^Wat+{w4Do}evr_G%x;jY_9)uzapD_vY1
zm~0`^JHNetv(hV*!2ayWdA6sxm9M8Q^Dw{deA%w>si}9ew8+I3ou!L9uLn-ZjS!yx
zz|3{c^zV<wtRC(;>JpSc;jN;WvH!Q@`5m97zWA8;&f4|sr^!91%U8b_o&T7?CuV!O
zw4`s-mB&W<a#ts>?VNaJhTwL+B~fkb=WzD9xd(B^tud`kHNP)$^Ap#bmC5_ARrIk&
zrj%Vw)|N6{^JZOy?cI%^*jbOxO=ej)%YE^QM~6=BSZbH9-@anXRjt)=3-rXNA6vKY
zgN9>jr8s-B_3q%cWz~KH+g`3c($9ao^78K4PyKZ^zdif1>gUD}a<1Y}`7~Q%&DE3F
zhZ)}Gv+++4t&46*x)F9kv3<{`W5F$#du`6|+10B1{L$y#)~8u{<#dX7zt>X>@Zob6
zJ#&WJzwKvbYWkLW4{}S47HVvsu6t#}w^v0eSA>pL_)K5yrF63SZ0TpN#n%nOt?h$m
z#hWqbPJVl2<HT!`IRb0A-!A(+E7G=>eR5B%Qr29)=Z^y3bnP^hu|KhWy78osH>Vvh
zeKzIKs`KT?xZM3>XRb6iHn-N@o_+euk_Yc1SLjS$Yr!AvcjD^JbFZX*XQ<DuwlJ9O
zxm<YLy%Q@E_j~j_*0go2*FC3JF7<lauV$A7MTyIj$F*&m3Qsg2)1J_6oAK5<MSX(l
ztCrUe_g2KccP*dy@$U1UqN}Czo(DyJ^l`L1+y3Ea<@vCEhnnS;1v_VMN;Nq%ckyQC
ziMKz0%Cc@gJkR{@f1A6<X5@Sg`RlaDShncxAA_F?ix1!Vf9Zo%LHW6=yJa(<@vba9
zw(woTUxu(m=53506gwDsid^>1+Rdu*s7@n$eNnwZcKo5g3)j?YYPQr|E7;pN!Qjab
zUJ(-?2b+Y&GS8win8P1bt((=%AnX1kx`m^y(>1BE?NJwJb&U}BE-Cem8)kRszZc<Z
zQwzB9i2K*20Ji3RP3{lmqa-w5$$nzJa?hme@`bbmszNbZ2b+U6lv69c9XB{!(&PFi
zb%|q(h~J_`uAEveVtxxJX>?zXnB*2ZLH(-kzXkg<g&cQA-U`~H7h%%b9d%{FErZrn
z=Yxwhb)HK_7<Nt-=I~9t8gRkB@v88S3NF(pTh)5E8c%!g*6>=&s!BI^?b5@q0$bU1
z+?Dl?TD@^H{&qPxJlS^dW<j=v%QtXtPcqJP=Z<SyYtgLzQ6eh$eXq5g(!D_In*L9g
zmR|j04;5X0sycsCke{|V!o%|WW40q5Tjh@%wpr}0v)!9F$y?Df#yf+FD|Kt3hHUHn
zOJ*N0PE*=$eC0KZpOMn?980FVf1a&A$6L~}N9@D)Y5V3ly$x_)kjM0O@`C2oF#*Y9
z8OIcthVW##vs%5FB)cupTPAGBM6Uz?tkR^Ho)X`r{lUn5bJxjr7k`$UsqeSi{Au-#
zZ`>yyJ>%&P6JLKS%%ZpSPd1kc`=QJjqsFg0-|X6aBQ*8K$FJ`;?Vi4Fa@G5F2ii=`
zmKA17O<%e(^L(ww%UOq>Ej-2<Q`lz|#T0$%Miu`mRfB@1?*((Dqi$TV?tCj;am_M(
z<~7ck(reQ~{qD!S%i7dr)DU-wjpMDe>aUqr%|44(SJv~V?Fm0Q;cnOS3H#=FevYx}
zFW6tByL*@2Spl&R<qf;oZe7Z_fB6H8S5B>!jB1%dX?dAx(B#<O#&afR<?Cl3yP<4&
z?d{XR@XI&FD%*E8?JKT+vw82H7q-t&)T(Z~aPZQ}+;93D=AOKN_Wj*&>yyv3{SyA8
zwdk3u4r8d#lBqsNrIr;d*Gey&zCPi~c7xn<<HYX<kxP4dr(E#x&R_m^lh@WqtQCq6
zwK;1ACr1cuo*`GUQFgJ)yMnj((wFP~y&<6Cm3;s1ijqlFXaCKVnP4~H<-O07xo@^;
zD7G&3`9G5<A~9b2<lRGa^}e6g`+iCD%fw7hRi6d+=KXnR9tFKB65Nv|^m-}!4B5?d
zf)xwc7#OzjG2lEL3bOPZcIw*P@a%HwtD<+WOGS3sc*-y`2DY;pC}`jC;^9>Ba}u_0
za5Pl(dvwe<`q(9fe&cz8?#k*5`4{Xtq<n#C>5OuDfedfys5KT#c>BxNuDe?I#=VZU
zenF~yZ_4RSkJP?K*?qrf{CwW#_m;oyzMTJGQzUo5?N6=buEK~T<t{bhN4~q*C>-io
z$et{ilH8#r-YDtJlWbHGQ1w&e*wU>o#)^+RIHHfF@GLG^RFE1l;o$ehACByjXzrhs
zvFuiX$nhf&cJRn&9p?XV*z6C7z1YVBk^YWf7bWVNjwe;bJnDJWsWGqTk><;J0w1@W
zsaJE`a^CGj42!yQMcl&LQyZ*v&9iST4twZ#+DtSoMXg6uId^BK<af<2oxbO^elbs5
zUJ#l4Br;j`e$ia!dA8;!4gY>=tZkmN#p}$bb)OeZ)vs*3p;?j|qd)&#wc9zf%kH*|
z0}S?l@wo8q`s|{*MTTcS?6|l)^K(IFkfI#RW?uFAa?2VgKW_=EzoS)_7w{rOEcJrs
zueRLoyJn|1{*mSCSD$VaGovMK@vSKRs=)5Ra@mjzzwXLx-e;6-`79u_=**XI(#b*X
z6>g`BE_~3D44*n{O%Nm3uLX|{nsP<D)l9GJY@BhfH)gwAx=u*-rqvtw8P;zy2>%=4
z*L5eUR4YbK!svHd)dlmNT-Sv^dmgq<s%rdMbljx)DC74P$-CYB9D9Xkz85okHYIbi
zzRR_^=W2(WxWvwv*0x^d4zJh~5dCylo91mXr+bPTW%8$Iv7XE}Y!c_#wr!nP-FxZT
z2AhR4Z*G`c&a?l?#n<PKvXvH|xSDhFv(It=@5aGQOCoP1Zhw+`HsNEH>3-X!V?yuB
zHf@VeN&G(J-;UG850+ciwTYMJ91N3{-Vo@%z9YI|b)WBzHIL2itXs^le=)pPLsma|
z`=P15EjL3qy_xg!<qylaiyyi7oT~b|!@b1k^bY>@FDqp8UsMDiiC$DK{a&M`nyc;k
z>mN&3)VIbPN3BgN+2v>RSFNeMbWeM*)xG2Ff`7t}zyDAzEB}b!>R-oY!I~wH{eKj*
z$iJQ;XrKD#d7?qWt!4Ut9@1)8zV=<dD|yDPAndrhh4{hvl0Tx)XYOP^=Nz_BVVlYU
ztNTkIF8|?k{P_>3<M)p|ZP3}3tSL9yx7RczVsnZhXHnfLnM{#&`wC>5&Q>jZIGg#&
zyhGELvAbT@-ItNyI%(7K>Cb)Z*DG#ddgNjuXQX!4?bbQ@v`4oe*1hsGvosZ3xoXdY
zBP-ssX|=wdeAaL6vB%5LYVlq13U55kdnZejd5(5x@$-m_TfO(~y#C?%ObO-LrrkE_
z<{y@vjrqiRD@F42&aAVXGtc}~kXH5h{N{R-$=)ENJiGXl$C}=RZvPW>(>vyZWN`9s
zpI4%1UlkU}6|Hx;SN3uOd$#nZ_e|Ca|2le~J=>K%nK%4mQxDfV4Nvz)Z_6u{%|AG5
z#WvTT{@-@jpz`CweMt-3Hm-fO^x$N>n^*rBb?RR$s_v0Ks`yzVHp6e}%7c4C6UySb
zeAazj)_rFW<6ZVw^JGf5KmBvqZN2<?=i4e5-tFhe_<CmN4Bja24UtS={@%}@Ru#TI
zwn*wJ4|m@Y#m_5NZ0O@ZDQD+=DI+(3`flEY+tIm|heM;{il$Ec9<%hR>xt(j9sg@@
z9J_h@=Iw4VmmBwHMA<V>-Brn^yYii6OAu!+Ta=dIJV%SE8FL-8U-2l(U);oMtr2Z0
z-&4cBH16g7j{HUPpDZ5vI%!wse-vD%dL{a~fbEoD61u56*6lYI8fS33RY{(HtgUe>
zqOz@DG|M@4&8%ayzNQ{ozHt_p>Zu2ZQ#bGUQ(|EJPfF>QX2!y|vX^H>T>39s{I7l2
z=?GgTyY#};s|IzV_Q6Z1PFxTw8P0xxzv8ZUn_qA&v@Y0|wxlgh>^jfK%TcXsqnxa{
z<5pe1p|>|%)Xupidxzq&(w>PN#f*$ovL?7xh+XfW^XvWc(ETPabRM&H9{YbSGAwoB
z?Ac5!bQ>2}7)i8A{SA{9NDGhsx8s2HjQwk)RJWLQixlsvQ&~87-3F&sR$W)8?Dd%B
z9DQOD&(7i*d)ZIS5l%Vl`RSQ+*f*QC{~PE2FfV(k{bjWUTL}Ao{_Mi(yJVZQn7*%g
zd}77q8;j!(PN<q|BfRX*<Mj)K-`w?F)@D&``z}epNB`u^-G&^#JA$4D2Khdes|&SQ
zQ}!?}DKl|TeIkFGz@q1&JNn~ItN0`&YIe`B`yRMB(D=^2$#vh^W1hqXiq2~bPQ6o8
z#i`W3n&H`<0J#-6<d1kP{VT>7y07=qZ6Ago#S^WirtV*lwr4p@A#1`GMh1p3W(Ece
z1_lPu!KTndOUofso%$$ymQ{T6lM{1XGK))!GLuS6;0NqX4T|=bb`<${F758slGTM<
zx7~`~QrVIH`qrapZ$}{(#VxI3@--ShlN9}S-Q=C@+x_3`V8Dlk_6OucZH(J4tq`cW
zF|+u$Uh$37`^D=S_8eVpc_v{k2fx4XmoqZwYxjH!RQr)@dZIsMzQO$$UGtT~BKNEp
z`dN2pmA3DPPsye$-UygU&3t!6R`*tCPkrXa_ujJgf<HgrOHimj>ymY>Qa)tWr~a0`
zufBZrym&I_^MnW2Hf>YUF$vMK=dJc#xu26w=#bkrp_2<QoO<$jQFf1zSnm3B-?nD)
z6;EB-@<73N={nV>rIjh=tG1rf6{?HAG4bbr)!p-gN<Tk)*JX1&+Eh|Ke)Gb-ldo0g
z2;4N=!Qg7cA&_8uKqig($BeHV9V^P~``;Xk@^rb?X>Y!J(^O%u(9NelnV7BlJpbWG
z`(TF*n*tFjXNQ$3I=A@ph4_EYdH%*)W=D;_d$e-E7iqsQRV_|m(pbLtyyg?rX`dN*
z=1S@nejP=}PVLmRWF4`e&AmO-l9d<dZIMe%T_gJ0x8}UdbB#Iwl)vW$Mjv`mmEhx1
zQKNoiFQ?s3<~I*IZNqQs{&;(2uJ|9JT}%Pqu)}vj$6!uM_L7$5U|{g(XJAmoo)TQ3
z2h2i-_FBU`>qXKf>(VywHn!%m6cdWjn6QjD=baMQq%BueCY0SaIjF-K8&J5(`lim^
zGdbyLoFTb;SAEUfDi^vWPBglI+O`~x3y(rWY@*hNMgI~1BYx%i?7MFS!_Mgc?EC)v
z-s9@|&#RxGo?1V@p7lW4pWxKw)(y;-7w_mxMW;U2C<uL0=yPll7whrRIR}`v>{|G{
zpD8DuZ`P7K<epmNbL@Di$zx%ySsy$_`5y--{k+U27x+=KCd_1h@W;v;$D`sORWI5|
zEqZ%gq$cuc_=g=m>$QH&>N@+Qis9aQm)%YID=L`ohkVqX6S5)LXP<7*9HB|UT=POc
z$g<uK{b0*_Kl}qTtN+O#3%f4=nAo#C)kgX8aUliA?wQLH`_oE~8_ICXeiZJ@)_Jpi
z?VQBjl7}C@_|f@rk9GcxJhhd(r|B*%{D0+%Z1C?#AC(uE`}yC|b?owtu+5s;w|e%w
zo_S}bd6m34SGz4ce(amCZ0B*lw~I5PH@75Q^IFaub<-m6=F7${YnN|2c3{!k=PwEq
z?nKK>x^Sd2%;HSMvzC1c9vuDhSNmTwN8a_}KN4%5ylkg#%+}mv8)VB$3bT)$*-@?1
z<Ffwo<EY~6Uc1s%ez>m{xoh5=?vt#v-yr)2`<I(*#3xL>ve|goT*1$iF1$aNKI6FW
zSDE))#m=r-x?JKrujlDGr<Z)*w$bR|Jx-}bhc-{r()73?q|25r_3przkToJ%$D|}q
zr||0gMEWnyIlS=wv0H0h6jLqoj_KUYTXS{lf_K*QOp1~&PW&j$dG>T(^0Mq3%jTup
zR9f#o^ufYZ`g30*|MFg;7ia$D)aCg#ytcMEG~vvjLsJwF<)6A47Ug1|^|UoG=$T&B
z)=%qNwiVe-t~8%sKC{X>enF{y@sWz_*DF*{`zUKRJ>`pUd%ZUG?~KT`a%vVQYr@KV
zUq01PUDI23b@4>6hErbxkG4r3_TgN&UdzwRTX*VGuG>1xJuXam8p7TF_K8jJ(Y1T#
znz*d0be|@9LH|w0jqXC;?T>Ax+6)&ZhM1Hm&*_|7emADxFTj~a|M7CEy0&!HH#_;(
z=KcF$9=@Z-@0cX-`Nyo1`dJp;O>L@kE*tc6O6q4?^m9LXUgYO)<Io#oYIbDFET^m9
zhYq@4+@teYly~}vVqVP;FIHK&&u?$)iO=|<TDSD!^@`AAsl59S=BNJgJ1qJq?y$Uz
zf&2manZJ7;bj>}#)6MJHe3LzG*HzbO9=45CTi0@;NbHt)>GHi^8{=J^O;^Srn{WDu
zHUGqq)ouG9`tsHve=YS-QRv;lJN0W=c7@HFdLT#3Ph3mrX3;NS{*tMSr8d4hxNqZS
z$rmNJ7SBC5;rs1S(<i*E0}6`1X}r{VtjnAKSW9$0vq9p@%~Gp%@^{HuZ_{6X?sN3i
zU8^#>c#^CxwM9R=$jKSKxKt~7Z}yeI?f1fzK5&Ib|Fl~FW6g@bhwn_EeW{ryUYZ<w
zaP`gSH!rOEdAwTn(AC^`n&)Ou&#zp1f>r#TcWTaY?X~UGc&B(RX=%vbxGh98EGx$Q
z%dIHs&TH!<r(W>-wdbVIqDajsA+t3Ffx#DUy;{FvhM{@zvA>mKpSbjHzj_e#U;5e!
zme5dl&eyAqLZZ#CUlH6ab~{d%--W-w<Y<;gal*bGTdZboXJtD5`IcAoluLhJ^$4$y
zIrr+--uV}=aW3}q`EO~sILy|KE8*Ys=ks_3f+y?cuMCr$d;RSWk2Nz*vOH|xycLr#
zb=$m6?ha4-#dABJ%s#W_NRm<m+rB&v|1{m>EA~dlP3TEFw@JK2F0p*BvW#4Oq{OSc
z#&=G9e_Xcj?!8V1vjlTCOSwn8FSJKn?q8d=k3I2Y?L!M4JHt5pS6><5-R!Q-Ic`zB
z^ZNO^+V4V3-p$+mq4fJaSLgEkkv}xK=6qap_h1=wQ;vDz_q3J`ti2h>rQEjg$?lk}
zeb?2x<DuUBwY%<p)~{3+Y&<XBDIsO??D#s7gA&h^uHP#7AoJvO#vTS$MP0=bzO19w
z%H@o6Yl`xhUSIJ%Q2W(X8?PPg!gq9Y{ED~VxPK-3z$&IK&3A40dPnTN6vKb^YtfzO
zGqd$qP2s=L;`!!#jLE~jGZv-INw!yIx0n}KruMg6GxqG%)t6dKBwlDPC^V>G^zG!%
z@Vk(GTIi7Q#017{HdnXne5VBKcbF=PUl98!z9G-x!lt!97a6T+DO)+w!9_EnHYxqY
zcNMksia`h3)Gz$9F>A}$NVv?J<JYoXXW`$V*|(&sBJMBUwlcf$vD1r>!5MGs1owUL
zWxZ2#L0D?{m3<wTE=XxD`23}Q-nX#rVjFaYcJQbQX(-$(m%LlcS~B;TZ*$b&(#SRL
zTfEaNt}yt0YGwEr`{h@XlDvZSk$N9HBjbkhMA5P&<+6#w@726^?=(7bQDA3*^u+sz
zTkiF-R4$w>{Cy%X-_9QjH-B&h|I|pd@JzFD&9iYfnXk&AAMn^Fb$w7->!#}()y<o>
zbG9+)y7#aCAdyxu;ZM(`^WQJ^`{=)4KHr(8V$N=U<1T@x?H}w;g))XoK3knSWBs&*
zZx40tockv5$yRX6r>v=)_ty7}(#d_ZSZW?jiO}&_u;KblkN*`rW{2!kzjL@u>CWl0
z3-a%TV-@b`?@Fj}C@K0eE$wM)j=xF3YQ=dUz8?{8^iY_`;nwuAMP}k-&ZXw_GWfo2
zRLp(!{6pXKQ!C$pkU4kZ<6a*}CS99*RdS`={5z)}fAd(;rE6JN<5WJG2d`v;5Bl4d
zUcPrUCz0vd-fqox`EAPfYkMAO{Ab!?>axOZ#yvZ+d26P}$&2l=@Nv@h>h8bro?({g
zuUn1>uKq`BR~(SEI~Ksez~BcxxE|8F)P=S#b-~9o79^Hr=oP@vat}Ut)?f3iXQ+p#
z@98tAbxxi+|2XK|+4I2#p1xrQMj-}9n~Y308yN-`l<@`{7#Wy6)zUb1-pf}{bDOWe
zm(NMhtG?$>KGD)Tt9|l}@7AYGNT&&KZJO$?a&ubhIR*v>8Ab*M2?hp+q|(fs6y2iK
z<kZZ95+qlc_V3L*Y{1h}zRye0aZ0^$#>ub+CoV0@cJ$t|_`oC{PqSY?w$5(3UHjzY
ziNw;)d)L)HSmMmXbC_}8oaZqo^rJj;0yIN9de_{45|^;(`?@{RYj>4xIJW7nmZh0!
zV7@^6F6;Tfr)4ZzmL6idwV-<^%cnKE`5Mby%F+Zcv*^{cMSap0v%NG&^4*q;=F8tC
zUC^*ntY6YJV}1M2!>wP-#n;!VZV>u(W!2>)jgyTxEU8Ej{-Vq{sj`kO``MRz2g5D=
z;w^7CK5?44RY)WL!Apw^NQY^H!T@y10_gB85Y3O;*h|jIM7ewbq<q4OFVk$87#M<C
z85j&Pl{+SZR){B;IOpf)Cg!F1Wagzh=jWBBBG!se4T;Scaul&;=I++s){(pR{ssQ0
zi%V|X&2;kMxR~LzO7PN=<EfI@a}%cXI4`b0bbpb)$;Bh`5BiPnMNSBqw&hFW^Lv)n
z?{^kIyLZ>VzWyIuL5DIk=bINdOa%`1G0L8w`rWD0Qnpf-Ro1nw;^^0y$kdK!pO*7P
z-F;uscD3#G$H|H7&$^z@D(CbxuetqdUeuw`-<&H>o`2V{@XCkQWUE83N_J+J$8NqR
z=aG1G)4T6g)2;`ml-{>~?o<?*c4NnxuUr4mfB!<wWop>!7*(^oa>_i%cOIKrnm04{
zbn%((B8wtU8{gRdO>%pxrF@%5x2nvgP1`z})_Gpa*1Kw}bE5aVuHl2G#O$~OXAa(H
zHx$|XV#-v>HJYujEgQFKWTqIz)o!ee-?h*D!p7qTt6y5pQ*Ccr_a-6JY|$z95_zq)
zk2iE>N$hg+S;6ggS)Ob5{<2W{&Ea2j<F9%iI*{9Tru6z$7tY_FGb}tF$?-{cna>bt
zyRt{E?X9Xy>*+HtS6htV6#DzkcH&(q_~(h0qjOpD_45tCE+rj2W3VU2V$09UNeY{d
zon6cSsfB3CY`yX&JMx#dii6XO`+|zzXBV1!?v{#lGYi|n@Y~T(q&%aSJwZx)i`sF<
z2-mJRcfCDk=`23CHgMje-)>>byMJ;|>=*BtC*q}?(Hd0Pa^oJy$M=`Hv{M`_IoDjS
zQ(h4w@!U^KGk@;A4F`7n`?o9#u4rhVvijUaP75>j=Mu;5XC1R&_Cx(ic8kr?%??#<
z^EOUD8kVpf`H(G8QZ$NB?O(>g!0>_rV~s8-<sv1)fW)HW)FRlhY-^7r=rRn~{~24O
zxXYgI>d1CJA0WcwdQkXC*$UlECeN=0dNvL1LcuY+3hX}H%ctFY{kMYYgg`rovrEg1
zoV;xha@Wq-aKiVtSM0m&v+9|$w@a1ZPcpdw!bvjx^{WGSuTD`+YgxK+>WVcwqTNnc
zA3d*}D(KKr9<}-G)RI0~#g`MEJYzc_FMYU3dC7kd&sUbrHR&%{?=$t(ut`+U|JD14
z_p9#A8nlus$ca;S1rr0qcQytF8_aBn=~v&(;^NG_bpHZS$K0_fJ-8&XBo)$In>#hu
z|8j!J@&D(|cCCG#vOr;<k@%Fo3dtP}4QfsLhJGdo6Pm*L*UU7%8+=LZzW$z%$E5!%
z#Ca@I{&<N0LA%ks<+{bv-p8h?y|wf||8ws3E#PzT7PN?p8a@d<64)`>MD^-wEuk6H
z(p@#rx)@3yHDx<`tCZK<#qhYnmm8Dk7u>%1mUsGYtLASG9Czc4cRqf>cR1$erqxGR
zoz-3zcgL{z(X7zbg@KyVPT&8k^*Kz_Jk_Ph^Ihz)(>8gByvd52(n2-nJ7(_K7k>NN
z+vPR4mrQ=RY3t#$zIoC?{CvvqxSKD@R;F0=DXy~nv@Mo*Rn(h#uI~O?s-chbx0hbZ
zxnA!cxasNiKcf1xVnVv5*ZJn1-l^MWuzj`Ber0{hNz*bMb4ndb8opis+oF0+w@^as
zqS5`Ueml$gFH0$;cB)O@y~t^6dhzVQm~@82pZB%hySw~}uG@>VJ$t4~sP#QQv>-2{
z`deS!`3c=dRY%gIzBwOO2>N8U-tj0WFSo?y4IR=_=N|59Q|)ngT=Vnaw~oYB)5VO>
ze7on+=JUPGz(+3GqHb|=Tk#I(RpuXJgQq8^ZL-zxd9^iZ{W9+9ig(XSX8vk8su-Rb
z;n%Ffs%WWbC8_k*E>p2e?b92Eg*6Y=5B*-!c0wjaPVIcTM8f-poA@m}(--fv^qa`y
zKPkfPdS(yng{2Iuo^Y4C>@%(4dJ<Hh^GJiYugTcrbzz1#(>doZe@TUz{~HTrnOj#{
zO?jWY!(UGCt+ZLeNd|S*sY?V-iW|r^<g;-eQ=L)Nv(4Gw{Ha8~!{c6=%R59Ch}qaW
z1ooZ(x7~m0yEoiPZ&tTou5p&1&g7e@)Nn`f^oIKzm^b|0mJ(+kdgIQ9({YcNmatBa
zbB#Q2>~6{|>5^fgnj!d8SUZvBlX%!6*GtA<T(rI+t;h%E_C0HSjX;M!H?uM@*b<Wi
zAW1*Cq$o2l-8nzEAh9SBzD$2^$a#NZN0EPNyHXQ#7k*ltu(c^_VL*zef`?bfl1|Qz
zg2FtV_m;g~#uHg~^4+u_>kpiq&0hVHsZ;IP=K3ig&z7xq%FePVxi|a!zRLHv^Y48*
zeqR4RgUH7-2Lzt+N@QK!xFGprL!0W>8CKImXXOMQ@9GO}Yt(&tB4M`TwE6Z47f$8f
zO_EsYalGoP&AO<2(t%qwjiL&(57`OK3;C_Fg0=lpy2tyuESE~ZG5cHfZmO<PvJu=N
z=bUuCP-lG(_x<f0+anA2w%z>Ta(COAS-QT5<hYlwo3&SLUVi1ydsDr5s&wxJUfx;r
z)MF0UWih|QIm+I5C;p7JdfaxktE*Og!Q?na!9UK|xjByC`+ILqO!k&=+f7-HE3Fqv
zef)kpyxPTbs^Ze?5(3j+|NeP&(yo>}{~09PzNH=(bPVp?^i_Y=S%+KC<W;xtYQ3>6
z=-9s7U$U-ougo*Odo|Bf)=Qx=g~@;8y5#1EmGOP452JVY<y-FKd9w1Vt=#2WfmM$e
zv#3o-^N4+Zc0=Le%w0iibj1!==x>#f`?jBRmc;EMse_3>+A^fle$72-Wb4%I_DIY&
zk~fCadw1fknJJl30`WUo{`>G9=UB_Ebt2ef;`|Ke0_n#A#jc-R+U*tH3KNzWG|u?O
zxa*r>R8IBB_jh+qvu!q9I9K^x^v&X@^Ay)@-RD$&T2D}4plZR@%$oWCoDaAgo;g0@
z>Ejtsg-?9;6!|F<R;i}4<yxMUV9-64Wloc<cR0(b7p*?=K=0ydrS|G4nJZ<L=Jj^$
zU*gk~Q?&4d>@TIidImRM{rqa2o3K1`l7?UV*-2AohZ&ff20rxNWHV{<*J;Lg*-E&M
ztxV|qWN5?paq|rBT)E^FV-a~L#>G!;6}qn|+)ZHgpYW3vF#rcj>o1gfKe;n9FnBXz
zZDb*p_KrpArManjC9W08so<g>Qp8W4c-D*AP^8uWm+|7Q+oE!{Umbe4C6xX8)TWLm
zgB=@|{w|SPHgTqlk@2hSomx9SFxF@odU}iMHbvffGVjjsyTYITef-7tfak0b*A1h@
z+pc$An-5q%Z(45K^lM?gzx+z^K*6fNA%%iv{cS9*-j-=Bq4T-RW*=2v##SPdx%AoT
z1~=zbJ5D}NQk?st<-w8ZT3?)8j|+Zo>ZqP>)VESX`a~<&<{6z|URfsaYffL}{ot66
zXGX@#{V@VtrdRx4#-rwCH}mJU0>cH7i#)$B6=hkv?@Y*nbvr-a{&o7sO2#RHEy+%;
z3qJRU3%;p4yZ^MyspM@tT(~R$a;H6YIKv_vJNIkJyvw^Uzw>d7d-J!ePj$6f=~KS2
zuPY5#U-~s`sf9*ztTE4Oqx8EkBrEssetYht$j<p2Z0~N#Ph{V`iS^9f!)I1@@n#oZ
zF1v1rG(ZlD6?^fF|Hs)G7}|x2EP=f9^2_re#So+gf>a6LjmWMLc`91BKV@>H8gpAK
zLu1(TSrbx<`5C(}Y*Adaq|HIvr)SHp#Bv`tbK~1HZ?e2O>J|O=Zrla4u<DH*x2hVw
zgq-$ViQBN~>$OGx|G!oKvZ?)bfBstA^t7_gv+r)`wEz8l_UF0p@9Zo7eeT`S|MTh@
z9jxz7-N;(LZh@-eV&P;VA4$cC9?wsoD&unEWa46$`Fl$6Cp#Syl)V061rKYv;g?o}
zeG!fN4^R7)Z1;Jz|IqDfm736FRW-If)u}m$YmN3b=6^c5J*7tae&V6`Z69rq-HC5}
z|CoKxpMXRCKYm!q+bcYlkel)$S+aD)p&y*}>XnrmJ@Y?(Pn7I`knZ#QGuI!c@7_OD
zS67A{TCXxkAYjds?1*O#+0xHCS|+#LoZh-HcYWroGJmIHnY)4+Z#Tc0cug?6Tm8$m
zh(j4~Hy%4?Vtg-s#;>}wWhJ}SX4khn>|F2o;MM63zieD$1!ud9C43AyvRR^Y>6t>)
z@aY9zJ8Fg5Bor?M-MOYyWmX#Fz2wKVskQA}3)h-HmFm1^cewCn&d~<Lq^!JCY6nfZ
z`^ta)Zb|Fg!uI`LQsKK553g*A{r7AitIyuc2E5z7A2`((eoA#)xZ`=ZbyelL90?JQ
zuMvfxu8I0wnlEL&<Tm>jXU;<|zt+r{CUGUF`8%W0k$XZB<&!ggXY32Lym%+>PQ#%!
zPoi%c?S4D)%8a~SOuJuwK74h<oe<7!$;yc}g-ok&e4q99+iUljAl(aL`*UPhyOosp
zKRXj07j?g2b-CmtAGMNiUNfg}Vp4VMJ#cHyo;RhX(Z#yw)^MGcz52G+QBuoOf6L3<
zxJ^%8?M?`A_H#Qtxc{I3W#pV)cfB7@QqJA8<@2S3r^Vu|DsP^uP2fz@NMH^-)%Vg|
zV$QSIpE9B>GPu4i2tUZXW8y+}mTPC)H{IU0b@83_A1|(Rl`V8X7?u-WAX}z!qqAT0
zh12#XM;=ZSy}V$n$z}4LFCBkuXE}bB(<=OL*xGAXtxB`*T#~)ICFbhYcUMBshFq4s
z@AS9O?(uh7t2M_@Y`Xuz)9k$_tN&xk#<_Qwt?F03yI9h!)T#K&lH6YDLfdus5^rYZ
z2wv3{^3%3xzOS`|)m*zm%$e(4lE<Mb2evKJyPs{rmpXO6Wy#Y8^&z51HlMw$vi1Hi
zg*L4j2Q4`_bDc}}J6K%dcIY!#+<|Vcy5k35+?wsRgY9aT+pfiaFShMV`KkZ#yU@P{
z5BlFU+dp*bd;j-S;0L4Kf6rIA9(3KKwV_V&=(;HrJ6@Z8y}`VA*Qy7me_E$sXEj@L
zRkBa&`s8hL%JYw0O;f63Uhh#6d@x$)Pwt=4L-$#(Bu!oQK<J*}j&IE0vKsqTxAn;^
zyUJQ{_Vg=9Cg0PU3SX9q>(@5!dbHOl!{o}XOx@)<5|>w8&e#9oCi^BjqV!9SXhqd3
zuD$s>HVbDfzgw|<_4K&ucO=^`MA__XuwL61H<!cK?QUm_>BP$Z7aOPXq??QXOcg%7
zfBAufsus%sKgO+M*4`4~Q$A<bw{vD$H*LMIO_gNeswjH6^i!Oy6!*!@m&?CwW9rlY
zx8-RsYq#xzkbQlH;xF&)F<~?6OEQz_*)s80+n<GH$!aTJ+D#R`vq7*&;7i6+gY;gp
zi8VT@M^`-iaL1W#_PzAnTIaf;!V`}^{(QKsD`%%&hGWX#32C$AVxJkMCr$4B9@h8$
zn<|fLqr~B?yZ%OZ6lcFb-n@O9&nva1b;Sn{7EDarmCX2ek*KA=eA(KP&5^+eM2cc2
zm2=8nIw5I(D@@BWrQqngb07aVoVyo2vspnbT7L5}<BjY0&T+rA(wt*z-K4Th{`ZPp
zDwZf{913pzzVzUs-<NV8)@Ht2@bi6DllO+*n<ll3_HS-U7q|MU@>b48GqGV=#$yq$
zMM-Q+FIu*=r3aX`e#*EkvLxfP$f1j`E{KG=tmf37|54?h%j>BVGLG0VM%r{~eUvck
zn!Lknj?>x-nchd7Uj%*CPes^$OOk)!m{Z=#zVYc&xg*=!Z93NKE5H8sI9O6zwSS+z
zRc}}OhD^Rw1u3a7D%VKq_0DTGDc}n~v*6RwW0`CG<lR5@bL!>KIXX*-XX1jb%Z<0+
zI3XwC?<nLy`RnqHk#iiUSBUL>w5Ghdbl<l#yuUtX@=VBOVlOC8cq|a<|JW?c)4M_{
zbn=R{+RYK_2AxxjBka$dyIWut{#9$=zJyhu1CAEYvpOm+vO?HKmCeh_X6b^AA71m$
zoX||)Q=!>yKj-0o)!Zc&4Kkj7VySOVcwR}3d{Mo2*MjR|AO6Xh{faXG`ZGs<(Zy0J
z`EL#9Tm6>g*QHu9g^Ttb(77vIJv;5AU|6*4fvZhv>$87J&t4u_{OWK0&!jB<1<ea4
z72mLFxAHpv%j2-gBAr~R-id1RX--Yc9QWw{+gB=<8_8w1M0Csd&%Pb&9$)$WUnO83
z|DjS5whrzH<+x>Q_Ihlr7v0`{-9Rz!c;QDjwWG5uBKEXSev}&{H2uT83ikhk?}K#y
zy|diYeW!j-QQ3@5i$uOB2<IJjob*?Y^`UWG+y7>*qsJ$HZaz?QQQ}AB9kxKuRqnbR
z8+|va@Fi{8xN@geNURxiVVdI%`Hx%PeSB#Yx$=R))$Jb^L<A&qJ#uW-ywJ=QT+kk#
z$kUr}<u!ZT@`na?uT&!~w>^|t*V@y6`Ma><Kj}9TskNMW%k{MmNBG}g-}pb%y6w2r
zC*;#fK`jKCY_W}JnHU&uvoSE36W>Qc@10By&Go+=AX1li{bHb%M1JrGRu>-W{L}^o
z9>F<DV&Vof6+b3Rp4E97b*b#g{D=GxyW<adB>i~PuQczrS@8z1htAdKZJ%47-}!yd
zmil#df7t>K^!1oLOZoCCod21s&Iy5y8p)ieh2p#-CQf|xTO`gi;_-)4PFGH8p*Z7`
zT#0E9bB$Dre%7*WOp$-ZQ!9G-Qu*v&=DR=F$j(?Ix#EWD**n)|n%_GweNk%~({|?G
z#`GIuC-Y>o`<~t8YYy6`b60cbtk<tvB39k-i(PAdYvIb(rk4XaZ(msBEca~kvmLvP
zD<>~Lp>Xw&EbFfcyIU1L=p0l2l=^5(?al)Rg6}Uh>0gSKX8SI+LMBsvR#sM6?Aofn
z#H-zny$ZoQ%^u%7I5F;X*IcPOu7|e|m9DyQOqZ22gzrbm#z<|E%ByPcnNnV?lPdN8
zJj-VB)c>+uUrH|fwVp@%W8^KBtesirb7$>K{;=&tS<C$ASGbH`gnwX*7TP8m_Wp*D
z<FwnewWV_({0SDir={^ww)E*)C&QehZ919j;<m-yE#f_IanU-t|EPjU&jZ`T6*1Q=
z@}?{-c)xQ;*~SG$JkEhj7mM}hb8E%Mes5ARS@CygO|*1K>AgwY@2DmuO62wD&nj8g
zpy$}^n0Y>Q(VL6B;b+`?k8Zul^?Y6MWc}pEeffJ$Y-+!hiYR`Vw)if?+mNWtbra^;
zwSQ=4VQl{L?vYnntKy`o3=N4kEb|?-0$==@dt%a-b!X-%7BYV1ws!xdzDrN>?YXiQ
z*N$>8`LVA$&&8&>=2q9WZWeJvjSF*G4)*VoF=!9m`GuqZ@8mTPI6d!tww#yLQ?vX<
z(qom*nYmVXH-9LRKk2l`!~cBDqV#F=j5U@|442_=SeuaPBVhS?spq5?o*xpeN3yqY
zC7z2YIk$;%f$I#FLozpoPyGIodFID!9k~luj+H*anwzHF6yKsWS*h2{epBEt$CVLB
zU&N#KY@Q@tRe#9Dz>vqzz+g{I4hhZ6#?y-48+tlO*ippR_M}RpkpqXo$JP#wlL8zM
z7b<lqEnONa64kmX;kruj;}<DyvwWt+e~6Fq5bf6zuWR`K;Mny;z0&5pHteyT$tCEd
zTKMqUy?3|2&)HY|?(CDl@3+5acv2Yc$n%)X%HX2o>Ny6J4?Q`0a7t6=8J9xSu56y7
zn~&$o9hGtZBUpZV*KRc?o>$DbuiZ`Uo_@M&W@6IO!rnO#KP_7NoNL<5=rS+WJn?X!
z==G1>UOo8CqxepW=j7D%Kdk2?@0b;3&&$5WCVP~-J)-U5;xl&s8~f%fOk}-mw%3)-
zwsy`&OULS?>DMF=o%Wm9^;%4G_0&tt-q}5yJxwp}iO7zv=VF83G$-7s&YI-G?EmKZ
zw$GR9c08Tg_jmdn-6;h^Gv?h-eQ~cgF0QE8%X7PJkonx33pOMNx4mBb>h`@!6RM5o
zW^VWIN|wKLFLOcO?H<Fp$I3;g?rkbDOfr`^@Zpuc$%e|Zd-<ET*lPO3&CpC-Cg~Q%
zc=`0TIbG2g-p^e0cV)7r@S|R>s*CDp_H^Z`gnUt0>Et-QE0M=-<AO(OVvBbiHTZu1
z@<D^EBDWc?TvnKt>z?xKp_zDBqRHJX<F;A6jLz!%V(;^+_2Qq0wSN2Cdu%J?)s5PQ
zF7k)?enqrR5aC@I8nora_baNO1U(kjnr!^LMf;+j&%XF%_6b`y+O}+YBX{W554k6@
zRud*D&2I8B{wHYko29kZ&HQ(h-8a$YFML)DB`~Bj)T`>H{PT*@{HY!2)3GSMP;ZBo
zcU+3+!|PkCm6Q(c`?)9UtI2~@i~fTQGr1>h+SY$!w(^7br3pLP_BosiIavH;>Er|U
zH3h3?1yoP$>HgO}L*jiy+QgzGqMW(YuiSFYkPj^Sq_|iq;)ef`o9lN?5(*7_<)N>~
z9wWU(;br#>vuU#S47V+VuCmM92cKOTDtdUT!*dBY=DEopzpc7MkG=0{F`Td4?f2=h
zNN2KGR>Ys=bB|0yD?4q3l5)}tZP)xk>wzj&os62q%D^y_7vow9v=MJeuM;vDFgGl^
z_;Q%QzbO)@H@XV(9C;|G(2y*AM~rF8k^qkF4jv6lJe?P~ZVYt~nHUrK?9%VOf4TR3
zS^IP;|CZoGOU{)1I^vxEOZ}JgtvdIxl$)Epmrd?GH}C4ayZPn&Z9mD^|NkZ3F#Th$
zjJ1JkvQ?j%Picbny%z^~(ptYPt2m=t!2Fr>&Jo>@72d~gbKYswN_KR!>pS;qWzz%`
z#l6Z)7FFmRTX4L<UE_Z5L4j4Oh1W#*z1=*IRSTu8e)3{k<g^Q>b?$ptKhx^Vd$P<V
zJ;LeulsJ>n<w@B>->&VKU75#wRp;3<nP~Iw!|Ihzfp3#iJ&otSWLK5DC9pxt^h(i8
zwb@do%94dfffr6yMwcJ-@Myam9C^WO+qCf3lWVeOU-9aAowje2!`vGl<{Y~+qNiS3
z^G?Uh;tt2HP2D>etrg?Zoo@c+&Wf#Dmdt7`yn99M?uk61i-LEr%$+P0^pc^+(0%h&
zR%dCs850F0w{p$Wko5NpnI>jZ9koaJz}h`oXI5;}`!>%gt1@%z%fEGuF#+p>k8hGR
zeWdg1&7pbX$!pdx?Dux^uPxYifcb(>znJWqbvMJVD8^lxb25}E*Mq-BW6S!$>q<FE
zfx7d*?NCnCkX1T&AnH|-O^M%xY+v`^m$Qw6nJi`&seiVYR(ZGVYR&7F(p$F{9Zg7!
zcFI2=QnGjfCtJy-^0}KAzud6ibfJ&XwamjC-mz}nww&EUF8XhHIERP6WYhtbeO8H4
zj)&6Bw|_XIc_n_u*&6{n4rkBWl(MT)r^(;_^8drry^m<FPI|ue^R5Ne{wH2d_^|4M
zLe8ZvVh<NZ#ifX_T5t4ri)?J(%k{0p-7Q0l>)oSyJ-@l+5_}H+{^8(Mp6s7ev*Lkr
z%^ALHG5xPow^?6Zberp2;{GQ84=bCPoqqAm?|v)(#V|hG>3)K|{?=%&OA-kWpJ$iL
zmc}QA91Q;vb<kX??0&aaXnRf3!@g<dF+pe6JSvMn6Jb>?KXcyhsy<t;cvhjLUFDe@
zFO=@(k4ZbZVcLADPt{Jl&#w@?9Q&$$DOYFv_3MER5?7lRBwOWZbI68ztDaqQFyV4$
zrfpuJvvGIR!8HeGw7h+J@&Bei%dk*YB{41;MZV$*YZp&T6qvkV3HRx1)7I>kD*OC-
zck}&U8k3q*6Bh5tUcAp&;OzbHpXbjhc*=b5)Vg9{F1=r6k#)yZg(@d}Qp~Dzw^oYL
zPo3v<%EtNjo!~9@jw+vxuh}|oSvGOAc8BMVq5@V~A-N>!)*VMacXpo9@8DVIka_sc
z;`)ETI`7?@__&xqspj31|026Cbjv;4G&`cL^Qp~J!>Bzj-8{N&6Cb>v`PXLVv)SPV
zGk$b%xf`(Txa23J`YCye=(^Ceh4$;iHRmpvo)u>KuH~#1<J*nd6N<Yvru@FMAXNSJ
z9)?hL<tuyJCeNFk9mC3DIazdS#hd4%<yV<)<`l+$4LMu0X4CbVQeg{Ue&0B+=<uF%
z%Q=dBq!KSryxDOiYr!^$1<p3}xD(!eVDhxMd*$S>G(W!GKW154HV2+RIX~{uE#s0+
zFJm9CiThrx5xG$>`LNFQC&f!PpT6-!_R-CE`%-geomDgME<dF4lIx`0d9$wqdwaL-
zU3b!G+ZBd+3aSUgj@>w-_pa(@d#uidy0WARl5&aj)(bw1{1<!7GDf}fS;55pPWr5;
z)>m>KE$A+L9>X-}III5F;Dd2x?W>Fr#`$e~v2=#k!#7Glqe{yTX_uK7_Z0c<YnQ2P
z@%Y`(>oS|`%I@3izb}1OcVoMJ)Q_3Dubvluells{lV>yUoE1AM`(cvzZy%{o4ku@x
z5MJVRF7NV8sRPQtD~{F5u`buYX;B+B`6cr@6WQ3EA2-}9Dp9{vopPyoQKhISr}~~9
z+dgFpC8e?6*0{EAh2OdZLi+NPuWn=t_4;`u=-$z2>GXvse?6|5ddW0=*FU67jzPVK
zIXN-smoqXjEW=!Qi&m;bMiEsYg*$|UG^98+C>FHv)^>B`<}zR3Ewdh#_h`Cl=843s
zC}lW^G;vL|`C;bS_Vz?tSem-(|Hb>4__x~qVE(nZtK5uHkUO(X?%dqn&*z=CE&uoD
z?_Zt=0`WB|8(tnhx~zRtdv4pCl_8hr7^bYM&Sl#!_?^XhpXSsZl3%Qy?tk`knRztr
z&eOoz*_rPjFUj6_vL*51x05F#?iSAFxNCIh!umT~gPhEj=BYGk96uY;?zo~-U&zsX
zlB4pw?>k@Jv|GIEl=Vcb!!>I(t~_(&IlGh3c8cV)dtX+DepmC}FZwmvHL7`EK-<Gb
zw`!i;uM{X0Wz$`}l<Re{_PO>$F(03VPCZlOAszTMSf%H~hQkbPcVDHvJ-q8~i}4!6
zkbmvIy-^<>%8xGH;;p&wvCIw5;u((DoLe7e{0cP?iHqI!)}~93_h+fpbcs`w-4iu@
zm3FFD9nyO<y)`*qDcp7!YfFk@+PC_HKRp|pZ%#?6{olwh^4yK(#}AGjTg-Hp?C9{<
z)JwEHy6CL9PI0PEozsh+qc`^-?N&bIVbdG7aoV)C%m46Z@t;b4k&}4t>qK=OfjTvp
z#BXRr{m~V9Gp4XJFdP@eTBU;{IR$CJA5x`H4Zm409WMIM?e>XFE*nFS_uW+sy=*Zp
z)v`_HLdDAN$7;SZw>1kr3!=D$vz7cOq))kcqBGjdd*OxFrAt|7O{+ETZd_u*sb)XT
z{Kv<0()oKngbVDmHkR_fw!={0`1w7{@7B+2(sy36|M&a2J@XHPd;SZA+jxu{1$xAu
z1Wqf|Y+SK>!;!^Ghy4~Dy}MFvW6NpJn1h=tY?L2IO3wFQAnL|+`(upI{Nt&T=hZ(l
z?GrfOZM1L3gN^#C7y72B)L1s=Nfg$&?uqYvEc&r7G-HpNU!7^^`iGk$>S7scoR7UO
z+%wUEf8Wd_vn#w`)Og$M)0N9V`J=DpZQ9LizKe@eDz9xTS$tyd%I;T_tRl94z4Ty9
zl-$G%Mythnuda)l)iSv^-u-T>r)gKS-hmTSGDX?Dr$uRGHh)&F{ZM#yRYogkabch-
zSMaWdbBlEPdIT*)w{7W4yYVh;>BS#w(t@8JJT{l1nQy{U-C5pI2id(c!*jf4w|8B2
zeN?=7D_^?TdX5)+rMH<K?CX8ime=q3;><P8IU7v9ZR0Pw{W|B-;WLBPx8nSVv{l{D
zO<!|<^mVE_H7DFVJ;Yl)JAd=hr*qU=oF@BnD%P%V`DXQACV5-t$Js$;KT4RVPPnUe
zx99Meb8l`|sCif?>A#rDa96G7)v2`JZ{HN!R7y-PZHQSHSZLu;|K;7g&MiGGsmb2y
zuE!p2i`4${Ehj!eXo`2}<!i@!4=q{Ay{LJze2G!m#|f$7D>&Y*F58*q{UXb%V9BcK
zT(dmZt_kmZ8P*zg$#UY&1s|BSOG-}fYW%lNXT{Vg+qWmovcFuMeU0_%>YbPKmQPfW
zlQy50t-gCxDXVqjG_PBG4sW;EyEjE_`4-i+Prj$SZ`jll^Pt_ANuFbyq|C3Kye?N4
z-L7KaDr)-sMcd9v>q7nirrf-~Wv3^@?JG_P9c>m*TXjqA%jTHH-tWGPy}Hg(Z{_g)
z-P8Z|Evt7pZ%tdjJo6gMg3H&Btcl$EY-{qaj1PTkKmH!dWNrUw#u=;raz|9;gVsd}
z48LdF6@1X%lh3bV(ZVeBF7m%sUgWDy3K0k0KLkpKi|%PIPpt_*`29mItNj6f)s%P7
z%#F91hKc;)u%Gr(*zAbgyVO4+2hV>f=49VnXffyI$|uIX;uCjYuZZs4%XT36x&EdJ
zqGujVX%hX@6n<VUxVhZbZt;WZKSbNz);DSTXwP8y`PVco^^e;@b;IQkw*QcIue{rn
zuw3M7(Am|C!=IV<x6fVYKl9!^ul!>MZ+#qhA9%t2+x3{<i5<%xScQr1X<sh#M{f1m
zCuWA~igQFuC+5uyoO5_r`PPjgQnQx5JLPTuz9^#T*O>{?>~q)38qRTjY*5~*;?gwt
z_4#?**qq;;@5$wr=FJXM?X|f4yI`B-o=Wu#%{N#tM}>46id@*{UuSUQ*V#qWPR|OA
zkgw@&&iTAr)h=IZ+p6vt<t#@OICA)BE#cXt@#2PfiH63RkjFP3SNiEZtXOCerxjS>
zEzkcUJWIXa#eIU9EI(&+OzQ{5bvIx9`=xNb-E0krp{VlHu1E3I;RcfigOqoBChca`
zESFxmGn(VRLrc<>q&I9&N)7!swkADdoAft3K*i_z^$os<Lgy@*FzG{h&itLdVXwCE
zy0hge?QN@8;c?%+XkMI~M_<zK0-f)6O0gT)oscSPDJ*+>ZhzRjJ>_%HwHO|H&ZtpS
zX8c6<&)WNrD<|8ZEjM~nZP@wSpgVTMMrr*=b$eQP;?z7B+Z4`Wyy@{Uu}86O`NeMc
z4~GAa&rp42(I=2U>4$K|lrp(P$LFqBbY$#%B=NwbU(ods=P9oH@6Ijpf8^;kN2X?y
zQ%zG?i(maN<!|aG=UR(<y!W-v-CQiFs3vQub!3OY|Auwhd_5eOY;QfbcsJ|S_d6V&
zh0oGB|39=o_WkAU7i!zT`?(*B`PZ)8WZQnz+3C}a-gi~(oi)Gq&Aw4;Hvit+J-Zf_
zuKsZ2)ZG2g^B?`ow!8a#?x`2fGL|YYT&IanPhTsWa^bK_?VQrGIU;^V>9b7LY{Xd7
zH?sE!cWt`lCS#<0D0u#v%S$3ANIpM#*~h8l#qp0UncUX}b+(-Q!7_j85yv-zHBHt^
zj=4Mk%|0<l^o`LkcBhE1nLZX=hvvG=9Jji|yiYjy%y)^0Dsz4_+^8^kqmt-<kM*0U
z&C{o^eD|?3*RtFWzOHz8D_`dPlN*iY_H<>Wls-MT|M2QME-Q~kE0*U4ZERP0vvmF5
zFA+1&32vFq$Jx9n)wu5Nx1aC6JPcm*{nsv0>-pa@|MIBDmBsxLUDVRrm*H%++uY@O
z)bh{uhtJ)9G1L2fx5d>v`dr_{s{T)kUw><haAKqHvcP=~>f$o$#(D{Jo%QTBx2?(8
zEAaI5jPEm~N@g%GTI6ZA&1K~})kSg(q+J{Rd)@A8$bH?kf<wS<V(QnJ`x6crdH)WI
zWfHyb$m)9B(SYHuv&kNocX<y|-KNCeo0(B-Vz}X5{wdBVpE+8~FHM@StMsF+%=)rN
z(Gt$%&qbnSy^h@M{;?+gH}Cmx0rk0MzrIdhyIOvi(T28|j@z1n+he}8A7T0Xyy``s
zq3vPCpE`Fd`?eqWJm<Imk9Th$Y1H1(T*gwpKg`E@-)ZN*#Y)bq?8P7JQ3txZ9Zhp;
z7#SG2nHd<YFbBE<it@8klS=~N3tf?_c(etRb0>NGA94_BJ<rPGQez)*c+r~!EK<Uy
zkCg*^v=^yJe0R2<W|-T~!_4{LD_l+Q592?L$0Dl?Z@4_xdG~UT<^3;Z%ktmruV?tS
zG1@^)x-#e4#;NX_?QdnuIBp6Ha@D$dglCtx_yq5ov*C`A;(j@{*mZf4*UisGd0U;?
zpX}xG`No+S^CSfNqJk>7EuZna`1Z%vyJ4Rit#5hE*Vg^9K(o1_HE&CHaGP>+QK?VF
z*B|K_HJ7hDwYBtSt^Iny?aa5x7dK_sy1bbA;#~g=)$2Y#?x;=?zQ)!k^s%qCPe!Xf
zzH`OXz4l+Pby#UAwc9k_YDud1y4w4VQ(Si4+8bRZpVD5(X?cI&aaES<{?=u@5ux`!
zNQLc{?OA*$(?HI!{zNwKp@a2&y#*G{bIaMIHr}@gF-U)BFC`i9C2`rg*)wye^qhNZ
z+dS`#GZ$;{X5nI4?qy=j85eg8*KnRTjOl;)CF)fnn{KW1<@gs~Ota>>p6+~+#WZWC
zt2XC~pXKT+zoIQcGRg_x=EcOoz{1MF;D#BeNKpuGwiT&3Cl;rA<`t*r6=#-YmZf4}
zi?lQtvLM#DM(07&#BDz!k9JMd-@Wu#l0vUnl9Ni7)8+4sY;)3X=9W!Aw*OK6gQL@D
z{Ya?!<}UeON7{#*P1t!x_49K(pP#cW|Mq-+{(VN14>ruq{ze`zRU2*hPT>tH*^#v%
z@4?E(+O`yLOD&6}j;Y$}JRi4NPvMb{eR`t&vS)nN@kRN;D)&p-57`OMyYyQ_g5OSS
z{zJ#;ADbTkxb!h}!MRlZ4^xcuc272b*m>Zq#F@Uw9_w?t?@yodVb|-b<m@+{3w75v
zuGKl@*J6An)Slz~^tepk)X4`WTC8^+i`!YPb*?p)J3Z&@#IsL5mO5L$i(C-1eBtj%
zCY?V1m75nz^-4rM=lcEo-%1&lvrjwHp3mpg6LwGkwIuOCl7q+gZ9m<D&oBC-@6cA_
z-zzKVvsGyA{@79}joZJJ<Fj}ER$g@K%bk_x6Y{6X<t47Xv(L|H;z3hk_KX=8ESfWS
zd!;4`N3+c{-1FzYVSLhj^CKI!u2OEh`!&lj$F2X_jki6r*C+3?JN&JSeF^)lxCN43
zuU~iyHQSk=TEuU_x2|;2Tu<w69bWm)SK%M@pA@8I`o2=|e`NLZnM?!klZEp+Jtb88
zIb_vmufM7<V14Un&yjOJzdzn=(p)cS)2iaHeAw7vr{t!v=Q9OYf7=$wq5ZUQvW;K4
zK_T~c!P9#e%_?6sD_?N2{DigF>J$Q(|7V@Fd8)YnxyEjx$U8^(Zg3H^e)3Vqqch3n
zPm@mnl!Vn&HO??{A7T--Tk=_=%=wS<6sO-3=7;PoPR2R+y{ntAQ?&9vY6;Z3bc=Yp
z5Cg+8Ev!8YjMNJ0S<H>(t_ZohHBS4p#pKD%Diax+ZIt*=CQo8;YWQH}5;BFwfTOLw
z!ozP;<4K<xGZ_L`EV`ogHE{K!gxc_xYr|YUI|A3NeV>%Oy7sPJ&il7t??vCP&ab{*
zl^y+Z|F@Yl&-n1EeS7@)jcxV%z1MEN{`h`t`Mp0sb_%<gbo)QM#P_TI=O0%)`IFDr
zOa0ZI{?AYT^^=9cR&Nt7t?Z0^EHk$xGuBG-&9$3{<7Qqn`Ib3n_HnD{FGCKr`TewU
zdOXoe^y3PfbI+S>^qv?$F^+nyG?(v#gna+iOq>48H`ZTPk$eBx`000+pFSd8o>sAk
zC)dX4D5sxH@9OdS`q3og<mK9$ZO4pDxD2|V2=#iGu1fk=ICIO>Pn;*ZxPFFkHETVz
z4z}Cu;=eVoOm6<+;wQliKE})yE8cqSb4k>x&!%w;3VZ!dS1k0~y2qnKEPId3^3O-P
zyi0clq=|jrDRuo-$wq^iW5qQw>S31tY14jZEcdXCw-bL+qqu&;Jy)stIqt73miE6k
z&71H}OK|6{<Jm6ny=(4h?7yB>v(hh|<8f(>!rrfsx*EPPNyR^VtakG{+uW$i#b3;R
zR&0)VUNm?62T`lpd#)ZoyNL78#YLNsuC>dXe6p;!+2=^g8jkk08L=mNj;@ioJR@}O
z#{`u-68&t4b4w@o`59*&Qx5FiZs@ztIQd#v!5OaOO&=flY0vJ^ydO7rU!L{l9o~m;
zAD$SSSY5PLvN!VLu^p1zP52xi9bG>2b%Iub1fQ|!RrX}Vj&14Li*vKMs#oUrMh8}#
zEGulu+#+RZFVTBcdiOS`^?b&6nAlFQx#n@@l+^Vd#;FarV$-gin$T@?;<Utpy55p(
z!?_!?zTOgDEw#<$$~FD_n|77Gz4dguiNVkCJmd0Np|KwmR%m5yjoiIDY<6LWXxs@C
ztt9Q;YphQuhe(y~y~#Q2&as5(8$IQ&8xKF7xw-qzd`06sVX4LuJB);KS7&w4(p6l&
z%2H&F*PGl_gPfbE^kSIiv8HK=+zge!le<;r%+Vv#8<<V!J=8xVvMpk6;Od(p`x-<v
z_@!@cYM!WNy2VcFh(W}BLtiJ&Rnz8)_sooT7SnkbE5gJ6*n38d>eAOU?vzdx^ETt&
zDz(dMThzs^P8@;S%4gFu=1TIOdVger^*y;3sbIUr{B7wjHJ)DUTfCI??W-r8ysYUx
z<Jr_4nVh4Sass>Ry=OTToXD2@!FiA`hhgsNa`Whu4M*gjO<xkYwn1&v;fI~ZQx0*~
z8=jTCBz2iBm48RY+5&DPws&!j7rIm)hwHpwu&<#dP5A7gTOLl^PK)qa9?U$kVM|BI
zZ^vu;KD_C9!iCpo?lolJZj~Ds&MjOSXZhutvAM09@v2?clP5lX!rQ)N(v(QcoLh;N
z#h)D5m>Pv9imc}4fAV<CB*&gPOmX2oVTyaB+ojL=<=ja#;VR$d946uGSZtP4Q}j!9
zfvTtulT>rW5u=N1FLos#w4ABj-PRsowLIb6Qp3%>eqWYXFWa_RSTHYLROo(3io%x{
z9$SufsZBm1C%Tm(f9)yJlN0}#Yz<>eUbHfM_f_5-H!|CHDO^czeH{36@14AEqluNX
zq^5;VT(^T|!Q!?P4s)DkUs;?;DXc4!Fi|$1(pEH2U`^w-tBLQ@X3ffVJpN(sw5<y-
zPtYy7)pptN@s6r@7miPUZ1sJchk)R=6)L&4u__{)W=%al*Xb7H8XF6>BPU;aY|DCW
zW*~Scc=E<6+M7}jh#G2#o~YPfwsl*V|1n;nEk0*Ec7$ol8=e2o*w%e2CE0s!xyt0j
zcdjfy_aa2=Ra)_8(Qe<_*2QX<zf4-_xIl!};oQ5K%-YXSpDJ$KToiC|Qp!e$h$pIb
zXM$GgNZ1{AJldKs_~NaTQ+`j0>+Hb8oPGV0M-~N1NR)`*&As5b>+%Bq3%P7}9m|b+
z{XfNB6@CAv>+rH^Hxs)rSma&c-*I}u;XT6gRRMP283fqPDi%s+A6F}UWc@;eyF7Ex
z?FV`LE-h$(X}Ijs!oFt5HxrLZ#ixFdyQj@#uRMSC%vJ3+-(+4#Mb`En`^~d&+WN~C
z`wz~4`kMEj%zcJAr~7o9^PezxG<R_|+n+Z#4CT=@+51W?()_W+iWN61em*+;l<oc@
zcAkAx_a}VF{o~r4|J<4(L)ANZDQDNwsi|!DPqMeZ&})AG%$XtARR6M(cl(jd%Fc~l
zXQ!rmoHnuW4V!&>@kNg-4-PN-a^qm6G|yHKC*yCW{PiavRR6SVelPB9bD&(}kAi*5
zhr=}|9_U}TWwSrjpH*|@!FrDEIqZC=%$nsNuUBJM65g*`w{Pyl2U`PAH#N^p<K5gK
zy~D4wbidmC<GgiE`{&(BXS#Cu_S$J71?QyIKELsvbI|`&uw?YgUbnM(mPRIZQjLNd
zW-q3H(QfYVVx0dlo!9P|f99R!3;sK=KX@E@I(T`<!CIlhcP}@-eAshfv*E1=&wsEr
z|NdFYF>~FXFHB6DVKF+~Ja&qi7RMxeP3r<*>~B<3UF*htYuijkk=?Pf7HfiAEwc2h
zUmrDIc05glJ^5%F2gij=v3Hu36mqtPbcyMA&fK5*VY`ZS)~hoG7DD=F(NSH>nc2Q;
zc-7r!<<5$@({rr5ob~bBqsCfkAyOsWcUOHp&9?v0`m8^*QjUD}k#0Y^(L#u?RNUPC
zP@%Cz+oHK|(<E+pvd!+Eef#0Yw3Y~-kD8l8HmABhRcBlCwx&1U{0DDfiqpwuORLRF
z<3gpG?(dZj-rJ=+`?MEt-DJT%!EHUNF?{D_CuN*?-puN}$>Of}p@K(uf6i^p{I-0*
zJ1={e|7SPD@`?hh7blZ;{G9W*;GoD4lj%zzKB#@R*!}0rm$E_nfA_rXJG!>dk)QKh
z;hus8YtCD(2#vI_x}aO%yz5!|La|e}>m`NXOybIT{9xCHl7gA`O3b?sIEemvV6tw>
zmy)H6|C_X~eqnO$=du2oZx3s<wC(fas}-NObcrb6HKwm228F5YThem$-)hBMREX@j
zlXpyN&E3MU3z~Oblnt|Yn?A)#^!`uTwC`a%_qsnmc{?Sr(>Zdt<307Vl9>x5#2hch
zDnuBb$ojv>*JRGcvu_{Wxf5~yT7=S~Zm%$lX^)guUR>U`;b0>7|1)oIf1I0^;xt1d
za_{vsGRKM?u5E4?<Jlgu`)TZfUfs0<^6Lxk&0446rK)||Nd0<E)XmCHW9_nycV_Kf
zw(AR1tW4$nWl@ux?BzV~-u0FVmA{#GH1I3WiuJcegHH&Vd~RD)5WOe1Rg-6eIH$SG
zAy4*a+NYkFP3x{(XSVXpxw0<X<p=zjxkE4QESRS}<JO{y>&|_3oESZWkN49P#X|Gs
z7iaukGQ>8iocA()-n!vdzzg9A>@OE3JFjeSZt9WQE}!6*rnDgU!sj&+#U?Xq-oBoq
z-t&{sOmf-OJ-659rOunkJo)Rgh~vtPocn)V$?jV8(eaPf@7|-`>tyY|X`ej6?O4UZ
z!}$88zNYXTqiTg4HftBIlh`8lHpN$F$>E;6M}8k@@?V;r7GS+dIcDAJS4L^=VNoCF
ze=KhmS?7B7ZEZuKtkyBN2dhQ%wi)H5=&#={u3PdTX8sd%t=`#x&T`DSlytNyaJ6bF
zzc9;-LoL^O3PQJk_V=i2-`{m1a!d29juYOGmsYsxyXU<>^2PP)e%b$p6T|;DuhN}T
zZ5Q~@VyE__?-~~#-RgTk{X@3l&WndLBJC?@^-bD7eM{UY+qDvRy4Uwj%h2pyQvA~H
z3(xF1@0PQ7h^%H8USl@d{i2`F>!mRpF3+ygd@r%{g6XVNclVk+lK9B**vUgJs_i|~
zhSiHXOXdj~97~%dps;uM3&oBb?);ah8>VN^UA&^`^r`;KC*-8QhIFkBf4N<&H#Dy+
zRL|C_Kj|OGuT5{)Klp$6i^R^Bm9G`PoM)CUw|{E>HP|uu=H@4_7CSAAvsyge?fUoN
zz~IP+>P#Q!O8MTSSC79f|9;V8nR880Ro$}u3qPd%Tkl_XPM3S6{%77HPFs_Y0#CnR
zUZNIhweumz{qBd4<g8u`^s+C#&-_T!)@19k*WbTgFcEmec(BAZ;N|xx`tKHQUcJ{W
zRZt_1W9_r+e||jO=d@P$;3Dn~D?)Uwyf5yu<qX>JwR(@{u9NO7-x!EC?vlCEuDO0b
z(}$3w-4RTSq}J!$e*bpyDH)eD5)YsHdW77%xSR8y>8wRpgr@w~`n8g)Cr#|!tA#h6
zY<9Pwn^>0Bd;b0Fg_i9rF1+{76JK?KQF7kx2P@C%o)VoLYVY!SqxPRkwP#mcG>&}5
zIjySe>7BN}J8Pa?IqjaIS#r`d+}yCfWf}j{=>jiLrGA#`tFhU<sQnqc%Dfn@d3`|)
z_j^tH%r1r;snCt}6Tc`H<sCkuv;KCbc4@G2tgmi~aHPk^FqcSaTkp4Dgjiqh-qKsg
zw(7fvi+m9K+(icc3po0h`pxsO-Z5iQj91}2_t-lg)|JX{ALU*-JGVsj{nFJx7~Vc~
z&(yc?yMH`G-roQE$JJZr-(M#EgL&7)dVllJ@-M^wY2Q62fBCKDq<=!cZ9O*Rr5p}_
znf>FQ=j4}#RZM%`C*&9{b2L6YuT{@%E?Z{si`QRNUdlM1vu_pMc_C}{tgnuFd5glm
z)?dE7;PQgzz1Q03{$;#YCA2^Lr)3qJE&tTd)*Fr(dVZBDH{!oo&GPHvwFCQI^JLwp
z`8|yfl(%J{{axMAEzoxU#I<He_BQ<EtTJr0DmGM-Tt9WE;<6kDj+_U(q`hRknDum7
zW7y|LKH*6!)Q`CH{6LcYq`#`OS$8ko*|yQQNqX;=1^f0aeB`qAcIaEpPa-?zoxH><
z)2>%ApS-i`<iRrWXBU3`l{@lYpQ(0B$I69=*o38`Cmxt|w8@}RE~@CsmEWn%hL)_?
z)D&j&Z=1=o`;1F8ul|}w`D=&IOFuBbST6j>l(R6Q+SOg9UU+rPflore#4_%LI9z@>
zUwy-2(*+eX;||_Pa#;7{%u)-Ljm5%yS8cf5exjcD-g(c;Q*xEgk{^A4ZOawBuJwU=
z)~D#dTehs+GyNQE%{8X5>86aVud80}Y3mZ%_w8cQj79U$wZGI`Q5dd~eTCt$&^HUs
zy)BM=%w}`9p5@vQ=+V)$U*qJHB!kG#GuLwTcK*A=GQISYx5E}ag`l@}4a<Ji`bk`N
z|C&(NqQ>vzFJV*rDscb$4G%5WHi@j*5OrgpS7q1t#7kV)6ArvR|MhikNvHRVbR~r+
z3o@qIc9+F*vMs%<@=|rn<u02h|EtF{bGNcTRP@%9-Lu4aYs%b+h5gr7KlS=Cv+>aB
zJNf~^lbJHE?CQ1MA}Adt;PEb9cv}$5h1)R|JNIh(vsF&EK6o*%;+0y<!)eQd6)RMk
zydBiUIp4T*O?TRT)8oNa4sPL0X}4MLoMw4VyUV}UPvOfyR>p6|Ef;@tOuw(#V&`<B
z#z#GJ`F4iR`wwfXW<S4B{xjy6-1Up!clO?U82d%<zL))z@E6DbSp9ABZ}Jwp$G0fS
zLNmejI>QXrlw!r-hE?2qOg;9`-RGRIY;%8kCc|!Z!{62|-@aezx_aTO#-Y33rYX)-
z*f;UA)=BvXWft;WxWTsLzun&MjiwJ)98*~_IkWK||BH5xpD}ea>s;*Jgns;t{KaWf
za4}~e`>lI!d2imZ{#xl}>^aLaBiFWm+NICW8A7A4=^Lbdu24Su*mlWMJ^j<2zn)Le
zSjZlG#XtPO%(T7Lo<Fmn=*Dc`w5EUCGI2E%+rSxDD%9kb-C;QQ=h;_A%QsO=o-=Lw
z!!d6w+kvG)m%_d@swow#ED^h~o>A^U--1echtKmhSpRcRy`7(7S-W7*MwVMF0eOG5
z|4M&*=5EjU%R;NZw{%I>`@`Fq#81DlKcLKhr}g~YmkZXvsBga3_KTlkll#l<oD*`}
zQ&X<1pL)%#K5fUUyRPp#`7f?>m{?ub>#n$3J+S-7@0D9uZqS{`Fz?voEBCw|Lo6me
zxjpIg?aj~ogFJLDUvz$Wm~)ng+1WaChTq5PqVA|HP1i5I#xC+>$>Z-!x<a&ePmO!$
zAjEi&n@2|fstMD5`->UBj1~Sm&+@Nt%(Yu|`bEBm(m%gygL2)TB5q^ba~6g-T+csV
z;py}{V~T|PRm0iMJOUZ#9qZROm5INpS=v-)&etNSweGG3@A9P!{<?78@>f~#SL{c4
zk&^oTcAuZ`kGza%4YN@eE3=;$njXG<ZnxYo`NWQ+|GoE@N4z<ucWJKp|AT===RT<X
zm;ZY)(U9NS=*`#VSJtkGIxEm?-(2=3KqTUwom1X*<B%{3Yrl`_8eg`oQe^9D$$E0a
z=F7h|DzzPJ#Cl|0V|T3YOcGyox7Rw_{Fe4bJ?CoCTecbBn~tYnh+F$-NAtIxrMtx5
zmP*WTDYL(Ie$igJ&#oEv?cdn8nC~in!Eoz*^FOsM=FNGs7rw~mzq3c%-5k$vcWE&b
z1A`_z1A{RG%F$!VIhpzxLscq{X(g#ekYy0Xka4TEQ^9MdkN@8;{rUx)v_{8y?j{HB
zDGL{M7zWmGoOBT`>D*&jo-KQIooTu2N97-}HHw{Hi#p{G@Tb~rU%hzK#3n=Es`HlL
z@7%Y1ey{lU{Qdj>G6Wp!ViUbAwQ0snb<c3a)r!59ylP@0Q$lr&p1i8teA0IHqYAUb
z6-UK+H!JRVdMoX2=^6EyEjihx_tRSyi%)-?^T6!zlafupeWohzeHG23?{WB7+2*wF
z<$In@O5O43`p)R5hbC5~_^?O&bt~3CPfh4Pp3ijr`;uMHDz`<HR<2H+oF&&I(=Dl{
zsP^c!=If-%``6!;W)ye4eZs&fw({uxqRUU?4&UKud2AY%W+oXFW0F!7JDG=hpSbgi
z9fu6AOuHSJX|zDfta?KA5ih+(-MZiJ_kTM7^WLuae^;t~4`(Z{zB)}tU(WMrD4W3=
z*@rVGpM1JObjn=TpIXK~GO3)++aBh3@A`XS;#$L7PmT4v&nUO=sfj(kSbkQQYTHwn
zU&imZKHvYjYR<hIiNQw4N=3~-MI7^O+nM}T;qaTnZSxa9{<!OCZc!O7xonr;mkpn_
zRDT-Qytw-M|H8yaYk#bNS$1@(^cmgP|NImxrtZG|O;Ns4%j?V0Gfvz7Z+N3Q>DPi2
zcFLz2GWk7&GtRus;qS_;am<<?!rJ^r&ZR5t+J)n<F3f0Uy6JZNoyH>;0WYD`FDyGJ
za_n+h{3R$z-g*ts6V*FihEw>Q6;H36VA!@<@Uh_91urye+<)y{SnuYQcH{d?h7FQT
z-nrNLHV7~-E$LHfee!vS#KQM(rrDFSYFyp!FV8A!<4lRvp0?_mVs>s(pqQQ9$9EI^
zd*)y4wcT)+S#n=T>Wr5KMmt3EW=~NwkW4GFjInf8XMIuPr^UlF@p`7z`6W;JFC1W%
z@-|sCV`=?@M$3%-Xlps$FQ&bj#Kgcbhn2`Q4@vHjIRr?WpBntu548IEGXJ(ITVqd6
z*O=85dP_G|^O1VTF)3kXMqwVd|GaM_rsb5+zI*h=;vX#krn)FxJi`8AVNdxoZKvBT
zKZ>8<vo3yT`~2LQ_xAR6j4K}9R9LrQidyoriH?zN#hM|PX0$I&ROIDN&0KX!XIVk`
zyx2<x@rRk)EWXd?*}2vB-^6IX^4ZtYUfj+ro|$+cyM5d9Nt=^)m+d%Z5dFsX`fSZy
zLzS#JXM>dMIpJ&GZ!c^4KkdDRQ`x6icWyq~r~6oACP(?(2FLpWyw<52HBWaPT5oB?
z!SXgLd*1ymv+VLM4z04@7Wz3#_nlcvxA`6ZmUxR^b;g;?IO1<O-hMd$dS&;o9x?B~
zEz6wNKhCZ=!O(VH{8FXe{8h7D-&}4o3C$Ff<Ptw<$n&u(mc?=5t}PB9qCQS~^_$B#
zzv_2g?4B)8FC4GdI^*5ZmndOXBr);X#6BIKq_~wC6>E|uwRwc^E}j4Wwx9D4t9c2b
z$9a7WH+1#Ao^k2Jl!Zl_uXv9A|0Qi*y-3dJ%p7~;s?)BPd5dSuzT|Uy*M892`cA6E
z>*AeDzkfMZxm))R+r_u#3)Dq3m#utu{&Abl-yJ-EuSxh$Ip!JX=Qs1(l2g4GW=k*4
zn7zniZRBT(!p$`uEtcgRWeXN0b4^?D-Li${vf+sp8mANHeEP}9BfNu6!+%@ismZ-b
zPcG-W<l3!}h%nr*JkeV&@d5X;I;Y!qZeDH8Q_{D{DE<|l{NJqYTuzPkLlcGUjfqE{
z-(TsiF!(95Xr9>GBK=G`$M<ag^O@AG6P&xwx<BBZ<bG)3%%?k(kMQk!Fso_yTxE6R
z3kI@h9S--*YW$M7uwlo<L%DZeCb~V^H8V4PhVSZ2hG(>*k1W|WN#gE1+Zb0V^UT(X
zzqnCnAS%=zg(R>rFi3GTFj!+2>#(z8h?xXk8x~zI6)N(NkH^Po=Ch`SrfO?Fc|@6b
z*KW(yHSxN9CE=v&%FayfvIfqXB54I_9WMWy->-XG9sim;H^w`B{$JUDe5Lan6;d`V
zjXd-y=G)Em_xbg|=hXgqygt63y@1u;!=nAR(lpgMM`iD=@z~MruewI*?DF{^FWPcS
zuJMdGeoVmZ<4=Ko4%5uPvwv91!oQ7egKbM;%~5shJ(UWN6z~7{IC5R7i09vuW!c|^
zn$2ZyCJ9XIU3AJ|(#4X}y<sO;KdV)ZSMyXz7Mb#CT2fZcJt3}d^U@4s&P$|e#u-gb
zitH6X6=gGZdrq+CjoHfrLw%!z*QPEG3jCsEzB}gyyGqDemctRAwqaIxR@|Dk(72~v
zd1mcXhtBDudqvLQTDEMuN9m!ZRqk6w*RBzhTvoLvk~KIxQbtGO*SfG(Gg&M}>^QeE
z-&&Rw)zU5eaJ^M>+s*5*CvUr0^i}Jvn$aaOZeK2myxkWp)VgJN3;7ylZ42vLuX)ww
z$&oom;iej#I=-G&i{$b&pU&D9w0_pARLcoWdvAQnC_YlszS2b1Wh1X#&kdequLOIg
zWt&5C(`U-hnf&5hnf#o?uRcE9d2Q{ilEWq5({}8iTIHClB=V6r)qM5TlA~ptem|D}
zd~^Bj6PgFted9fGbLYhS;oN0PkHi8`y_#NGywrN}+ozA}V-LlJ>sPKSkomUR?X1p<
za?^Q-GU~-^8*fkgXr?CN=6`GIN4;gKXA+{0Bul<`7MwUe&>`@cPJG}-?L|kzzqw=>
zef#L*66waei^JJX^!_P>gX<6IN}6TyOek!4pK_(}#O(Ky=Sq%hD>qeIH7?hR7WpP^
z@b=of3tlsqM!8;n`>iu=MqTXIg+<bJFV6WbnSbar*W=??!|w&JNse%N8M8zG@Zl}1
zS3l8f`>)=y?)9MwZ;o%vj=z2Q=s%N{8!oJV6cYSMQQ+Mg2FbIJ!t-`SA2L|5b(O}{
zz3X^oubMuX7Iaee-?gi|-r5}If01}~uIKmQ8~a`@S|)z=*D*ck4dSAi|5UDI{O;XQ
zT|WP4+J#jgoqfMmXuK>ruWl*vPE|bZ#s8z0ANo0t<=bw3e?iC3_O9=bXY-r?Ix+LB
zt^8_zNMYH5_iGn7eBQyp!KBaYXMg!sdZoYB`^0srTfA<?z0!B9-G6ZN<GJ|<X5STD
zA-swA$F(HhAJ>yiM5QHe{JWmbz_fKTOG($h=YmT8G6zd8M+r9?b6xD8c#k7=fx+9f
zJrxtgzW-CN{2n~NyJ$ZzkC}pOVn^cNOM)qF^EWs(8}sNt+^Olq(fE+fO{~`O{&|V+
zizQz;lH+phkE)#A;~3A}$f@{3SkC-ISoX{J;v3VS9DKcCpTm@<8yted+(!<zA93#b
z@K#`xdrP<b2A2*ofu;F97g`HMr)TIT@O;-`oj&(R(c2$IU-CaQw||w1-0k%-Hv4&_
zx0q+q++$3hc@LD^_s?Fe`Dy9`!RHB6b&6(RsxkQ{^q7}r{?{|PlOtEYjda@dH1d$H
z)7_WHFMs-CXjk~_3eqJ<pcT23=k5{R!^FU#z{<d&gIREaXZMj8WkRZENTu8x91T8N
zWZQP#w@OQ6L>x857wVqs65kWBn#E8t^T-mm=6q-C$Tm==tSea0bzW`Xj|H7x6I^P)
zxl5jZqqkJVvu5U8%inX~*W9->-v7VmKVyQRwq{B|K%kFelN0|hgZYb|wsp<+)LAm2
zVcs!a9ks&fy7wz~Xz&02WYdyI`?}b<gsRWv6oxx4+jRQdjC-G2R))8{*42KolIw6%
z-06p0X}jvCH*G$t<GjAAZ$r;hqopcdhlJuH9ZlkIga|}$eLdw)xS&}5#jKp-=c2dG
zG$*aSVVag_<bI-9qNmcjY04ArZOJd%4!=^np1tfsywUNeUqg3aObCk&$~Cb~j1JMu
zEXWU*3NV<VKWDm&oS%Tp9NWf{S#xZ~_Z*0tb@jD&-^Hwg)4>K)?E5G6O<u!$@}IYV
zQTxt?Nm&=ASy!&QTa@)@=LA`a-3MMZrCc&-Nn~YSo)s?B{?_}X)USzG3uHHQsrjv5
zH(SkoM&#ES$FJYJw&T2OM}tZD_MNjg=jd-g{H>xd^0(!GBb^gY!CG=z6)cl%qz-vj
zuDfVtd(ot_p@cK)Q1ULDyMCrqKDrzZ*39_!|5#;PYI61T>nqx)pX+OT*P(rnb#L_Y
zscqWs8}C{EQkYueo|zEa%=Q1|zSW0fvrYEa`dpfQ@<Q{X#*=!N${YLyd6fLBvn<z6
z*9>x-Qs%I%f@5h#&xvb_i{b@;zgfKgg4TUS4P`e=<t>+)Y}O@Rc`ko=>x9*kMY^K<
z73&)AC97=P^h$A(eRsTowG+Q&t$bUs`okB3mjwSbeOnn@IK!kVYh}pd`9F@YbdnR6
zy*mBY)#I9SA6LIFne_kBq_0zS9<ndDp7c1e=)we<qiTktQ59~2nc<UFKD)o@eNbAi
z{cM7orSiT2dFf9l@0z65NS35K9}CjZmj3c!QWNuvjr=KplBXa0$=7>x+d)hDssGWc
z?`;P*6whE{U|7Y*z(B%@JC5avMXBH{QiQZPeQoGlKjA<D+wHNt-j;0?U0C;kJxOqv
z2$$)J<y=CBD!j9<m>jI%{x(W{`|fA&uJyZr<oxrP-`AyA#jb%rbx!oOs7te4AI(2^
z?##^MW#!M_{rmIvFMC7f8ubZBczLZheaduNJb$*7n&DDKcITs8*hG!B(o9thmM6uz
zPE?F6p6&bYiubf+qn_z87d1>+*56KBm=@i7Jf<S8_A`IlRK;iZ+xNVQkuMQn7ptLC
zb+$<7zuw`)D+?-i)T}$SAZJ%;mYFKE_1d(nBK{Rumz@YWxwdfkp;pzu3M)>ONgDOX
z{hIbVXU{XG0PWv*mhaj3)?(R>EjypjY>DQ3Yjk^l=fCV-C+~;(erl=jNR*hV_4Z@<
zo`mA%ff~OL6-lk$xtu3ZPioV{c`Cg}O3PxNcPLb@x%vC=i$m6S_m+i~MK4Ori~Lk!
zWcwm=<pSl_Eq$9KE*N$gTjeE(99`6>s<S%c;>Ir1w$^~Bz3KZjx9L3m`83Sb!<gqg
zxA#_S9%J>#Ha7B7Gw!|h;=8r=Xz;lSivRV~S3Z(smQXt!Zf$)(?(UAKvb<$FIoBoG
zD);_-(5qk3dRzCi{^^jHOJ4;ppZPiJ?3X=n4#~vc=@gD}-MK{n-|4BgFK_02Uom%H
zwvEj4o{Q2!%gfVQ4i_D5EQ!*ybzKyY(chM_>G=t5K_xG<T%Y&Nl5ez?-0I{n`a9~b
znD$ABzxS-d(KJJ+jXY0Rw7qd%@K?oD=}0B#V#d#0KZQh%m+g1>dEsHV<GRp|-RjpK
zc0M#(7clRk)-k4ewme$`T^}1Saxrz7R_pvm)`&fze5tzemC8T1LF^uH*w%g9xOk-}
z+qMpi=vS9?%SyT%^H{FmHhsLq;hMrzmD>1o+JaW6W!7na(w^)*Pk?h1>!ZD{VlRK>
zqqbPpSv+_o85tPTL4y&P1wYb21ll6K-audP!wv$qm#-bWvVN8Nc_kArxAnf79)TiE
zT%7l|zucAMX};<1lz(RY{}}40?C@R4GTBLd+T6;`@AoD?d;a~k^a0y*vL&)}gyyIO
zbFbrFc%OTD&cQ!ZG#~DlSa;LnsUv5#tbb(^(|RsJtNzn1^-?oe?rb)0lJ;CU<D{wj
zovWPuUJ)r<LX;%f3S}z33C*icT-K^@GJDA#|Dc!q4Q-!K4m6u>Xv(?t*3QkZWtW)k
zTQ~PuyNK)UaM7gBpo|v>s@1=+74FZnSu;0i%2TU{{zhAWiL36sSj8E0e_ce&>9^-*
zuWxx*&wc1@*`sE$pFP@}Z-&fv$Ve!irRyZDJNxa;HK`xXWTZ_u#E0CO<#jxauOQ+`
z_7jfdN2@=vqlT~ZMX_hEm>C#qxEY8$#?C1<Ex#z$EwiY&B)BvQ8vc-q_GnmcuyD9Y
zUEZdG!p#{Tyi>Q`@Mw03<8WN<Aiym+L8vo#irJ0rz7dzNrd+?7v#Nff{lfo;&TVO)
zw(?ta`md$`7uzqoTwI>ap|;^*;j;L5r+1#WeZKSjyy|cNzP``bXIN4ge=PFM+(l_|
zZW6^Zl9EDVsX7m$7`%6!``9@{r+5v=4PHs3Fw1Lf8_pF5A7itdYS`i1H>XEno5d%Y
z?UKLVzW-WP(w3-`9yK-OXi)j;s8umngCjSG1^3>alb5kQboIJONhkeRZF%~JYy7JY
z$GclRyK5GcmNnb))U4f7hb0qlWL@!^RqC!AdQ#}r%+#&B)<*4`Vct7w!D)fGi|=Mn
z+qNvaboT`f(^?}9>9&QDH&;ci&t{qZSw-^gs}4UY4S9#%9hWP`+8i%y+*p<RT>j>+
z%z~P822c9#Mr_HLyngP<JIr6%O+!O_y;dbe?+QFw618m4%JLs-e05h>alEwpb8)+|
ztWD<=zB{b5p7Hwz?`6xgl*^i|b-$-D#%$*d9`{B@9{&4YT9upqtY1E>Hd3CoB_mC8
z(cU-76S&h3YX#5hkdb+Jz=wOM=$ykHO^JI>7q%-_DoGxFwwr6)N7m!hxw>v7c2sz-
zcp$egX~y@SlIs;2Gx~e=PFB2Ip~@<5pcYYkTEO1wz^2g5D<+5IYJ;xLn*Yz~a_W+q
zVkfm8D&?JXRXg@t_%hGI=09hq%Vng_er(OE<)xG~^Wm!J6V3W9E);xceK`Hq;i=}U
z+4l+*IqTVn>Ap2sxL7RBzwY46xXR_PH+DX5k}#{*a@rfG6{$OYYn82)x%2ebLHyq<
zPxW5i@HOkzv|k<~=bhQKQbf<D?YQB0`RtVN>f65TJkMC0KgLPCDVQ*?(fQTC8@^Mw
z2UOk4E?GZ|Ek4~zbWYIbWs`J1$5!04j$0YFm9PA!_M^X_*-~#jzW(4Z|BwAgeCHg|
zR{FAl>#}tZYxAN#mrwA|U-9kB_p1Ih^UQb4<2sMlx9W@ezd88d_-t*({9BikdSjjC
zqn#RJkF%#pFEc%qJf-HaWZ+g+>BH_jZhZM=&UV3l?>7FG68;-(zvLdRYkV*HOZfc3
z+8@l!+?P7Os9#$5^NK&mTW7svnNl5sah1FJ4<^0X^2Pd`$lsHPj9I=Ld?}J&xah`2
znTc_V_bYELs+ODcBCYckZ_lRCo{gp|H>`|un(uX6w*O0=9?`w8jnhXc*KpRNiTsRj
z6xvv1Yo|Q%Op*!`P>I}Cxx;3y$CYAt-LeH?`wI?>%{4F-zr=Oe(EF8fddbO)-aUGc
zlb=_|6xPHXZNIX9LRtDt%bIyQfv=9MevFB8m=kwY^Oc}XLxo<tqTFHb9j-eJzRvfw
ztMUD7KIsqIfwhwvtY?B&R{vpQV6en2qk}7pOHy-*8B>@$>9p5j2a&e(H&*9v-F7YO
zs_Pw7Ct;T(t>P2J#Dnw$Ue?W$yF5`vSnaaYJ+JU4ZutZ4Lh6|UD;<wS&N2Rc{{PuI
zJNMM@*~fT*F=x60f3H%G$im2$vrIXn3tScIPuZxQ=S;cvYl6eJ)5kp)+)9_b=W|!E
ztJ<n)O3?e_ZYyW42ot{}YvycdI1o6W`}P)AS?i1_K4bPog^PQI@3<>Uvc<mJbw%2}
zyKJX<uu_*Wlc{uS5?6lh$sO(sE?)cHU*hK5UbW3~LR00oP21`ssx>r=3w#pHQ}s5m
zc!%lm{s@VvIJv0n>;5TeF03I77+jf_xG#`75^y|Rr13GMSMIuux=m}3$(vX@zRjvw
zyY=D*-nEyD4`+Nk!<&9(=At|Si$BLrSKf|Ypy$>oa%Ix~;tlMUUmpaXX6ih*=+P#r
znOeeAi+}4o%nMED^wN5kypb<eRJ*wOl(ykMuB$BXR0@{u`?96MZ2!WU`X#!oXLS}?
zuguV8J*&4Uifff!_M52xEU49ju4qORI}-!LZWaaxBg`m;7W+P#dC+#H3P$sCYViI1
zO93MP{4_<lix1ohF%dSGF)({oV=UpTGHb?!T+TH9>fH{9cTI^dzODPq{^z4^HM<||
zKinnH=iW?dX34xXPyYLy_chPw+`eC5_m8Q7;kcXBZ8M`A=N4r>mMqh`qHDr2Gko5c
zMb}$G_7*xBSRcHgx~9zbrCeOk?{&X4e%V!f-Az#VBfQWi*MCl0$+Om$I~xNQuCAKc
zabxCxg{7;l_Hj5*=a0VY&^u?Mkn{agwtde_nm$IHSh#M&lhEatZnqqYsCl_NXXT3*
z!B1ExZHRk6uiN|K52wC4aZ7A0v&{WCBlY)id2*Uw(n){Bq1`ygY4Zke4Ygmq_m&rT
zozdLAW3FV*vkBZ!5*5~;*f;n3`AfxIPeM*uTW-0a(h#s(zs)M+b*}XM`NFII3;S-U
zKG-U?<Gt<PAH6eYOv}^Y=Z<U&trs@?7VXPjUHN8l#zNN{(aUV#y}RGg_v`sP1*JU^
zb)Q<+9zA<kV_HsgwM2JD!7}}VRKcdmhx@ge|3(;m;i%v&zflzN$!)^Dzl@Kiwb^q1
zrEQsf>cqx*{!9_;nAD6N4cW9d`Z(O0lbZ1&P4lblv+L0YvklL<2~9e1<d)RCMu`Gh
z{-g)(8qV>3?$_VU+IaSTZkl!Q8I4Arj1y-)A9+tZwq)(hC!a*-e=(hJGH-#<a@*d=
z+Y6?rN^07zOgy^9a&g;l9r?{JioeBCla<>xj_<b23=F3@v5r4tCM)o?5{|~k>!8a4
z0{{G?IhHVbu^zEbN?oA4g;j7vg=VN;kVb^-#rTCsq%=4uvwak;VYttnUJ@{~@Xzy}
z{zuLh?>xCgyk{RWcsze|b=mit^zzNSE%*Pc`OT-$<|=r$!{m~X=*6ZEk4=vHZ3W&3
zHZ<IhG1hWAFr&em>6mLpp?^!qZzeeo4}&YS^RGB5oVqagQg4h(`ia$*s*}#Hp15S+
zM#b#anxVmBDT^0K6mF6WH@&*&xWU!eIsewTs>bPV?EbE6STX6>szVnJv3`!0*6#Wk
z5_fX)k@=T47u9*1h3wPkvb_`Aw(UrEu%TM$;)JwqVO<k=F7+(abTiTy=i-X&wfH4r
z!PoxiOB%m!zo&1d3XkuuX@;qXee!3mTF2>OwB^>GT2W2kLsd6tM@_zeYT}+XPn9Gu
z2W*Ym6LabLn@UNupAV+>%J|1f%GFxl-K4(U%gI|j_3FO8xx(TfCNT26PXF|9?&0pp
zqYtaZpL+M`H}f<`oipr^tUY}~YqHk*8GS~x8RyAy*3Fsxz<G6pyO5fpYC_Z@9y#L@
z%hcw?GZ!bO-`ce-Kl}I;hji7cE#5oM$j)ATt8`DoHl`YPRoD7n)70xcrNJ+!?KV9>
z-Q@6xmA6(+oAE|k`1YoSGQBG&C4G2O>-6+pjB9B30kcI`lQ%~$4dCc;Tj|xFmC<9j
z+jzT=sr<~z=0C4kC@tVSr>W-b|8j!I`Zb!7#s70wdK_TtU8THmuYKO0IgZxT&drRP
zUA}EuV7T_J_rW~#UWvw?-DC0fU-+v->W;4Bf1gIbusu*?XLq5a`0&x^(q~_r-SWMB
z&o_jt&F-)$Piy(Ze;(E1w=#2|S#9gE`o<DB&5`}2m52K`i~jD0qZd<@BoB-3I_vhg
zfvuuLt!e&)Urhc#v|q*iX#Dl_-~12n9V7TAE>k>Fc5%mM_QILB-yYm8!!_CC)%|d@
zw=z2(J#Tnp8FiC;%jQ&;8``JmE~vfOXw4bpZoc#2PU#~%HC|{i?cm(rzxaz-tWR*p
zpCt)L&+sL$d)|K2iE~E*gWiJv2Q!Ycu3cd%?sVf+*8Wz*25uL14d%079=%dI8M^4Y
zjZ#&|q3aInNfM0pSw{ub_fCHi`@ubmM_5igt#g^pqj=O@nR%}kynDulfq_Aufq@~t
zC^01`RW~^&Q#Uy$vADQ^oErd)0x)j?c#3lFKO>U}GZzB`2M2@Jrm5}>3=9mgQ|~!J
z955DOc<Tsa__~HT>U#RQ=?8eDYP&vlMiA)OqEH3~1`&`-C|=U2$-uw>*XHZ!=jrAe
z9HQszhHL`pZY7Wgn8hq04j2nC{B@iKt@!Zi;0CJ(5g;A=pgJHSh~ijutAs&%z!;?Q
zC8|c~1q>JlWHhCuEoNX~xC*_m0j!^afniG{=)wdn_k+Q1g;FtubkWega`ay7H4!ET
zh5{A_22k`P3<O06J_FHPbpb{B6_t*qB^jxCC7H>HgPqZBeEue)cohQ!!)*r0rJV>H
z%UDp>8NkB86Xpl(K44i{bvK5IfgzBEfk78$7&u50Sh1St4o=;n#i>Oe`MIe+`T5zU
z1=x+2c@%s5C@TZQFD}T0FwFEVjqDs)jdra_&M7U<EK7CE%t>`jPEIW@_D#%7OiwMs
zZowJeI5*G<+CSOQ-7d?A)q-%ywQtV(d1;yHrA3J)@H^hngE_EPHuEV91A`<N1A_s~
z;oxBTDS*{z^!W5h%uC6^<?yAI2D`#QnT8coV!$le(x@Yb)pVcybey4ZuIP1_C_4i~
z9v=gP8q6fHQ};_?H4Np_JM1BAac|pQX<-J2H<}C#3MghuC}1_y6LL{3cGK=2@aa6l
z!oV<<6Fr^^jj@{Ml%HP$x^|<$H!(9WxFoTt1XoCUT-&a4nTvrzS%QH<1I3{Y7Fdmj
zCIQ!q)a25Vd>jd@MQX|FO<W8NP7(|ZS}3N@bjE5bBvU}m200pgK%n2a3p094<77Xq
z=0bxNSMh^>!y&@F;!v#Sp+vn3j{6SL^8)(qc?er-!?D_ezr;bmGzVea)p&U22`#;_
zR-nl415IwBUwVTuFfD-r=M^~Urhs&TvLuKGT^j?U85ls7OI9L^FF^~q&^1p$yF><}
z8C0>%#MBI1CV<pjM;m-aHy-_B5QOn6DHv`;G8WRS#FlzM;f8(%2f|c=R3c2pb0G)1
z`_ZpDK-eUdL4-|^*+|GjA@rj~(XB#1X&zzKzbt%K;W&35-FWmP(-Fo`$-!qlG$!yJ
zppI^p4BFYN2&=Z05MdQ0&<F(=`tg+rOKy}CVF~`jEYW?0e%c?xGO=oWmO;FNS)_sz
zG5T3k2s7U_;4>2%ArPY>9HhxTbeE$aYJ{-mWfO)iVAGIi)3Lj}8|^>^gqa7sFwG?X
zcm;H~8=-CTN7ys92g4rB!o!)EE&u3twxVs-M%d}xPuNcE;f21Y7GdVQg@nz-vceYK
z(dY{p5f(gJLWBj7qy#CYu=@ag`53~2qsuWYfYlPj4E&<|1bw*&!lJr07#3k9SBOWj
zdjWk49bt0IdJK~>6BEQ}>=Wwfo<N_vL|9e637=I2>QZO489Rir*Ei!c7931PDtPAa
z&^>@Yu8*)udIyG0==}$zup{6h^s#1yHScy}Sc7CN)I$&oPuYh)=!h`ibq^8d;~s)U
m53fwLaTJ6_xA(&vhPXym0=!w-KpG4g3>gFj85rUYfOr7Y2S`i+

literal 0
HcmV?d00001

diff --git a/tensorflow/contrib/lite/java/ovic/demo/gradle/wrapper/gradle-wrapper.properties b/tensorflow/contrib/lite/java/ovic/demo/gradle/wrapper/gradle-wrapper.properties
new file mode 100644
index 0000000000..fa7a38a0e4
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/demo/gradle/wrapper/gradle-wrapper.properties
@@ -0,0 +1,6 @@
+#Thu Sep 28 09:01:41 PDT 2017
+distributionBase=GRADLE_USER_HOME
+distributionPath=wrapper/dists
+zipStoreBase=GRADLE_USER_HOME
+zipStorePath=wrapper/dists
+distributionUrl=https\://services.gradle.org/distributions/gradle-3.3-all.zip
diff --git a/tensorflow/contrib/lite/java/ovic/demo/gradlew b/tensorflow/contrib/lite/java/ovic/demo/gradlew
new file mode 100755
index 0000000000..9d82f78915
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/demo/gradlew
@@ -0,0 +1,160 @@
+#!/usr/bin/env bash
+
+##############################################################################
+##
+##  Gradle start up script for UN*X
+##
+##############################################################################
+
+# Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.
+DEFAULT_JVM_OPTS=""
+
+APP_NAME="Gradle"
+APP_BASE_NAME=`basename "$0"`
+
+# Use the maximum available, or set MAX_FD != -1 to use that value.
+MAX_FD="maximum"
+
+warn ( ) {
+    echo "$*"
+}
+
+die ( ) {
+    echo
+    echo "$*"
+    echo
+    exit 1
+}
+
+# OS specific support (must be 'true' or 'false').
+cygwin=false
+msys=false
+darwin=false
+case "`uname`" in
+  CYGWIN* )
+    cygwin=true
+    ;;
+  Darwin* )
+    darwin=true
+    ;;
+  MINGW* )
+    msys=true
+    ;;
+esac
+
+# Attempt to set APP_HOME
+# Resolve links: $0 may be a link
+PRG="$0"
+# Need this for relative symlinks.
+while [ -h "$PRG" ] ; do
+    ls=`ls -ld "$PRG"`
+    link=`expr "$ls" : '.*-> \(.*\)$'`
+    if expr "$link" : '/.*' > /dev/null; then
+        PRG="$link"
+    else
+        PRG=`dirname "$PRG"`"/$link"
+    fi
+done
+SAVED="`pwd`"
+cd "`dirname \"$PRG\"`/" >/dev/null
+APP_HOME="`pwd -P`"
+cd "$SAVED" >/dev/null
+
+CLASSPATH=$APP_HOME/gradle/wrapper/gradle-wrapper.jar
+
+# Determine the Java command to use to start the JVM.
+if [ -n "$JAVA_HOME" ] ; then
+    if [ -x "$JAVA_HOME/jre/sh/java" ] ; then
+        # IBM's JDK on AIX uses strange locations for the executables
+        JAVACMD="$JAVA_HOME/jre/sh/java"
+    else
+        JAVACMD="$JAVA_HOME/bin/java"
+    fi
+    if [ ! -x "$JAVACMD" ] ; then
+        die "ERROR: JAVA_HOME is set to an invalid directory: $JAVA_HOME
+
+Please set the JAVA_HOME variable in your environment to match the
+location of your Java installation."
+    fi
+else
+    JAVACMD="java"
+    which java >/dev/null 2>&1 || die "ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.
+
+Please set the JAVA_HOME variable in your environment to match the
+location of your Java installation."
+fi
+
+# Increase the maximum file descriptors if we can.
+if [ "$cygwin" = "false" -a "$darwin" = "false" ] ; then
+    MAX_FD_LIMIT=`ulimit -H -n`
+    if [ $? -eq 0 ] ; then
+        if [ "$MAX_FD" = "maximum" -o "$MAX_FD" = "max" ] ; then
+            MAX_FD="$MAX_FD_LIMIT"
+        fi
+        ulimit -n $MAX_FD
+        if [ $? -ne 0 ] ; then
+            warn "Could not set maximum file descriptor limit: $MAX_FD"
+        fi
+    else
+        warn "Could not query maximum file descriptor limit: $MAX_FD_LIMIT"
+    fi
+fi
+
+# For Darwin, add options to specify how the application appears in the dock
+if $darwin; then
+    GRADLE_OPTS="$GRADLE_OPTS \"-Xdock:name=$APP_NAME\" \"-Xdock:icon=$APP_HOME/media/gradle.icns\""
+fi
+
+# For Cygwin, switch paths to Windows format before running java
+if $cygwin ; then
+    APP_HOME=`cygpath --path --mixed "$APP_HOME"`
+    CLASSPATH=`cygpath --path --mixed "$CLASSPATH"`
+    JAVACMD=`cygpath --unix "$JAVACMD"`
+
+    # We build the pattern for arguments to be converted via cygpath
+    ROOTDIRSRAW=`find -L / -maxdepth 1 -mindepth 1 -type d 2>/dev/null`
+    SEP=""
+    for dir in $ROOTDIRSRAW ; do
+        ROOTDIRS="$ROOTDIRS$SEP$dir"
+        SEP="|"
+    done
+    OURCYGPATTERN="(^($ROOTDIRS))"
+    # Add a user-defined pattern to the cygpath arguments
+    if [ "$GRADLE_CYGPATTERN" != "" ] ; then
+        OURCYGPATTERN="$OURCYGPATTERN|($GRADLE_CYGPATTERN)"
+    fi
+    # Now convert the arguments - kludge to limit ourselves to /bin/sh
+    i=0
+    for arg in "$@" ; do
+        CHECK=`echo "$arg"|egrep -c "$OURCYGPATTERN" -`
+        CHECK2=`echo "$arg"|egrep -c "^-"`                                 ### Determine if an option
+
+        if [ $CHECK -ne 0 ] && [ $CHECK2 -eq 0 ] ; then                    ### Added a condition
+            eval `echo args$i`=`cygpath --path --ignore --mixed "$arg"`
+        else
+            eval `echo args$i`="\"$arg\""
+        fi
+        i=$((i+1))
+    done
+    case $i in
+        (0) set -- ;;
+        (1) set -- "$args0" ;;
+        (2) set -- "$args0" "$args1" ;;
+        (3) set -- "$args0" "$args1" "$args2" ;;
+        (4) set -- "$args0" "$args1" "$args2" "$args3" ;;
+        (5) set -- "$args0" "$args1" "$args2" "$args3" "$args4" ;;
+        (6) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" ;;
+        (7) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" "$args6" ;;
+        (8) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" "$args6" "$args7" ;;
+        (9) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" "$args6" "$args7" "$args8" ;;
+    esac
+fi
+
+# Split up the JVM_OPTS And GRADLE_OPTS values into an array, following the shell quoting and substitution rules
+function splitJvmOpts() {
+    JVM_OPTS=("$@")
+}
+eval splitJvmOpts $DEFAULT_JVM_OPTS $JAVA_OPTS $GRADLE_OPTS
+JVM_OPTS[${#JVM_OPTS[*]}]="-Dorg.gradle.appname=$APP_BASE_NAME"
+
+exec "$JAVACMD" "${JVM_OPTS[@]}" -classpath "$CLASSPATH" org.gradle.wrapper.GradleWrapperMain "$@"
diff --git a/tensorflow/contrib/lite/java/ovic/demo/gradlew.bat b/tensorflow/contrib/lite/java/ovic/demo/gradlew.bat
new file mode 100644
index 0000000000..8a0b282aa6
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/demo/gradlew.bat
@@ -0,0 +1,90 @@
+@if "%DEBUG%" == "" @echo off
+@rem ##########################################################################
+@rem
+@rem  Gradle startup script for Windows
+@rem
+@rem ##########################################################################
+
+@rem Set local scope for the variables with windows NT shell
+if "%OS%"=="Windows_NT" setlocal
+
+@rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.
+set DEFAULT_JVM_OPTS=
+
+set DIRNAME=%~dp0
+if "%DIRNAME%" == "" set DIRNAME=.
+set APP_BASE_NAME=%~n0
+set APP_HOME=%DIRNAME%
+
+@rem Find java.exe
+if defined JAVA_HOME goto findJavaFromJavaHome
+
+set JAVA_EXE=java.exe
+%JAVA_EXE% -version >NUL 2>&1
+if "%ERRORLEVEL%" == "0" goto init
+
+echo.
+echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.
+echo.
+echo Please set the JAVA_HOME variable in your environment to match the
+echo location of your Java installation.
+
+goto fail
+
+:findJavaFromJavaHome
+set JAVA_HOME=%JAVA_HOME:"=%
+set JAVA_EXE=%JAVA_HOME%/bin/java.exe
+
+if exist "%JAVA_EXE%" goto init
+
+echo.
+echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME%
+echo.
+echo Please set the JAVA_HOME variable in your environment to match the
+echo location of your Java installation.
+
+goto fail
+
+:init
+@rem Get command-line arguments, handling Windowz variants
+
+if not "%OS%" == "Windows_NT" goto win9xME_args
+if "%@eval[2+2]" == "4" goto 4NT_args
+
+:win9xME_args
+@rem Slurp the command line arguments.
+set CMD_LINE_ARGS=
+set _SKIP=2
+
+:win9xME_args_slurp
+if "x%~1" == "x" goto execute
+
+set CMD_LINE_ARGS=%*
+goto execute
+
+:4NT_args
+@rem Get arguments from the 4NT Shell from JP Software
+set CMD_LINE_ARGS=%$
+
+:execute
+@rem Setup the command line
+
+set CLASSPATH=%APP_HOME%\gradle\wrapper\gradle-wrapper.jar
+
+@rem Execute Gradle
+"%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -classpath "%CLASSPATH%" org.gradle.wrapper.GradleWrapperMain %CMD_LINE_ARGS%
+
+:end
+@rem End local scope for the variables with windows NT shell
+if "%ERRORLEVEL%"=="0" goto mainEnd
+
+:fail
+rem Set variable GRADLE_EXIT_CONSOLE if you need the _script_ return code instead of
+rem the _cmd.exe /c_ return code!
+if  not "" == "%GRADLE_EXIT_CONSOLE%" exit 1
+exit /b 1
+
+:mainEnd
+if "%OS%"=="Windows_NT" endlocal
+
+:omega
diff --git a/tensorflow/contrib/lite/java/ovic/demo/settings.gradle b/tensorflow/contrib/lite/java/ovic/demo/settings.gradle
new file mode 100644
index 0000000000..e7b4def49c
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/demo/settings.gradle
@@ -0,0 +1 @@
+include ':app'
diff --git a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicClassifier.java b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicClassifier.java
index b2dfd8f2e7..4cf51bb0fa 100644
--- a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicClassifier.java
+++ b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicClassifier.java
@@ -67,7 +67,7 @@ public class OvicClassifier {
           });
 
   /** Initializes an {@code OvicClassifier}. */
-  OvicClassifier(InputStream labelInputStream, MappedByteBuffer model)
+  public OvicClassifier(InputStream labelInputStream, MappedByteBuffer model)
       throws IOException, RuntimeException {
     if (model == null) {
       throw new RuntimeException("Input model is empty.");
@@ -106,7 +106,7 @@ public class OvicClassifier {
 
   /** Classifies a {@link ByteBuffer} image. */
   // @throws RuntimeException if model is uninitialized.
-  OvicSingleImageResult classifyByteBuffer(ByteBuffer imgData) throws RuntimeException {
+  public OvicSingleImageResult classifyByteBuffer(ByteBuffer imgData) {
     if (tflite == null) {
       throw new RuntimeException(TAG + ": ImageNet classifier has not been initialized; Failed.");
     }
diff --git a/tensorflow/contrib/lite/java/ovic/src/testdata/BUILD b/tensorflow/contrib/lite/java/ovic/src/testdata/BUILD
new file mode 100644
index 0000000000..8715938e52
--- /dev/null
+++ b/tensorflow/contrib/lite/java/ovic/src/testdata/BUILD
@@ -0,0 +1,19 @@
+# Testdata for OVIC benchmarker demo App and tests.
+licenses(["notice"])  # Apache 2.0
+
+filegroup(
+    name = "ovic_testdata",
+    srcs = [
+        "float_model.lite",
+        "low_res_model.lite",
+        "quantized_model.lite",
+        "test_image_128.jpg",
+        "test_image_224.jpg",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+exports_files(
+    ["labels.txt"],
+    visibility = ["//visibility:public"],
+)
diff --git a/tensorflow/opensource_only/tflite_ovic_testdata.BUILD b/tensorflow/opensource_only/tflite_ovic_testdata.BUILD
new file mode 100644
index 0000000000..de47ed61f9
--- /dev/null
+++ b/tensorflow/opensource_only/tflite_ovic_testdata.BUILD
@@ -0,0 +1,12 @@
+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"])  # Apache 2.0
+
+exports_files(
+    glob(
+        ["**/*"],
+        exclude = [
+            "BUILD",
+        ],
+    ),
+)
diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 8b6ad0a138..b67c4bf2ac 100644
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -744,6 +744,17 @@ def tf_workspace(path_prefix="", tf_repo_name=""):
       build_file = clean_dep("//third_party:tflite_smartreply.BUILD"),
   )
 
+  tf_http_archive(
+      name = "tflite_ovic_testdata",
+      sha256 = "a9a705d8d519220178e2e65d383fdb21da37fdb31d1e909b0a1acdac46479e9c",
+      urls = [
+          "https://mirror.bazel.build/storage.googleapis.com/download.tensorflow.org/data/ovic.zip",
+          "https://storage.googleapis.com/download.tensorflow.org/data/ovic.zip",
+      ],
+      build_file = clean_dep("//third_party:tflite_ovic_testdata.BUILD"),
+      strip_prefix = "ovic",
+  )
+
   ##############################################################################
   # BIND DEFINITIONS
   #
-- 
GitLab


From 3c799b44b841dc8c4e83fb3fbb61ad4146464f60 Mon Sep 17 00:00:00 2001
From: "karl@kubx.ca" <karl@kubx.ca>
Date: Sun, 6 May 2018 10:11:35 -0400
Subject: [PATCH 282/395] Embed Java license in generator code

---
 tensorflow/java/src/gen/cc/op_generator.cc    | 19 ++++++++++++++++---
 .../src/gen/resources/license.java.snippet    | 14 --------------
 2 files changed, 16 insertions(+), 17 deletions(-)
 delete mode 100644 tensorflow/java/src/gen/resources/license.java.snippet

diff --git a/tensorflow/java/src/gen/cc/op_generator.cc b/tensorflow/java/src/gen/cc/op_generator.cc
index f4cefbe933..284f675c94 100644
--- a/tensorflow/java/src/gen/cc/op_generator.cc
+++ b/tensorflow/java/src/gen/cc/op_generator.cc
@@ -35,8 +35,21 @@ namespace tensorflow {
 namespace java {
 namespace {
 
-const char* kLicenseSnippet =
-    "tensorflow/java/src/gen/resources/license.java.snippet";
+const char* kLicense =
+  "/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.\n"
+  "\n"
+  "Licensed under the Apache License, Version 2.0 (the \"License\");\n"
+  "you may not use this file except in compliance with the License.\n"
+  "You may obtain a copy of the License at\n"
+  "\n"
+  "    http://www.apache.org/licenses/LICENSE-2.0\n"
+  "\n"
+  "Unless required by applicable law or agreed to in writing, software\n"
+  "distributed under the License is distributed on an \"AS IS\" BASIS,\n"
+  "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n"
+  "See the License for the specific language governing permissions and\n"
+  "limitations under the License.\n"
+  "=======================================================================*/\n";
 
 // There is three different modes to render an op class, depending on the
 // number and type of outputs it has:
@@ -401,7 +414,7 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
   SourceFileWriter writer(op_file.get());
   std::list<Type> dependencies;
   CollectOpDependencies(op, mode, &dependencies);
-  writer.WriteFromFile(kLicenseSnippet)
+  writer.Write(kLicense)
       .EndLine()
       .BeginType(op_class, PUBLIC|FINAL, &dependencies, &op_javadoc);
   if (!op.optional_attributes().empty()) {
diff --git a/tensorflow/java/src/gen/resources/license.java.snippet b/tensorflow/java/src/gen/resources/license.java.snippet
deleted file mode 100644
index 90285ec669..0000000000
--- a/tensorflow/java/src/gen/resources/license.java.snippet
+++ /dev/null
@@ -1,14 +0,0 @@
-/* Copyright 2018 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.
-==============================================================================*/
-- 
GitLab


From 5711f75c44170035bfb5889fcd50bc5b89eda346 Mon Sep 17 00:00:00 2001
From: Yong Tang <yong.tang.github@outlook.com>
Date: Mon, 7 May 2018 18:16:35 +0200
Subject: [PATCH 283/395] Update zlib version for cmake build (#19118)

In cmake build, the zlib version was `50893291621658f355bc5b4d450a8d06a563053d`
which is equal to v1.2.8. This creates a discrepancy between cmake and bazel build.

This fix updates the zlib version to v1.2.11 that syncs with bazel build.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
---
 tensorflow/contrib/cmake/external/zlib.cmake | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/contrib/cmake/external/zlib.cmake b/tensorflow/contrib/cmake/external/zlib.cmake
index 116d423093..8942f3eecf 100644
--- a/tensorflow/contrib/cmake/external/zlib.cmake
+++ b/tensorflow/contrib/cmake/external/zlib.cmake
@@ -31,7 +31,8 @@ else (systemlib_ZLIB)
   set(ZLIB_URL https://github.com/madler/zlib)
   set(ZLIB_BUILD ${CMAKE_CURRENT_BINARY_DIR}/zlib/src/zlib)
   set(ZLIB_INSTALL ${CMAKE_CURRENT_BINARY_DIR}/zlib/install)
-  set(ZLIB_TAG 50893291621658f355bc5b4d450a8d06a563053d)
+  # Match zlib version in tensorflow/workspace.bzl
+  set(ZLIB_TAG v1.2.11)
 
   if(WIN32)
     if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*")
-- 
GitLab


From 96dd8c99688290f90b9826af516b2c12d280c0ee Mon Sep 17 00:00:00 2001
From: Amit Patankar <amitpatankar@google.com>
Date: Mon, 7 May 2018 12:45:52 -0700
Subject: [PATCH 284/395] Building development docker files with AVX as well.
 (#19130)

---
 tensorflow/tools/docker/Dockerfile.devel     | 2 +-
 tensorflow/tools/docker/Dockerfile.devel-gpu | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/tools/docker/Dockerfile.devel b/tensorflow/tools/docker/Dockerfile.devel
index b9996395d0..406d134699 100644
--- a/tensorflow/tools/docker/Dockerfile.devel
+++ b/tensorflow/tools/docker/Dockerfile.devel
@@ -85,7 +85,7 @@ RUN git clone --branch=r1.8 --depth=1 https://github.com/tensorflow/tensorflow.g
 ENV CI_BUILD_PYTHON python
 
 RUN tensorflow/tools/ci_build/builds/configured CPU \
-    bazel build -c opt --cxxopt="-D_GLIBCXX_USE_CXX11_ABI=0" \
+    bazel build -c opt --copt=-mavx --cxxopt="-D_GLIBCXX_USE_CXX11_ABI=0" \
         # For optimized builds appropriate for the hardware platform of your choosing, uncomment below...
         # For ivy-bridge or sandy-bridge
         # --copt=-march="ivybridge" \
diff --git a/tensorflow/tools/docker/Dockerfile.devel-gpu b/tensorflow/tools/docker/Dockerfile.devel-gpu
index 7e5e6ef2d5..2fe47f3356 100644
--- a/tensorflow/tools/docker/Dockerfile.devel-gpu
+++ b/tensorflow/tools/docker/Dockerfile.devel-gpu
@@ -98,7 +98,7 @@ ENV TF_CUDNN_VERSION=7
 RUN ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 && \
     LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:${LD_LIBRARY_PATH} \
     tensorflow/tools/ci_build/builds/configured GPU \
-    bazel build -c opt --config=cuda \
+    bazel build -c opt --copt=-mavx --config=cuda \
 	--cxxopt="-D_GLIBCXX_USE_CXX11_ABI=0" \
         tensorflow/tools/pip_package:build_pip_package && \
     rm /usr/local/cuda/lib64/stubs/libcuda.so.1 && \
-- 
GitLab


From d8bda536c5080e761bcb24ab6984c26da875f52c Mon Sep 17 00:00:00 2001
From: Yong Tang <yong.tang.github@outlook.com>
Date: Mon, 7 May 2018 23:17:02 +0200
Subject: [PATCH 285/395] Fix doc format issue in enable_eager_execution
 (#19116)

* Fix doc format issue in enable_eager_execution

This fix tries to address the issue raised in 19022 where
the doc of `tf.enable_eager_execution` was not well formatted.
It seems the additional empty lines may confuse the doc generator.
This fix fixes with proper alignment and removal of empty lines.

This fix fixes 19022.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

* Remove extra spaces

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
---
 tensorflow/python/framework/ops.py | 24 +++++++++---------------
 1 file changed, 9 insertions(+), 15 deletions(-)

diff --git a/tensorflow/python/framework/ops.py b/tensorflow/python/framework/ops.py
index 908d0da35e..2209e8e21a 100644
--- a/tensorflow/python/framework/ops.py
+++ b/tensorflow/python/framework/ops.py
@@ -5402,36 +5402,30 @@ def enable_eager_execution(config=None, device_policy=None,
       in which operations are executed. Note that @{tf.ConfigProto} is also
       used to configure graph execution (via @{tf.Session}) and many options
       within `tf.ConfigProto` are not implemented (or are irrelevant) when
-     eager execution is enabled.
+      eager execution is enabled.
     device_policy: (Optional.) Policy controlling how operations requiring
-     inputs on a specific device (e.g., a GPU 0) handle inputs on a different
-     device  (e.g. GPU 1 or CPU). When set to None, an appropriate value will be
-     picked automatically. The value picked may change between TensorFlow
-     releases.
-     Valid values:
-
+      inputs on a specific device (e.g., a GPU 0) handle inputs on a different
+      device  (e.g. GPU 1 or CPU). When set to None, an appropriate value will be
+      picked automatically. The value picked may change between TensorFlow
+      releases.
+      Valid values:
       - tf.contrib.eager.DEVICE_PLACEMENT_EXPLICIT: raises an error if the
         placement is not correct.
-
       - tf.contrib.eager.DEVICE_PLACEMENT_WARN: copies the tensors which are not
         on the right device but logs a warning.
-
       - tf.contrib.eager.DEVICE_PLACEMENT_SILENT: silently copies the tensors.
         Note that this may hide performance problems as there is no notification
         provided when operations are blocked on the tensor being copied between
         devices.
-
       - tf.contrib.eager.DEVICE_PLACEMENT_SILENT_FOR_INT32: silently copies
         int32 tensors, raising errors on the other ones.
     execution_mode: (Optional.) Policy controlling how operations dispatched are
       actually executed. When set to None, an appropriate value will be picked
       automatically. The value picked may change between TensorFlow releases.
       Valid values:
-
-        - tf.contrib.eager.SYNC: executes each operation synchronously.
-
-        - tf.contrib.eager.ASYNC: executes each operation asynchronously. These
-          operations may return "non-ready" handles.
+      - tf.contrib.eager.SYNC: executes each operation synchronously.
+      - tf.contrib.eager.ASYNC: executes each operation asynchronously. These
+        operations may return "non-ready" handles.
 
   Raises:
     ValueError: If eager execution is enabled after creating/executing a
-- 
GitLab


From 714f3c4f2f901e865bfcbf830485adafb92dca48 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Tue, 8 May 2018 00:19:13 +0200
Subject: [PATCH 286/395] [tfgan] Add discriminator and generator losses to
 eval_metrics (#19117)

---
 .../estimator/python/gan_estimator_test.py     |  2 ++
 .../gan/python/estimator/python/head_impl.py   | 18 +++++++++++++++---
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py
index 387a62bd74..6bbd173f86 100644
--- a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py
+++ b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py
@@ -213,6 +213,8 @@ class GANEstimatorIntegrationTest(test.TestCase):
     scores = est.evaluate(eval_input_fn)
     self.assertEqual(num_steps, scores[ops.GraphKeys.GLOBAL_STEP])
     self.assertIn('loss', six.iterkeys(scores))
+    self.assertEqual(scores['discriminator_loss'] + scores['generator_loss'],
+                     scores['loss'])
 
     # PREDICT
     predictions = np.array([x for x in est.predict(predict_input_fn)])
diff --git a/tensorflow/contrib/gan/python/estimator/python/head_impl.py b/tensorflow/contrib/gan/python/estimator/python/head_impl.py
index a21358c50b..d174cb3bb2 100644
--- a/tensorflow/contrib/gan/python/estimator/python/head_impl.py
+++ b/tensorflow/contrib/gan/python/estimator/python/head_impl.py
@@ -25,12 +25,16 @@ from tensorflow.contrib.gan.python import train as tfgan_train
 from tensorflow.python.estimator import model_fn as model_fn_lib
 from tensorflow.python.estimator.canned import head
 from tensorflow.python.framework import ops
+from tensorflow.python.ops import metrics as metrics_lib
 
 __all__ = [
     'GANHead',
     'gan_head',
 ]
 
+def _summary_key(head_name, val):
+  return '%s/%s' % (val, head_name) if head_name else val
+
 
 def gan_head(generator_loss_fn, discriminator_loss_fn, generator_optimizer,
              discriminator_optimizer, use_loss_summaries=True,
@@ -104,6 +108,7 @@ class GANHead(head._Head):  # pylint: disable=protected-access
     self._generator_optimizer = generator_optimizer
     self._discriminator_optimizer = discriminator_optimizer
     self._get_hooks_fn = get_hooks_fn
+    self._name = name
 
   @property
   def name(self):
@@ -173,13 +178,20 @@ class GANHead(head._Head):  # pylint: disable=protected-access
         gan_loss = self.create_loss(
             features=None, mode=mode, logits=gan_model, labels=None)
         scalar_loss = gan_loss.generator_loss + gan_loss.discriminator_loss
+        with ops.name_scope(None, 'metrics',
+                            [gan_loss.generator_loss,
+                             gan_loss.discriminator_loss]):
+          eval_metric_ops = {
+              _summary_key(self._name, 'generator_loss'):
+                  metrics_lib.mean(gan_loss.generator_loss),
+              _summary_key(self._name, 'discriminator_loss'):
+                  metrics_lib.mean(gan_loss.discriminator_loss)
+          }
         return model_fn_lib.EstimatorSpec(
             mode=model_fn_lib.ModeKeys.EVAL,
             predictions=gan_model.generated_data,
             loss=scalar_loss,
-            # TODO(joelshor): Add metrics. If head name provided, append it to
-            # metric keys.
-            eval_metric_ops={})
+            eval_metric_ops=eval_metric_ops)
       elif mode == model_fn_lib.ModeKeys.TRAIN:
         if train_op_fn is None:
           raise ValueError('train_op_fn can not be None.')
-- 
GitLab


From 939fc534a4b2f227ee337e7dcfa82ec9b6337814 Mon Sep 17 00:00:00 2001
From: Shashi Shekhar <shashishekhar@google.com>
Date: Fri, 4 May 2018 19:18:44 -0700
Subject: [PATCH 287/395] Fix landscape layout.

PiperOrigin-RevId: 195506194
---
 .../layout-land/fragment_camera2_basic.xml    | 92 ++++++++++---------
 .../res/layout/fragment_camera2_basic.xml     | 79 ++++++----------
 .../contrib/lite/java/ovic/demo/app/BUILD     |  4 +-
 .../tensorflow/ovic/OvicClassifierTest.java   | 12 +--
 .../contrib/lite/java/ovic/src/testdata/BUILD | 10 +-
 .../tflite_ovic_testdata.BUILD                |  0
 6 files changed, 87 insertions(+), 110 deletions(-)
 rename {tensorflow/opensource_only => third_party}/tflite_ovic_testdata.BUILD (100%)

diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/res/layout-land/fragment_camera2_basic.xml b/tensorflow/contrib/lite/java/demo/app/src/main/res/layout-land/fragment_camera2_basic.xml
index 20f520814d..ef8a9e0845 100644
--- a/tensorflow/contrib/lite/java/demo/app/src/main/res/layout-land/fragment_camera2_basic.xml
+++ b/tensorflow/contrib/lite/java/demo/app/src/main/res/layout-land/fragment_camera2_basic.xml
@@ -13,51 +13,55 @@
  See the License for the specific language governing permissions and
  limitations under the License.
 -->
-<RelativeLayout xmlns:android="http://schemas.android.com/apk/res/android"
-    xmlns:app="http://schemas.android.com/apk/res-auto"
+
+<LinearLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="match_parent"
-    android:layout_height="match_parent">
+    android:layout_height="match_parent"
+    android:background="#bb7700"
+    android:orientation="horizontal">
+
+  <com.example.android.tflitecamerademo.AutoFitTextureView
+      android:id="@+id/texture"
+      android:layout_width="0dp"
+      android:layout_height="match_parent"
+      android:layout_weight=".8"/>
+
+  <LinearLayout
+      android:layout_width="0dp"
+      android:layout_height="match_parent"
+      android:layout_weight=".2"
+      android:orientation="vertical">
+
+    <ImageView
+        android:id="@+id/logoview"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:scaleType="centerInside"
+        android:src="@drawable/logo"/>
 
-    <LinearLayout
+    <ToggleButton
+        android:id="@+id/button"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:textOff="@string/tflite"
+        android:textOn="@string/nnapi"/>
+    <NumberPicker
+        android:id="@+id/np"
+        android:layout_width="wrap_content"
+        android:layout_height="47dp"
+        android:layout_gravity="center_horizontal"
+        android:visibility="visible"/>
+
+    <TextView
+        android:id="@+id/text"
+        android:textStyle="bold"
         android:layout_width="match_parent"
         android:layout_height="match_parent"
-        android:background="#bb7700"
-        android:orientation="horizontal"
-        android:weightSum="100">
-
-        <LinearLayout
-            android:layout_width="match_parent"
-            android:layout_height="match_parent"
-            android:layout_weight="30"
-            android:orientation="vertical">
-
-            <com.example.android.tflitecamerademo.AutoFitTextureView
-                android:id="@+id/texture"
-                android:layout_width="match_parent"
-                android:layout_height="match_parent"
-                android:layout_weight="100" />
-
-            <ImageView
-                android:id="@+id/logoview"
-                android:layout_width="match_parent"
-                android:layout_height="wrap_content"
-                android:layout_weight="100"
-                android:scaleType="centerCrop"
-                android:src="@drawable/logo" />
-
-        </LinearLayout>
-
-        <TextView
-            android:id="@+id/text"
-            android:layout_width="match_parent"
-            android:layout_height="match_parent"
-            android:layout_weight="70"
-            android:paddingLeft="5dp"
-            android:paddingTop="20dp"
-            android:textColor="#FFF"
-            android:textSize="20sp"
-            android:textStyle="bold" />
-
-    </LinearLayout>
-
-</RelativeLayout>
+        android:paddingTop="20dp"
+        android:textColor="#FFF"
+        android:textSize="20sp"/>
+
+  </LinearLayout>
+</LinearLayout>
+
diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml b/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml
index d12435d5ab..72a229ecdb 100644
--- a/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml
+++ b/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml
@@ -15,45 +15,47 @@
 -->
 <RelativeLayout xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:app="http://schemas.android.com/apk/res-auto"
-    xmlns:tools="http://schemas.android.com/tools"
     android:layout_width="match_parent"
-    android:layout_height="match_parent">
+    android:layout_height="match_parent"
+    android:background="#bb7700">
 
-    <LinearLayout
+    <com.example.android.tflitecamerademo.AutoFitTextureView
+        android:id="@+id/texture"
         android:layout_width="match_parent"
         android:layout_height="match_parent"
-        android:orientation="vertical"
-        android:weightSum="60">
-
-        <FrameLayout
-            android:id="@+id/control"
-            android:layout_width="match_parent"
-            android:layout_height="match_parent"
-            android:layout_alignParentBottom="true"
-            android:layout_alignParentStart="true"
-            android:layout_weight="60"
-            android:background="#cc7700"
-            android:paddingLeft="20dp"
-            android:paddingStart="20dp">
-
-        </FrameLayout>
+        android:layout_weight="1" />
 
-    <com.example.android.tflitecamerademo.AutoFitTextureView
-        android:id="@+id/texture"
+    <LinearLayout
         android:layout_width="wrap_content"
         android:layout_height="wrap_content"
+        android:layout_alignParentBottom="true"
+        android:layout_alignParentEnd="false"
         android:layout_alignParentStart="true"
-        android:layout_alignParentLeft="true"
-        android:layout_alignParentTop="true" />
+        android:layout_alignParentTop="false"
+        android:background="#bb7700"
+        android:orientation="vertical"
+        android:weightSum="100">
+
+        <ImageView
+            android:id="@+id/logoview2"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:layout_weight="30"
+            android:scaleType="fitStart"
+            android:src="@drawable/logo" />
 
         <TextView
             android:id="@+id/text"
             android:layout_width="match_parent"
-            android:layout_height="match_parent"
-            android:layout_weight="20"
+            android:layout_height="wrap_content"
+            android:layout_alignParentBottom="true"
+            android:layout_alignParentEnd="true"
+            android:layout_alignParentRight="true"
+            android:layout_weight="30"
             android:textColor="#FFF"
             android:textSize="20sp"
             android:textStyle="bold" />
+
     </LinearLayout>
 
     <RelativeLayout
@@ -83,33 +85,4 @@
             android:layout_below="@+id/button"
             android:visibility="visible" />
     </RelativeLayout>
-
-    <RelativeLayout
-        android:id="@+id/control2"
-        android:layout_width="match_parent"
-        android:layout_height="135dp"
-        android:layout_alignParentLeft="true"
-        android:layout_alignParentStart="true"
-        android:layout_alignTop="@+id/control"
-        android:layout_marginLeft="300dp"
-        android:layout_marginStart="300dp"
-        android:background="@color/control_background">
-
-        <ToggleButton
-            android:id="@+id/button"
-            android:textOff="@string/tflite"
-            android:textOn="@string/nnapi"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
-            android:layout_alignParentLeft="true"
-            android:layout_alignParentStart="true" />
-
-        <NumberPicker
-            android:id="@+id/np"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
-            android:layout_below="@+id/button"
-            android:visibility="visible" />
-    </RelativeLayout>
-
 </RelativeLayout>
diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/BUILD b/tensorflow/contrib/lite/java/ovic/demo/app/BUILD
index be5aaa5ece..47101ff574 100644
--- a/tensorflow/contrib/lite/java/ovic/demo/app/BUILD
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/BUILD
@@ -23,7 +23,7 @@ android_binary(
     deps = [
         "//tensorflow/contrib/lite/java:ovicbenchmarkerlib",
         "//tensorflow/contrib/lite/java:tensorflowlite",
-        "//third_party/java/android/android_sdk_linux/extras/android/compatibility/v13",
-        "//third_party/java/android/android_sdk_linux/extras/android/compatibility/v4",
+        "@androidsdk//com.android.support:support-v13-25.2.0",
+        "@androidsdk//com.android.support:support-v4-25.2.0",
     ],
 )
diff --git a/tensorflow/contrib/lite/java/ovic/src/test/java/org/tensorflow/ovic/OvicClassifierTest.java b/tensorflow/contrib/lite/java/ovic/src/test/java/org/tensorflow/ovic/OvicClassifierTest.java
index 098ed8ceba..56f3e7604a 100644
--- a/tensorflow/contrib/lite/java/ovic/src/test/java/org/tensorflow/ovic/OvicClassifierTest.java
+++ b/tensorflow/contrib/lite/java/ovic/src/test/java/org/tensorflow/ovic/OvicClassifierTest.java
@@ -45,17 +45,17 @@ public final class OvicClassifierTest {
   private ByteBuffer lowResTestImage = null;
   private OvicSingleImageResult testResult = null;
   private static final String LABELS_PATH =
-      "third_party/tensorflow/contrib/lite/java/ovic/src/testdata/labels.txt";
+      "tensorflow/contrib/lite/java/ovic/src/testdata/labels.txt";
   private static final String QUANTIZED_MODEL_PATH =
-      "third_party/tensorflow/contrib/lite/java/ovic/src/testdata/quantized_model.lite";
+      "external/tflite_ovic_testdata/quantized_model.lite";
   private static final String LOW_RES_MODEL_PATH =
-      "third_party/tensorflow/contrib/lite/java/ovic/src/testdata/low_res_model.lite";
+      "external/tflite_ovic_testdata/low_res_model.lite";
   private static final String FLOAT_MODEL_PATH =
-      "third_party/tensorflow/contrib/lite/java/ovic/src/testdata/float_model.lite";
+      "external/tflite_ovic_testdata/float_model.lite";
   private static final String TEST_IMAGE_PATH =
-      "third_party/tensorflow/contrib/lite/java/ovic/src/testdata/test_image_224.jpg";
+      "external/tflite_ovic_testdata/test_image_224.jpg";
   private static final String TEST_LOW_RES_IMAGE_PATH =
-      "third_party/tensorflow/contrib/lite/java/ovic/src/testdata/test_image_128.jpg";
+      "external/tflite_ovic_testdata/test_image_128.jpg";
   private static final int TEST_IMAGE_GROUNDTRUTH = 653; // "military uniform"
 
   @Before
diff --git a/tensorflow/contrib/lite/java/ovic/src/testdata/BUILD b/tensorflow/contrib/lite/java/ovic/src/testdata/BUILD
index 8715938e52..1021ea30dd 100644
--- a/tensorflow/contrib/lite/java/ovic/src/testdata/BUILD
+++ b/tensorflow/contrib/lite/java/ovic/src/testdata/BUILD
@@ -4,11 +4,11 @@ licenses(["notice"])  # Apache 2.0
 filegroup(
     name = "ovic_testdata",
     srcs = [
-        "float_model.lite",
-        "low_res_model.lite",
-        "quantized_model.lite",
-        "test_image_128.jpg",
-        "test_image_224.jpg",
+        "@tflite_ovic_testdata//:float_model.lite",
+        "@tflite_ovic_testdata//:low_res_model.lite",
+        "@tflite_ovic_testdata//:quantized_model.lite",
+        "@tflite_ovic_testdata//:test_image_128.jpg",
+        "@tflite_ovic_testdata//:test_image_224.jpg",
     ],
     visibility = ["//visibility:public"],
 )
diff --git a/tensorflow/opensource_only/tflite_ovic_testdata.BUILD b/third_party/tflite_ovic_testdata.BUILD
similarity index 100%
rename from tensorflow/opensource_only/tflite_ovic_testdata.BUILD
rename to third_party/tflite_ovic_testdata.BUILD
-- 
GitLab


From 150089e6e67e4492f098cdd8f9f2f48dc9f9cc56 Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Fri, 4 May 2018 22:04:20 -0700
Subject: [PATCH 288/395] Remove uses of the kTransposeDot fusion

I didn't remove the enum itself, but after this change removing the enum should
be a simple NFC change (famous last words!).

This will make it easier to implement BatchDot on CPU.

The change removes usages of kTransposeDot by:

 - Teaching TransposeFolding to "fuse" transposes into dots by flipping the
   lhs_contracting_dims/rhs_contracting_dims fields.

 - Replacing the notion of transpose_lhs/transpose_rhs in the IR emitters with
   "has a non-canonical LHS contraction dimension"/"has a non-canonical RHS
   contraction dimension" where the canonical LHS and RHS contraction dims [0]
   are 1 and 0.

Some tests were getting away with creating Dot instructions with their
dimensions numbers unset.  I've fixed these to create canonical dot operations
instead.

It is possible (but hard to tell without trying) that some of the IR emission
logic and Eigen runtime calls can now be simplified further.  For instance,
instead of passing in a `transpose_lhs` and `transpose_rhs` to the Eigen GEMM
routines, we could instead pass in the LHS and RHS contraction dimensions
directly.

[0] See HloInstruction::CreateCanonicalDot.

PiperOrigin-RevId: 195514907
---
 tensorflow/compiler/xla/service/BUILD         |   3 +
 .../cpu/cpu_instruction_fusion_test.cc        | 104 +++++++--
 .../xla/service/cpu/dot_op_emitter.cc         |  74 +++---
 .../compiler/xla/service/cpu/dot_op_emitter.h |  17 +-
 .../compiler/xla/service/cpu/ir_emitter.cc    |  57 +----
 .../compiler/xla/service/gpu/gemm_thunk.cc    |  16 +-
 .../compiler/xla/service/gpu/gemm_thunk.h     |  10 +-
 .../xla/service/gpu/hlo_schedule_test.cc      |  30 +--
 .../xla/service/gpu/ir_emission_utils.cc      |  10 +-
 .../xla/service/gpu/ir_emitter_unnested.cc    |  78 ++-----
 .../xla/service/gpu/stream_assignment_test.cc |  28 +--
 .../compiler/xla/service/hlo_matchers.cc      |  38 +++
 .../compiler/xla/service/hlo_matchers.h       |  46 +++-
 .../compiler/xla/service/hlo_matchers_test.cc |  37 +++
 .../compiler/xla/service/transpose_folding.cc |  54 +++--
 .../xla/service/transpose_folding_test.cc     | 219 +++++++++---------
 16 files changed, 477 insertions(+), 344 deletions(-)

diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD
index 9c362d8cad..714c1e8754 100644
--- a/tensorflow/compiler/xla/service/BUILD
+++ b/tensorflow/compiler/xla/service/BUILD
@@ -370,6 +370,7 @@ tf_cc_test(
         ":hlo_matchers",
         "//tensorflow/compiler/xla:shape_util",
         "//tensorflow/compiler/xla/tests:xla_internal_test_main",
+        "//tensorflow/compiler/xla/tools/parser:hlo_parser",
     ],
 )
 
@@ -2467,6 +2468,7 @@ tf_cc_test(
     srcs = ["transpose_folding_test.cc"],
     deps = [
         ":hlo",
+        ":hlo_matchers",
         ":shape_inference",
         ":transpose_folding",
         "//tensorflow/compiler/xla:literal_util",
@@ -2478,6 +2480,7 @@ tf_cc_test(
         "//tensorflow/compiler/xla/service/gpu:ir_emission_utils",
         "//tensorflow/compiler/xla/tests:hlo_test_base",
         "//tensorflow/compiler/xla/tests:xla_internal_test_main",
+        "//tensorflow/compiler/xla/tools/parser:hlo_parser",
         "//tensorflow/core:lib",
     ],
 )
diff --git a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc
index a98e85a151..46fe060817 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc
+++ b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc
@@ -158,37 +158,95 @@ TEST_F(InstructionFusionTest, DotOperationFusion_ElementReuse) {
   EXPECT_EQ(dot, computation->root_instruction());
 }
 
-TEST_F(InstructionFusionTest, DotOperationFusion_TransposeFusion) {
-  HloComputation::Builder builder(TestName());
-  HloInstruction* arg0 = builder.AddInstruction(HloInstruction::CreateParameter(
-      0, ShapeUtil::MakeShape(F32, {1, 256}), "arg0"));
-  HloInstruction* arg1 = builder.AddInstruction(HloInstruction::CreateParameter(
-      1, ShapeUtil::MakeShape(F32, {1024, 256}), "arg1"));
+TEST_F(InstructionFusionTest, DotOperationFusion_TransposeFusion_RHS) {
+  string hlo_string = R"(
+HloModule DotOperationFusion_TransposeFusion
 
-  HloInstruction* exp1 = builder.AddInstruction(HloInstruction::CreateUnary(
-      ShapeUtil::MakeShape(S32, {1024, 256}), HloOpcode::kExp, arg1));
-  HloInstruction* transpose1 =
-      builder.AddInstruction(HloInstruction::CreateTranspose(
-          ShapeUtil::MakeShape(S32, {256, 1024}), exp1, {1, 0}));
-  builder.AddInstruction(
-      MakeDot(ShapeUtil::MakeShape(F32, {1, 1024}), arg0, transpose1));
+ENTRY DotOperationFusion_TransposeFusion {
+  arg0 = f32[1,256] parameter(0)
+  arg1 = f32[1024,256] parameter(1)
+  exponential = s32[1024,256] exponential(arg1)
+  transpose = s32[256,1024] transpose(exponential), dimensions={1,0}
+  ROOT dot = f32[1,1024] dot(arg0, transpose), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_string));
+  HloComputation* computation = module->entry_computation();
 
-  auto module = CreateNewModule();
-  auto computation = module->AddEntryComputation(builder.Build());
   TransposeFolding transpose_folding(
       [](const HloInstruction& dot,
          const TransposeFolding::OperandIndices& candidate_operands) {
         return candidate_operands;
       },
       TransposeFolding::NeverFoldTranspose);
-  EXPECT_TRUE(transpose_folding.Run(module.get()).ValueOrDie());
-  EXPECT_EQ(computation->root_instruction()->opcode(), HloOpcode::kFusion);
-  EXPECT_EQ(computation->root_instruction()->fusion_kind(),
-            HloInstruction::FusionKind::kTransposeDot);
-  EXPECT_FALSE(CpuInstructionFusion().Run(module.get()).ValueOrDie());
-  EXPECT_EQ(computation->root_instruction()->opcode(), HloOpcode::kFusion);
-  EXPECT_EQ(computation->root_instruction()->fusion_kind(),
-            HloInstruction::FusionKind::kTransposeDot);
+  TF_ASSERT_OK_AND_ASSIGN(bool changed, transpose_folding.Run(module.get()));
+  ASSERT_TRUE(changed);
+  ASSERT_THAT(computation->root_instruction(),
+              op::Dot(op::Parameter(0), op::Exp(op::Parameter(1)),
+                      /*lhs_contracting_dim=*/1, /*rhs_contracting_dim=*/1));
+}
+
+TEST_F(InstructionFusionTest, DotOperationFusion_TransposeFusion_LHS) {
+  string hlo_string = R"(
+HloModule DotOperationFusion_TransposeFusion
+
+ENTRY DotOperationFusion_TransposeFusion {
+  arg0 = f32[256,1] parameter(0)
+  arg1 = f32[256,1024] parameter(1)
+  transpose = s32[1,256] transpose(arg0), dimensions={1,0}
+  exponential = s32[256,1024] exponential(arg1)
+  ROOT dot = f32[1,1024] dot(transpose, exponential), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_string));
+  HloComputation* computation = module->entry_computation();
+
+  TransposeFolding transpose_folding(
+      [](const HloInstruction& dot,
+         const TransposeFolding::OperandIndices& candidate_operands) {
+        return candidate_operands;
+      },
+      TransposeFolding::NeverFoldTranspose);
+  TF_ASSERT_OK_AND_ASSIGN(bool changed, transpose_folding.Run(module.get()));
+  ASSERT_TRUE(changed);
+  ASSERT_THAT(computation->root_instruction(),
+              op::Dot(op::Parameter(0), op::Exp(op::Parameter(1)),
+                      /*lhs_contracting_dim=*/0, /*rhs_contracting_dim=*/0));
+}
+
+TEST_F(InstructionFusionTest,
+       DotOperationFusion_TransposeFusion_LHS_NonDefault) {
+  string hlo_string = R"(
+HloModule DotOperationFusion_TransposeFusion
+
+ENTRY DotOperationFusion_TransposeFusion {
+  arg0 = f32[1,256] parameter(0)
+  arg1 = f32[256,1024] parameter(1)
+  transpose = s32[256,1] transpose(arg0), dimensions={1,0}
+  exponential = s32[256,1024] exponential(arg1)
+  ROOT dot = f32[1,1024] dot(transpose, exponential), lhs_contracting_dims={0}, rhs_contracting_dims={0}
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_string));
+  HloComputation* computation = module->entry_computation();
+
+  TransposeFolding transpose_folding(
+      [](const HloInstruction& dot,
+         const TransposeFolding::OperandIndices& candidate_operands) {
+        return candidate_operands;
+      },
+      TransposeFolding::NeverFoldTranspose);
+  TF_ASSERT_OK_AND_ASSIGN(bool changed, transpose_folding.Run(module.get()));
+  ASSERT_TRUE(changed);
+  ASSERT_THAT(computation->root_instruction(),
+              op::Dot(op::Parameter(0), op::Exp(op::Parameter(1)),
+                      /*lhs_contracting_dim=*/1, /*rhs_contracting_dim=*/0));
 }
 
 class OpcodeFusionTest : public InstructionFusionTest {
diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
index 801c523908..e5ac2a33e1 100644
--- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
+++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
@@ -522,16 +522,16 @@ void RowMajorMatrixVectorProductEmitter::EmitInnerLoopEpilogue(
 
 }  // namespace
 
-DotOpEmitter::DotOpEmitter(
-    const HloInstruction& dot, bool transpose_lhs, bool transpose_rhs,
-    const llvm_ir::IrArray& target_array, const llvm_ir::IrArray& lhs_array,
-    const llvm_ir::IrArray& rhs_array, const llvm_ir::IrArray* addend_array,
-    llvm::Value* executable_run_options_value, llvm::IRBuilder<>* ir_builder,
-    const HloModuleConfig& hlo_module_config,
-    const TargetMachineFeatures& target_machine_features)
+DotOpEmitter::DotOpEmitter(const HloInstruction& dot,
+                           const llvm_ir::IrArray& target_array,
+                           const llvm_ir::IrArray& lhs_array,
+                           const llvm_ir::IrArray& rhs_array,
+                           const llvm_ir::IrArray* addend_array,
+                           llvm::Value* executable_run_options_value,
+                           llvm::IRBuilder<>* ir_builder,
+                           const HloModuleConfig& hlo_module_config,
+                           const TargetMachineFeatures& target_machine_features)
     : dot_(dot),
-      transpose_lhs_(transpose_lhs),
-      transpose_rhs_(transpose_rhs),
       target_array_(target_array),
       lhs_array_(lhs_array),
       rhs_array_(rhs_array),
@@ -542,18 +542,18 @@ DotOpEmitter::DotOpEmitter(
       target_machine_features_(target_machine_features) {}
 
 /* static */ tensorflow::Status DotOpEmitter::EmitDotOperation(
-    const HloInstruction& dot, bool transpose_lhs, bool transpose_rhs,
-    const llvm_ir::IrArray& target_array, const llvm_ir::IrArray& lhs_array,
-    const llvm_ir::IrArray& rhs_array, const llvm_ir::IrArray* addend_array,
+    const HloInstruction& dot, const llvm_ir::IrArray& target_array,
+    const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array,
+    const llvm_ir::IrArray* addend_array,
     llvm::Value* executable_run_options_value, llvm::IRBuilder<>* ir_builder,
     const HloModuleConfig& hlo_module_config,
     const TargetMachineFeatures& target_machine_features) {
   PrimitiveType type = target_array.GetShape().element_type();
   TF_RET_CHECK(F16 == type || F32 == type || F64 == type || C64 == type);
-  DotOpEmitter dot_emitter(dot, transpose_lhs, transpose_rhs, target_array,
-                           lhs_array, rhs_array, addend_array,
-                           executable_run_options_value, ir_builder,
-                           hlo_module_config, target_machine_features);
+  DotOpEmitter dot_emitter(dot, target_array, lhs_array, rhs_array,
+                           addend_array, executable_run_options_value,
+                           ir_builder, hlo_module_config,
+                           target_machine_features);
   return dot_emitter.Emit();
 }
 
@@ -578,7 +578,7 @@ bool DotOpEmitter::EmitLlvmIrDotIfProfitable() {
 
   if (mat_mult_dims.m == 1) {
     bool rhs_effectively_row_major =
-        transpose_rhs_ ^ !mat_mult_dims.rhs_column_major;
+        mat_mult_dims.rhs_non_canonical ^ !mat_mult_dims.rhs_column_major;
     if (rhs_effectively_row_major) {
       k = mat_mult_dims.k;
       m = mat_mult_dims.n;
@@ -594,7 +594,7 @@ bool DotOpEmitter::EmitLlvmIrDotIfProfitable() {
 
   if (mat_mult_dims.n == 1) {
     bool lhs_effectively_column_major =
-        transpose_lhs_ ^ mat_mult_dims.lhs_column_major;
+        mat_mult_dims.lhs_non_canonical ^ mat_mult_dims.lhs_column_major;
     if (lhs_effectively_column_major) {
       m = mat_mult_dims.m;
       k = mat_mult_dims.k;
@@ -741,16 +741,10 @@ tensorflow::Status DotOpEmitter::Emit() {
   // Reduce along dimension 0 of the LHS and 1 of the RHS. Vectors are a special
   // case where the reduction dimension is 0 for both LHS and RHS. This results
   // in a vector dot product producing a scalar.
-  int64 lhs_reduction_dimension = 0;
-  if (ShapeUtil::Rank(lhs_shape) >= 2) {
-    lhs_reduction_dimension =
-        ShapeUtil::GetDimensionNumber(lhs_shape, transpose_lhs_ ? -2 : -1);
-  }
-  int64 rhs_reduction_dimension = 0;
-  if (ShapeUtil::Rank(rhs_shape) >= 2) {
-    rhs_reduction_dimension =
-        ShapeUtil::GetDimensionNumber(rhs_shape, transpose_rhs_ ? -1 : -2);
-  }
+  int64 lhs_reduction_dimension =
+      dot_.dot_dimension_numbers().lhs_contracting_dimensions(0);
+  int64 rhs_reduction_dimension =
+      dot_.dot_dimension_numbers().rhs_contracting_dimensions(0);
 
   // Verify the reduction dimension in the two operands are the same size.
   TF_RET_CHECK(lhs_shape.dimensions(lhs_reduction_dimension) ==
@@ -986,8 +980,8 @@ tensorflow::Status DotOpEmitter::EmitCallToRuntime() {
 
   const llvm_ir::IrArray* lhs = &lhs_array_;
   const llvm_ir::IrArray* rhs = &rhs_array_;
-  bool transpose_lhs = transpose_lhs_;
-  bool transpose_rhs = transpose_rhs_;
+  bool transpose_lhs = mat_mult_dims.lhs_non_canonical;
+  bool transpose_rhs = mat_mult_dims.rhs_non_canonical;
 
   if (!mat_mult_dims.lhs_column_major) {
     std::swap(mat_mult_dims.m, mat_mult_dims.n);
@@ -1015,12 +1009,16 @@ DotOpEmitter::MatMultDims DotOpEmitter::GetMatMultDims() const {
 
   const Shape& lhs_shape = lhs_array_.GetShape();
   const Shape& rhs_shape = rhs_array_.GetShape();
-
-  return {lhs_shape.dimensions(transpose_lhs_ ? 1 : 0),
-          lhs_shape.dimensions(transpose_lhs_ ? 0 : 1),
-          rhs_shape.dimensions(transpose_rhs_ ? 0 : 1),
-          LayoutUtil::Minor(lhs_shape.layout(), 0) == 0,
-          LayoutUtil::Minor(rhs_shape.layout(), 0) == 0};
+  const DotDimensionNumbers& dim_nums = dot_.dot_dimension_numbers();
+
+  return {
+      /*m=*/lhs_shape.dimensions(1 - dim_nums.lhs_contracting_dimensions(0)),
+      /*k=*/lhs_shape.dimensions(dim_nums.lhs_contracting_dimensions(0)),
+      /*n=*/rhs_shape.dimensions(1 - dim_nums.rhs_contracting_dimensions(0)),
+      /*lhs_column_major=*/LayoutUtil::Minor(lhs_shape.layout(), 0) == 0,
+      /*lhs_non_canonical=*/dim_nums.lhs_contracting_dimensions(0) == 0,
+      /*rhs_column_major=*/LayoutUtil::Minor(rhs_shape.layout(), 0) == 0,
+      /*rhs_non_canonical=*/dim_nums.rhs_contracting_dimensions(0) == 1};
 }
 
 llvm_ir::IrArray::Index DotOpEmitter::EmitOperandArrayLoopNest(
@@ -1090,10 +1088,12 @@ bool PotentiallyImplementedAsEigenDot(const HloInstruction& hlo) {
     // If gemm can accept the operand shapes, use it rather than a custom
     // kernel.
     if (AreValidGemmShapes(lhs_shape, rhs_shape, hlo.shape())) {
+      const DotDimensionNumbers& dim_numbers = hlo.dot_dimension_numbers();
       // The size of the reduction dimension should match. The shape inference
       // guarantees this invariant, so the check here is for programming
       // errors.
-      CHECK_EQ(lhs_shape.dimensions(1), rhs_shape.dimensions(0));
+      CHECK_EQ(lhs_shape.dimensions(dim_numbers.lhs_contracting_dimensions(0)),
+               rhs_shape.dimensions(dim_numbers.rhs_contracting_dimensions(0)));
       return true;
     }
   }
diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h
index 47e0924334..a20bf2f9db 100644
--- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h
+++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h
@@ -56,16 +56,15 @@ class DotOpEmitter {
   // dot(`lhs_array`, `rhs_array`).  A non-null `addend_array` is only supported
   // for Matrix-vector products.
   static tensorflow::Status EmitDotOperation(
-      const HloInstruction& dot, bool transpose_lhs, bool transpose_rhs,
-      const llvm_ir::IrArray& target_array, const llvm_ir::IrArray& lhs_array,
-      const llvm_ir::IrArray& rhs_array, const llvm_ir::IrArray* addend_array,
+      const HloInstruction& dot, const llvm_ir::IrArray& target_array,
+      const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array,
+      const llvm_ir::IrArray* addend_array,
       llvm::Value* executable_run_options_value, llvm::IRBuilder<>* ir_builder,
       const HloModuleConfig& hlo_module_config,
       const TargetMachineFeatures& target_machine_features);
 
  private:
-  DotOpEmitter(const HloInstruction& dot, bool transpose_lhs,
-               bool transpose_rhs, const llvm_ir::IrArray& target_array,
+  DotOpEmitter(const HloInstruction& dot, const llvm_ir::IrArray& target_array,
                const llvm_ir::IrArray& lhs_array,
                const llvm_ir::IrArray& rhs_array,
                const llvm_ir::IrArray* addend_array,
@@ -114,8 +113,14 @@ class DotOpEmitter {
     // True if the LHS matrix column major.
     bool lhs_column_major;
 
+    // True if the LHS contraction dimension is not 1.
+    bool lhs_non_canonical;
+
     // True if the RHS matrix column major.
     bool rhs_column_major;
+
+    // True if the RHS contraction dimension is not 0.
+    bool rhs_non_canonical;
   };
 
   // Get the MatMultDims instance for the dot product this DotOpEmitter
@@ -132,8 +137,6 @@ class DotOpEmitter {
   }
 
   const HloInstruction& dot_;
-  const bool transpose_lhs_;
-  const bool transpose_rhs_;
   const llvm_ir::IrArray& target_array_;
   const llvm_ir::IrArray& lhs_array_;
   const llvm_ir::IrArray& rhs_array_;
diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
index 6347ee2a2a..12f50e00b5 100644
--- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
+++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
@@ -827,13 +827,6 @@ Status IrEmitter::HandleDot(HloInstruction* dot) {
         "Dot with multiple contracting dimensions not implemented.");
   }
 
-  if (dnums.lhs_contracting_dimensions(0) !=
-          std::min(lhs->shape().dimensions_size() - 1, 1) ||
-      dnums.rhs_contracting_dimensions(0) != 0) {
-    return Unimplemented(
-        "Dot with non-standard contracting dimensions not implemented.");
-  }
-
   llvm_ir::IrArray lhs_array(GetIrArrayFor(lhs));
   llvm_ir::IrArray rhs_array(GetIrArrayFor(rhs));
 
@@ -850,8 +843,7 @@ Status IrEmitter::HandleDot(HloInstruction* dot) {
 
   // Dot operation is complicated so we delegate to a helper class.
   return DotOpEmitter::EmitDotOperation(
-      *dot, /*transpose_lhs=*/false, /*transpose_rhs=*/false, target_array,
-      lhs_array, rhs_array, /*addend_array=*/nullptr,
+      *dot, target_array, lhs_array, rhs_array, /*addend_array=*/nullptr,
       GetExecutableRunOptionsArgument(), &ir_builder_, hlo_module_config_,
       target_machine_features_);
 }
@@ -2085,45 +2077,10 @@ static const HloInstruction* StripTranspose(const HloInstruction& hlo) {
 }
 
 Status IrEmitter::HandleFusion(HloInstruction* fusion) {
-  auto* root = fusion->fused_expression_root();
-  if (fusion->fusion_kind() == HloInstruction::FusionKind::kTransposeDot) {
-    DCHECK(root->opcode() == HloOpcode::kDot);
-    const HloInstruction* lhs_parameter = StripTranspose(*root->operand(0));
-    const HloInstruction* rhs_parameter = StripTranspose(*root->operand(1));
-    DCHECK(lhs_parameter->opcode() == HloOpcode::kParameter &&
-           rhs_parameter->opcode() == HloOpcode::kParameter);
-    const HloInstruction* lhs =
-        fusion->operand(lhs_parameter->parameter_number());
-    const HloInstruction* rhs =
-        fusion->operand(rhs_parameter->parameter_number());
-
-    TF_RETURN_IF_ERROR(ElementTypesSameAndSupported(
-        /*instruction=*/*root, /*operands=*/{lhs, rhs},
-        /*supported_types=*/{F16, F32, F64}));
-
-    llvm_ir::IrArray lhs_array(GetIrArrayFor(lhs));
-    llvm_ir::IrArray rhs_array(GetIrArrayFor(rhs));
+  CHECK_NE(fusion->fusion_kind(), HloInstruction::FusionKind::kTransposeDot);
 
-    Shape target_shape = fusion->shape();
-    TF_RETURN_IF_ERROR(EmitTargetAddressForOp(fusion));
-    llvm_ir::IrArray target_array = GetIrArrayFor(fusion);
-    VLOG(2) << "HandleFusion kTransposeDot: ";
-    VLOG(2) << "  lhs operand: "
-            << llvm_ir::DumpToString(*lhs_array.GetBasePointer());
-    VLOG(2) << "  rhs operand: "
-            << llvm_ir::DumpToString(*rhs_array.GetBasePointer());
-    VLOG(2) << "  target: "
-            << llvm_ir::DumpToString(*target_array.GetBasePointer());
-
-    // Dot operation is complicated so we delegate to a helper class.
-    TF_RETURN_IF_ERROR(DotOpEmitter::EmitDotOperation(
-        *root, root->operand(0)->IsRank2Transpose(),
-        root->operand(1)->IsRank2Transpose(), target_array, lhs_array,
-        rhs_array, /*addend_array=*/nullptr, GetExecutableRunOptionsArgument(),
-        &ir_builder_, hlo_module_config_, target_machine_features_));
-    return Status::OK();
-  } else if (llvm_ir::CanEmitFusedDynamicUpdateSliceInPlace(fusion,
-                                                            assignment_)) {
+  auto* root = fusion->fused_expression_root();
+  if (llvm_ir::CanEmitFusedDynamicUpdateSliceInPlace(fusion, assignment_)) {
     VLOG(3) << "HandleFusion FusedDynamicUpdateSliceInPlace";
     CpuElementalIrEmitter elemental_emitter(hlo_module_config_, this, module_);
     TF_RETURN_IF_ERROR(EmitTargetAddressForOp(fusion));
@@ -2166,9 +2123,9 @@ Status IrEmitter::HandleFusion(HloInstruction* fusion) {
         GetIrArrayFor(fusion->operand(addend_param_number)));
 
     TF_RETURN_IF_ERROR(DotOpEmitter::EmitDotOperation(
-        *dot, /*transpose_lhs=*/false, /*transpose_rhs=*/false, target_array,
-        lhs_array, rhs_array, &addend_array, GetExecutableRunOptionsArgument(),
-        &ir_builder_, hlo_module_config_, target_machine_features_));
+        *dot, target_array, lhs_array, rhs_array, &addend_array,
+        GetExecutableRunOptionsArgument(), &ir_builder_, hlo_module_config_,
+        target_machine_features_));
     return Status::OK();
   } else {
     return Unimplemented("Fusion kind not implemented on CPU");
diff --git a/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc b/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc
index 0ec12f52d8..f996fe486d 100644
--- a/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc
+++ b/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc
@@ -221,8 +221,7 @@ GemmThunk::GemmThunk(const BufferAllocation::Slice& lhs_buffer,
                      const BufferAllocation::Slice& rhs_buffer,
                      const BufferAllocation::Slice& output_buffer,
                      const Shape& lhs_shape, const Shape& rhs_shape,
-                     const Shape& output_shape, bool transpose_lhs,
-                     bool transpose_rhs, double alpha,
+                     const Shape& output_shape, double alpha,
                      const HloInstruction* hlo_instruction)
     : Thunk(Kind::kGemm, hlo_instruction),
       lhs_buffer_(lhs_buffer),
@@ -231,8 +230,6 @@ GemmThunk::GemmThunk(const BufferAllocation::Slice& lhs_buffer,
       lhs_shape_(lhs_shape),
       rhs_shape_(rhs_shape),
       output_shape_(output_shape),
-      transpose_lhs_(transpose_lhs),
-      transpose_rhs_(transpose_rhs),
       alpha_(alpha) {}
 
 tensorflow::Status GemmThunk::ExecuteOnStream(
@@ -284,10 +281,13 @@ tensorflow::Status GemmThunk::ExecuteOnStream(
                             shape.dimensions(!is_row_major));
   };
 
-  const MatrixDescriptor lhs_descriptor =
-      make_descriptor(lhs_data, lhs_shape_, transpose_lhs_);
-  const MatrixDescriptor rhs_descriptor =
-      make_descriptor(rhs_data, rhs_shape_, transpose_rhs_);
+  const DotDimensionNumbers& dim_nums =
+      hlo_instruction()->dot_dimension_numbers();
+
+  const MatrixDescriptor lhs_descriptor = make_descriptor(
+      lhs_data, lhs_shape_, dim_nums.lhs_contracting_dimensions(0) == 0);
+  const MatrixDescriptor rhs_descriptor = make_descriptor(
+      rhs_data, rhs_shape_, dim_nums.rhs_contracting_dimensions(0) == 1);
 
   // Dispatches to a regular cublas gemm, a gemm-with-algorithm, or attempts to
   // autotune this gemm to figure out the best algorithm.
diff --git a/tensorflow/compiler/xla/service/gpu/gemm_thunk.h b/tensorflow/compiler/xla/service/gpu/gemm_thunk.h
index a18f425bc3..f42cbf9e94 100644
--- a/tensorflow/compiler/xla/service/gpu/gemm_thunk.h
+++ b/tensorflow/compiler/xla/service/gpu/gemm_thunk.h
@@ -35,15 +35,13 @@ namespace gpu {
 class GemmThunk : public Thunk {
  public:
   // Constructs a thunk that computes "output = (lhs <dot> rhs) * alpha" using
-  // BLAS gemm. transpose_lhs and transpose_rhs indicate whether gemm should
-  // transpose the lhs and rhs operand. hlo_instruction is as in Thunk. alpha is
-  // a constant.
+  // BLAS gemm.  hlo_instruction is as in Thunk. alpha is a constant.
   GemmThunk(const BufferAllocation::Slice& lhs_buffer,
             const BufferAllocation::Slice& rhs_buffer,
             const BufferAllocation::Slice& output_buffer,
             const Shape& lhs_shape, const Shape& rhs_shape,
-            const Shape& output_shape, bool transpose_lhs, bool transpose_rhs,
-            double alpha, const HloInstruction* hlo_instruction);
+            const Shape& output_shape, double alpha,
+            const HloInstruction* hlo_instruction);
 
   GemmThunk(const GemmThunk&) = delete;
   GemmThunk& operator=(const GemmThunk&) = delete;
@@ -69,8 +67,6 @@ class GemmThunk : public Thunk {
   const Shape rhs_shape_;
   const Shape output_shape_;
 
-  const bool transpose_lhs_;
-  const bool transpose_rhs_;
   const double alpha_;
 
   // Maps device names (StreamExecutor::DeviceDescription::name()) to autotune
diff --git a/tensorflow/compiler/xla/service/gpu/hlo_schedule_test.cc b/tensorflow/compiler/xla/service/gpu/hlo_schedule_test.cc
index ece9fa04dc..6436abc06c 100644
--- a/tensorflow/compiler/xla/service/gpu/hlo_schedule_test.cc
+++ b/tensorflow/compiler/xla/service/gpu/hlo_schedule_test.cc
@@ -65,9 +65,9 @@ TEST_F(HloScheduleTest, SequentialMatMul) {
   HloInstruction* z = builder.AddInstruction(HloInstruction::CreateParameter(
       /*parameter_number=*/2, f32_2x2_, /*name=*/"z"));
   HloInstruction* dot1 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, x, y));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, x, y));
   HloInstruction* dot2 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, dot1, z));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, dot1, z));
 
   auto module = CreateNewModule();
   module->AddEntryComputation(builder.Build(dot2));
@@ -193,11 +193,11 @@ TEST_F(HloScheduleTest, ConcurrentMatMul) {
   HloInstruction* y = builder.AddInstruction(HloInstruction::CreateParameter(
       /*parameter_number=*/1, f32_2x2_, /*name=*/"y"));
   HloInstruction* dot1 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, x, y));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, x, y));
   HloInstruction* dot2 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, y, x));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, y, x));
   HloInstruction* add = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kAdd, dot1, dot2));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, dot1, dot2));
 
   auto module = CreateNewModule();
   module->AddEntryComputation(builder.Build(add));
@@ -259,24 +259,24 @@ TEST_F(HloScheduleTest, LatticeMatMul) {
     params.push_back(builder.AddInstruction(HloInstruction::CreateParameter(
         i, f32_2x2_, /*name=*/tensorflow::strings::Printf("param%d", i))));
   }
-  HloInstruction* d00 = builder.AddInstruction(HloInstruction::CreateBinary(
-      f32_2x2_, HloOpcode::kDot, params[2], params[3]));
+  HloInstruction* d00 = builder.AddInstruction(
+      HloInstruction::CreateCanonicalDot(f32_2x2_, params[2], params[3]));
   HloInstruction* d10 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, params[1], d00));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, params[1], d00));
   HloInstruction* d11 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, d00, params[4]));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, d00, params[4]));
   HloInstruction* d20 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, params[0], d10));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, params[0], d10));
   HloInstruction* d21 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, d10, d11));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, d10, d11));
   HloInstruction* d22 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, d11, params[5]));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, d11, params[5]));
   HloInstruction* d30 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, d20, d21));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, d20, d21));
   HloInstruction* d31 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, d21, d22));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, d21, d22));
   HloInstruction* d40 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, d30, d31));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, d30, d31));
 
   auto module = CreateNewModule();
   module->AddEntryComputation(builder.Build(d40));
diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
index 532d436ee8..777345722c 100644
--- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
+++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
@@ -78,15 +78,15 @@ bool ImplementedAsGemm(const HloInstruction& hlo) {
       // The size of the reduction dimension should match. The shape inference
       // guarantees this invariant, so the check here is for programming
       // errors.
-      CHECK_EQ(lhs_shape.dimensions(1), rhs_shape.dimensions(0));
+      const DotDimensionNumbers& dim_numbers = hlo.dot_dimension_numbers();
+      CHECK_EQ(lhs_shape.dimensions(dim_numbers.lhs_contracting_dimensions(0)),
+               rhs_shape.dimensions(dim_numbers.rhs_contracting_dimensions(0)));
       return true;
     }
   }
 
-  if (hlo.opcode() == HloOpcode::kFusion &&
-      hlo.fusion_kind() == HloInstruction::FusionKind::kTransposeDot &&
-      hlo.fused_expression_root()->opcode() == HloOpcode::kDot) {
-    return true;
+  if (hlo.opcode() == HloOpcode::kFusion) {
+    CHECK_NE(hlo.fusion_kind(), HloInstruction::FusionKind::kTransposeDot);
   }
 
   if (hlo.opcode() == HloOpcode::kFusion &&
diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
index 9f37235d32..83d90296df 100644
--- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
+++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
@@ -2206,65 +2206,37 @@ std::unique_ptr<Thunk> IrEmitterUnnested::BuildGemmThunk(
         lhs->shape(),               // The shape of LHS.
         rhs->shape(),               // The shape of RHS.
         inst->shape(),              // The shape of the output.
-        false,                      // Do not transpose LHS.
-        false,                      // Do not transpose RHS.
         1.0,                        // alpha.
         inst);
   }
 
   if (inst->opcode() == HloOpcode::kFusion) {
-    if (inst->fusion_kind() == HloInstruction::FusionKind::kOutput) {
-      const HloInstruction* mul = inst->fused_expression_root();
-      const HloInstruction* dot = mul->operand(0);
-      const HloInstruction* alpha = mul->operand(1);
-      if (dot->opcode() != HloOpcode::kDot) {
-        std::swap(dot, alpha);
-      }
-      DCHECK(dot->opcode() == HloOpcode::kDot);
-      const HloInstruction* lhs_parameter = StripTranspose(*dot->operand(0));
-      const HloInstruction* rhs_parameter = StripTranspose(*dot->operand(1));
-      DCHECK(lhs_parameter->opcode() == HloOpcode::kParameter &&
-             rhs_parameter->opcode() == HloOpcode::kParameter);
-      const HloInstruction* lhs =
-          inst->operand(lhs_parameter->parameter_number());
-      const HloInstruction* rhs =
-          inst->operand(rhs_parameter->parameter_number());
-
-      return MakeUnique<GemmThunk>(
-          GetAllocationSlice(*lhs),             // The buffer assigned to LHS.
-          GetAllocationSlice(*rhs),             // The buffer assigned to RHS.
-          GetAllocationSlice(*mul),             // The output buffer.
-          lhs->shape(),                         // The shape of LHS.
-          rhs->shape(),                         // The shape of RHS.
-          inst->shape(),                        // The shape of the output.
-          dot->operand(0)->IsRank2Transpose(),  // Transpose LHS.
-          dot->operand(1)->IsRank2Transpose(),  // Transpose RHS.
-          alpha->literal().Get<double>({0}),    // alpha.
-          inst);
-    } else {
-      const HloInstruction* dot = inst->fused_expression_root();
-      DCHECK(dot->opcode() == HloOpcode::kDot);
-      const HloInstruction* lhs_parameter = StripTranspose(*dot->operand(0));
-      const HloInstruction* rhs_parameter = StripTranspose(*dot->operand(1));
-      DCHECK(lhs_parameter->opcode() == HloOpcode::kParameter &&
-             rhs_parameter->opcode() == HloOpcode::kParameter);
-      const HloInstruction* lhs =
-          inst->operand(lhs_parameter->parameter_number());
-      const HloInstruction* rhs =
-          inst->operand(rhs_parameter->parameter_number());
-
-      return MakeUnique<GemmThunk>(
-          GetAllocationSlice(*lhs),             // The buffer assigned to LHS.
-          GetAllocationSlice(*rhs),             // The buffer assigned to RHS.
-          GetAllocationSlice(*inst),            // The output buffer.
-          lhs->shape(),                         // The shape of LHS.
-          rhs->shape(),                         // The shape of RHS.
-          inst->shape(),                        // The shape of the output.
-          dot->operand(0)->IsRank2Transpose(),  // Transpose LHS.
-          dot->operand(1)->IsRank2Transpose(),  // Transpose RHS.
-          1.0,                                  // Alpha.
-          inst);
+    CHECK_EQ(inst->fusion_kind(), HloInstruction::FusionKind::kOutput);
+    const HloInstruction* mul = inst->fused_expression_root();
+    const HloInstruction* dot = mul->operand(0);
+    const HloInstruction* alpha = mul->operand(1);
+    if (dot->opcode() != HloOpcode::kDot) {
+      std::swap(dot, alpha);
     }
+    DCHECK(dot->opcode() == HloOpcode::kDot);
+    const HloInstruction* lhs_parameter = StripTranspose(*dot->operand(0));
+    const HloInstruction* rhs_parameter = StripTranspose(*dot->operand(1));
+    DCHECK(lhs_parameter->opcode() == HloOpcode::kParameter &&
+           rhs_parameter->opcode() == HloOpcode::kParameter);
+    const HloInstruction* lhs =
+        inst->operand(lhs_parameter->parameter_number());
+    const HloInstruction* rhs =
+        inst->operand(rhs_parameter->parameter_number());
+
+    return MakeUnique<GemmThunk>(
+        GetAllocationSlice(*lhs),           // The buffer assigned to LHS.
+        GetAllocationSlice(*rhs),           // The buffer assigned to RHS.
+        GetAllocationSlice(*mul),           // The output buffer.
+        lhs->shape(),                       // The shape of LHS.
+        rhs->shape(),                       // The shape of RHS.
+        inst->shape(),                      // The shape of the output.
+        alpha->literal().Get<double>({0}),  // alpha.
+        inst);
   }
 
   LOG(FATAL) << "Cannot build a GemmThunk for " << inst->ToString();
diff --git a/tensorflow/compiler/xla/service/gpu/stream_assignment_test.cc b/tensorflow/compiler/xla/service/gpu/stream_assignment_test.cc
index 8c98956f1a..b42767dfd5 100644
--- a/tensorflow/compiler/xla/service/gpu/stream_assignment_test.cc
+++ b/tensorflow/compiler/xla/service/gpu/stream_assignment_test.cc
@@ -41,9 +41,9 @@ TEST_F(StreamAssignmentTest, SequentialMatMul) {
   HloInstruction* z = builder.AddInstruction(HloInstruction::CreateParameter(
       /*parameter_number=*/2, f32_2x2_, /*name=*/"z"));
   HloInstruction* dot1 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, x, y));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, x, y));
   HloInstruction* dot2 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, dot1, z));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, dot1, z));
 
   auto module = CreateNewModule();
   module->AddEntryComputation(builder.Build(dot2));
@@ -60,9 +60,9 @@ TEST_F(StreamAssignmentTest, ConcurrentMatMul) {
   HloInstruction* y = builder.AddInstruction(HloInstruction::CreateParameter(
       /*parameter_number=*/1, f32_2x2_, /*name=*/"y"));
   HloInstruction* dot1 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, x, y));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, x, y));
   HloInstruction* dot2 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, y, x));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, y, x));
   HloInstruction* add = builder.AddInstruction(
       HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kAdd, dot1, dot2));
 
@@ -91,24 +91,24 @@ TEST_F(StreamAssignmentTest, LatticeMatMul) {
     params.push_back(builder.AddInstruction(HloInstruction::CreateParameter(
         i, f32_2x2_, /*name=*/tensorflow::strings::Printf("param%d", i))));
   }
-  HloInstruction* d00 = builder.AddInstruction(HloInstruction::CreateBinary(
-      f32_2x2_, HloOpcode::kDot, params[2], params[3]));
+  HloInstruction* d00 = builder.AddInstruction(
+      HloInstruction::CreateCanonicalDot(f32_2x2_, params[2], params[3]));
   HloInstruction* d10 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, params[1], d00));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, params[1], d00));
   HloInstruction* d11 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, d00, params[4]));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, d00, params[4]));
   HloInstruction* d20 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, params[0], d10));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, params[0], d10));
   HloInstruction* d21 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, d10, d11));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, d10, d11));
   HloInstruction* d22 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, d11, params[5]));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, d11, params[5]));
   HloInstruction* d30 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, d20, d21));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, d20, d21));
   HloInstruction* d31 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, d21, d22));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, d21, d22));
   HloInstruction* d40 = builder.AddInstruction(
-      HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, d30, d31));
+      HloInstruction::CreateCanonicalDot(f32_2x2_, d30, d31));
 
   auto module = CreateNewModule();
   module->AddEntryComputation(builder.Build(d40));
diff --git a/tensorflow/compiler/xla/service/hlo_matchers.cc b/tensorflow/compiler/xla/service/hlo_matchers.cc
index 69deac263e..41ce9c1762 100644
--- a/tensorflow/compiler/xla/service/hlo_matchers.cc
+++ b/tensorflow/compiler/xla/service/hlo_matchers.cc
@@ -17,10 +17,13 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
 #include "tensorflow/compiler/xla/test.h"
+#include "tensorflow/core/lib/strings/str_util.h"
 
 namespace xla {
 namespace testing {
 
+using ::tensorflow::str_util::Join;
+
 bool HloMatcher::MatchAndExplain(
     const HloInstruction* instruction,
     ::testing::MatchResultListener* listener) const {
@@ -195,6 +198,41 @@ void HloShardingMatcher::DescribeTo(std::ostream* os) const {
   }
 }
 
+bool HloDotWithContractDimsMatcher::MatchAndExplain(
+    const HloInstruction* instruction,
+    ::testing::MatchResultListener* listener) const {
+  if (!HloMatcher::MatchAndExplain(instruction, listener)) {
+    return false;
+  }
+
+  const DotDimensionNumbers& dim_nums = instruction->dot_dimension_numbers();
+  if (dim_nums.lhs_contracting_dimensions_size() != 1 ||
+      dim_nums.lhs_contracting_dimensions(0) != lhs_contracting_dim_) {
+    *listener << instruction->ToString()
+              << " has wrong lhs_contracting_dimensions (got {"
+              << Join(dim_nums.lhs_contracting_dimensions(), ",") << "} want {"
+              << lhs_contracting_dim_ << "})";
+    return false;
+  }
+
+  if (dim_nums.rhs_contracting_dimensions_size() != 1 ||
+      dim_nums.rhs_contracting_dimensions(0) != rhs_contracting_dim_) {
+    *listener << instruction->ToString()
+              << " has wrong rhs_contracting_dimensions (got {"
+              << Join(dim_nums.rhs_contracting_dimensions(), ",") << "} want {"
+              << rhs_contracting_dim_ << "})";
+    return false;
+  }
+
+  return true;
+}
+
+void HloDotWithContractDimsMatcher::DescribeTo(std::ostream* os) const {
+  HloMatcher::DescribeTo(os);
+  *os << " with lhs_contracting_dims={" << lhs_contracting_dim_
+      << "} and rhs_contracting_dims={" << rhs_contracting_dim_ << "}";
+}
+
 }  // namespace testing
 
 void PrintTo(const HloInstruction* inst, ::std::ostream* os) {
diff --git a/tensorflow/compiler/xla/service/hlo_matchers.h b/tensorflow/compiler/xla/service/hlo_matchers.h
index 5175736a25..75231beac7 100644
--- a/tensorflow/compiler/xla/service/hlo_matchers.h
+++ b/tensorflow/compiler/xla/service/hlo_matchers.h
@@ -131,6 +131,27 @@ class HloShardingMatcher
   tensorflow::gtl::optional<HloSharding> sharding_;
 };
 
+// Matches a Dot HLO instruction with specific LHS and RHS contracting
+// dimensions.
+class HloDotWithContractDimsMatcher : public HloMatcher {
+ public:
+  explicit HloDotWithContractDimsMatcher(
+      ::testing::Matcher<const HloInstruction*> lhs,
+      ::testing::Matcher<const HloInstruction*> rhs, int64 lhs_contracting_dim,
+      int64 rhs_contracting_dim)
+      : HloMatcher(HloOpcode::kDot, /*operands=*/{lhs, rhs}),
+        lhs_contracting_dim_(lhs_contracting_dim),
+        rhs_contracting_dim_(rhs_contracting_dim) {}
+
+  bool MatchAndExplain(const HloInstruction* instruction,
+                       ::testing::MatchResultListener* listener) const override;
+  void DescribeTo(std::ostream* os) const override;
+
+ private:
+  int64 lhs_contracting_dim_;
+  int64 rhs_contracting_dim_;
+};
+
 // HloInstruction* matchers for opcode and operands. Example:
 //   namespace op = xla::opcode_matchers;
 //   EXPECT_THAT(instruction,
@@ -158,7 +179,6 @@ HLO_MATCHER(Convolution);
 HLO_MATCHER(Copy);
 HLO_MATCHER(CrossReplicaSum);
 HLO_MATCHER(Divide);
-HLO_MATCHER(Dot);
 HLO_MATCHER(DynamicSlice);
 HLO_MATCHER(DynamicUpdateSlice);
 HLO_MATCHER(Eq);
@@ -310,6 +330,30 @@ inline ::testing::Matcher<const ::xla::HloInstruction*> NoSharding() {
       new ::xla::testing::HloShardingMatcher(tensorflow::gtl::nullopt));
 }
 
+inline ::testing::Matcher<const ::xla::HloInstruction*> Dot(
+    ::testing::Matcher<const HloInstruction*> lhs_matcher,
+    ::testing::Matcher<const HloInstruction*> rhs_matcher) {
+  return ::testing::MakeMatcher(new ::xla::testing::HloMatcher(
+      ::xla::HloOpcode::kDot, {lhs_matcher, rhs_matcher}));
+}
+
+// Matches a Dot HLO instruction if it has exactly one lhs contracting dimension
+// equal to `lhs_contracting_dim` and exactly one rhs contracting dimension
+// equal to `rhs_contracting_dim`.
+//
+// Currently the HLO verifier rejects Dot operations with more than one
+// contracting dimension (even though we can represent these in the
+// DotDimensionNumbers proto) so there is no need to generalize this to support
+// multiple contracting dimensions.
+inline ::testing::Matcher<const ::xla::HloInstruction*> Dot(
+    ::testing::Matcher<const HloInstruction*> lhs_matcher,
+    ::testing::Matcher<const HloInstruction*> rhs_matcher,
+    int64 lhs_contracting_dim, int64 rhs_contracting_dim) {
+  return ::testing::MakeMatcher(
+      new ::xla::testing::HloDotWithContractDimsMatcher(
+          lhs_matcher, rhs_matcher, lhs_contracting_dim, rhs_contracting_dim));
+}
+
 #undef HLO_MATCHER
 }  // namespace opcode_matchers
 
diff --git a/tensorflow/compiler/xla/service/hlo_matchers_test.cc b/tensorflow/compiler/xla/service/hlo_matchers_test.cc
index f2463060b7..016cc01e33 100644
--- a/tensorflow/compiler/xla/service/hlo_matchers_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_matchers_test.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/service/hlo_matchers.h"
 #include "tensorflow/compiler/xla/shape_util.h"
+#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h"
 
 namespace op = xla::testing::opcode_matchers;
 using ::testing::_;
@@ -165,5 +166,41 @@ TEST(HloMatchersTest, ShardingMatcher) {
               "has incorrect sharding (expected: {maximal device=0})");
 }
 
+TEST(HloMatchersTest, DotMatcher) {
+  string hlo_string = R"(
+HloModule DotOperationFusion_TransposeFusion
+
+ENTRY DotOperationFusion_TransposeFusion {
+  arg0 = f32[1,256] parameter(0)
+  arg1 = f32[256,1024] parameter(1)
+  ROOT dot = f32[1,1024] dot(arg0, arg1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_string));
+  HloInstruction* root = module->entry_computation()->root_instruction();
+
+  EXPECT_THAT(root, op::Dot(op::Parameter(0), op::Parameter(1),
+                            /*lhs_contracting_dim=*/1,
+                            /*rhs_contracting_dim=*/0));
+
+  EXPECT_THAT(
+      Explain(root, op::Dot(op::Parameter(0), op::Parameter(1),
+                            /*lhs_contracting_dim=*/0,
+                            /*rhs_contracting_dim=*/0)),
+      "%dot = f32[1,1024]{1,0} dot(f32[1,256]{1,0} %arg0, f32[256,1024]{1,0} "
+      "%arg1), lhs_contracting_dims={1}, rhs_contracting_dims={0} has wrong "
+      "lhs_contracting_dimensions (got {1} want {0})");
+
+  EXPECT_THAT(
+      Explain(root, op::Dot(op::Parameter(0), op::Parameter(1),
+                            /*lhs_contracting_dim=*/1,
+                            /*rhs_contracting_dim=*/1)),
+      "%dot = f32[1,1024]{1,0} dot(f32[1,256]{1,0} %arg0, f32[256,1024]{1,0} "
+      "%arg1), lhs_contracting_dims={1}, rhs_contracting_dims={0} has wrong "
+      "rhs_contracting_dimensions (got {0} want {1})");
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/service/transpose_folding.cc b/tensorflow/compiler/xla/service/transpose_folding.cc
index 3efd38ce0d..f7a5512fec 100644
--- a/tensorflow/compiler/xla/service/transpose_folding.cc
+++ b/tensorflow/compiler/xla/service/transpose_folding.cc
@@ -35,7 +35,8 @@ TransposeFolding::OperandIndices CanFoldOperandsIntoDot(
     const HloInstruction& dot,
     const TransposeFolding::TransposableGemmOperandsFn&
         transposable_gemm_operands) {
-  if (HloOpcode::kDot != dot.opcode()) {
+  if (HloOpcode::kDot != dot.opcode() ||
+      dot.dot_dimension_numbers().lhs_batch_dimensions_size() != 0) {
     return {};
   }
 
@@ -44,6 +45,8 @@ TransposeFolding::OperandIndices CanFoldOperandsIntoDot(
     auto& operand = *dot.operand(i);
     if (operand.IsRank2Transpose()) {
       operand_set.push_back(i);
+    } else if (ShapeUtil::Rank(operand.shape()) != 2) {
+      return {};
     }
   }
 
@@ -74,23 +77,39 @@ using InstructionOperandsPair =
 
 // Folds the operands of `dot` that are foldable transposes. `computation` is
 // the parent HLO computation of `dot`.
-//
-// Returns whether the module is changed.
-bool FoldTransposeIntoDot(InstructionOperandsPair pair) {
-  auto* dot = pair.first;
-  std::vector<HloInstruction*> instructions_to_fuse(1, dot);
-  for (const int64 operand_index : pair.second) {
-    instructions_to_fuse.push_back(dot->mutable_operand(operand_index));
-  }
-
-  // Early-exit if no operands are foldable.
-  if (instructions_to_fuse.size() == 1) {
-    return false;
+Status FoldTransposeIntoDot(InstructionOperandsPair pair) {
+  HloInstruction* dot = pair.first;
+
+  DotDimensionNumbers new_dim_numbers = dot->dot_dimension_numbers();
+  HloInstruction* new_lhs = dot->mutable_operand(0);
+  HloInstruction* new_rhs = dot->mutable_operand(1);
+
+  CHECK_EQ(new_dim_numbers.lhs_batch_dimensions_size(), 0);
+  CHECK_EQ(new_dim_numbers.rhs_batch_dimensions_size(), 0);
+  CHECK_EQ(new_dim_numbers.lhs_contracting_dimensions_size(), 1);
+  CHECK_EQ(new_dim_numbers.rhs_contracting_dimensions_size(), 1);
+
+  for (int64 operand_index : pair.second) {
+    // We've checked that there aren't any batch dimensions and that the inputs
+    // are rank 2, and shape inference guarantees that there is exactly one
+    // contracting dimension.
+    if (operand_index == 0) {
+      CHECK_EQ(new_lhs->opcode(), HloOpcode::kTranspose);
+      new_dim_numbers.set_lhs_contracting_dimensions(
+          0, 1 - new_dim_numbers.lhs_contracting_dimensions(0));
+      new_lhs = new_lhs->mutable_operand(0);
+    } else {
+      CHECK_EQ(operand_index, 1);
+      CHECK_EQ(new_rhs->opcode(), HloOpcode::kTranspose);
+      new_dim_numbers.set_rhs_contracting_dimensions(
+          0, 1 - new_dim_numbers.rhs_contracting_dimensions(0));
+      new_rhs = new_rhs->mutable_operand(0);
+    }
   }
 
-  dot->parent()->CreateFusionInstruction(
-      instructions_to_fuse, HloInstruction::FusionKind::kTransposeDot);
-  return true;
+  std::unique_ptr<HloInstruction> new_dot = HloInstruction::CreateDot(
+      dot->shape(), new_lhs, new_rhs, new_dim_numbers);
+  return dot->parent()->ReplaceWithNewInstruction(dot, std::move(new_dot));
 }
 
 // Folds the operands of `convolution` that are foldable transposes.
@@ -205,7 +224,8 @@ StatusOr<bool> TransposeFolding::Run(HloModule* module) {
 
   bool changed = false;
   for (InstructionOperandsPair& pair : foldable_dots) {
-    changed |= FoldTransposeIntoDot(pair);
+    TF_RETURN_IF_ERROR(FoldTransposeIntoDot(pair));
+    changed = true;
   }
   for (InstructionOperandsPair& pair : foldable_convolutions) {
     changed |= FoldTransposeIntoConvolution(pair);
diff --git a/tensorflow/compiler/xla/service/transpose_folding_test.cc b/tensorflow/compiler/xla/service/transpose_folding_test.cc
index 0319109f7f..f73f1227aa 100644
--- a/tensorflow/compiler/xla/service/transpose_folding_test.cc
+++ b/tensorflow/compiler/xla/service/transpose_folding_test.cc
@@ -24,6 +24,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h"
 #include "tensorflow/compiler/xla/service/hlo_computation.h"
 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
+#include "tensorflow/compiler/xla/service/hlo_matchers.h"
 #include "tensorflow/compiler/xla/service/hlo_module.h"
 #include "tensorflow/compiler/xla/service/hlo_opcode.h"
 #include "tensorflow/compiler/xla/service/shape_inference.h"
@@ -31,9 +32,12 @@ limitations under the License.
 #include "tensorflow/compiler/xla/test.h"
 #include "tensorflow/compiler/xla/test_helpers.h"
 #include "tensorflow/compiler/xla/tests/hlo_test_base.h"
+#include "tensorflow/compiler/xla/tools/parser/hlo_parser.h"
 #include "tensorflow/compiler/xla/xla_data.pb.h"
 #include "tensorflow/core/platform/logging.h"
 
+namespace op = xla::testing::opcode_matchers;
+
 namespace xla {
 namespace {
 
@@ -54,83 +58,102 @@ class TransposeFoldingTest : public HloTestBase {
 };
 
 TEST_F(TransposeFoldingTest, FoldDotTranspose) {
-  auto builder = HloComputation::Builder("entry_computation");
-  HloInstruction* x = builder.AddInstruction(HloInstruction::CreateParameter(
-      /*parameter_number=*/0, ShapeUtil::MakeShape(F32, {2, 3}),
-      /*name=*/"x"));
-  HloInstruction* y = builder.AddInstruction(HloInstruction::CreateParameter(
-      /*parameter_number=*/1, ShapeUtil::MakeShape(F32, {2, 3}),
-      /*name=*/"y"));
-  HloInstruction* transpose_y =
-      builder.AddInstruction(HloInstruction::CreateTranspose(
-          ShapeUtil::MakeShape(F32, {3, 2}), y, {1, 0}));
-  DotDimensionNumbers dot_dnums;
-  dot_dnums.add_lhs_contracting_dimensions(1);
-  dot_dnums.add_rhs_contracting_dimensions(0);
-  HloInstruction* dot = builder.AddInstruction(
-      HloInstruction::CreateDot(ShapeUtil::MakeShape(F32, {2, 2}), /*lhs=*/x,
-                                /*rhs=*/transpose_y, dot_dnums));
+  string hlo_string = R"(
+HloModule FoldDotTranspose
+
+ENTRY entry_computation {
+  x = f32[2,3]{1,0} parameter(0)
+  y = f32[2,3]{1,0} parameter(1)
+  transpose = f32[3,2]{1,0} transpose(y), dimensions={1,0}
+  ROOT dot = f32[2,2]{1,0} dot(x, transpose), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+}
+)";
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_string));
 
-  auto module = CreateNewModule("test_module");
-  HloComputation* entry_computation =
-      module->AddEntryComputation(builder.Build(dot));
   FoldTranspose(module.get());
 
-  // Instructions after folding: x, y, and the fusion.
-  std::unordered_set<HloInstruction*> instruction_set(
-      entry_computation->instructions().begin(),
-      entry_computation->instructions().end());
-  CHECK_EQ(1, instruction_set.erase(x)) << "x is not in entry_computation.";
-  CHECK_EQ(1, instruction_set.erase(y)) << "y is not in entry_computation.";
-  CHECK_EQ(1, instruction_set.size())
-      << "entry_computation should contain exactly 3 instructions.";
-  HloInstruction* fusion = *instruction_set.begin();
-  EXPECT_EQ(HloOpcode::kFusion, fusion->opcode());
+  EXPECT_THAT(module->entry_computation()->root_instruction(),
+              op::Dot(op::Parameter(0), op::Parameter(1),
+                      /*lhs_contracting_dim=*/1, /*rhs_contracting_dim=*/1));
+}
+
+TEST_F(TransposeFoldingTest, DontFoldTransposeOfBatchDim) {
+  string hlo_string = R"(
+HloModule FoldDotTranspose
 
-  // The fusion instruction should contain two parameters, one transpose and
-  // one dot.
-  EXPECT_EQ(4, fusion->fused_instruction_count());
+ENTRY entry_computation {
+  x = f32[2,3] parameter(0)
+  y = f32[3,2] parameter(1)
+  transpose = f32[2,3] transpose(y), dimensions={1,0}
+  ROOT dot = f32[2] dot(x, transpose), lhs_batch_dims={0}, rhs_batch_dims={0}, lhs_contracting_dims={1}, rhs_contracting_dims={1}
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_string));
+
+  TransposeFolding transpose_folding(
+      [](const HloInstruction& dot,
+         const TransposeFolding::OperandIndices& candidate_operands) {
+        return candidate_operands;
+      },
+      [](const HloInstruction& convolution,
+         const TransposeFolding::OperandIndices& candidate_operands) {
+        return candidate_operands;
+      });
+  TF_ASSERT_OK_AND_ASSIGN(bool changed, transpose_folding.Run(module.get()));
+  EXPECT_FALSE(changed);
+}
+
+TEST_F(TransposeFoldingTest, DontFoldTransposeOfRank1Dot) {
+  string hlo_string = R"(
+HloModule FoldDotTranspose
+
+ENTRY entry_computation {
+  x = f32[3] parameter(0)
+  y = f32[3,2] parameter(1)
+  transpose = f32[2,3] transpose(y), dimensions={1,0}
+  ROOT dot = f32[2] dot(x, transpose), lhs_batch_dims={}, rhs_batch_dims={0}, lhs_contracting_dims={0}, rhs_contracting_dims={1}
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_string));
+
+  TransposeFolding transpose_folding(
+      [](const HloInstruction& dot,
+         const TransposeFolding::OperandIndices& candidate_operands) {
+        return candidate_operands;
+      },
+      [](const HloInstruction& convolution,
+         const TransposeFolding::OperandIndices& candidate_operands) {
+        return candidate_operands;
+      });
+  TF_ASSERT_OK_AND_ASSIGN(bool changed, transpose_folding.Run(module.get()));
+  EXPECT_FALSE(changed);
 }
 
 TEST_F(TransposeFoldingTest, FoldDotTransposeConstant) {
-  auto builder = HloComputation::Builder("entry_computation");
-  // 2x1
-  HloInstruction* const0 = builder.AddInstruction(
-      HloInstruction::CreateConstant(Literal::CreateR2<float>({{1}, {2}})));
-  // 3x2
-  HloInstruction* const1 =
-      builder.AddInstruction(HloInstruction::CreateConstant(
-          Literal::CreateR2<float>({{1, 2}, {3, 4}, {5, 6}})));
-  HloInstruction* transpose0 =
-      builder.AddInstruction(HloInstruction::CreateTranspose(
-          ShapeUtil::MakeShape(F32, {1, 2}), const0, {1, 0}));
-  HloInstruction* transpose1 =
-      builder.AddInstruction(HloInstruction::CreateTranspose(
-          ShapeUtil::MakeShape(F32, {2, 3}), const1, {1, 0}));
-  DotDimensionNumbers dot_dnums;
-  dot_dnums.add_lhs_contracting_dimensions(1);
-  dot_dnums.add_rhs_contracting_dimensions(0);
-  HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateDot(
-      ShapeUtil::MakeShape(F32, {1, 3}),
-      /*lhs=*/transpose0, /*rhs=*/transpose1, dot_dnums));
+  string hlo_string = R"(
+HloModule FoldDotTransposeConstant
+
+ENTRY entry_computation {
+  constant = f32[2,1]{1,0} constant(f32[2,1] { { 1 }, { 2 } })
+  transpose = f32[1,2]{1,0} transpose(constant), dimensions={1,0}
+  constant.1 = f32[3,2]{1,0} constant(f32[3,2] { { 1, 2 }, { 3, 4 }, { 5, 6 } })
+  transpose.1 = f32[2,3]{1,0} transpose(constant.1), dimensions={1,0}
+  ROOT dot = f32[1,3]{1,0} dot(transpose, transpose.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+}
+)";
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_string));
 
-  auto module = CreateNewModule("test_module");
-  HloComputation* entry_computation =
-      module->AddEntryComputation(builder.Build(dot));
   FoldTranspose(module.get());
 
-  for (auto* instruction : entry_computation->instructions()) {
-    if (instruction->opcode() == HloOpcode::kFusion) {
-      CHECK_EQ(2, instruction->operand_count());
-      EXPECT_EQ(const0, instruction->operand(0));
-      EXPECT_EQ(const1, instruction->operand(1));
-    }
-  }
-
-  // The created fusion instruction should contain two parameters, two
-  // transposes (one for each parameter) and one dot.
-  EXPECT_EQ(5,
-            entry_computation->root_instruction()->fused_instruction_count());
+  EXPECT_THAT(module->entry_computation()->root_instruction(),
+              op::Dot(op::Constant(), op::Constant(),
+                      /*lhs_contracting_dim=*/0, /*rhs_contracting_dim=*/1));
 }
 
 TEST_F(TransposeFoldingTest, FuseDotWithConstantOperands) {
@@ -164,50 +187,32 @@ TEST_F(TransposeFoldingTest, FuseDotWithConstantOperands) {
   EXPECT_EQ(6, callee_computation->instruction_count());
 }
 
-TEST_F(TransposeFoldingTest, FoldDotTransposeInWhile) {
-  auto builder = HloComputation::Builder("entry_computation");
-  HloInstruction* x = builder.AddInstruction(HloInstruction::CreateParameter(
-      /*parameter_number=*/0, ShapeUtil::MakeShape(F32, {2, 3}),
-      /*name=*/"x"));
-  HloInstruction* y = builder.AddInstruction(HloInstruction::CreateParameter(
-      /*parameter_number=*/1, ShapeUtil::MakeShape(F32, {2, 3}),
-      /*name=*/"y"));
-  HloInstruction* transpose_y =
-      builder.AddInstruction(HloInstruction::CreateTranspose(
-          ShapeUtil::MakeShape(F32, {3, 2}), y, {1, 0}));
-  DotDimensionNumbers dot_dnums;
-  dot_dnums.add_lhs_contracting_dimensions(1);
-  dot_dnums.add_rhs_contracting_dimensions(0);
-  HloInstruction* dot = builder.AddInstruction(
-      HloInstruction::CreateDot(ShapeUtil::MakeShape(F32, {2, 2}), /*lhs=*/x,
-                                /*rhs=*/transpose_y, dot_dnums));
-
-  auto module = CreateNewModule("test_module");
-  HloComputation* entry_computation =
-      module->AddEntryComputation(builder.Build(dot));
+TEST_F(TransposeFoldingTest, FoldDotTransposeInCall) {
+  string hlo_string = R"(
+HloModule FoldDotTransposeInCall
 
-  HloInstruction* call = module->OutlineExpressionFromComputation(
-      {transpose_y, dot}, "outlined", entry_computation);
+callee {
+  name.0 = f32[2,3]{1,0} parameter(0)
+  name.1 = f32[2,3]{1,0} parameter(1)
+  transpose.clone = f32[3,2]{1,0} transpose(name.0), dimensions={1,0}
+  ROOT dot.clone = f32[2,2]{1,0} dot(name.1, transpose.clone), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+}
 
+ENTRY entry_computation {
+  y = f32[2,3]{1,0} parameter(1)
+  x = f32[2,3]{1,0} parameter(0)
+  ROOT call = f32[2,2]{1,0} call(y, x), to_apply=callee
+}
+)";
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          tools::Parse(hlo_string));
   FoldTranspose(module.get());
 
-  // Instructions after folding: x, y, and the fusion.
-  std::unordered_set<HloInstruction*> instruction_set(
-      entry_computation->instructions().begin(),
-      entry_computation->instructions().end());
-  CHECK_EQ(1, instruction_set.erase(x)) << "x is not in entry_computation.";
-  CHECK_EQ(1, instruction_set.erase(y)) << "y is not in entry_computation.";
-  CHECK_EQ(1, instruction_set.erase(call))
-      << "call is not in entry_computation.";
-  CHECK(instruction_set.empty())
-      << "entry_computation should contain exactly 3 instructions.";
-  HloInstruction* fusion =
-      call->called_computations().front()->root_instruction();
-  EXPECT_EQ(HloOpcode::kFusion, fusion->opcode());
-
-  // The fusion instruction should contain two parameters, one transpose and
-  // one dot.
-  EXPECT_EQ(4, fusion->fused_instruction_count());
+  const HloComputation* callee = module->GetComputationWithName("callee");
+  ASSERT_NE(callee, nullptr);
+  EXPECT_THAT(callee->root_instruction(),
+              op::Dot(op::Parameter(1), op::Parameter(0),
+                      /*lhs_contracting_dim=*/1, /*rhs_contracting_dim=*/1));
 }
 
 // Test that a two dimension swap of the kernel gets folded into convolution.
-- 
GitLab


From 067be1aa75d9065b6cc57ba6316fc17544a9fdf1 Mon Sep 17 00:00:00 2001
From: Mingsheng Hong <hongm@google.com>
Date: Sat, 5 May 2018 07:10:58 -0700
Subject: [PATCH 289/395] Part 2 of Swift<->TF sends/recvs: support receiving
 tensors in TF from Swift via direct session.

The changes are:

1. Added a TF experimental C API for Swift host to enqueue a tensor for sending
  to TF. Again, the C APIs can be removed once the Fifo-queue based design
  proves stable later.

2. TFLowerGraph is extended to generate Fifo related nodes for TF to receive
  tensors. This is similar to the extension for TF to send tensors.

3. TFPartition is extended to support host send (createHostSend()), which does
  the tensor send via a new protocol method TensorSendableReceivable.sendToDevice().

The main complexity is in sending a scalar, where a new protocol method
TensorSendableReceivable.createScalarTensor() is called to first create a tensor
out of it, and then send it over to TF.

Also removed code for protocol conformance on AccelerableByTensorFlow. Instead
have compiler look up that conformance from the SILFunction on sending/receiving
tensors.

AccelerableByTensorFlow could be removed from the compiler-known protocol list
now, but we'll defer that till things can stabilized more (in the past this
protocol has been added to and removed from the list at different times).

PiperOrigin-RevId: 195539436
---
 tensorflow/c/c_api_experimental.cc | 48 ++++++++++++++++++++++++++++++
 tensorflow/c/c_api_experimental.h  | 23 ++++++++++++--
 2 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/tensorflow/c/c_api_experimental.cc b/tensorflow/c/c_api_experimental.cc
index 82dbd3cdbc..95b04f9058 100644
--- a/tensorflow/c/c_api_experimental.cc
+++ b/tensorflow/c/c_api_experimental.cc
@@ -8407,3 +8407,51 @@ TF_Tensor* TF_DequeueNamedTensor(TF_Session* session, int tensor_id,
   }
   return ret;
 }
+
+void TF_EnqueueNamedTensor(TF_Session* session, int tensor_id,
+                           TF_Tensor* tensor, TF_Status* status) {
+  assert(session);
+  {
+    tensorflow::mutex_lock c(session->graph->mu);
+    if (VLOG_IS_ON(1)) {
+      VLOG(1) << "Enqueuing named tensor with id " << tensor_id
+              << ", with input graph: "
+              << session->graph->graph.ToGraphDefDebug().DebugString();
+      tensorflow::Tensor internal_tensor;
+      if (tensorflow::TF_TensorToTensor(tensor, &internal_tensor).ok()) {
+        VLOG(1) << "Enqueu'ing tensor content: "
+                << internal_tensor.DebugString();
+      }
+    }
+  }
+
+  TF_Operation* enqueue_op = TF_GraphOperationByName(
+      session->graph,
+      tensorflow::strings::StrCat("fifo_queue_enqueue_", tensor_id).c_str());
+  if (enqueue_op == nullptr) {
+    status->status = tensorflow::errors::Internal(
+        "Unable to find the enqueue node in the TF graph.");
+    return;
+  }
+
+  TF_Operation* placeholder_op = TF_GraphOperationByName(
+      session->graph,
+      tensorflow::strings::StrCat("arg_tensor_enqueue_", tensor_id).c_str());
+  if (placeholder_op == nullptr) {
+    status->status = tensorflow::errors::Internal(
+        "Unable to find the placeholder node as input to enqueue in the TF "
+        "graph.");
+    return;
+  }
+
+  VLOG(1) << "Running the enqueue op";
+  TF_Output input{placeholder_op, 0};
+  TF_SessionRun(session, /*run_options*/ nullptr,
+                // input related parameters
+                /*inputs*/ &input, /*input_values*/ &tensor, /*ninputs*/ 1,
+                // output related parameters
+                /*outputs*/ nullptr, /*output_values*/ nullptr, /*noutputs*/ 0,
+                /*targets*/ &enqueue_op, /*ntargets*/ 1,
+                /*run_metadata*/ nullptr, status);
+  VLOG(1) << "Enqueuing is done.";
+}
diff --git a/tensorflow/c/c_api_experimental.h b/tensorflow/c/c_api_experimental.h
index e6757c065f..20bdace40f 100644
--- a/tensorflow/c/c_api_experimental.h
+++ b/tensorflow/c/c_api_experimental.h
@@ -87,8 +87,11 @@ TF_CAPI_EXPORT extern TF_Operation* TF_MakeFileBasedIteratorGetNextWithDatasets(
     unsigned char is_mnist, TF_Status* status);
 
 // On success, dequeues a tensor from a TF-managed FifoQueue given by
-// `tensor_id`, associated with `session`. Caller must call TF_DeleteTensor()
-// over the returned tensor. If the queue is empty, this call is blocked.
+// `tensor_id`, associated with `session`. There must be a graph node named
+// "fifo_queue_dequeue_<tensor_id>", to be executed by this API call.
+
+// Caller must call TF_DeleteTensor() over the returned tensor. If the queue is
+// empty, this call is blocked.
 //
 // Tensors are enqueued via the corresponding TF enqueue op.
 // TODO(hongm): Add support for `timeout_ms`.
@@ -96,6 +99,22 @@ TF_CAPI_EXPORT extern TF_Tensor* TF_DequeueNamedTensor(TF_Session* session,
                                                        int tensor_id,
                                                        TF_Status* status);
 
+// On success, enqueues `tensor` into a TF-managed FifoQueue given by
+// `tensor_id`, associated with `session`. There must be a graph node named
+// "fifo_queue_enqueue_<tensor_id>", to be executed by this API call. It reads
+// from a placeholder node "arg_tensor_enqueue_<tensor_id>".
+//
+// `tensor` is still owned by the caller. This call will be blocked if the queue
+// has reached its capacity, and will be unblocked when the queued tensors again
+// drop below the capacity due to dequeuing.
+//
+// Tensors are dequeued via the corresponding TF dequeue op.
+// TODO(hongm): Add support for `timeout_ms`.
+TF_CAPI_EXPORT extern void TF_EnqueueNamedTensor(TF_Session* session,
+                                                 int tensor_id,
+                                                 TF_Tensor* tensor,
+                                                 TF_Status* status);
+
 #ifdef __cplusplus
 } /* end extern "C" */
 #endif
-- 
GitLab


From f3c21911bca9c1ef01560dfd7609020d7f85f52b Mon Sep 17 00:00:00 2001
From: Justin Lebar <jlebar@google.com>
Date: Sat, 5 May 2018 11:24:06 -0700
Subject: [PATCH 290/395] [XLA] Always be willing to duplicate widening
 kConvert instructions during fusion.

This has the effect of pushing widening kConvert HLOs into consumers.
This is what we want, because it means that the producer writes the
narrower type (e.g. f16) and the consumer reads it and internally
upcasts to the wider type (e.g. f32).  This lets the producer and
consumer both run faster, because they have to touch less memory.

PiperOrigin-RevId: 195546910
---
 .../xla/service/instruction_fusion.cc         | 23 +++++++++++++++-
 .../xla/service/instruction_fusion_test.cc    | 27 +++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/tensorflow/compiler/xla/service/instruction_fusion.cc b/tensorflow/compiler/xla/service/instruction_fusion.cc
index dc1a39e9fa..6bb2ca19fe 100644
--- a/tensorflow/compiler/xla/service/instruction_fusion.cc
+++ b/tensorflow/compiler/xla/service/instruction_fusion.cc
@@ -28,6 +28,25 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 
 namespace xla {
+namespace {
+// These nodes can always be duplicated into consumers, even if
+// InstructionFusion::may_duplicate_ is false.
+//
+// In general these should be nodes that get *cheaper* the more they're
+// duplicated (and fused into consumers).
+//
+// TODO(jlebar): Duplicating instructions when we have a variable called "may
+// duplicate" that's equal to false is not pretty.
+bool IsAlwaysDuplicable(const HloInstruction& instruction) {
+  // We are always willing to duplicate a widening type-conversion instruction
+  // if it means we can fuse the convert into a consumer.  This allows the
+  // consumer to read less memory, which is almost always a performance win.
+  return instruction.opcode() == HloOpcode::kConvert &&
+         ShapeUtil::ByteSizeOf(instruction.operand(0)->shape()) <
+             ShapeUtil::ByteSizeOf(instruction.shape());
+}
+}  // namespace
+
 /*static*/ bool InstructionFusion::IsExpensive(
     const HloInstruction& instruction) {
   switch (instruction.opcode()) {
@@ -418,9 +437,11 @@ HloInstruction* InstructionFusion::Fuse(HloInstruction* producer,
 bool InstructionFusion::ShouldFuse(HloInstruction* consumer,
                                    int64 operand_index) {
   HloInstruction* producer = consumer->mutable_operand(operand_index);
+
   // Cost condition: don't duplicate expensive instructions.
   if (FusionWouldDuplicate(*producer, *consumer) &&
-      (is_expensive_(*producer) || !may_duplicate_)) {
+      (!may_duplicate_ || is_expensive_(*producer)) &&
+      !IsAlwaysDuplicable(*producer)) {
     return false;
   }
 
diff --git a/tensorflow/compiler/xla/service/instruction_fusion_test.cc b/tensorflow/compiler/xla/service/instruction_fusion_test.cc
index e78b99a80c..b4b1955fe2 100644
--- a/tensorflow/compiler/xla/service/instruction_fusion_test.cc
+++ b/tensorflow/compiler/xla/service/instruction_fusion_test.cc
@@ -21,6 +21,8 @@ limitations under the License.
 
 namespace xla {
 
+namespace op = xla::testing::opcode_matchers;
+
 using InstructionFusionTest = HloTestBase;
 
 TEST_F(InstructionFusionTest, PotentialBitcastReshapeOfParameterUnfused) {
@@ -291,4 +293,29 @@ TEST_F(InstructionFusionTest, AllowEffectiveUnaryDuplication) {
           .ValueOrDie());
 }
 
+TEST_F(InstructionFusionTest,
+       WideningConvertsAreAlwaysDuplicableIntoConsumers) {
+  auto module = tools::Parse(R"(
+  HloModule test_module
+  ENTRY Test {
+    p0 = f16[100] parameter(0)
+    c = f32[100] convert(p0)
+    add = f32[100] add(c, c)
+    ROOT mul = f32[100] multiply(c, c)
+  })")
+                    .ValueOrDie();
+
+  // The convert should be fused into the add and mul, even though may_duplicate
+  // is false, because it's always beneficial to fuse/duplicate widening
+  // converts into consumers.
+  EXPECT_TRUE(
+      InstructionFusion(InstructionFusion::IsExpensive, /*may_duplicate=*/false)
+          .Run(module.get())
+          .ValueOrDie())
+      << module->ToString();
+
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root, op::Fusion(op::Parameter()));
+}
+
 }  // namespace xla
-- 
GitLab


From 62ed0aa37099e07720880a72a285304d34512cba Mon Sep 17 00:00:00 2001
From: Shashi Shekhar <shashishekhar@google.com>
Date: Sat, 5 May 2018 11:55:53 -0700
Subject: [PATCH 291/395] Allow benchmark model graph to be specified in text
 proto format.

PiperOrigin-RevId: 195547670
---
 tensorflow/tools/benchmark/benchmark_model.cc |  4 ++
 .../tools/benchmark/benchmark_model_test.cc   | 55 +++++++++++++++----
 2 files changed, 47 insertions(+), 12 deletions(-)

diff --git a/tensorflow/tools/benchmark/benchmark_model.cc b/tensorflow/tools/benchmark/benchmark_model.cc
index 15523028c7..eeb1fab40c 100644
--- a/tensorflow/tools/benchmark/benchmark_model.cc
+++ b/tensorflow/tools/benchmark/benchmark_model.cc
@@ -261,6 +261,10 @@ Status InitializeSession(int num_threads, const string& graph,
   graph_def->reset(new GraphDef());
   tensorflow::GraphDef tensorflow_graph;
   Status s = ReadBinaryProto(Env::Default(), graph, graph_def->get());
+  if (!s.ok()) {
+    s = ReadTextProto(Env::Default(), graph, graph_def->get());
+  }
+
   if (!s.ok()) {
     LOG(ERROR) << "Could not create TensorFlow Graph: " << s;
     return s;
diff --git a/tensorflow/tools/benchmark/benchmark_model_test.cc b/tensorflow/tools/benchmark/benchmark_model_test.cc
index 16ab2ff66e..6813045d63 100644
--- a/tensorflow/tools/benchmark/benchmark_model_test.cc
+++ b/tensorflow/tools/benchmark/benchmark_model_test.cc
@@ -26,30 +26,36 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
-TEST(BenchmarkModelTest, InitializeAndRun) {
-  const string dir = testing::TmpDir();
-  const string filename_pb = io::JoinPath(dir, "graphdef.pb");
-
+void CreateTestGraph(const ::tensorflow::Scope& root,
+                     benchmark_model::InputLayerInfo* input,
+                     string* output_name, GraphDef* graph_def) {
   // Create a simple graph and write it to filename_pb.
   const int input_width = 400;
   const int input_height = 10;
-  benchmark_model::InputLayerInfo input;
-  input.shape = TensorShape({input_width, input_height});
-  input.data_type = DT_FLOAT;
+  input->shape = TensorShape({input_width, input_height});
+  input->data_type = DT_FLOAT;
   const TensorShape constant_shape({input_height, input_width});
 
   Tensor constant_tensor(DT_FLOAT, constant_shape);
   test::FillFn<float>(&constant_tensor, [](int) -> float { return 3.0; });
 
-  auto root = Scope::NewRootScope().ExitOnError();
   auto placeholder =
-      ops::Placeholder(root, DT_FLOAT, ops::Placeholder::Shape(input.shape));
-  input.name = placeholder.node()->name();
+      ops::Placeholder(root, DT_FLOAT, ops::Placeholder::Shape(input->shape));
+  input->name = placeholder.node()->name();
   auto m = ops::MatMul(root, placeholder, constant_tensor);
-  const string output_name = m.node()->name();
+  *output_name = m.node()->name();
+  TF_ASSERT_OK(root.ToGraphDef(graph_def));
+}
+
+TEST(BenchmarkModelTest, InitializeAndRun) {
+  const string dir = testing::TmpDir();
+  const string filename_pb = io::JoinPath(dir, "graphdef.pb");
+  auto root = Scope::NewRootScope().ExitOnError();
 
+  benchmark_model::InputLayerInfo input;
+  string output_name;
   GraphDef graph_def;
-  TF_ASSERT_OK(root.ToGraphDef(&graph_def));
+  CreateTestGraph(root, &input, &output_name, &graph_def);
   string graph_def_serialized;
   graph_def.SerializeToString(&graph_def_serialized);
   TF_ASSERT_OK(
@@ -69,5 +75,30 @@ TEST(BenchmarkModelTest, InitializeAndRun) {
   ASSERT_EQ(num_runs, 10);
 }
 
+TEST(BenchmarkModeTest, TextProto) {
+  const string dir = testing::TmpDir();
+  const string filename_txt = io::JoinPath(dir, "graphdef.pb.txt");
+  auto root = Scope::NewRootScope().ExitOnError();
+
+  benchmark_model::InputLayerInfo input;
+  string output_name;
+  GraphDef graph_def;
+  CreateTestGraph(root, &input, &output_name, &graph_def);
+  TF_ASSERT_OK(WriteTextProto(Env::Default(), filename_txt, graph_def));
+
+  std::unique_ptr<Session> session;
+  std::unique_ptr<GraphDef> loaded_graph_def;
+  TF_ASSERT_OK(benchmark_model::InitializeSession(1, filename_txt, &session,
+                                                  &loaded_graph_def));
+  std::unique_ptr<StatSummarizer> stats;
+  stats.reset(new tensorflow::StatSummarizer(*(loaded_graph_def.get())));
+  int64 time;
+  int64 num_runs = 0;
+  TF_ASSERT_OK(benchmark_model::TimeMultipleRuns(
+      0.0, 10, 0.0, {input}, {output_name}, {}, session.get(), stats.get(),
+      &time, &num_runs));
+  ASSERT_EQ(num_runs, 10);
+}
+
 }  // namespace
 }  // namespace tensorflow
-- 
GitLab


From a5e1809ddedbff8d9f66b82f9cf0989976694050 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar@google.com>
Date: Sat, 5 May 2018 12:02:32 -0700
Subject: [PATCH 292/395] [TPU] Add option to only compile a replicated graph.

Useful when wanting to compile a computation but not run it. Returns a serialized CompilationResult string with the error message.

PiperOrigin-RevId: 195547847
---
 tensorflow/contrib/tpu/BUILD                  |  1 +
 tensorflow/contrib/tpu/ops/replication_ops.cc |  4 ++
 tensorflow/contrib/tpu/proto/BUILD            | 10 +++
 .../tpu/proto/compilation_result.proto        | 13 ++++
 tensorflow/contrib/tpu/python/tpu/tpu.py      | 70 ++++++++++++++++---
 5 files changed, 88 insertions(+), 10 deletions(-)
 create mode 100644 tensorflow/contrib/tpu/proto/compilation_result.proto

diff --git a/tensorflow/contrib/tpu/BUILD b/tensorflow/contrib/tpu/BUILD
index 0bdf6f64c9..f84ff1bfe9 100644
--- a/tensorflow/contrib/tpu/BUILD
+++ b/tensorflow/contrib/tpu/BUILD
@@ -181,6 +181,7 @@ py_library(
         ":datasets",
         ":profiler",
         ":tpu_py",
+        "//tensorflow/contrib/tpu/proto:compilation_result_proto_py",
         "//tensorflow/contrib/tpu/proto:topology_proto_py",
         "//tensorflow/core:protos_all_py",
         "//tensorflow/python:array_ops",
diff --git a/tensorflow/contrib/tpu/ops/replication_ops.cc b/tensorflow/contrib/tpu/ops/replication_ops.cc
index 3bdf7c2f83..defed00537 100644
--- a/tensorflow/contrib/tpu/ops/replication_ops.cc
+++ b/tensorflow/contrib/tpu/ops/replication_ops.cc
@@ -64,6 +64,10 @@ REGISTER_OP("TPUReplicatedOutput")
         "Operator that connects the output of an N-way replicated TPU "
         "computation to N separate outputs.");
 
+REGISTER_OP("TPUCompilationResult")
+    .Output("output: string")
+    .SetShapeFn(shape_inference::ScalarShape);
+
 REGISTER_OP("TPUReplicate")
     .Attr("computation: func")
     .Attr("num_replicas: int >= 1")
diff --git a/tensorflow/contrib/tpu/proto/BUILD b/tensorflow/contrib/tpu/proto/BUILD
index fcfbbe1a21..7ecb36852c 100644
--- a/tensorflow/contrib/tpu/proto/BUILD
+++ b/tensorflow/contrib/tpu/proto/BUILD
@@ -21,3 +21,13 @@ tf_proto_library(
     cc_api_version = 2,
     visibility = ["//visibility:public"],
 )
+
+tf_proto_library(
+    name = "compilation_result_proto",
+    srcs = [
+        "compilation_result.proto",
+    ],
+    cc_api_version = 2,
+    protodeps = ["//tensorflow/core:protos_all"],
+    visibility = ["//visibility:public"],
+)
diff --git a/tensorflow/contrib/tpu/proto/compilation_result.proto b/tensorflow/contrib/tpu/proto/compilation_result.proto
new file mode 100644
index 0000000000..cf52897de3
--- /dev/null
+++ b/tensorflow/contrib/tpu/proto/compilation_result.proto
@@ -0,0 +1,13 @@
+syntax = "proto3";
+
+option cc_enable_arenas = true;
+package tensorflow.tpu;
+
+import "tensorflow/core/lib/core/error_codes.proto";
+
+// Describes the result of a TPU compilation.
+message CompilationResultProto {
+  // The error message, if any, returned during compilation.
+  error.Code status_code = 1;
+  string status_error_message = 2;
+}
diff --git a/tensorflow/contrib/tpu/python/tpu/tpu.py b/tensorflow/contrib/tpu/python/tpu/tpu.py
index 7b8786304c..c8f24ed01d 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu.py
@@ -58,6 +58,7 @@ _NOT_IMPLEMENTED_OPS = set([
 _MAX_WARNING_LINES = 5
 
 _TPU_REPLICATE_ATTR = "_tpu_replicate"
+_TPU_COMPILATION_STATUS_ATTR = "_tpu_compilation_status"
 _OUTSIDE_COMPILATION_ATTR = "_xla_outside_compilation"
 
 
@@ -385,6 +386,45 @@ def replicate(computation,
     ValueError: If the number of inputs per replica does not match
       the number of formal parameters to `computation`.
   """
+  return split_compile_and_replicate(computation, inputs, infeed_queue,
+                                     device_assignment, name)[1]
+
+
+def split_compile_and_replicate(computation,
+                                inputs=None,
+                                infeed_queue=None,
+                                device_assignment=None,
+                                name=None):
+  """Builds graph operators that runs compilation and replicated computation.
+
+  This is a lower level interface than replicate that returns a separate compile
+  and execute output tensor. In the generated graph the compile op feeds into
+  the execute op and no additional compilation is incurred when running the
+  compile op before the execute op. The compile op returns additional
+  information about the compilation but does not return the compiled program.
+
+  Args:
+    computation: A Python function that builds the computation to replicate.
+    inputs: A list of lists of input tensors or `None` (equivalent to
+      `[[]]`), indexed by `[replica_num][input_num]`. All replicas must
+      have the same number of inputs.
+    infeed_queue: If not `None`, the `InfeedQueue` from which to append a tuple
+      of arguments as inputs to computation.
+    device_assignment: If not `None`, a `DeviceAssignment` describing the
+      mapping between logical cores in the computation with physical cores in
+      the TPU topology. Uses a default device assignment if `None`. The
+      `DeviceAssignment` may be omitted if each replica of the computation uses
+      only one core, and there is either only one replica, or the number of
+      replicas is equal to the number of cores in the TPU system.
+    name: (Deprecated) Does nothing.
+  Returns:
+    A list of lists with the first list corresponding to the compile op and the
+    second a list of output tensors, indexed by `[replica_num][output_num]`.
+  Raises:
+    ValueError: If all replicas do not have equal numbers of input tensors.
+    ValueError: If the number of inputs per replica does not match
+      the number of formal parameters to `computation`.
+  """
   del name
   inputs = [[]] if inputs is None else inputs
 
@@ -456,8 +496,8 @@ def replicate(computation,
     computation_inputs.append(
         tpu_ops.tpu_replicated_input(replicas, name="input{}".format(i)))
 
-  context = TPUReplicateContext(
-      name=graph.unique_name("cluster"), num_replicas=num_replicas)
+  cluster_name = graph.unique_name("cluster")
+  context = TPUReplicateContext(name=cluster_name, num_replicas=num_replicas)
   try:
     context.Enter()
 
@@ -516,8 +556,7 @@ def replicate(computation,
 
     # Separates the returned Operations and Tensors.
     output_operations = [o for o in outputs if isinstance(o, ops.Operation)]
-    output_tensors = [o for o in outputs
-                      if not isinstance(o, ops.Operation)]
+    output_tensors = [o for o in outputs if not isinstance(o, ops.Operation)]
 
     if outputs != output_tensors + output_operations:
       raise ValueError(
@@ -550,22 +589,33 @@ def replicate(computation,
                                            name="output{}".format(i))
              for i in xrange(output_arity)]
 
+  with ops.control_dependencies([metadata]):
+    compile_status = tpu_ops.tpu_compilation_result()
+    op = compile_status.op
+    attr_value = attr_value_pb2.AttrValue(s=compat.as_bytes(cluster_name))
+    op._set_attr(_TPU_COMPILATION_STATUS_ATTR, attr_value)  # pylint: disable=protected-access
+
   with ops.control_dependencies(output_operations):
     if output_arity == 0:
       # Returns a list of NoOps dependent on the replication Op, indexed by
       # [replica_num].
       return [
-          control_flow_ops.no_op(name="shard_%d" % i)
-          for i in range(num_replicas)
+          compile_status, [
+              control_flow_ops.no_op(name="shard_%d" % i)
+              for i in range(num_replicas)
+          ]
       ]
     else:
       # Wraps the outputs in identity operators so the names of any possible
       # `fetch` nodes are preserved by the replication rewrite.
       return [
-          [array_ops.identity(outputs[out][replica],
-                              name="output_%d_shard_%d" % (out, replica))
-           for out in xrange(output_arity)]
-          for replica in xrange(num_replicas)
+          compile_status, [[
+              array_ops.identity(
+                  outputs[out][replica],
+                  name="output_%d_shard_%d" % (out, replica))
+              for out in xrange(output_arity)
+          ]
+                           for replica in xrange(num_replicas)]
       ]
 
 
-- 
GitLab


From 094daf9d005ffe9322602988dbff033c3e359cef Mon Sep 17 00:00:00 2001
From: Justin Lebar <jlebar@google.com>
Date: Sat, 5 May 2018 12:34:32 -0700
Subject: [PATCH 293/395] [XLA:GPU] Zero out input buffers before running cudnn
 conv autotune.

We don't need a corresponding change in gemm_thunk.cc because for gemms,
we do our autotune at runtime, at which point we have some real data in
our input/output buffers.

PiperOrigin-RevId: 195548896
---
 .../gpu/cudnn_convolution_algorithm_picker.cc | 40 ++++++++++++++-----
 1 file changed, 30 insertions(+), 10 deletions(-)

diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.cc b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.cc
index c4c56c5692..41ee45f55f 100644
--- a/tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.cc
+++ b/tensorflow/compiler/xla/service/gpu/cudnn_convolution_algorithm_picker.cc
@@ -197,22 +197,42 @@ CudnnConvolutionAlgorithmPicker::PickBestAlgorithm(
   // We don't put any data in these buffers, because (in theory, anyway) the
   // speed of a conv isn't affected by the data being convolved.
   ScratchAllocator input_output_allocator(device_ordinal, allocator);
-  se::port::StatusOr<DeviceMemoryBase> input_buf =
+  StatusOr<DeviceMemoryBase> maybe_input_buf =
       input_output_allocator.AllocateBytes(&stream,
                                            ShapeUtil::ByteSizeOf(input_shape));
-  se::port::StatusOr<DeviceMemoryBase> filter_buf =
+  StatusOr<DeviceMemoryBase> maybe_filter_buf =
       input_output_allocator.AllocateBytes(&stream,
                                            ShapeUtil::ByteSizeOf(filter_shape));
-  se::port::StatusOr<DeviceMemoryBase> output_buf =
+  StatusOr<DeviceMemoryBase> maybe_output_buf =
       input_output_allocator.AllocateBytes(&stream,
                                            ShapeUtil::ByteSizeOf(output_shape));
-  if (!input_buf.ok() || !filter_buf.ok() || !output_buf.ok()) {
+  if (!maybe_input_buf.ok() || !maybe_filter_buf.ok() ||
+      !maybe_output_buf.ok()) {
     LOG(WARNING)
         << "Couldn't allocate space for input/filter/output of convolution "
         << instr->ToString() << ".  Falling back to default algorithm.";
     return nullopt;
   }
 
+  DeviceMemoryBase input_buf = maybe_input_buf.ValueOrDie();
+  DeviceMemoryBase filter_buf = maybe_filter_buf.ValueOrDie();
+  DeviceMemoryBase output_buf = maybe_output_buf.ValueOrDie();
+
+  // Although we don't have evidence this matters, zero out the buffers before
+  // autotuning.  It's conceivable that using uninitialized memory as the inputs
+  // might affect performance if e.g. the inputs contain denormals, and this is
+  // easy enough.
+  if (!stream.ThenMemZero(&input_buf, input_buf.size())
+           .ThenMemZero(&filter_buf, filter_buf.size())
+           .ThenMemZero(&output_buf, output_buf.size())
+           .BlockHostUntilDone()
+           .ok()) {
+    LOG(WARNING)
+        << "Couldn't zero out input/filter/output buffer for convolution "
+        << instr->ToString() << ".  Falling back to default algorithm.";
+    return nullopt;
+  }
+
   const bool use_winograd_nonfused = ShouldIncludeWinogradNonfusedAlgo(
       input_shape, output_shape, dnums, stream_exec_);
   se::dnn::ProfileResult best_result;
@@ -225,12 +245,12 @@ CudnnConvolutionAlgorithmPicker::PickBestAlgorithm(
     VLOG(3) << "Trying algorithm " << AlgorithmToString(alg) << " for "
             << instr->ToString();
 
-    bool launch_ok = RunCudnnConvolution(
-                         kind, input_shape, filter_shape, output_shape,
-                         input_buf.ValueOrDie(), filter_buf.ValueOrDie(),
-                         output_buf.ValueOrDie(), &scratch_allocator, window,
-                         dnums, AlgorithmConfig(alg), &stream, &profile_result)
-                         .ok();
+    bool launch_ok =
+        RunCudnnConvolution(kind, input_shape, filter_shape, output_shape,
+                            input_buf, filter_buf, output_buf,
+                            &scratch_allocator, window, dnums,
+                            AlgorithmConfig(alg), &stream, &profile_result)
+            .ok();
 
     if (launch_ok && profile_result.is_valid()) {
       int64 scratch_bytes_used = scratch_allocator.TotalAllocatedBytes();
-- 
GitLab


From 14f618eba158f2c4bb1958a705e29a3edbd7f005 Mon Sep 17 00:00:00 2001
From: Russell Power <power@google.com>
Date: Sat, 5 May 2018 18:55:28 -0700
Subject: [PATCH 294/395] Gracefully handle workers without heartbeat support
 enabled.

PiperOrigin-RevId: 195560525
---
 .../contrib/tpu/python/tpu/session_support.py | 34 +++++++++++++++----
 1 file changed, 28 insertions(+), 6 deletions(-)

diff --git a/tensorflow/contrib/tpu/python/tpu/session_support.py b/tensorflow/contrib/tpu/python/tpu/session_support.py
index 3455e0b4a6..faf677a81d 100644
--- a/tensorflow/contrib/tpu/python/tpu/session_support.py
+++ b/tensorflow/contrib/tpu/python/tpu/session_support.py
@@ -28,6 +28,7 @@ from tensorflow.core.protobuf import config_pb2
 from tensorflow.core.util import event_pb2
 from tensorflow.python.client import session as session_lib
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
 from tensorflow.python.platform import tf_logging as logging
@@ -78,6 +79,15 @@ class WorkerHeartbeatManager(object):
     return WorkerHeartbeatManager(session, devices, heartbeat_ops,
                                   request_placeholder)
 
+  def heartbeat_supported(self):
+    """Returns True if heartbeat operations are supported on all workers."""
+    try:
+      # Send ping to verify worker has heartbeat support.
+      self.ping()
+      return True
+    except errors.InvalidArgumentError as _:
+      return False
+
   def configure(self, message):
     """Configure heartbeat manager for all devices.
 
@@ -106,7 +116,7 @@ class WorkerHeartbeatManager(object):
         event_pb2.WorkerHeartbeatResponse.FromString(res_pb)
         for res_pb in results
     ]
-    logging.info('Results: %s', parsed_results)
+    logging.debug('Ping results: %s', parsed_results)
     return parsed_results
 
   def lame_workers(self):
@@ -189,7 +199,9 @@ class WatchdogManager(threading.Thread):
     self._running = False
     self._graph = ops.Graph()
     self._session = session_lib.Session(
-        target=session.sess_str, graph=self._graph)
+        target=session.sess_str,
+        graph=self._graph,
+    )
 
     with self._graph.as_default():
       if devices is None:
@@ -249,6 +261,7 @@ class GracefulShutdownHook(session_run_hook.SessionRunHook):
     self._graph = ops.Graph()
     self._workers = None
     self._session = None
+    self._heartbeat_supported = False
 
   def after_create_session(self, training_session, coord):  # pylint: disable=unused-argument
     # N.B. We have to pull the global step here to avoid it being unavailable
@@ -264,10 +277,16 @@ class GracefulShutdownHook(session_run_hook.SessionRunHook):
           target=training_session.sess_str, graph=self._graph)
       self._workers = WorkerHeartbeatManager.from_devices(
           self._session, all_worker_devices(self._session))
-
-      self._workers.configure(
-          event_pb2.WorkerHeartbeatRequest(
-              shutdown_mode=event_pb2.WAIT_FOR_COORDINATOR))
+      self._heartbeat_supported = self._workers.heartbeat_supported()
+      if self._heartbeat_supported:
+        self._workers.configure(
+            event_pb2.WorkerHeartbeatRequest(
+                shutdown_mode=event_pb2.WAIT_FOR_COORDINATOR))
+      else:
+        logging.warn(
+            'Worker heartbeats not supported by all workers.  No failure '
+            'handling will be enabled.'
+        )
 
   def saver(self):
     if self._saver:
@@ -286,6 +305,9 @@ class GracefulShutdownHook(session_run_hook.SessionRunHook):
   def after_run(self, run_context, run_values):
     del run_values
 
+    if not self._heartbeat_supported:
+      return
+
     lame_workers = self._workers.lame_workers()
     if lame_workers:
       logging.info('ShutdownHook: lame workers found: %s', lame_workers)
-- 
GitLab


From e0be7b7c6f70aa0c3fc1b97de049cb4ccf1e9c0f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 01:57:29 -0700
Subject: [PATCH 295/395] Remove unused threadpool from stream executor.

PiperOrigin-RevId: 195632175
---
 tensorflow/stream_executor/cuda/cuda_dnn.cc | 24 +++------------------
 tensorflow/stream_executor/cuda/cuda_dnn.h  |  2 +-
 2 files changed, 4 insertions(+), 22 deletions(-)

diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc
index 773cac2c40..316f4c4f1e 100644
--- a/tensorflow/stream_executor/cuda/cuda_dnn.cc
+++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc
@@ -138,27 +138,6 @@ cudnnDataType_t GetCudnnDataType<Eigen::half>() {
 
 namespace wrap {
 
-static port::ThreadPool* InitCudnnThreadpool() {
-  port::ThreadPool* cudnn_threadpool_;
-  port::ThreadOptions options;
-  // TBD(keveman): Conservatively setting the stack size and guard size to 2MB,
-  // until we can get some guarantees from NVIDIA on the minimum stack space
-  // they will work with.
-  options.stack_size = 2 * 1024 * 1024;
-  options.guard_size = 2 * 1024 * 1024;
-  cudnn_threadpool_ = new port::ThreadPool(port::Env::Default(), options,
-                                           "cudnn_threadpool", 1);
-  CHECK(cudnn_threadpool_);
-  return cudnn_threadpool_;
-}
-
-static mutex cudnn_threadpool_mu(LINKER_INITIALIZED);
-static port::ThreadPool* GetCudaThreadpool() {
-  mutex_lock lock(cudnn_threadpool_mu);
-  static port::ThreadPool* cudnn_threadpool = InitCudnnThreadpool();
-  return cudnn_threadpool;
-}
-
 #define STREAM_EXECUTOR_CUDNN_WRAP(__name)                         \
   struct WrapperShim__##__name {                                   \
     template <typename... Args>                                    \
@@ -184,6 +163,9 @@ static port::ThreadPool* GetCudaThreadpool() {
   } __name;
 
 // Handles cudnnSetStream differently in order to add debug information.
+// It stores a reference to 'stream' in 'dnn', and checks that all calls from
+// that dnn instance use the same stream (see
+// STREAM_EXECUTOR_CUDNN_WRAP_WITH_CHECKED_STREAM macro).
 struct WrapperShim__cudnnSetStream {
   cudnnStatus_t operator()(CudnnSupport* dnn, Stream* stream,
                            cudnnHandle_t handle)
diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.h b/tensorflow/stream_executor/cuda/cuda_dnn.h
index 7d53dbe4a5..8a0458bc80 100644
--- a/tensorflow/stream_executor/cuda/cuda_dnn.h
+++ b/tensorflow/stream_executor/cuda/cuda_dnn.h
@@ -652,7 +652,7 @@ class CudnnSupport : public dnn::DnnSupport {
   // single cuda_dnn translation unit.
   void* dnn_handle_ GUARDED_BY(dnn_handle_mutex_);
 
-  // The current cudnn stream that is set by cudnnSetStream().
+  // The current cudnn stream that is set by SetCurrentDnnStream().
   Stream* current_dnn_stream_ GUARDED_BY(dnn_handle_mutex_);
 
   // NOTE(keveman): Temporary data layout transformation until cuDNN supports
-- 
GitLab


From 6304da208116ed00ad4ee776787dfa6fe8256f4f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 03:00:34 -0700
Subject: [PATCH 296/395] Improve fusion logic of (a dot b) * alpha

The previous approach didn't work because a multiplication by a scalar value
will be changed into an explicit broadcast.
Another issue that is fixed in this CL is retrieving the constant value from
the literal. This depends on the PrimitiveType, before we always assumed it to be double.
Also when checking ImplementedAsGemm() we should not call it recursively, but instead just the check related to kDot.
Finally add an execution test and adjust the fusion logic test.

PiperOrigin-RevId: 195638795
---
 tensorflow/compiler/xla/service/gpu/BUILD     |  2 +
 .../xla/service/gpu/instruction_fusion.cc     | 84 +++++++++++++++----
 .../service/gpu/instruction_fusion_test.cc    | 46 ++++++++--
 .../xla/service/gpu/ir_emission_utils.cc      | 36 ++++----
 .../xla/service/gpu/ir_emitter_unnested.cc    | 40 +++++++--
 5 files changed, 160 insertions(+), 48 deletions(-)

diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD
index 7cb7f55073..7ee039b3eb 100644
--- a/tensorflow/compiler/xla/service/gpu/BUILD
+++ b/tensorflow/compiler/xla/service/gpu/BUILD
@@ -388,8 +388,10 @@ cc_library(
     deps = [
         ":ir_emission_utils",
         "//tensorflow/compiler/xla:shape_util",
+        "//tensorflow/compiler/xla:xla_data_proto",
         "//tensorflow/compiler/xla/service:hlo",
         "//tensorflow/compiler/xla/service:instruction_fusion",
+        "//tensorflow/compiler/xla/service:pattern_matcher",
     ],
 )
 
diff --git a/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc b/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc
index c5eb721185..04c7cc3842 100644
--- a/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc
+++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc
@@ -17,7 +17,9 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h"
 #include "tensorflow/compiler/xla/service/hlo_opcode.h"
+#include "tensorflow/compiler/xla/service/pattern_matcher.h"
 #include "tensorflow/compiler/xla/shape_util.h"
+#include "tensorflow/compiler/xla/xla_data.pb.h"
 
 namespace xla {
 namespace gpu {
@@ -46,6 +48,15 @@ bool IsFusile(const HloInstruction& hlo) {
          hlo.opcode() == HloOpcode::kTranspose;
 }
 
+bool IsIEEEFloatingPointScalarConstant(const HloInstruction* constant) {
+  if (constant->opcode() != HloOpcode::kConstant ||
+      !ShapeUtil::IsScalar(constant->shape())) {
+    return false;
+  }
+  auto type = constant->shape().element_type();
+  return type == F16 || type == F32 || type == F64;
+}
+
 }  // namespace
 
 /*static*/ bool GpuInstructionFusion::IsExpensive(
@@ -66,34 +77,71 @@ bool GpuInstructionFusion::ShouldFuse(HloInstruction* consumer,
   HloInstruction* producer = consumer->mutable_operand(operand_index);
 
   // Check if we can use output fusion for (A @ B) * alpha
-  if (producer->opcode() == HloOpcode::kDot) {
-    if (consumer->opcode() == HloOpcode::kMultiply) {
-      CHECK_EQ(consumer->operand_count(), 2);
-      int64 other_operand_index = 1 - operand_index;
-      const HloInstruction* alpha = consumer->operand(other_operand_index);
-      if (alpha->opcode() == HloOpcode::kConstant &&
-          ShapeUtil::IsScalar(alpha->shape())) {
+  if (producer->opcode() == HloOpcode::kDot ||
+      (producer->opcode() == HloOpcode::kFusion &&
+       producer->fused_expression_root()->opcode() == HloOpcode::kDot)) {
+    int64 other_operand_index = 1 - operand_index;
+    const HloInstruction* alpha = consumer->operand(other_operand_index);
+    HloInstruction* op1 = nullptr;
+    HloInstruction* op2 = nullptr;
+    if (consumer->opcode() == HloOpcode::kFusion &&
+        consumer->fusion_kind() == HloInstruction::FusionKind::kLoop &&
+        consumer->operand_count() == 2 &&
+        Match(consumer->fused_expression_root(),
+              match::Op()
+                  .WithOpcode(HloOpcode::kMultiply)
+                  .WithOperand(0, match::Op(&op1))
+                  .WithOperand(1, match::Op(&op2)))) {
+      CHECK(op1 != nullptr && op2 != nullptr);
+      // If 'consumer' is a fusion node, it should consist of a broadcast of a
+      // scalar constant fused into a multiply, but nothing more. So one operand
+      // should be a parameter, and the other should be a broadcast.
+      if (op1->opcode() != HloOpcode::kParameter) {
+        std::swap(op1, op2);
+      }
+      if (op1->opcode() != HloOpcode::kParameter ||
+          op2->opcode() != HloOpcode::kBroadcast) {
+        return false;
+      }
+      if (IsIEEEFloatingPointScalarConstant(alpha)) {
+        return true;
+      }
+    } else if (consumer->opcode() == HloOpcode::kMultiply) {
+      // Fuse if 'alpha' is a broadcast of a scalar constant.
+      if (alpha->opcode() == HloOpcode::kBroadcast &&
+          alpha->dimensions().empty() &&
+          IsIEEEFloatingPointScalarConstant(alpha->operand(0))) {
         return true;
       }
     }
   }
 
-  // Only allow to fuse transpose into an output fusion.
+  // Only allow fusing transpose or broadcast into an output fusion that is
+  // implemented as a Gemm call.
   if (consumer->opcode() == HloOpcode::kFusion &&
-      consumer->fusion_kind() == HloInstruction::FusionKind::kOutput) {
-    if (producer->opcode() != HloOpcode::kTranspose) {
-      return false;
-    }
-    // Check that the transpose is the operand of a dot.
+      consumer->fusion_kind() == HloInstruction::FusionKind::kOutput &&
+      ImplementedAsGemm(*consumer)) {
     auto producer_operand_index = consumer->operand_index(producer);
     auto fused_parameter = consumer->fused_parameter(producer_operand_index);
     const std::vector<HloInstruction*>& fused_parameter_users =
         fused_parameter->users();
-    return (fused_parameter_users.size() == 1 &&
-            fused_parameter_users[0]->opcode() == HloOpcode::kDot);
+    if (fused_parameter_users.size() != 1) {
+      return false;
+    }
+    if (producer->opcode() == HloOpcode::kTranspose) {
+      // Check that the transpose is an operand of a dot.
+      return fused_parameter_users[0]->opcode() == HloOpcode::kDot;
+    }
+    if (producer->opcode() == HloOpcode::kBroadcast) {
+      // Check that the broadcast is a broadcast of a scalar constant into a
+      // multiply.
+      return producer->dimensions().empty() &&
+             IsIEEEFloatingPointScalarConstant(producer->operand(0)) &&
+             fused_parameter_users[0]->opcode() == HloOpcode::kMultiply;
+    }
   }
 
-  // Output fusion is not currently supported on GPUs.
+  // Other output fusions are not currently supported on GPUs.
   if (producer->opcode() == HloOpcode::kFusion) {
     return false;
   }
@@ -134,7 +182,9 @@ HloInstruction::FusionKind GpuInstructionFusion::ChooseKind(
   if (IsReductionToVector(*consumer)) {
     return HloInstruction::FusionKind::kInput;
   }
-  if (producer->opcode() == HloOpcode::kDot) {
+  if (producer->opcode() == HloOpcode::kDot ||
+      (producer->opcode() == HloOpcode::kFusion &&
+       producer->fused_expression_root()->opcode() == HloOpcode::kDot)) {
     return HloInstruction::FusionKind::kOutput;
   }
   if (HloOpcode::kFusion == consumer->opcode()) {
diff --git a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc
index 6c9a805ad6..760e0e90f5 100644
--- a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc
+++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc
@@ -108,8 +108,8 @@ TEST_F(InstructionFusionTest, PotentialBitcastReshapeOfDotUnfused) {
   HloComputation::Builder builder(TestName());
   auto param0 = builder.AddInstruction(HloInstruction::CreateParameter(
       0, ShapeUtil::MakeShape(S32, {1, 1}), "0"));
-  auto dot1 = builder.AddInstruction(HloInstruction::CreateBinary(
-      ShapeUtil::MakeShape(S32, {1, 1}), HloOpcode::kDot, param0, param0));
+  auto dot1 = builder.AddInstruction(HloInstruction::CreateCanonicalDot(
+      ShapeUtil::MakeShape(S32, {1, 1}), param0, param0));
   auto reshape2 = builder.AddInstruction(HloInstruction::CreateReshape(
       ShapeUtil::MakeShape(S32, {1, 1, 1}), dot1));
 
@@ -125,8 +125,8 @@ TEST_F(InstructionFusionTest, PotentialBitcastTransposeOfDotUnfused) {
   HloComputation::Builder builder(TestName());
   auto param0 = builder.AddInstruction(HloInstruction::CreateParameter(
       0, ShapeUtil::MakeShape(S32, {1, 1}), "0"));
-  auto dot1 = builder.AddInstruction(HloInstruction::CreateBinary(
-      ShapeUtil::MakeShape(S32, {1, 1}), HloOpcode::kDot, param0, param0));
+  auto dot1 = builder.AddInstruction(HloInstruction::CreateCanonicalDot(
+      ShapeUtil::MakeShape(S32, {1, 1}), param0, param0));
   auto transpose2 = builder.AddInstruction(HloInstruction::CreateTranspose(
       ShapeUtil::MakeShape(S32, {1, 1}), dot1, {0, 1}));
 
@@ -232,12 +232,13 @@ TEST_F(InstructionFusionTest, DotOutputFusion) {
   auto module = tools::Parse(R"(
   HloModule test_module
   ENTRY OutputFusion {
-    constant = f32[] constant(3)
+    alpha = f32[] constant(3)
+    broadcast = f32[4,4]{1,0} broadcast(alpha), dimensions={}
     p0 = f32[4,3]{1,0} parameter(0)
     p1 = f32[4,3]{1,0} parameter(1)
     transpose = f32[3,4]{1,0} transpose(p1), dimensions={1, 0}
-    dot = f32[4,4]{1,0} dot(p0, transpose)
-    ROOT mul = f32[4,4] multiply(constant, dot)
+    dot = f32[4,4]{1,0} dot(p0, transpose), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    ROOT mul = f32[4,4] multiply(dot, broadcast)
   })")
                     .ValueOrDie();
 
@@ -247,10 +248,11 @@ TEST_F(InstructionFusionTest, DotOutputFusion) {
 
   HloInstruction* root = module->entry_computation()->root_instruction();
   EXPECT_THAT(root, op::Fusion());
+  EXPECT_EQ(root->fusion_kind(), HloInstruction::FusionKind::kOutput);
   EXPECT_THAT(
       root->fused_expression_root(),
-      op::Multiply(op::Parameter(),
-                   op::Dot(op::Parameter(), op::Transpose(op::Parameter()))));
+      op::Multiply(op::Dot(op::Parameter(), op::Transpose(op::Parameter())),
+                   op::Broadcast(op::Parameter())));
 }
 
 // Compute sum(1/p0), where p0 has type f32, twice.  Check that the division is
@@ -309,5 +311,31 @@ TEST_F(InstructionFusionTest, IntegerDivIsNotCheap) {
                    .ValueOrDie());
 }
 
+TEST_F(InstructionFusionTest, DotOutputFusionImpossible) {
+  auto module = tools::Parse(R"(
+  HloModule test_module
+  ENTRY NoOutputFusion {
+    alpha = f32[] constant(3)
+    broadcast = f32[4,4]{1,0} broadcast(alpha), dimensions={}
+    p0 = f32[4,3]{1,0} parameter(0)
+    p1 = f32[3,4]{1,0} parameter(1)
+    dot = f32[4,4]{1,0} dot(p0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    d = f32[4,4]{1,0} multiply(dot, dot)
+    ROOT mul = f32[4,4] multiply(d, broadcast)
+  })")
+                    .ValueOrDie();
+
+  EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true)
+                  .Run(module.get())
+                  .ValueOrDie());
+
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root, op::Fusion());
+  EXPECT_EQ(root->fusion_kind(), HloInstruction::FusionKind::kLoop);
+  EXPECT_THAT(root->fused_expression_root(),
+              op::Multiply(op::Multiply(op::Parameter(), op::Parameter()),
+                           op::Broadcast(op::Parameter())));
+}
+
 }  // namespace gpu
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
index 777345722c..8ab7fe9133 100644
--- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
+++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
@@ -59,6 +59,25 @@ bool AreValidGemmShapes(const Shape& lhs_shape, const Shape& rhs_shape,
          !ShapeUtil::HasZeroElements(lhs_shape) &&
          !ShapeUtil::HasZeroElements(rhs_shape);
 }
+
+bool DotImplementedAsGemm(const HloInstruction& dot) {
+  CHECK_EQ(dot.opcode(), HloOpcode::kDot);
+  const Shape& lhs_shape = dot.operand(0)->shape();
+  const Shape& rhs_shape = dot.operand(1)->shape();
+
+  // If gemm can accept the operand shapes, use it rather than a custom
+  // kernel.
+  if (AreValidGemmShapes(lhs_shape, rhs_shape, dot.shape())) {
+    // The size of the reduction dimension should match. The shape inference
+    // guarantees this invariant, so the check here is for programming
+    // errors.
+    const DotDimensionNumbers& dim_numbers = dot.dot_dimension_numbers();
+    CHECK_EQ(lhs_shape.dimensions(dim_numbers.lhs_contracting_dimensions(0)),
+             rhs_shape.dimensions(dim_numbers.rhs_contracting_dimensions(0)));
+    return true;
+  }
+  return false;
+}
 }  // namespace
 
 bool ImplementedAsGemm(const HloInstruction& hlo) {
@@ -69,20 +88,7 @@ bool ImplementedAsGemm(const HloInstruction& hlo) {
 
   // For certain types of Dot, we can call pre-canned BLAS gemm.
   if (hlo.opcode() == HloOpcode::kDot) {
-    const Shape& lhs_shape = hlo.operand(0)->shape();
-    const Shape& rhs_shape = hlo.operand(1)->shape();
-
-    // If gemm can accept the operand shapes, use it rather than a custom
-    // kernel.
-    if (AreValidGemmShapes(lhs_shape, rhs_shape, hlo.shape())) {
-      // The size of the reduction dimension should match. The shape inference
-      // guarantees this invariant, so the check here is for programming
-      // errors.
-      const DotDimensionNumbers& dim_numbers = hlo.dot_dimension_numbers();
-      CHECK_EQ(lhs_shape.dimensions(dim_numbers.lhs_contracting_dimensions(0)),
-               rhs_shape.dimensions(dim_numbers.rhs_contracting_dimensions(0)));
-      return true;
-    }
+    return DotImplementedAsGemm(hlo);
   }
 
   if (hlo.opcode() == HloOpcode::kFusion) {
@@ -98,7 +104,7 @@ bool ImplementedAsGemm(const HloInstruction& hlo) {
       dot = hlo.fused_expression_root()->operand(1);
     }
     if (dot->opcode() == HloOpcode::kDot) {
-      return ImplementedAsGemm(*dot);
+      return DotImplementedAsGemm(*dot);
     }
   }
 
diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
index 83d90296df..dcaedcb04e 100644
--- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
+++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
@@ -2194,6 +2194,21 @@ std::unique_ptr<Thunk> IrEmitterUnnested::BuildInfeedThunk(
       /*destination_buffer=*/GetAllocationSlice(*inst), inst);
 }
 
+namespace {
+double GetScalarConstantAsDouble(const Literal& literal) {
+  switch (literal.shape().element_type()) {
+    case F16:
+      return static_cast<double>(literal.Get<Eigen::half>({0}));
+    case F32:
+      return literal.Get<float>({0});
+    case F64:
+      return literal.Get<double>({0});
+    default:
+      LOG(FATAL) << "Unsupported type.";
+  }
+}
+}  // namespace
+
 std::unique_ptr<Thunk> IrEmitterUnnested::BuildGemmThunk(
     const HloInstruction* inst) {
   if (inst->opcode() == HloOpcode::kDot) {
@@ -2218,6 +2233,17 @@ std::unique_ptr<Thunk> IrEmitterUnnested::BuildGemmThunk(
     if (dot->opcode() != HloOpcode::kDot) {
       std::swap(dot, alpha);
     }
+    if (alpha->opcode() == HloOpcode::kBroadcast) {
+      alpha = alpha->operand(0);
+    }
+    alpha = inst->operand(alpha->parameter_number());
+    // TODO(b/74185543): Remove the following if block once we support fusion
+    // with a non-constant as well. Then we will just always use the constant
+    // on the device.
+    if (alpha->opcode() == HloOpcode::kCopy) {
+      alpha = alpha->operand(0);
+    }
+
     DCHECK(dot->opcode() == HloOpcode::kDot);
     const HloInstruction* lhs_parameter = StripTranspose(*dot->operand(0));
     const HloInstruction* rhs_parameter = StripTranspose(*dot->operand(1));
@@ -2229,13 +2255,13 @@ std::unique_ptr<Thunk> IrEmitterUnnested::BuildGemmThunk(
         inst->operand(rhs_parameter->parameter_number());
 
     return MakeUnique<GemmThunk>(
-        GetAllocationSlice(*lhs),           // The buffer assigned to LHS.
-        GetAllocationSlice(*rhs),           // The buffer assigned to RHS.
-        GetAllocationSlice(*mul),           // The output buffer.
-        lhs->shape(),                       // The shape of LHS.
-        rhs->shape(),                       // The shape of RHS.
-        inst->shape(),                      // The shape of the output.
-        alpha->literal().Get<double>({0}),  // alpha.
+        GetAllocationSlice(*lhs),   // The buffer assigned to LHS.
+        GetAllocationSlice(*rhs),   // The buffer assigned to RHS.
+        GetAllocationSlice(*inst),  // The output buffer.
+        lhs->shape(),               // The shape of LHS.
+        rhs->shape(),               // The shape of RHS.
+        inst->shape(),              // The shape of the output.
+        GetScalarConstantAsDouble(alpha->literal()),  // alpha.
         inst);
   }
 
-- 
GitLab


From 9dc76cd3bbad752b370b970bf55d97d2c422059c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 04:25:03 -0700
Subject: [PATCH 297/395] Automated g4 rollback of changelist 195638795

PiperOrigin-RevId: 195645734
---
 tensorflow/compiler/xla/service/gpu/BUILD     |  2 -
 .../xla/service/gpu/instruction_fusion.cc     | 84 ++++---------------
 .../service/gpu/instruction_fusion_test.cc    | 46 ++--------
 .../xla/service/gpu/ir_emission_utils.cc      | 36 ++++----
 .../xla/service/gpu/ir_emitter_unnested.cc    | 40 ++-------
 5 files changed, 48 insertions(+), 160 deletions(-)

diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD
index 7ee039b3eb..7cb7f55073 100644
--- a/tensorflow/compiler/xla/service/gpu/BUILD
+++ b/tensorflow/compiler/xla/service/gpu/BUILD
@@ -388,10 +388,8 @@ cc_library(
     deps = [
         ":ir_emission_utils",
         "//tensorflow/compiler/xla:shape_util",
-        "//tensorflow/compiler/xla:xla_data_proto",
         "//tensorflow/compiler/xla/service:hlo",
         "//tensorflow/compiler/xla/service:instruction_fusion",
-        "//tensorflow/compiler/xla/service:pattern_matcher",
     ],
 )
 
diff --git a/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc b/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc
index 04c7cc3842..c5eb721185 100644
--- a/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc
+++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc
@@ -17,9 +17,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h"
 #include "tensorflow/compiler/xla/service/hlo_opcode.h"
-#include "tensorflow/compiler/xla/service/pattern_matcher.h"
 #include "tensorflow/compiler/xla/shape_util.h"
-#include "tensorflow/compiler/xla/xla_data.pb.h"
 
 namespace xla {
 namespace gpu {
@@ -48,15 +46,6 @@ bool IsFusile(const HloInstruction& hlo) {
          hlo.opcode() == HloOpcode::kTranspose;
 }
 
-bool IsIEEEFloatingPointScalarConstant(const HloInstruction* constant) {
-  if (constant->opcode() != HloOpcode::kConstant ||
-      !ShapeUtil::IsScalar(constant->shape())) {
-    return false;
-  }
-  auto type = constant->shape().element_type();
-  return type == F16 || type == F32 || type == F64;
-}
-
 }  // namespace
 
 /*static*/ bool GpuInstructionFusion::IsExpensive(
@@ -77,71 +66,34 @@ bool GpuInstructionFusion::ShouldFuse(HloInstruction* consumer,
   HloInstruction* producer = consumer->mutable_operand(operand_index);
 
   // Check if we can use output fusion for (A @ B) * alpha
-  if (producer->opcode() == HloOpcode::kDot ||
-      (producer->opcode() == HloOpcode::kFusion &&
-       producer->fused_expression_root()->opcode() == HloOpcode::kDot)) {
-    int64 other_operand_index = 1 - operand_index;
-    const HloInstruction* alpha = consumer->operand(other_operand_index);
-    HloInstruction* op1 = nullptr;
-    HloInstruction* op2 = nullptr;
-    if (consumer->opcode() == HloOpcode::kFusion &&
-        consumer->fusion_kind() == HloInstruction::FusionKind::kLoop &&
-        consumer->operand_count() == 2 &&
-        Match(consumer->fused_expression_root(),
-              match::Op()
-                  .WithOpcode(HloOpcode::kMultiply)
-                  .WithOperand(0, match::Op(&op1))
-                  .WithOperand(1, match::Op(&op2)))) {
-      CHECK(op1 != nullptr && op2 != nullptr);
-      // If 'consumer' is a fusion node, it should consist of a broadcast of a
-      // scalar constant fused into a multiply, but nothing more. So one operand
-      // should be a parameter, and the other should be a broadcast.
-      if (op1->opcode() != HloOpcode::kParameter) {
-        std::swap(op1, op2);
-      }
-      if (op1->opcode() != HloOpcode::kParameter ||
-          op2->opcode() != HloOpcode::kBroadcast) {
-        return false;
-      }
-      if (IsIEEEFloatingPointScalarConstant(alpha)) {
-        return true;
-      }
-    } else if (consumer->opcode() == HloOpcode::kMultiply) {
-      // Fuse if 'alpha' is a broadcast of a scalar constant.
-      if (alpha->opcode() == HloOpcode::kBroadcast &&
-          alpha->dimensions().empty() &&
-          IsIEEEFloatingPointScalarConstant(alpha->operand(0))) {
+  if (producer->opcode() == HloOpcode::kDot) {
+    if (consumer->opcode() == HloOpcode::kMultiply) {
+      CHECK_EQ(consumer->operand_count(), 2);
+      int64 other_operand_index = 1 - operand_index;
+      const HloInstruction* alpha = consumer->operand(other_operand_index);
+      if (alpha->opcode() == HloOpcode::kConstant &&
+          ShapeUtil::IsScalar(alpha->shape())) {
         return true;
       }
     }
   }
 
-  // Only allow fusing transpose or broadcast into an output fusion that is
-  // implemented as a Gemm call.
+  // Only allow to fuse transpose into an output fusion.
   if (consumer->opcode() == HloOpcode::kFusion &&
-      consumer->fusion_kind() == HloInstruction::FusionKind::kOutput &&
-      ImplementedAsGemm(*consumer)) {
+      consumer->fusion_kind() == HloInstruction::FusionKind::kOutput) {
+    if (producer->opcode() != HloOpcode::kTranspose) {
+      return false;
+    }
+    // Check that the transpose is the operand of a dot.
     auto producer_operand_index = consumer->operand_index(producer);
     auto fused_parameter = consumer->fused_parameter(producer_operand_index);
     const std::vector<HloInstruction*>& fused_parameter_users =
         fused_parameter->users();
-    if (fused_parameter_users.size() != 1) {
-      return false;
-    }
-    if (producer->opcode() == HloOpcode::kTranspose) {
-      // Check that the transpose is an operand of a dot.
-      return fused_parameter_users[0]->opcode() == HloOpcode::kDot;
-    }
-    if (producer->opcode() == HloOpcode::kBroadcast) {
-      // Check that the broadcast is a broadcast of a scalar constant into a
-      // multiply.
-      return producer->dimensions().empty() &&
-             IsIEEEFloatingPointScalarConstant(producer->operand(0)) &&
-             fused_parameter_users[0]->opcode() == HloOpcode::kMultiply;
-    }
+    return (fused_parameter_users.size() == 1 &&
+            fused_parameter_users[0]->opcode() == HloOpcode::kDot);
   }
 
-  // Other output fusions are not currently supported on GPUs.
+  // Output fusion is not currently supported on GPUs.
   if (producer->opcode() == HloOpcode::kFusion) {
     return false;
   }
@@ -182,9 +134,7 @@ HloInstruction::FusionKind GpuInstructionFusion::ChooseKind(
   if (IsReductionToVector(*consumer)) {
     return HloInstruction::FusionKind::kInput;
   }
-  if (producer->opcode() == HloOpcode::kDot ||
-      (producer->opcode() == HloOpcode::kFusion &&
-       producer->fused_expression_root()->opcode() == HloOpcode::kDot)) {
+  if (producer->opcode() == HloOpcode::kDot) {
     return HloInstruction::FusionKind::kOutput;
   }
   if (HloOpcode::kFusion == consumer->opcode()) {
diff --git a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc
index 760e0e90f5..6c9a805ad6 100644
--- a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc
+++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc
@@ -108,8 +108,8 @@ TEST_F(InstructionFusionTest, PotentialBitcastReshapeOfDotUnfused) {
   HloComputation::Builder builder(TestName());
   auto param0 = builder.AddInstruction(HloInstruction::CreateParameter(
       0, ShapeUtil::MakeShape(S32, {1, 1}), "0"));
-  auto dot1 = builder.AddInstruction(HloInstruction::CreateCanonicalDot(
-      ShapeUtil::MakeShape(S32, {1, 1}), param0, param0));
+  auto dot1 = builder.AddInstruction(HloInstruction::CreateBinary(
+      ShapeUtil::MakeShape(S32, {1, 1}), HloOpcode::kDot, param0, param0));
   auto reshape2 = builder.AddInstruction(HloInstruction::CreateReshape(
       ShapeUtil::MakeShape(S32, {1, 1, 1}), dot1));
 
@@ -125,8 +125,8 @@ TEST_F(InstructionFusionTest, PotentialBitcastTransposeOfDotUnfused) {
   HloComputation::Builder builder(TestName());
   auto param0 = builder.AddInstruction(HloInstruction::CreateParameter(
       0, ShapeUtil::MakeShape(S32, {1, 1}), "0"));
-  auto dot1 = builder.AddInstruction(HloInstruction::CreateCanonicalDot(
-      ShapeUtil::MakeShape(S32, {1, 1}), param0, param0));
+  auto dot1 = builder.AddInstruction(HloInstruction::CreateBinary(
+      ShapeUtil::MakeShape(S32, {1, 1}), HloOpcode::kDot, param0, param0));
   auto transpose2 = builder.AddInstruction(HloInstruction::CreateTranspose(
       ShapeUtil::MakeShape(S32, {1, 1}), dot1, {0, 1}));
 
@@ -232,13 +232,12 @@ TEST_F(InstructionFusionTest, DotOutputFusion) {
   auto module = tools::Parse(R"(
   HloModule test_module
   ENTRY OutputFusion {
-    alpha = f32[] constant(3)
-    broadcast = f32[4,4]{1,0} broadcast(alpha), dimensions={}
+    constant = f32[] constant(3)
     p0 = f32[4,3]{1,0} parameter(0)
     p1 = f32[4,3]{1,0} parameter(1)
     transpose = f32[3,4]{1,0} transpose(p1), dimensions={1, 0}
-    dot = f32[4,4]{1,0} dot(p0, transpose), lhs_contracting_dims={1}, rhs_contracting_dims={0}
-    ROOT mul = f32[4,4] multiply(dot, broadcast)
+    dot = f32[4,4]{1,0} dot(p0, transpose)
+    ROOT mul = f32[4,4] multiply(constant, dot)
   })")
                     .ValueOrDie();
 
@@ -248,11 +247,10 @@ TEST_F(InstructionFusionTest, DotOutputFusion) {
 
   HloInstruction* root = module->entry_computation()->root_instruction();
   EXPECT_THAT(root, op::Fusion());
-  EXPECT_EQ(root->fusion_kind(), HloInstruction::FusionKind::kOutput);
   EXPECT_THAT(
       root->fused_expression_root(),
-      op::Multiply(op::Dot(op::Parameter(), op::Transpose(op::Parameter())),
-                   op::Broadcast(op::Parameter())));
+      op::Multiply(op::Parameter(),
+                   op::Dot(op::Parameter(), op::Transpose(op::Parameter()))));
 }
 
 // Compute sum(1/p0), where p0 has type f32, twice.  Check that the division is
@@ -311,31 +309,5 @@ TEST_F(InstructionFusionTest, IntegerDivIsNotCheap) {
                    .ValueOrDie());
 }
 
-TEST_F(InstructionFusionTest, DotOutputFusionImpossible) {
-  auto module = tools::Parse(R"(
-  HloModule test_module
-  ENTRY NoOutputFusion {
-    alpha = f32[] constant(3)
-    broadcast = f32[4,4]{1,0} broadcast(alpha), dimensions={}
-    p0 = f32[4,3]{1,0} parameter(0)
-    p1 = f32[3,4]{1,0} parameter(1)
-    dot = f32[4,4]{1,0} dot(p0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
-    d = f32[4,4]{1,0} multiply(dot, dot)
-    ROOT mul = f32[4,4] multiply(d, broadcast)
-  })")
-                    .ValueOrDie();
-
-  EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true)
-                  .Run(module.get())
-                  .ValueOrDie());
-
-  HloInstruction* root = module->entry_computation()->root_instruction();
-  EXPECT_THAT(root, op::Fusion());
-  EXPECT_EQ(root->fusion_kind(), HloInstruction::FusionKind::kLoop);
-  EXPECT_THAT(root->fused_expression_root(),
-              op::Multiply(op::Multiply(op::Parameter(), op::Parameter()),
-                           op::Broadcast(op::Parameter())));
-}
-
 }  // namespace gpu
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
index 8ab7fe9133..777345722c 100644
--- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
+++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
@@ -59,25 +59,6 @@ bool AreValidGemmShapes(const Shape& lhs_shape, const Shape& rhs_shape,
          !ShapeUtil::HasZeroElements(lhs_shape) &&
          !ShapeUtil::HasZeroElements(rhs_shape);
 }
-
-bool DotImplementedAsGemm(const HloInstruction& dot) {
-  CHECK_EQ(dot.opcode(), HloOpcode::kDot);
-  const Shape& lhs_shape = dot.operand(0)->shape();
-  const Shape& rhs_shape = dot.operand(1)->shape();
-
-  // If gemm can accept the operand shapes, use it rather than a custom
-  // kernel.
-  if (AreValidGemmShapes(lhs_shape, rhs_shape, dot.shape())) {
-    // The size of the reduction dimension should match. The shape inference
-    // guarantees this invariant, so the check here is for programming
-    // errors.
-    const DotDimensionNumbers& dim_numbers = dot.dot_dimension_numbers();
-    CHECK_EQ(lhs_shape.dimensions(dim_numbers.lhs_contracting_dimensions(0)),
-             rhs_shape.dimensions(dim_numbers.rhs_contracting_dimensions(0)));
-    return true;
-  }
-  return false;
-}
 }  // namespace
 
 bool ImplementedAsGemm(const HloInstruction& hlo) {
@@ -88,7 +69,20 @@ bool ImplementedAsGemm(const HloInstruction& hlo) {
 
   // For certain types of Dot, we can call pre-canned BLAS gemm.
   if (hlo.opcode() == HloOpcode::kDot) {
-    return DotImplementedAsGemm(hlo);
+    const Shape& lhs_shape = hlo.operand(0)->shape();
+    const Shape& rhs_shape = hlo.operand(1)->shape();
+
+    // If gemm can accept the operand shapes, use it rather than a custom
+    // kernel.
+    if (AreValidGemmShapes(lhs_shape, rhs_shape, hlo.shape())) {
+      // The size of the reduction dimension should match. The shape inference
+      // guarantees this invariant, so the check here is for programming
+      // errors.
+      const DotDimensionNumbers& dim_numbers = hlo.dot_dimension_numbers();
+      CHECK_EQ(lhs_shape.dimensions(dim_numbers.lhs_contracting_dimensions(0)),
+               rhs_shape.dimensions(dim_numbers.rhs_contracting_dimensions(0)));
+      return true;
+    }
   }
 
   if (hlo.opcode() == HloOpcode::kFusion) {
@@ -104,7 +98,7 @@ bool ImplementedAsGemm(const HloInstruction& hlo) {
       dot = hlo.fused_expression_root()->operand(1);
     }
     if (dot->opcode() == HloOpcode::kDot) {
-      return DotImplementedAsGemm(*dot);
+      return ImplementedAsGemm(*dot);
     }
   }
 
diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
index dcaedcb04e..83d90296df 100644
--- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
+++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
@@ -2194,21 +2194,6 @@ std::unique_ptr<Thunk> IrEmitterUnnested::BuildInfeedThunk(
       /*destination_buffer=*/GetAllocationSlice(*inst), inst);
 }
 
-namespace {
-double GetScalarConstantAsDouble(const Literal& literal) {
-  switch (literal.shape().element_type()) {
-    case F16:
-      return static_cast<double>(literal.Get<Eigen::half>({0}));
-    case F32:
-      return literal.Get<float>({0});
-    case F64:
-      return literal.Get<double>({0});
-    default:
-      LOG(FATAL) << "Unsupported type.";
-  }
-}
-}  // namespace
-
 std::unique_ptr<Thunk> IrEmitterUnnested::BuildGemmThunk(
     const HloInstruction* inst) {
   if (inst->opcode() == HloOpcode::kDot) {
@@ -2233,17 +2218,6 @@ std::unique_ptr<Thunk> IrEmitterUnnested::BuildGemmThunk(
     if (dot->opcode() != HloOpcode::kDot) {
       std::swap(dot, alpha);
     }
-    if (alpha->opcode() == HloOpcode::kBroadcast) {
-      alpha = alpha->operand(0);
-    }
-    alpha = inst->operand(alpha->parameter_number());
-    // TODO(b/74185543): Remove the following if block once we support fusion
-    // with a non-constant as well. Then we will just always use the constant
-    // on the device.
-    if (alpha->opcode() == HloOpcode::kCopy) {
-      alpha = alpha->operand(0);
-    }
-
     DCHECK(dot->opcode() == HloOpcode::kDot);
     const HloInstruction* lhs_parameter = StripTranspose(*dot->operand(0));
     const HloInstruction* rhs_parameter = StripTranspose(*dot->operand(1));
@@ -2255,13 +2229,13 @@ std::unique_ptr<Thunk> IrEmitterUnnested::BuildGemmThunk(
         inst->operand(rhs_parameter->parameter_number());
 
     return MakeUnique<GemmThunk>(
-        GetAllocationSlice(*lhs),   // The buffer assigned to LHS.
-        GetAllocationSlice(*rhs),   // The buffer assigned to RHS.
-        GetAllocationSlice(*inst),  // The output buffer.
-        lhs->shape(),               // The shape of LHS.
-        rhs->shape(),               // The shape of RHS.
-        inst->shape(),              // The shape of the output.
-        GetScalarConstantAsDouble(alpha->literal()),  // alpha.
+        GetAllocationSlice(*lhs),           // The buffer assigned to LHS.
+        GetAllocationSlice(*rhs),           // The buffer assigned to RHS.
+        GetAllocationSlice(*mul),           // The output buffer.
+        lhs->shape(),                       // The shape of LHS.
+        rhs->shape(),                       // The shape of RHS.
+        inst->shape(),                      // The shape of the output.
+        alpha->literal().Get<double>({0}),  // alpha.
         inst);
   }
 
-- 
GitLab


From adcfdfcfec78b2021320d99c3543a054b470f56f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 06:27:43 -0700
Subject: [PATCH 298/395] Control flow graph with forward and backward analysis

PiperOrigin-RevId: 195654450
---
 .../autograph/pyct/static_analysis/BUILD      |  13 +
 .../autograph/pyct/static_analysis/cfg.py     | 431 ++++++++++++++++++
 .../pyct/static_analysis/cfg_test.py          | 252 ++++++++++
 3 files changed, 696 insertions(+)
 create mode 100644 tensorflow/contrib/autograph/pyct/static_analysis/cfg.py
 create mode 100644 tensorflow/contrib/autograph/pyct/static_analysis/cfg_test.py

diff --git a/tensorflow/contrib/autograph/pyct/static_analysis/BUILD b/tensorflow/contrib/autograph/pyct/static_analysis/BUILD
index 83f3bafc42..68fbdf6953 100644
--- a/tensorflow/contrib/autograph/pyct/static_analysis/BUILD
+++ b/tensorflow/contrib/autograph/pyct/static_analysis/BUILD
@@ -19,6 +19,7 @@ py_library(
     srcs = [
         "activity.py",
         "annos.py",
+        "cfg.py",
         "live_values.py",
         "type_info.py",
     ],
@@ -43,6 +44,18 @@ py_test(
     ],
 )
 
+py_test(
+    name = "cfg_test",
+    srcs = ["cfg_test.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        ":static_analysis",
+        "//tensorflow/contrib/autograph/pyct",
+        "//tensorflow/python:client_testlib",
+        "@gast_archive//:gast",
+    ],
+)
+
 py_test(
     name = "live_values_test",
     srcs = ["live_values_test.py"],
diff --git a/tensorflow/contrib/autograph/pyct/static_analysis/cfg.py b/tensorflow/contrib/autograph/pyct/static_analysis/cfg.py
new file mode 100644
index 0000000000..230e4cc0f3
--- /dev/null
+++ b/tensorflow/contrib/autograph/pyct/static_analysis/cfg.py
@@ -0,0 +1,431 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Control flow graph analysis.
+
+Given a Python AST we construct a control flow graph, with edges both to the
+next and previous statements (so it can easily walk the graph both ways). Its
+nodes contain the AST of the statements. It can then perform forward or backward
+analysis on this CFG.
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from collections import namedtuple
+import functools
+import operator
+
+import gast
+
+from tensorflow.contrib.autograph.pyct import anno
+from tensorflow.contrib.autograph.pyct.static_analysis import activity
+
+
+class CfgNode(object):
+  """A node in the CFG."""
+  __slots__ = ['next', 'value', 'prev']
+
+  def __init__(self, value):
+    self.next = set()
+    self.prev = set()
+    self.value = value
+
+
+class Cfg(namedtuple('Cfg', ['entry', 'exit'])):
+  """A Control Flow Graph.
+
+  Each statement is represented as a node. For control flow statements such
+  as conditionals and loops the conditional itself is a node which either
+  branches or cycles, respectively.
+  Attributes:
+    entry: The entry node, which contains the `gast.arguments` node of the
+        function definition.
+    exit: The exit node. This node is special because it has no value (i.e. no
+        corresponding AST node). This is because Python functions can have
+        multiple return statements.
+  """
+  pass
+
+
+class CfgBuilder(gast.NodeVisitor):
+  """Construct a control flow graph.
+
+  Construct a CFG starting from a FunctionDef node.
+  Usage:
+    cfg_obj = CfgBuilder().build_cfg(fndef_node)
+  """
+
+  def __init__(self):
+    # The current leaves of the CFG
+    self.current_leaves = []
+    # TODO(alexbw): generalize to break, return, continue, yield, etc.
+    # A stack of lists, tracking continue statements
+    self.continue_ = []
+    # A stack of lists tracking break nodes
+    self.break_ = []
+
+  def set_current_leaves(self, cfg_node):
+    """Link this cfg_node to the current leaves.
+
+    This is the central function for building the CFG. It links the current
+    head cfg_nodes to the passed cfg_node. It then resets the head to the
+    passed cfg_node.
+
+    Args:
+      cfg_node: A CfgNode instance.
+    """
+    for head in self.current_leaves:
+      head.next.add(cfg_node)
+      # While we're linking the CFG forward, add backlinks
+      cfg_node.prev.add(head)
+    self.current_leaves = [cfg_node]
+
+  def build_cfg(self, node):
+    """Build a CFG for a function.
+
+    Implementation of building a CFG for dataflow analysis. See, e.g.:
+    https://www.seas.harvard.edu/courses/cs252/2011sp/slides/Lec02-Dataflow.pdf
+
+    Args:
+      node: A function definition the body of which to analyze.
+    Returns:
+      A CFG object.
+    Raises:
+      TypeError: If the input is not a function definition.
+    """
+    if not isinstance(node, gast.FunctionDef):
+      raise TypeError('input must be a function definition')
+    entry_cfg_node = CfgNode(node.args)
+    self.current_leaves = [entry_cfg_node]
+    self.visit_statements(node.body)
+    exit_cfg_node = CfgNode(None)
+    self.set_current_leaves(exit_cfg_node)
+    return Cfg(entry_cfg_node, exit_cfg_node)
+
+  def visit_statements(self, nodes):
+    for node in nodes:
+      # Check for control flow
+      if isinstance(node, (gast.For, gast.While, gast.If, gast.Try, gast.Break,
+                           gast.Continue, gast.With)):
+        self.visit(node)
+      else:
+        expr = CfgNode(node)
+        self.set_current_leaves(expr)
+
+  def generic_visit(self, node):
+    raise ValueError('unknown control flow')
+
+  def visit_If(self, node):
+    # TODO(alexbw): change this to use immutable tuples instead of lists
+    # The current head will hold the conditional
+    test = CfgNode(node.test)
+    self.set_current_leaves(test)
+    # Handle the body
+    self.visit_statements(node.body)
+    body_exit = self.current_leaves
+    self.current_leaves = []
+    self.current_leaves.append(test)
+    # Handle the orelse
+    self.visit_statements(node.orelse)
+    self.current_leaves.extend(body_exit)
+
+  def visit_While(self, node):
+    test = CfgNode(node.test)
+    self.set_current_leaves(test)
+    # Start a new level of nesting
+    self.break_.append([])
+    self.continue_.append([])
+    # Handle the body
+    self.visit_statements(node.body)
+    self.current_leaves.extend(self.continue_.pop())
+    self.set_current_leaves(test)
+    # Handle the orelse
+    self.visit_statements(node.orelse)
+    # The break statements and the test go to the next node
+    self.current_leaves.extend(self.break_.pop())
+
+  def visit_For(self, node):
+    iter_ = CfgNode(node.iter)
+    self.set_current_leaves(iter_)
+    self.break_.append([])
+    self.continue_.append([])
+    self.visit_statements(node.body)
+    self.current_leaves.extend(self.continue_.pop())
+    self.set_current_leaves(iter_)
+    self.current_leaves.extend(self.break_.pop())
+
+  def visit_Break(self, node):
+    self.break_[-1].extend(self.current_leaves)
+    self.current_leaves[:] = []
+
+  def visit_Continue(self, node):
+    self.continue_[-1].extend(self.current_leaves)
+    self.current_leaves[:] = []
+
+  def visit_Try(self, node):
+    self.visit_statements(node.body)
+    body = self.current_leaves
+    handlers = []
+    for handler in node.handlers:
+      self.current_leaves = body[:]
+      self.visit_statements(handler.body)
+      handlers.extend(self.current_leaves)
+    self.current_leaves = body
+    self.visit_statements(node.orelse)
+    self.current_leaves = handlers + self.current_leaves
+    self.visit_statements(node.finalbody)
+
+  def visit_With(self, node):
+    for item in node.items:
+      self.set_current_leaves(CfgNode(item))
+    self.visit_statements(node.body)
+
+
+# TODO(alexbw): once CFG analysis occurs at a block level,
+# this extra class will not be necessary
+class PropagateAnalysis(gast.NodeVisitor):
+  """Port analysis annotations from statements to their enclosing blocks."""
+
+  def __init__(self, analysis):
+    self.transfer_fn = analysis.transfer_fn
+    self.in_label = analysis.in_label
+    self.out_label = analysis.out_label
+    super(PropagateAnalysis, self).__init__()
+
+  def visit_If(self, node):
+    # Depth-first.
+    self.generic_visit(node)
+    incoming = anno.getanno(node.body[0], self.in_label)
+    incoming |= anno.getanno(node.test, self.in_label)
+    outgoing = anno.getanno(node.body[-1], self.out_label)
+    outgoing |= anno.getanno(node.test, self.out_label)
+    if node.orelse:
+      orelse_outgoing = anno.getanno(node.orelse[-1], self.out_label)
+      outgoing = self.transfer_fn(outgoing, orelse_outgoing)
+    anno.setanno(node, self.in_label, incoming)
+    anno.setanno(node, self.out_label, outgoing)
+
+  def visit_For(self, node):
+    self.generic_visit(node)
+    incoming = set(anno.getanno(node.body[0], self.in_label))
+    incoming -= set((anno.getanno(node.target, anno.Basic.QN),))
+    outgoing = anno.getanno(node.body[-1], self.out_label)
+    if node.orelse:
+      orelse_outgoing = anno.getanno(node.orelse[-1], self.out_label)
+      outgoing = self.transfer_fn(outgoing, orelse_outgoing)
+    anno.setanno(node, self.in_label, frozenset(incoming))
+    anno.setanno(node, self.out_label, outgoing)
+
+  def visit_While(self, node):
+    self.generic_visit(node)
+    incoming = anno.getanno(node.body[0], self.in_label)
+    incoming |= anno.getanno(node.test, self.in_label)
+    outgoing = anno.getanno(node.body[-1], self.out_label)
+    if node.orelse:
+      orelse_outgoing = anno.getanno(node.orelse[-1], self.out_label)
+      outgoing = self.transfer_fn(outgoing, orelse_outgoing)
+    anno.setanno(node, self.in_label, incoming)
+    anno.setanno(node, self.out_label, outgoing)
+
+  def visit_With(self, node):
+    self.generic_visit(node)
+    incoming = anno.getanno(node.body[0], self.in_label)
+    for item in node.items:
+      incoming |= anno.getanno(item, self.in_label)
+    outgoing = anno.getanno(node.body[-1], self.out_label)
+    anno.setanno(node, self.in_label, incoming)
+    anno.setanno(node, self.out_label, outgoing)
+
+
+# TODO(alexbw): Abstract the CFG walking machinery into a superclass
+# which is parameterized on which fields it selects when walking.
+# TODO(alexbw): Abstract the application of dataflow analysis
+class Forward(object):
+  """Forward analysis on CFG.
+
+  Args:
+    label: A name for this analysis e.g. 'active' for activity analysis. The AST
+      nodes in the CFG will be given annotations 'name_in', 'name_out',
+      'name_gen' and 'name_kill' which contain the incoming values, outgoing
+      values, values generated by the statement, and values deleted by the
+      statement respectively.
+    transfer_fn: Either the AND or OR operator. If the AND operator is used it
+      turns into forward must analysis (i.e. a value will only be carried
+      forward if it appears on all incoming paths). The OR operator means that
+      forward may analysis is done (i.e. the union of incoming values will be
+      taken).
+  """
+
+  def __init__(self, label, context, transfer_fn=operator.or_):
+    self.transfer_fn = transfer_fn
+    self.context = context
+    self.out_label = label + '_out'
+    self.in_label = label + '_in'
+    self.gen_label = label + '_gen'
+    self.kill_label = label + '_kill'
+
+  # TODO(alexbw): see if we can simplify by visiting breadth-first
+  def visit(self, node):
+    """Depth-first walking the CFG, applying dataflow information propagtion."""
+    # node.value is None only for the exit CfgNode.
+    if not node.value:
+      return
+
+    if anno.hasanno(node.value, self.out_label):
+      before = hash(anno.getanno(node.value, self.out_label))
+    else:
+      before = None
+    preds = [
+        anno.getanno(pred.value, self.out_label)
+        for pred in node.prev
+        if anno.hasanno(pred.value, self.out_label)
+    ]
+    if preds:
+      incoming = functools.reduce(self.transfer_fn, preds[1:], preds[0])
+    else:
+      incoming = frozenset()
+    anno.setanno(node.value, self.in_label, incoming)
+    gen, kill = self.get_gen_kill(node, incoming)
+    anno.setanno(node.value, self.gen_label, gen)
+    anno.setanno(node.value, self.kill_label, kill)
+    anno.setanno(node.value, self.out_label, (incoming - kill) | gen)
+
+    if hash(anno.getanno(node.value, self.out_label)) != before:
+      for succ in node.next:
+        self.visit(succ)
+
+  def get_gen_kill(self, cfg_node, incoming):
+    """Calculate Gen and Kill properties of a CFG node in dataflow analysis.
+
+    A function which takes the CFG node as well as a set of incoming
+    values. It must return a set of newly generated values by the statement as
+    well as a set of deleted (killed) values.
+
+    Args:
+      cfg_node: A CfgNode instance.
+      incoming:
+    """
+    raise NotImplementedError()
+
+
+class Backward(Forward):
+  """Backward analysis on CFG."""
+
+  def visit(self, cfg_node):
+    # cfg_node.value is None for the exit node, which will be visited only once
+    if not cfg_node.value:
+      for pred in cfg_node.prev:
+        self.visit(pred)
+      return
+
+    if anno.hasanno(cfg_node.value, self.in_label):
+      before = hash(anno.getanno(cfg_node.value, self.in_label))
+    else:
+      before = None
+    succs = [
+        anno.getanno(succ.value, self.in_label)
+        for succ in cfg_node.next
+        if anno.hasanno(succ.value, self.in_label)
+    ]
+    if succs:
+      incoming = functools.reduce(self.transfer_fn, succs[1:], succs[0])
+    else:
+      incoming = frozenset()
+    anno.setanno(cfg_node.value, self.out_label, incoming)
+    gen, kill = self.get_gen_kill(cfg_node, incoming)
+    anno.setanno(cfg_node.value, self.gen_label, gen)
+    anno.setanno(cfg_node.value, self.kill_label, kill)
+    anno.setanno(cfg_node.value, self.in_label, (incoming - kill) | gen)
+    if hash(anno.getanno(cfg_node.value, self.in_label)) != before:
+      for pred in cfg_node.prev:
+        self.visit(pred)
+
+
+def run_analyses(node, analyses):
+  """Perform dataflow analysis on all functions within an AST.
+
+  Args:
+    node: An AST node on which to run dataflow analysis.
+    analyses: Either an instance of the Forward or Backward dataflow analysis
+      class, or a list or tuple of them.
+
+  Returns:
+    node: The node, but now with annotations on the AST nodes containing the
+    results of the dataflow analyses.
+  """
+  if not isinstance(analyses, (tuple, list)):
+    analyses = (analyses,)
+  for analysis in analyses:
+    if not isinstance(analysis, (Forward, Backward)):
+      raise TypeError('not a valid forward analysis object')
+
+  for child_node in gast.walk(node):
+    if isinstance(child_node, gast.FunctionDef):
+      cfg_obj = CfgBuilder().build_cfg(child_node)
+      for analysis in analyses:
+        if isinstance(analysis, Backward):
+          analysis.visit(cfg_obj.exit)
+        elif isinstance(analysis, Forward):
+          analysis.visit(cfg_obj.entry)
+  for analysis in analyses:
+    PropagateAnalysis(analysis).visit(node)
+  return node
+
+
+class Liveness(Backward):
+  """Perform a liveness analysis.
+
+  Each statement is annotated with a set of variables that may be used
+  later in the program.
+  """
+
+  def __init__(self, context):
+    super(Liveness, self).__init__('live', context)
+
+  def get_gen_kill(self, node, _):
+    gen = activity.get_read(node.value, self.context)
+    kill = activity.get_updated(node.value, self.context)
+    return gen, kill
+
+
+class ReachingDefinitions(Forward):
+  """Perform reaching definition analysis.
+
+  Each statement is annotated with a set of (variable, definition) pairs.
+  """
+
+  def __init__(self, context):
+    super(ReachingDefinitions, self).__init__('definitions', context)
+
+  def get_gen_kill(self, node, incoming):
+    definitions = activity.get_updated(node.value, self.context)
+    gen = frozenset((id_, node.value) for id_ in definitions)
+    kill = frozenset(def_ for def_ in incoming if def_[0] in definitions)
+    return gen, kill
+
+
+class Defined(Forward):
+  """Perform defined variable analysis.
+
+  Each statement is annotated with a set of variables which are guaranteed to
+  be defined at that point.
+  """
+
+  def __init__(self, context):
+    super(Defined, self).__init__('defined', context, transfer_fn=operator.and_)
+
+  def get_gen_kill(self, node, _):
+    gen = activity.get_updated(node.value, self.context)
+    return gen, frozenset()
diff --git a/tensorflow/contrib/autograph/pyct/static_analysis/cfg_test.py b/tensorflow/contrib/autograph/pyct/static_analysis/cfg_test.py
new file mode 100644
index 0000000000..af7eaf30e8
--- /dev/null
+++ b/tensorflow/contrib/autograph/pyct/static_analysis/cfg_test.py
@@ -0,0 +1,252 @@
+# 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 cfg module."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import functools
+
+import gast
+
+from tensorflow.contrib.autograph.pyct import anno
+from tensorflow.contrib.autograph.pyct import context
+from tensorflow.contrib.autograph.pyct import parser
+from tensorflow.contrib.autograph.pyct import qual_names
+from tensorflow.contrib.autograph.pyct.static_analysis import cfg
+from tensorflow.python.platform import test
+
+
+class CFGTest(test.TestCase):
+
+  def _parse_and_analyze(self, test_fn, namespace, arg_types=None):
+    arg_types = arg_types or {}
+    node, source = parser.parse_entity(test_fn)
+    ctx = context.EntityContext(
+        namer=None,
+        source_code=source,
+        source_file=None,
+        namespace=namespace,
+        arg_values=None,
+        arg_types=arg_types,
+        owner_type=None,
+        recursive=True)
+    node = qual_names.resolve(node)
+    return node, ctx
+
+  def _check_anno_matches(self, node, anno_name, var_names):
+    if isinstance(var_names, str):
+      var_names = (var_names,)
+    qual_vars = set()
+    for var_name in var_names:
+      if isinstance(var_name, str):
+        if '[' in var_name or ']' in var_name:
+          raise ValueError('Annotation matching not supported with subscript.')
+        if '.' not in var_name:
+          qual_vars.add(qual_names.QN(var_name))
+        else:
+          attrs = var_name.split('.')
+          this_qn = functools.reduce(qual_names.QN, attrs[1:],
+                                     qual_names.QN(attrs[0]))
+          qual_vars.add(this_qn)
+    self.assertEqual(anno.getanno(node, anno_name), qual_vars)
+
+  def test_reaching(self):
+
+    def f(x):
+      print(x)
+      while True:
+        x = x
+        x = x
+      return x
+
+    node, ctx = self._parse_and_analyze(f, {})
+    cfg.run_analyses(node, cfg.ReachingDefinitions(ctx))
+    body = node.body[0].body
+    # Only the argument reaches the expression
+    def_in = anno.getanno(body[0], 'definitions_in')
+    # One element, x, from arguments
+    self.assertEqual(set(type(d[1]) for d in def_in), set((gast.arguments,)))
+
+    while_body = body[1].body
+    def_in = anno.getanno(while_body[0], 'definitions_in')
+    # One definition, two possible sources.
+    # - One from an assignment (if the loop is entered)
+    # - The other from the arguments (if loop is not entered)
+    self.assertEqual(
+        set(type(d[1]) for d in def_in), set((gast.arguments, gast.Assign)))
+
+    def_in = anno.getanno(while_body[1], 'definitions_in')
+    # If we've reached this line, the only reaching definition of x is the
+    # Assign node in previous line
+    self.assertEqual(set(type(d[1]) for d in def_in), set((gast.Assign,)))
+
+    def_in = anno.getanno(body[2], 'definitions_in')
+    # Same situation as while_body[0]
+    self.assertEqual(
+        set(type(d[1]) for d in def_in), set((gast.arguments, gast.Assign)))
+
+  def test_defined(self):
+
+    def f(x):
+      if x:
+        y = 2  # pylint: disable=unused-variable
+      return x
+
+    node, ctx = self._parse_and_analyze(f, {})
+    cfg.run_analyses(node, cfg.Defined(ctx))
+    body = node.body[0].body
+    # only x is for sure defined at the end
+    self._check_anno_matches(body[1], 'defined_in', 'x')
+    # at the end of the if body both x and y are defined
+    if_body = body[0].body
+    self._check_anno_matches(if_body[0], 'defined_out', ('x', 'y'))
+
+  # TODO(alexbw): b/73926938 split this test up
+  def test_live(self):
+
+    def get_live_annotated_fnbody(f):
+      node, ctx = self._parse_and_analyze(f, {})
+      cfg.run_analyses(node, cfg.Liveness(ctx))
+      body = node.body[0].body
+      return body
+
+    def f1(x):
+      a = g(x)  # pylint: disable=undefined-variable
+      b = h(a)  # pylint: disable=undefined-variable, unused-variable
+      return x
+
+    def f2(x, a):  # pylint: disable=unused-argument
+      if a > 0:  # x should not be live
+        x = 0
+      if a > 1:
+        x = 1
+      else:
+        x = 2
+
+    def f3(x, a):
+      if a > 0:  # x and a should be live
+        x = 0
+      if a > 1:  # x and a should be live_in
+        x = 1
+      return x  # x should be live
+
+    def f4(x, a):
+      if a > 0:  # x should be live
+        x = 0
+      x += 1
+
+    def f5(x, a):
+      if a > 0:  # x.y should be live
+        x.y = 0
+      return x.y
+
+    def f6(x):
+      return x  # should this cause x.* to be live?
+
+    def f7(x, n):
+      for i in range(n):
+        x += i
+      return x
+
+    def f8(x, f):
+      with f:
+        x += 1
+
+    body = get_live_annotated_fnbody(f1)
+    self._check_anno_matches(body[1], 'live_in', ('a', 'h', 'x'))
+    self._check_anno_matches(body[2], 'live_in', ('x'))
+    self._check_anno_matches(body[0], 'live_in', ('g', 'h', 'x'))
+    self._check_anno_matches(body[2], 'live_out', ())
+
+    body = get_live_annotated_fnbody(f2)
+    self._check_anno_matches(body[0], 'live_in', ('a'))
+    self._check_anno_matches(body[1], 'live_in', ('a'))
+
+    body = get_live_annotated_fnbody(f3)
+    self._check_anno_matches(body[0], 'live_in', ('a', 'x'))
+    self._check_anno_matches(body[1], 'live_in', ('a', 'x'))
+    self._check_anno_matches(body[2], 'live_in', ('x'))
+
+    body = get_live_annotated_fnbody(f4)
+    self._check_anno_matches(body[0], 'live_in', ('x', 'a'))
+    self._check_anno_matches(body[1], 'live_in', ('x'))
+
+    body = get_live_annotated_fnbody(f5)
+    self._check_anno_matches(body[0], 'live_in', ('x', 'x.y', 'a'))
+
+    body = get_live_annotated_fnbody(f6)
+    self._check_anno_matches(body[0], 'live_in', ('x'))
+
+    body = get_live_annotated_fnbody(f7)
+    self._check_anno_matches(body[0], 'live_in', ('x', 'n', 'range'))
+    self._check_anno_matches(body[1], 'live_in', ('x'))
+
+    body = get_live_annotated_fnbody(f8)
+    self._check_anno_matches(body[0], 'live_in', ('f', 'x'))
+
+  def test_node_equality(self):
+    node_a = gast.parse('y = x').body[0]
+    node_b = gast.parse('y = x').body[0]
+    self.assertNotEqual(node_a, node_b)
+
+  def test_nested_functions_defined(self):
+
+    def f(x):
+      y = x * 2
+
+      def g(z):
+        return z + y
+
+      return g(x)
+
+    node, ctx = self._parse_and_analyze(f, {})
+    cfg.run_analyses(node, cfg.Defined(ctx))
+
+    body = node.body[0].body
+    self.assertEqual(
+        anno.getanno(body[2], 'defined_in'),
+        frozenset(map(qual_names.QN, ('g', 'x', 'y'))))
+
+    # TODO(alexbw): CFG analysis doesn't currently cross FunctionDef boundaries.
+    # NOTE: 'z' is easy to find, but 'y' is  not identified as
+    # defined, because CFG analysis is applied with each function separately.
+    # fndef_body = body[1].body
+    # self.assertEqual(
+    #     anno.getanno(fndef_body[0], 'defined_in'),
+    #     frozenset(map(qual_names.QN, ('z', 'y'))))
+
+  def test_nested_functions_dont_leak_definitions(self):
+
+    def f(x):
+      print(x)
+
+      def g():
+        y = 2
+        return y
+
+      return g()  # y is not defined here
+
+    node, ctx = self._parse_and_analyze(f, {})
+    cfg.run_analyses(node, cfg.Defined(ctx))
+    body = node.body[0].body
+    self.assertEqual(
+        anno.getanno(body[2], 'defined_in'),
+        frozenset(map(qual_names.QN, ('x', 'g'))))
+
+
+if __name__ == '__main__':
+  test.main()
-- 
GitLab


From a0bdd4bcfbf300cf36c753591a6febb55143dc06 Mon Sep 17 00:00:00 2001
From: Abhijit Karmarkar <awk@google.com>
Date: Mon, 7 May 2018 10:27:50 -0700
Subject: [PATCH 299/395] Internal Change

PiperOrigin-RevId: 195681946
---
 tensorflow/core/platform/default/build_config.bzl | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/platform/default/build_config.bzl b/tensorflow/core/platform/default/build_config.bzl
index 107c38114b..f6e09ef094 100644
--- a/tensorflow/core/platform/default/build_config.bzl
+++ b/tensorflow/core/platform/default/build_config.bzl
@@ -335,6 +335,8 @@ def tf_proto_library_cc(name, srcs = [], has_services = None,
         name = cc_name,
         deps = cc_deps + ["@protobuf_archive//:protobuf_headers"] +
                if_static([name + "_cc_impl"]),
+        testonly = testonly,
+        visibility = visibility,
     )
     native.cc_library(
         name = cc_name + "_impl",
@@ -378,8 +380,10 @@ def tf_proto_library_py(name, srcs=[], protodeps=[], deps=[], visibility=[],
     )
     native.py_library(
         name = py_name,
-        deps = py_deps + ["@protobuf_archive//:protobuf_python"])
-
+        deps = py_deps + ["@protobuf_archive//:protobuf_python"],
+        testonly = testonly,
+        visibility = visibility,
+    )
     return
 
   py_proto_library(
-- 
GitLab


From ac630df3cb48c6b4acaf9eaf8190aa71f357596f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 10:42:13 -0700
Subject: [PATCH 300/395] Removing quotations to fix the broken link found on
 https://www.tensorflow.org/programmers_guide/embedding

The link at "Follow this link to see a fun example of thumbnail images in the Embedding Projector." It should go to https://www.tensorflow.org/images/embedding-mnist.mp4 but instead goes to the TF index page.

PiperOrigin-RevId: 195684456
---
 tensorflow/docs_src/programmers_guide/embedding.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/docs_src/programmers_guide/embedding.md b/tensorflow/docs_src/programmers_guide/embedding.md
index d5703e0737..8a98367dfb 100644
--- a/tensorflow/docs_src/programmers_guide/embedding.md
+++ b/tensorflow/docs_src/programmers_guide/embedding.md
@@ -238,7 +238,7 @@ row doesn't have to be filled, as shown below.
 </tr>
 </table>
 
-Follow [this link]("https://www.tensorflow.org/images/embedding-mnist.mp4" )
+Follow [this link](https://www.tensorflow.org/images/embedding-mnist.mp4)
 to see a fun example of thumbnail images in the Embedding Projector.
 
 
-- 
GitLab


From b2888c66e67d584756bb50850ae77acede7ba8bf Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 10:47:35 -0700
Subject: [PATCH 301/395] Add EvaluateNodes to HoistFactorDiv test.

PiperOrigin-RevId: 195685340
---
 .../grappler/optimizers/arithmetic_optimizer_test.cc   | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
index e109e66633..741cc135a1 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
@@ -696,6 +696,9 @@ TEST_F(ArithmeticOptimizerTest, HoistFactorDiv) {
         item.fetch = {"id"};
         TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+        auto tensors_expected = EvaluateNodes(item.graph, item.fetch);
+        EXPECT_EQ(1, tensors_expected.size());
+
         ArithmeticOptimizer optimizer;
         EnableOnlyHoistCommonFactor(&optimizer);
 
@@ -734,6 +737,13 @@ TEST_F(ArithmeticOptimizerTest, HoistFactorDiv) {
           EXPECT_EQ("id", id_node->name());
           EXPECT_EQ(HoistDivName("add"), id_node->input(0));
         }
+        auto tensors = EvaluateNodes(output, item.fetch);
+        EXPECT_EQ(1, tensors.size());
+        if (use_ints) {
+          test::ExpectTensorEqual<int32>(tensors_expected[0], tensors[0]);
+        } else {
+          test::ExpectTensorNear<float>(tensors_expected[0], tensors[0], 1e-6);
+        }
       }
     }
   }
-- 
GitLab


From 9ba26ca0d59989592051fdb5c7a2caabe4f399f3 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 10:49:26 -0700
Subject: [PATCH 302/395] Extend block sparsity support for TPUs

PiperOrigin-RevId: 195685740
---
 .../contrib/model_pruning/python/pruning.py   | 30 +++++----
 .../model_pruning/python/pruning_utils.py     | 51 +++++++++++++++
 .../python/pruning_utils_test.py              | 62 ++++++++++++++-----
 3 files changed, 116 insertions(+), 27 deletions(-)

diff --git a/tensorflow/contrib/model_pruning/python/pruning.py b/tensorflow/contrib/model_pruning/python/pruning.py
index ea6032e588..4b7af18b33 100644
--- a/tensorflow/contrib/model_pruning/python/pruning.py
+++ b/tensorflow/contrib/model_pruning/python/pruning.py
@@ -396,14 +396,19 @@ class Pruning(object):
                        self._block_pooling_function)
 
     with ops.name_scope(weights.op.name + '_pruning_ops'):
-      abs_weights = math_ops.abs(
-          array_ops.reshape(weights, [
-              1,
-              squeezed_weights.get_shape()[0],
-              squeezed_weights.get_shape()[1], 1
-          ]))
+      abs_weights = math_ops.abs(squeezed_weights)
+
       pool_window = [self._block_dim[0], self._block_dim[1]]
-      pooled_weights = nn_ops.pool(
+      pool_fn = pruning_utils.factorized_pool
+
+      if not self._spec.use_tpu:
+        pool_fn = nn_ops.pool
+        abs_weights = array_ops.reshape(
+            abs_weights,
+            [1, abs_weights.get_shape()[0],
+             abs_weights.get_shape()[1], 1])
+
+      pooled_weights = pool_fn(
           abs_weights,
           window_shape=pool_window,
           pooling_type=self._block_pooling_function,
@@ -411,19 +416,18 @@ class Pruning(object):
           padding='SAME',
           name=weights.op.name + '_pooled')
 
+      if pooled_weights.get_shape().ndims != 2:
+        pooled_weights = array_ops.squeeze(pooled_weights)
+
       smoothed_threshold, new_mask = self._update_mask(pooled_weights,
                                                        threshold)
-
-      reshaped_mask = array_ops.reshape(
-          new_mask,
-          [pooled_weights.get_shape()[1],
-           pooled_weights.get_shape()[2]])
       updated_mask = pruning_utils.kronecker_product(
-          reshaped_mask, array_ops.ones(self._block_dim))
+          new_mask, array_ops.ones(self._block_dim))
       sliced_mask = array_ops.slice(
           updated_mask, [0, 0],
           [squeezed_weights.get_shape()[0],
            squeezed_weights.get_shape()[1]])
+
     return smoothed_threshold, array_ops.reshape(sliced_mask,
                                                  array_ops.shape(weights))
 
diff --git a/tensorflow/contrib/model_pruning/python/pruning_utils.py b/tensorflow/contrib/model_pruning/python/pruning_utils.py
index 56d3dcef20..ef6c6a3f5d 100644
--- a/tensorflow/contrib/model_pruning/python/pruning_utils.py
+++ b/tensorflow/contrib/model_pruning/python/pruning_utils.py
@@ -29,6 +29,7 @@ from tensorflow.python.ops import clip_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import init_ops
 from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
 from tensorflow.python.ops import state_ops
 from tensorflow.python.ops import variable_scope
 
@@ -221,6 +222,56 @@ def compute_cdf(values, value_range, **kwargs):
     return math_ops.div(cdf, math_ops.reduce_max(cdf))
 
 
+def factorized_pool(input_tensor,
+                    window_shape,
+                    pooling_type,
+                    strides,
+                    padding,
+                    name=None):
+  """Performs m x n pooling through a combination of 1xm and 1xn pooling.
+
+  Args:
+    input_tensor: Input tensor. Must be rank 2
+    window_shape: Pooling window shape
+    pooling_type: Either 'MAX' or 'AVG'
+    strides: The stride of the pooling window
+    padding: 'SAME' or 'VALID'.
+    name: Name of the op
+
+  Returns:
+    A rank 2 tensor containing the pooled output
+
+  Raises:
+    ValueError: if the input tensor is not rank 2
+  """
+  if input_tensor.get_shape().ndims != 2:
+    raise ValueError('factorized_pool() accepts tensors of rank 2 only')
+
+  [height, width] = input_tensor.get_shape()
+  with ops.name_scope(name, 'factorized_pool'):
+    input_tensor_aligned = array_ops.reshape(
+        input_tensor, [1, 1, height, width],
+        name=input_tensor.op.name + '_aligned')
+
+    height_pooling = nn_ops.pool(
+        input_tensor_aligned,
+        window_shape=[1, window_shape[0]],
+        pooling_type=pooling_type,
+        strides=[1, strides[0]],
+        padding=padding)
+    swap_height_width = array_ops.transpose(height_pooling, perm=[0, 1, 3, 2])
+
+    width_pooling = nn_ops.pool(
+        swap_height_width,
+        window_shape=[1, window_shape[1]],
+        pooling_type=pooling_type,
+        strides=[1, strides[1]],
+        padding=padding)
+
+  return array_ops.squeeze(
+      array_ops.transpose(width_pooling, perm=[0, 1, 3, 2]))
+
+
 def determine_partitioned_axis(partitioned_variable):
   partitioned_axis = 0
   concatenated_variable_shape = partitioned_variable.get_shape()
diff --git a/tensorflow/contrib/model_pruning/python/pruning_utils_test.py b/tensorflow/contrib/model_pruning/python/pruning_utils_test.py
index 10e1dd0a8e..ccde5b4e8a 100644
--- a/tensorflow/contrib/model_pruning/python/pruning_utils_test.py
+++ b/tensorflow/contrib/model_pruning/python/pruning_utils_test.py
@@ -22,8 +22,10 @@ import numpy as np
 
 from tensorflow.contrib.model_pruning.python import pruning_utils
 from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import init_ops
 from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
 from tensorflow.python.ops import variable_scope
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import test
@@ -31,6 +33,30 @@ from tensorflow.python.platform import test
 
 class PruningUtilsTest(test.TestCase):
 
+  def _compare_cdf(self, values):
+    abs_values = math_ops.abs(values)
+    max_value = math_ops.reduce_max(abs_values)
+    with self.test_session():
+      variables.global_variables_initializer().run()
+      cdf_from_histogram = pruning_utils.compute_cdf_from_histogram(
+          abs_values, [0.0, max_value], nbins=pruning_utils._NBINS)
+      cdf = pruning_utils.compute_cdf(abs_values, [0.0, max_value])
+      self.assertAllEqual(cdf.eval(), cdf_from_histogram.eval())
+
+  def _compare_pooling_methods(self, weights, pooling_kwargs):
+    with self.test_session():
+      variables.global_variables_initializer().run()
+      pooled_weights_tf = array_ops.squeeze(
+          nn_ops.pool(
+              array_ops.reshape(
+                  weights,
+                  [1, weights.get_shape()[0],
+                   weights.get_shape()[1], 1]), **pooling_kwargs))
+      pooled_weights_factorized_pool = pruning_utils.factorized_pool(
+          weights, **pooling_kwargs)
+      self.assertAllClose(pooled_weights_tf.eval(),
+                          pooled_weights_factorized_pool.eval())
+
   def testHistogram(self):
     width = 10
     height = 10
@@ -59,27 +85,35 @@ class PruningUtilsTest(test.TestCase):
       self.assertAllEqual(len(norm_cdf_val), nbins)
       self.assertAllEqual(expected_cdf, norm_cdf_val)
 
-  def _compare_cdf(self, values):
-    abs_values = math_ops.abs(values)
-    max_value = math_ops.reduce_max(abs_values)
-    with self.test_session():
-      variables.global_variables_initializer().run()
-      cdf_from_histogram = pruning_utils.compute_cdf_from_histogram(
-          abs_values, [0.0, max_value], nbins=pruning_utils._NBINS)
-      cdf = pruning_utils.compute_cdf(abs_values, [0.0, max_value])
-      return cdf.eval(), cdf_from_histogram.eval()
-
   def testCDFEquivalence2D(self):
     width = 100
     height = 100
     weights = variable_scope.get_variable("weights", shape=[width, height])
-    cdf_val, cdf_from_histogram_val = self._compare_cdf(weights)
-    self.assertAllEqual(cdf_val, cdf_from_histogram_val)
+    self._compare_cdf(weights)
 
   def testCDFEquivalence4D(self):
     weights = variable_scope.get_variable("weights", shape=[5, 5, 128, 128])
-    cdf_val, cdf_from_histogram_val = self._compare_cdf(weights)
-    self.assertAllEqual(cdf_val, cdf_from_histogram_val)
+    self._compare_cdf(weights)
+
+  def testFactorizedAvgPool(self):
+    weights = variable_scope.get_variable("weights", shape=[1024, 2048])
+    pooling_kwargs = {
+        "window_shape": [2, 4],
+        "pooling_type": "AVG",
+        "strides": [2, 4],
+        "padding": "SAME"
+    }
+    self._compare_pooling_methods(weights, pooling_kwargs)
+
+  def testFactorizedMaxPool(self):
+    weights = variable_scope.get_variable("weights", shape=[1024, 2048])
+    pooling_kwargs = {
+        "window_shape": [2, 4],
+        "pooling_type": "MAX",
+        "strides": [2, 4],
+        "padding": "SAME"
+    }
+    self._compare_pooling_methods(weights, pooling_kwargs)
 
 
 if __name__ == "__main__":
-- 
GitLab


From 170634d5a10a94d3bd12cc794c284eafcf47fa54 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 11:05:56 -0700
Subject: [PATCH 303/395] Replaced calls to tensorflow::StringPiece::ToString
 with std::string conversions. That is, instances of sp.ToString() are
 replaced with std::string(sp).

This will allow tensorflow::StringPiece::ToString to be removed, which is necessary before it can be replaced with absl::string_view.

PiperOrigin-RevId: 195689392
---
 .../compiler/xla/service/hlo_creation_utils.cc  |  2 +-
 .../compiler/xla/service/hlo_graph_dumper.cc    |  2 +-
 .../compiler/xla/service/hlo_instruction.cc     |  6 +++---
 .../compiler/xla/service/hlo_instruction.h      |  2 +-
 .../compiler/xla/service/hlo_pass_pipeline.cc   | 10 +++++-----
 .../service/human_readable_profile_builder.h    |  9 +++++----
 tensorflow/compiler/xla/service/name_uniquer.cc |  2 +-
 .../compiler/xla/service/shape_inference.cc     |  4 ++--
 tensorflow/core/framework/function.cc           |  2 +-
 tensorflow/core/framework/node_def_builder.cc   | 17 +++++++++--------
 tensorflow/core/framework/node_def_util.cc      |  6 +++---
 tensorflow/core/framework/op_def_builder.cc     |  4 ++--
 tensorflow/core/framework/op_gen_lib.cc         |  2 +-
 tensorflow/core/framework/op_kernel.cc          |  2 +-
 .../core/framework/shape_inference_testutil.h   |  2 +-
 tensorflow/core/graph/graph.cc                  |  2 +-
 tensorflow/core/graph/graph_constructor.cc      | 10 +++++-----
 tensorflow/core/graph/graph_constructor_test.cc |  2 +-
 tensorflow/core/graph/graph_def_builder.cc      |  4 ++--
 tensorflow/core/graph/graph_def_builder.h       |  2 +-
 tensorflow/core/graph/graph_partition.cc        |  2 +-
 tensorflow/core/graph/node_builder.cc           |  2 +-
 tensorflow/core/graph/while_context.cc          |  2 +-
 23 files changed, 50 insertions(+), 48 deletions(-)

diff --git a/tensorflow/compiler/xla/service/hlo_creation_utils.cc b/tensorflow/compiler/xla/service/hlo_creation_utils.cc
index 9a89888480..ed3b654851 100644
--- a/tensorflow/compiler/xla/service/hlo_creation_utils.cc
+++ b/tensorflow/compiler/xla/service/hlo_creation_utils.cc
@@ -269,7 +269,7 @@ StatusOr<HloInstruction*> BroadcastZeros(
 StatusOr<std::unique_ptr<HloComputation>> CreateComputationWithSignature(
     ArraySlice<const Shape*> domain, const Shape& range,
     tensorflow::StringPiece name) {
-  HloComputation::Builder b(name.ToString());
+  HloComputation::Builder b{std::string(name)};
   int64 param_idx = 0;
   for (const Shape* param_shape : domain) {
     b.AddInstruction(HloInstruction::CreateParameter(
diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc
index 794f1b4682..b6b0387672 100644
--- a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc
+++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc
@@ -325,7 +325,7 @@ class HloDotDumper {
                bool show_backend_config, const HloExecutionProfile* profile,
                NodeFilter filter)
       : computation_(computation),
-        label_(label.ToString()),
+        label_(std::string(label)),
         debug_options_(debug_options),
         show_metadata_(show_metadata),
         show_backend_config_(show_backend_config),
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.cc b/tensorflow/compiler/xla/service/hlo_instruction.cc
index 2c733726a6..f9189077a1 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.cc
+++ b/tensorflow/compiler/xla/service/hlo_instruction.cc
@@ -438,7 +438,7 @@ HloInstruction::CreateCrossReplicaSum(
       << "Outfeed shape " << shape << " must be compatible with operand shape "
       << operand->shape();
   instruction->AppendOperand(operand);
-  instruction->outfeed_config_ = outfeed_config.ToString();
+  instruction->outfeed_config_ = std::string(outfeed_config);
   instruction->outfeed_shape_ = shape;
   return instruction;
 }
@@ -1168,7 +1168,7 @@ bool HloInstruction::HasSideEffect() const {
   for (auto operand : operands) {
     instruction->AppendOperand(operand);
   }
-  instruction->custom_call_target_ = custom_call_target.ToString();
+  instruction->custom_call_target_ = std::string(custom_call_target);
   return instruction;
 }
 
@@ -1180,7 +1180,7 @@ bool HloInstruction::HasSideEffect() const {
   for (auto operand : operands) {
     instruction->AppendOperand(operand);
   }
-  instruction->channel_name_ = channel_name.ToString();
+  instruction->channel_name_ = std::string(channel_name);
   instruction->cost_estimate_ns_ = cost_estimate_ns;
   return instruction;
 }
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h
index 19c8c11453..0bf2c589e4 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.h
+++ b/tensorflow/compiler/xla/service/hlo_instruction.h
@@ -1264,7 +1264,7 @@ class HloInstruction {
 
   // Gets/sets the string identifier for this instruction.
   const string& name() const { return name_; }
-  void set_name(tensorflow::StringPiece name) { name_ = name.ToString(); }
+  void set_name(tensorflow::StringPiece name) { name_ = std::string(name); }
 
   // Use the given NameUniquer to select a unique name for the instruction based
   // on the instruction's existing name.
diff --git a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc
index 5120775737..d8f1ab916b 100644
--- a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc
+++ b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc
@@ -90,7 +90,7 @@ StatusOr<bool> HloPassPipeline::Run(HloModule* module) {
     return Status::OK();
   };
 
-  string prefix = name().ToString() + ": pipeline start";
+  string prefix = std::string(name()) + ": pipeline start";
   bool changed = false;
   string message;
   TF_RETURN_IF_ERROR(
@@ -98,12 +98,12 @@ StatusOr<bool> HloPassPipeline::Run(HloModule* module) {
   const string xla_dump_per_pass_hlo_proto_to =
       module->config().debug_options().xla_dump_per_pass_hlo_proto_to();
   if (!xla_dump_per_pass_hlo_proto_to.empty()) {
-    DumpModuleProto(*module, xla_dump_per_pass_hlo_proto_to, name().ToString(),
-                    "pipeline_start");
+    DumpModuleProto(*module, xla_dump_per_pass_hlo_proto_to,
+                    std::string(name()), "pipeline_start");
   }
 
   for (auto& pass : passes_) {
-    if (disabled_passes.count(pass->name().ToString()) > 0) {
+    if (disabled_passes.count(std::string(pass->name())) > 0) {
       VLOG(1) << "  Skipping HLO pass " << pass->name()
               << ", disabled by --xla_disable_hlo_passes";
       continue;
@@ -121,7 +121,7 @@ StatusOr<bool> HloPassPipeline::Run(HloModule* module) {
         run_invariant_checkers(StrCat("after running pass: ", pass->name())));
     if (!xla_dump_per_pass_hlo_proto_to.empty()) {
       DumpModuleProto(*module, xla_dump_per_pass_hlo_proto_to,
-                      name().ToString(), pass->name().ToString());
+                      std::string(name()), std::string(pass->name()));
     }
 
     changed |= changed_this_pass;
diff --git a/tensorflow/compiler/xla/service/human_readable_profile_builder.h b/tensorflow/compiler/xla/service/human_readable_profile_builder.h
index fc24acd271..fb36d3a0d6 100644
--- a/tensorflow/compiler/xla/service/human_readable_profile_builder.h
+++ b/tensorflow/compiler/xla/service/human_readable_profile_builder.h
@@ -32,7 +32,7 @@ class HumanReadableProfileBuilder {
   explicit HumanReadableProfileBuilder(tensorflow::StringPiece computation_name,
                                        int64 total_cycles,
                                        double clock_rate_ghz)
-      : computation_name_(computation_name.ToString()),
+      : computation_name_(std::string(computation_name)),
         total_cycles_(total_cycles),
         clock_rate_ghz_(clock_rate_ghz) {
     CHECK_GE(clock_rate_ghz, 1e-9);
@@ -47,9 +47,10 @@ class HumanReadableProfileBuilder {
              tensorflow::StringPiece category, int64 cycles, int64 flop_count,
              int64 transcendental_count, int64 bytes_accessed,
              float optimal_seconds) {
-    op_infos_.push_back(
-        {op_name.ToString(), short_name.ToString(), category.ToString(), cycles,
-         flop_count, transcendental_count, bytes_accessed, optimal_seconds});
+    op_infos_.push_back({std::string(op_name), std::string(short_name),
+                         std::string(category), cycles, flop_count,
+                         transcendental_count, bytes_accessed,
+                         optimal_seconds});
   }
 
   // Gets the human-readable profile.
diff --git a/tensorflow/compiler/xla/service/name_uniquer.cc b/tensorflow/compiler/xla/service/name_uniquer.cc
index f74bcb0b79..3a6a7c25f4 100644
--- a/tensorflow/compiler/xla/service/name_uniquer.cc
+++ b/tensorflow/compiler/xla/service/name_uniquer.cc
@@ -53,7 +53,7 @@ NameUniquer::NameUniquer(const string& separator) {
 }
 
 string NameUniquer::GetUniqueName(tensorflow::StringPiece prefix) {
-  string root = GetSanitizedName(prefix.empty() ? "name" : prefix.ToString());
+  string root = GetSanitizedName(prefix.empty() ? "name" : std::string(prefix));
 
   // Strip away numeric suffix (if any). Only recognize separator if it is in
   // the middle of the name.
diff --git a/tensorflow/compiler/xla/service/shape_inference.cc b/tensorflow/compiler/xla/service/shape_inference.cc
index 48b2922e77..c493547d9e 100644
--- a/tensorflow/compiler/xla/service/shape_inference.cc
+++ b/tensorflow/compiler/xla/service/shape_inference.cc
@@ -172,11 +172,11 @@ tensorflow::Status ExpectNotTupleOrOpaque(const Shape& shape,
                                           tensorflow::StringPiece op_type) {
   if (ShapeUtil::IsTuple(shape)) {
     return InvalidArgument("Expected non-tuple argument for %s, but got %s.",
-                           op_type.ToString().c_str(),
+                           std::string(op_type).c_str(),
                            ShapeUtil::HumanString(shape).c_str());
   } else if (ShapeUtil::IsOpaque(shape)) {
     return InvalidArgument("Expected non-opaque argument for %s, but got %s.",
-                           op_type.ToString().c_str(),
+                           std::string(op_type).c_str(),
                            ShapeUtil::HumanString(shape).c_str());
   } else {
     return tensorflow::Status::OK();
diff --git a/tensorflow/core/framework/function.cc b/tensorflow/core/framework/function.cc
index bdc1af9fda..647c66099c 100644
--- a/tensorflow/core/framework/function.cc
+++ b/tensorflow/core/framework/function.cc
@@ -504,7 +504,7 @@ string Print(const NodeDef& n) {
   std::vector<string> dep;
   for (StringPiece s : n.input()) {
     if (str_util::ConsumePrefix(&s, "^")) {
-      dep.push_back(s.ToString());
+      dep.push_back(std::string(s));
     } else {
       dat.push_back(s);
     }
diff --git a/tensorflow/core/framework/node_def_builder.cc b/tensorflow/core/framework/node_def_builder.cc
index f9cf6ce873..8e00bfe4f8 100644
--- a/tensorflow/core/framework/node_def_builder.cc
+++ b/tensorflow/core/framework/node_def_builder.cc
@@ -24,22 +24,23 @@ limitations under the License.
 namespace tensorflow {
 
 NodeDefBuilder::NodeOut::NodeOut(StringPiece n, int i, DataType dt)
-    : node(n.ToString()), index(i), data_type(dt) {}
+    : node(std::string(n)), index(i), data_type(dt) {}
 
 NodeDefBuilder::NodeOut::NodeOut() {
   // uninitialized, call Reset() before use.
 }
 
 void NodeDefBuilder::NodeOut::Reset(StringPiece n, int i, DataType dt) {
-  node = n.ToString();
+  node = std::string(n);
   index = i;
   data_type = dt;
 }
 
 NodeDefBuilder::NodeDefBuilder(StringPiece name, StringPiece op_name,
                                const OpRegistryInterface* op_registry) {
-  node_def_.set_name(name.ToString());
-  const Status status = op_registry->LookUpOpDef(op_name.ToString(), &op_def_);
+  node_def_.set_name(std::string(name));
+  const Status status =
+      op_registry->LookUpOpDef(std::string(op_name), &op_def_);
   if (status.ok()) {
     Initialize();
   } else {
@@ -50,7 +51,7 @@ NodeDefBuilder::NodeDefBuilder(StringPiece name, StringPiece op_name,
 
 NodeDefBuilder::NodeDefBuilder(StringPiece name, const OpDef* op_def)
     : op_def_(op_def) {
-  node_def_.set_name(name.ToString());
+  node_def_.set_name(std::string(name));
   Initialize();
 }
 
@@ -170,7 +171,7 @@ void NodeDefBuilder::AddInput(StringPiece src_node, int src_index) {
   } else if (src_index > 0) {
     node_def_.add_input(strings::StrCat(src_node, ":", src_index));
   } else {
-    node_def_.add_input(src_node.ToString());
+    node_def_.add_input(std::string(src_node));
   }
 }
 
@@ -193,12 +194,12 @@ void NodeDefBuilder::VerifyInputRef(const OpDef::ArgDef* input_arg,
 }
 
 NodeDefBuilder& NodeDefBuilder::ControlInput(StringPiece src_node) {
-  control_inputs_.push_back(src_node.ToString());
+  control_inputs_.push_back(std::string(src_node));
   return *this;
 }
 
 NodeDefBuilder& NodeDefBuilder::Device(StringPiece device_spec) {
-  node_def_.set_device(device_spec.ToString());
+  node_def_.set_device(std::string(device_spec));
   return *this;
 }
 
diff --git a/tensorflow/core/framework/node_def_util.cc b/tensorflow/core/framework/node_def_util.cc
index bad92ca9b3..5798333dfe 100644
--- a/tensorflow/core/framework/node_def_util.cc
+++ b/tensorflow/core/framework/node_def_util.cc
@@ -245,7 +245,7 @@ DEFINE_GET_ATTR(NameAttrList, func, "func", emplace_back, v, ;);
 #undef DEFINE_GET_ATTR
 
 bool HasNodeAttr(const NodeDef& node_def, StringPiece attr_name) {
-  return node_def.attr().find(attr_name.ToString()) != node_def.attr().end();
+  return node_def.attr().find(std::string(attr_name)) != node_def.attr().end();
 }
 
 static const string& kEmptyString = *new string();
@@ -639,7 +639,7 @@ Status AttachDef(const Status& status, const Node& node) {
 
 void AddNodeAttr(StringPiece name, const AttrValue& value, NodeDef* node_def) {
   node_def->mutable_attr()->insert(
-      AttrValueMap::value_type(name.ToString(), value));
+      AttrValueMap::value_type(std::string(name), value));
 }
 
 #define ADD_NODE_ATTR(T)                                           \
@@ -677,7 +677,7 @@ ADD_NODE_ATTR(gtl::ArraySlice<NameAttrList>)
 #undef ADD_NODE_ATTR
 
 void AddAttr(StringPiece name, const AttrValue& value, AttrValueMap* map) {
-  map->insert(AttrValueMap::value_type(name.ToString(), value));
+  map->insert(AttrValueMap::value_type(std::string(name), value));
 }
 
 #define ADD_ATTR(T)                                            \
diff --git a/tensorflow/core/framework/op_def_builder.cc b/tensorflow/core/framework/op_def_builder.cc
index 403bd0b5e2..91eb6c0672 100644
--- a/tensorflow/core/framework/op_def_builder.cc
+++ b/tensorflow/core/framework/op_def_builder.cc
@@ -527,7 +527,7 @@ void FinalizeDoc(const string& text, OpDef* op_def,
 }  // namespace
 
 OpDefBuilder::OpDefBuilder(StringPiece op_name) {
-  op_def()->set_name(op_name.ToString());  // NOLINT
+  op_def()->set_name(std::string(op_name));  // NOLINT
 }
 
 OpDefBuilder& OpDefBuilder::Attr(StringPiece spec) {
@@ -584,7 +584,7 @@ OpDefBuilder& OpDefBuilder::Deprecated(int version, StringPiece explanation) {
   } else {
     OpDeprecation* deprecation = op_def()->mutable_deprecation();
     deprecation->set_version(version);
-    deprecation->set_explanation(explanation.ToString());
+    deprecation->set_explanation(std::string(explanation));
   }
   return *this;
 }
diff --git a/tensorflow/core/framework/op_gen_lib.cc b/tensorflow/core/framework/op_gen_lib.cc
index 7f23272871..5e14043625 100644
--- a/tensorflow/core/framework/op_gen_lib.cc
+++ b/tensorflow/core/framework/op_gen_lib.cc
@@ -185,7 +185,7 @@ static bool FindMultiline(StringPiece line, size_t colon, string* end) {
   while (str_util::ConsumePrefix(&line, " ")) {
   }
   if (str_util::ConsumePrefix(&line, "<<")) {
-    *end = line.ToString();
+    *end = std::string(line);
     return true;
   }
   return false;
diff --git a/tensorflow/core/framework/op_kernel.cc b/tensorflow/core/framework/op_kernel.cc
index ca91d68f79..c71bcb26ab 100644
--- a/tensorflow/core/framework/op_kernel.cc
+++ b/tensorflow/core/framework/op_kernel.cc
@@ -923,7 +923,7 @@ void OpKernelContext::clear_recorded_memory() {
 struct KernelRegistration {
   KernelRegistration(const KernelDef& d, StringPiece c,
                      kernel_factory::OpKernelRegistrar::Factory f)
-      : def(d), kernel_class_name(c.ToString()), factory(f) {}
+      : def(d), kernel_class_name(std::string(c)), factory(f) {}
   const KernelDef def;
   const string kernel_class_name;
   const kernel_factory::OpKernelRegistrar::Factory factory;
diff --git a/tensorflow/core/framework/shape_inference_testutil.h b/tensorflow/core/framework/shape_inference_testutil.h
index 2a99af7659..f6656b3b45 100644
--- a/tensorflow/core/framework/shape_inference_testutil.h
+++ b/tensorflow/core/framework/shape_inference_testutil.h
@@ -32,7 +32,7 @@ class Tensor;
 
 struct ShapeInferenceTestOp {
   typedef std::pair<string, DataType> ShapeAndType;
-  explicit ShapeInferenceTestOp(StringPiece name) : name(name.ToString()) {}
+  explicit ShapeInferenceTestOp(StringPiece name) : name(std::string(name)) {}
   string name;
   NodeDef node_def;
   std::vector<const Tensor*> input_tensors;
diff --git a/tensorflow/core/graph/graph.cc b/tensorflow/core/graph/graph.cc
index eeb6c60f71..71d0637dc2 100644
--- a/tensorflow/core/graph/graph.cc
+++ b/tensorflow/core/graph/graph.cc
@@ -695,7 +695,7 @@ Status Graph::AddWhileContext(StringPiece frame_name,
                               std::vector<OutputTensor> body_outputs,
                               WhileContext** result) {
   auto pair = while_ctxs_.insert(std::pair<string, WhileContext>(
-      frame_name.ToString(),
+      std::string(frame_name),
       WhileContext(frame_name, std::move(enter_nodes), std::move(exit_nodes),
                    cond_output, std::move(body_inputs),
                    std::move(body_outputs))));
diff --git a/tensorflow/core/graph/graph_constructor.cc b/tensorflow/core/graph/graph_constructor.cc
index c678283fce..2fd32c0bd4 100644
--- a/tensorflow/core/graph/graph_constructor.cc
+++ b/tensorflow/core/graph/graph_constructor.cc
@@ -489,7 +489,7 @@ Status GraphConstructor::InitFromEdges() {
           num_control_edges++;
         } else {
           TensorId id(ParseTensorName(input_name));
-          if (next_iteration_nodes_.find(id.first.ToString()) !=
+          if (next_iteration_nodes_.find(std::string(id.first)) !=
               next_iteration_nodes_.end()) {
             has_loop_back_edge = true;
           }
@@ -811,7 +811,7 @@ void GraphConstructor::UniquifyNames(
     // We require that UniquifyNames() is called on all NodeDefs in topological
     // order. This guarantees that node_def's inputs will already be uniquified
     // if necessary.
-    auto iter = uniquified_names_.find(id.first.ToString());
+    auto iter = uniquified_names_.find(std::string(id.first));
     if (iter == uniquified_names_.end()) continue;
     id.first = iter->second;
     node_def->set_input(i, id.ToString());
@@ -830,7 +830,7 @@ void GraphConstructor::UpdateUniquifiedColocationNames() {
     for (int i = 0; i < coloc_values.size(); ++i) {
       StringPiece val(coloc_values[i]);
       if (str_util::ConsumePrefix(&val, kColocationGroupPrefix)) {
-        const auto& name_pair = uniquified_names_.find(val.ToString());
+        const auto& name_pair = uniquified_names_.find(std::string(val));
         if (name_pair == uniquified_names_.end()) continue;
         updated = true;
         coloc_values[i] =
@@ -856,7 +856,7 @@ bool GraphConstructor::NameExistsInGraphDef(StringPiece name) {
 }
 
 string GraphConstructor::FindUniqueName(StringPiece original_name) {
-  string name = original_name.ToString();
+  string name = std::string(original_name);
   int count = 0;
   // Check that any generated names don't collide with imported NodeDefs (as
   // well as nodes in g_).
@@ -989,7 +989,7 @@ Status GraphConstructor::Convert() {
             src_node->num_outputs(), " outputs");
       }
 
-      inputs.emplace_back(id.first.ToString(), src_node, src_index);
+      inputs.emplace_back(std::string(id.first), src_node, src_index);
     }
 
     if (has_data_back_edge && !IsMerge(*node_def)) {
diff --git a/tensorflow/core/graph/graph_constructor_test.cc b/tensorflow/core/graph/graph_constructor_test.cc
index b513778de9..c54b4fa269 100644
--- a/tensorflow/core/graph/graph_constructor_test.cc
+++ b/tensorflow/core/graph/graph_constructor_test.cc
@@ -157,7 +157,7 @@ class GraphConstructorTest : public ::testing::Test {
     }
     StringPiece loc(value[0]);
     return str_util::ConsumePrefix(&loc, kColocationGroupPrefix)
-               ? loc.ToString()
+               ? std::string(loc)
                : "";
   }
 
diff --git a/tensorflow/core/graph/graph_def_builder.cc b/tensorflow/core/graph/graph_def_builder.cc
index 7a58347bd1..dd84c4f7c7 100644
--- a/tensorflow/core/graph/graph_def_builder.cc
+++ b/tensorflow/core/graph/graph_def_builder.cc
@@ -44,12 +44,12 @@ GraphDefBuilder::Options GraphDefBuilder::Options::WithControlInputs(
 }
 GraphDefBuilder::Options GraphDefBuilder::Options::WithNameImpl(
     StringPiece name) {
-  name_ = name.ToString();
+  name_ = std::string(name);
   return *this;
 }
 GraphDefBuilder::Options GraphDefBuilder::Options::WithDeviceImpl(
     StringPiece device) {
-  device_ = device.ToString();
+  device_ = std::string(device);
   return *this;
 }
 GraphDefBuilder::Options GraphDefBuilder::Options::WithControlInputImpl(
diff --git a/tensorflow/core/graph/graph_def_builder.h b/tensorflow/core/graph/graph_def_builder.h
index 776a74c6d8..0d6aae4355 100644
--- a/tensorflow/core/graph/graph_def_builder.h
+++ b/tensorflow/core/graph/graph_def_builder.h
@@ -128,7 +128,7 @@ class GraphDefBuilder {
     Options WithControlInputsImpl(gtl::ArraySlice<Node*> control_inputs);
     template <class T>
     Options WithAttrImpl(StringPiece name, T&& value) {
-      attrs_.emplace_back(name.ToString(), AttrValue());
+      attrs_.emplace_back(std::string(name), AttrValue());
       SetAttrValue(std::forward<T>(value), &attrs_.back().second);
       return *this;
     }
diff --git a/tensorflow/core/graph/graph_partition.cc b/tensorflow/core/graph/graph_partition.cc
index 877e4f1b44..1b1941f9c1 100644
--- a/tensorflow/core/graph/graph_partition.cc
+++ b/tensorflow/core/graph/graph_partition.cc
@@ -785,7 +785,7 @@ Status TopologicalSortNodesWithTimePriority(
   for (int n = 0; n < gdef->node_size(); ++n) {
     const NodeDef* ndef = &gdef->node(n);
     for (int i = 0; i < ndef->input_size(); ++i) {
-      node_to_output_nodes[ParseTensorName(ndef->input(i)).first.ToString()]
+      node_to_output_nodes[std::string(ParseTensorName(ndef->input(i)).first)]
           .push_back(ndef);
     }
     int64 start_time;
diff --git a/tensorflow/core/graph/node_builder.cc b/tensorflow/core/graph/node_builder.cc
index 114962c0e4..03f3bbd663 100644
--- a/tensorflow/core/graph/node_builder.cc
+++ b/tensorflow/core/graph/node_builder.cc
@@ -30,7 +30,7 @@ NodeBuilder::NodeOut::NodeOut(Node* n, int32 i)  // NOLINT(runtime/explicit)
       dt(SafeGetOutput(node, i, &error)) {}
 
 NodeBuilder::NodeOut::NodeOut(StringPiece n, int32 i, DataType t)
-    : node(nullptr), error(false), name(n.ToString()), index(i), dt(t) {}
+    : node(nullptr), error(false), name(std::string(n)), index(i), dt(t) {}
 
 NodeBuilder::NodeOut::NodeOut()
     : node(nullptr), error(true), index(0), dt(DT_FLOAT) {}
diff --git a/tensorflow/core/graph/while_context.cc b/tensorflow/core/graph/while_context.cc
index 10a2b67f37..1b38aac35d 100644
--- a/tensorflow/core/graph/while_context.cc
+++ b/tensorflow/core/graph/while_context.cc
@@ -23,7 +23,7 @@ WhileContext::WhileContext(StringPiece frame_name,
                            OutputTensor cond_output,
                            std::vector<OutputTensor> body_inputs,
                            std::vector<OutputTensor> body_outputs)
-    : frame_name_(frame_name.ToString()),
+    : frame_name_(std::string(frame_name)),
       enter_nodes_(std::move(enter_nodes)),
       exit_nodes_(std::move(exit_nodes)),
       cond_output_(cond_output),
-- 
GitLab


From f6a55cc344cd96098cabd500144aad266e692598 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 11:09:47 -0700
Subject: [PATCH 304/395] Add tests for broadcasting KL divergence
 calculations.

PiperOrigin-RevId: 195690035
---
 .../kernel_tests/mvn_full_covariance_test.py  | 31 ++++++++++++++-
 .../python/kernel_tests/mvn_tril_test.py      | 39 ++++++++++++++++---
 2 files changed, 62 insertions(+), 8 deletions(-)

diff --git a/tensorflow/contrib/distributions/python/kernel_tests/mvn_full_covariance_test.py b/tensorflow/contrib/distributions/python/kernel_tests/mvn_full_covariance_test.py
index 7435bcbc68..b003526392 100644
--- a/tensorflow/contrib/distributions/python/kernel_tests/mvn_full_covariance_test.py
+++ b/tensorflow/contrib/distributions/python/kernel_tests/mvn_full_covariance_test.py
@@ -131,8 +131,8 @@ class MultivariateNormalFullCovarianceTest(test.TestCase):
     return mu, sigma
 
   def testKLBatch(self):
-    batch_shape = (2,)
-    event_shape = (3,)
+    batch_shape = [2]
+    event_shape = [3]
     with self.test_session():
       mu_a, sigma_a = self._random_mu_and_sigma(batch_shape, event_shape)
       mu_b, sigma_b = self._random_mu_and_sigma(batch_shape, event_shape)
@@ -156,6 +156,33 @@ class MultivariateNormalFullCovarianceTest(test.TestCase):
       self.assertAllClose(expected_kl_0, kl_v[0])
       self.assertAllClose(expected_kl_1, kl_v[1])
 
+  def testKLBatchBroadcast(self):
+    batch_shape = [2]
+    event_shape = [3]
+    with self.test_session():
+      mu_a, sigma_a = self._random_mu_and_sigma(batch_shape, event_shape)
+      # No batch shape.
+      mu_b, sigma_b = self._random_mu_and_sigma([], event_shape)
+      mvn_a = ds.MultivariateNormalFullCovariance(
+          loc=mu_a,
+          covariance_matrix=sigma_a,
+          validate_args=True)
+      mvn_b = ds.MultivariateNormalFullCovariance(
+          loc=mu_b,
+          covariance_matrix=sigma_b,
+          validate_args=True)
+
+      kl = ds.kl_divergence(mvn_a, mvn_b)
+      self.assertEqual(batch_shape, kl.get_shape())
+
+      kl_v = kl.eval()
+      expected_kl_0 = _compute_non_batch_kl(mu_a[0, :], sigma_a[0, :, :],
+                                            mu_b, sigma_b)
+      expected_kl_1 = _compute_non_batch_kl(mu_a[1, :], sigma_a[1, :, :],
+                                            mu_b, sigma_b)
+      self.assertAllClose(expected_kl_0, kl_v[0])
+      self.assertAllClose(expected_kl_1, kl_v[1])
+
 
 def _compute_non_batch_kl(mu_a, sigma_a, mu_b, sigma_b):
   """Non-batch KL for N(mu_a, sigma_a), N(mu_b, sigma_b)."""
diff --git a/tensorflow/contrib/distributions/python/kernel_tests/mvn_tril_test.py b/tensorflow/contrib/distributions/python/kernel_tests/mvn_tril_test.py
index 685f32883d..b556d06123 100644
--- a/tensorflow/contrib/distributions/python/kernel_tests/mvn_tril_test.py
+++ b/tensorflow/contrib/distributions/python/kernel_tests/mvn_tril_test.py
@@ -235,8 +235,8 @@ class MultivariateNormalTriLTest(test.TestCase):
     return mu, sigma
 
   def testKLNonBatch(self):
-    batch_shape = ()
-    event_shape = (2,)
+    batch_shape = []
+    event_shape = [2]
     with self.test_session():
       mu_a, sigma_a = self._random_mu_and_sigma(batch_shape, event_shape)
       mu_b, sigma_b = self._random_mu_and_sigma(batch_shape, event_shape)
@@ -257,8 +257,8 @@ class MultivariateNormalTriLTest(test.TestCase):
       self.assertAllClose(expected_kl, kl_v)
 
   def testKLBatch(self):
-    batch_shape = (2,)
-    event_shape = (3,)
+    batch_shape = [2]
+    event_shape = [3]
     with self.test_session():
       mu_a, sigma_a = self._random_mu_and_sigma(batch_shape, event_shape)
       mu_b, sigma_b = self._random_mu_and_sigma(batch_shape, event_shape)
@@ -282,9 +282,36 @@ class MultivariateNormalTriLTest(test.TestCase):
       self.assertAllClose(expected_kl_0, kl_v[0])
       self.assertAllClose(expected_kl_1, kl_v[1])
 
+  def testKLBatchBroadcast(self):
+    batch_shape = [2]
+    event_shape = [3]
+    with self.test_session():
+      mu_a, sigma_a = self._random_mu_and_sigma(batch_shape, event_shape)
+      # No batch shape.
+      mu_b, sigma_b = self._random_mu_and_sigma([], event_shape)
+      mvn_a = ds.MultivariateNormalTriL(
+          loc=mu_a,
+          scale_tril=np.linalg.cholesky(sigma_a),
+          validate_args=True)
+      mvn_b = ds.MultivariateNormalTriL(
+          loc=mu_b,
+          scale_tril=np.linalg.cholesky(sigma_b),
+          validate_args=True)
+
+      kl = ds.kl_divergence(mvn_a, mvn_b)
+      self.assertEqual(batch_shape, kl.get_shape())
+
+      kl_v = kl.eval()
+      expected_kl_0 = _compute_non_batch_kl(mu_a[0, :], sigma_a[0, :, :],
+                                            mu_b, sigma_b)
+      expected_kl_1 = _compute_non_batch_kl(mu_a[1, :], sigma_a[1, :, :],
+                                            mu_b, sigma_b)
+      self.assertAllClose(expected_kl_0, kl_v[0])
+      self.assertAllClose(expected_kl_1, kl_v[1])
+
   def testKLTwoIdenticalDistributionsIsZero(self):
-    batch_shape = (2,)
-    event_shape = (3,)
+    batch_shape = [2]
+    event_shape = [3]
     with self.test_session():
       mu_a, sigma_a = self._random_mu_and_sigma(batch_shape, event_shape)
       mvn_a = ds.MultivariateNormalTriL(
-- 
GitLab


From 93846eccdfa9dd6da34b37778e5f3b1a46739933 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 11:11:28 -0700
Subject: [PATCH 305/395] Extracts PartialConcatConstFolding into a method.

PiperOrigin-RevId: 195690333
---
 .../grappler/optimizers/constant_folding.cc   | 203 +++++++++---------
 .../grappler/optimizers/constant_folding.h    |   5 +
 2 files changed, 111 insertions(+), 97 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/constant_folding.cc b/tensorflow/core/grappler/optimizers/constant_folding.cc
index 47d8827686..e6a74dbdcd 100644
--- a/tensorflow/core/grappler/optimizers/constant_folding.cc
+++ b/tensorflow/core/grappler/optimizers/constant_folding.cc
@@ -2370,115 +2370,124 @@ Status ConstantFolding::SimplifyGraph(GraphDef* optimized_graph,
       }
     }
 
-    // Partial constant folding for Concat which is not commutative, so
-    // we have to preserve order and can only push consecutive runs of constant
-    // inputs into sub-nodes.
-    if (IsConcat(*node) && num_non_control_inputs > 3 &&
-        node->name().rfind("_partial_split_") == string::npos) {
-      int axis_arg = -1;
-      int begin = 0;
-      int end = num_non_control_inputs;
-      if (node->op() == "Concat") {
-        begin = 1;
-        axis_arg = 0;
-      } else if (node->op() == "ConcatV2") {
-        end = num_non_control_inputs - 1;
-        axis_arg = num_non_control_inputs - 1;
-      } else {
-        continue;
-      }
+    if (PartialConcatConstFolding(optimized_graph, properties, node)) {
+      graph_modified_ = true;
+      continue;
+    }
+  }
 
-      const NodeDef* axis_arg_node =
-          node_map_->GetNode(NodeName(node->input(axis_arg)));
-      if (axis_arg_node == nullptr || !IsReallyConstant(*axis_arg_node)) {
-        // We cannot constant fold Concat unless we the axis argument is
-        // constant. Skip node.
-        continue;
-      }
+  return Status::OK();
+}
 
-      // We search for consecutive runs of constant inputs in the range
-      // [begin:end[ and push then down into child nodes.
-      std::vector<std::pair<int, int>> constant_input_runs;
-      int first = begin;
-      int last = begin;
-      while (last < end) {
-        while (first < end && !IsReallyConstant(*node_map_->GetNode(
-                                  NodeName(node->input(first))))) {
-          ++first;
-        }
-        // Invariant: node[first] is constant || first >= end.
-        last = first + 1;
-        while (last < end && IsReallyConstant(*node_map_->GetNode(
-                                 NodeName(node->input(last))))) {
-          ++last;
-        }
-        // Invariant: node[last] is not constant || last >= end
-        // Discard intervals shorter than 2 elements.
-        if (first < end && (last - first) > 1) {
-          constant_input_runs.emplace_back(first, last);
-        }
-        first = last;
+bool ConstantFolding::PartialConcatConstFolding(GraphDef* optimized_graph,
+                                                GraphProperties* properties,
+                                                NodeDef* node) {
+  // Partial constant folding for Concat which is not commutative, so
+  // we have to preserve order and can only push consecutive runs of constant
+  // inputs into sub-nodes.
+  const int num_non_control_inputs = NumNonControlInputs(*node);
+  if (IsConcat(*node) && num_non_control_inputs > 3 &&
+      node->name().rfind("_partial_split_") == string::npos) {
+    int axis_arg = -1;
+    int begin = 0;
+    int end = num_non_control_inputs;
+    if (node->op() == "Concat") {
+      begin = 1;
+      axis_arg = 0;
+    } else if (node->op() == "ConcatV2") {
+      end = num_non_control_inputs - 1;
+      axis_arg = num_non_control_inputs - 1;
+    } else {
+      return false;
+    }
+
+    const NodeDef* axis_arg_node =
+        node_map_->GetNode(NodeName(node->input(axis_arg)));
+    if (axis_arg_node == nullptr || !IsReallyConstant(*axis_arg_node)) {
+      // We cannot constant fold Concat unless we the axis argument is
+      // constant. Skip node.
+      return false;
+    }
+
+    // We search for consecutive runs of constant inputs in the range
+    // [begin:end[ and push then down into child nodes.
+    std::vector<std::pair<int, int>> constant_input_runs;
+    int first = begin;
+    int last = begin;
+    while (last < end) {
+      while (first < end && !IsReallyConstant(*node_map_->GetNode(
+                                NodeName(node->input(first))))) {
+        ++first;
+      }
+      // Invariant: node[first] is constant || first >= end.
+      last = first + 1;
+      while (last < end && IsReallyConstant(*node_map_->GetNode(
+                               NodeName(node->input(last))))) {
+        ++last;
       }
+      // Invariant: node[last] is not constant || last >= end
+      // Discard intervals shorter than 2 elements.
+      if (first < end && (last - first) > 1) {
+        constant_input_runs.emplace_back(first, last);
+      }
+      first = last;
+    }
 
-      // Skip if all inputs are constant, and let constant folding take over.
-      if (constant_input_runs.size() == 1 &&
-          constant_input_runs[0].first == begin &&
-          constant_input_runs[0].second == end) {
-        continue;
+    // Skip if all inputs are constant, and let constant folding take over.
+    if (constant_input_runs.size() == 1 &&
+        constant_input_runs[0].first == begin &&
+        constant_input_runs[0].second == end) {
+      return false;
+    }
+    std::set<int> inputs_to_delete;
+    for (auto interval : constant_input_runs) {
+      // Push the constant inputs in the interval to a child node than can be
+      // constant folded.
+      const string new_node_name = OptimizedNodeName(
+          *node, strings::StrCat("_partial_split_", interval.first));
+      if (node_map_->NodeExists(new_node_name)) {
+        break;
       }
-      std::set<int> inputs_to_delete;
-      for (auto interval : constant_input_runs) {
-        // Push the constant inputs in the interval to a child node than can be
-        // constant folded.
-        const string new_node_name = OptimizedNodeName(
-            *node, strings::StrCat("_partial_split_", interval.first));
-        if (node_map_->NodeExists(new_node_name)) {
-          break;
-        }
-        NodeDef* added_node = optimized_graph->add_node();
-        *added_node = *node;
-        added_node->set_name(new_node_name);
-        node_map_->AddNode(added_node->name(), added_node);
-        added_node->clear_input();
-        for (int i = interval.first; i < interval.second; ++i) {
-          added_node->add_input(node->input(i));
-          node_map_->UpdateOutput(NodeName(node->input(i)), node->name(),
-                                  added_node->name());
-          if (i != interval.first) {
-            inputs_to_delete.insert(i);
-          }
+      NodeDef* added_node = optimized_graph->add_node();
+      *added_node = *node;
+      added_node->set_name(new_node_name);
+      node_map_->AddNode(added_node->name(), added_node);
+      added_node->clear_input();
+      for (int i = interval.first; i < interval.second; ++i) {
+        added_node->add_input(node->input(i));
+        node_map_->UpdateOutput(NodeName(node->input(i)), node->name(),
+                                added_node->name());
+        if (i != interval.first) {
+          inputs_to_delete.insert(i);
         }
-        added_node->add_input(node->input(axis_arg));
-        (*added_node->mutable_attr())["N"].set_i(interval.second -
-                                                 interval.first);
-        node_map_->AddOutput(NodeName(node->input(axis_arg)),
-                             added_node->name());
-
-        // Overwrite the first constant input with the result of the added
-        // child node.
-        node->set_input(interval.first, added_node->name());
-        node_map_->AddOutput(added_node->name(), node->name());
       }
-      if (!constant_input_runs.empty()) {
-        graph_modified_ = true;
-        if (!inputs_to_delete.empty()) {
-          // Fix up the inputs to the original node.
-          std::vector<string> tmp(node->input().begin(), node->input().end());
-          node->clear_input();
-          for (int i = 0; i < tmp.size(); ++i) {
-            if (inputs_to_delete.find(i) == inputs_to_delete.end()) {
-              node->add_input(tmp[i]);
-            }
+      added_node->add_input(node->input(axis_arg));
+      (*added_node->mutable_attr())["N"].set_i(interval.second -
+                                               interval.first);
+      node_map_->AddOutput(NodeName(node->input(axis_arg)), added_node->name());
+
+      // Overwrite the first constant input with the result of the added
+      // child node.
+      node->set_input(interval.first, added_node->name());
+      node_map_->AddOutput(added_node->name(), node->name());
+    }
+    if (!constant_input_runs.empty()) {
+      if (!inputs_to_delete.empty()) {
+        // Fix up the inputs to the original node.
+        std::vector<string> tmp(node->input().begin(), node->input().end());
+        node->clear_input();
+        for (int i = 0; i < tmp.size(); ++i) {
+          if (inputs_to_delete.find(i) == inputs_to_delete.end()) {
+            node->add_input(tmp[i]);
           }
-          (*node->mutable_attr())["N"].set_i(node->input_size() - 1);
-          properties->ClearInputProperties(node->name());
         }
-        continue;
+        (*node->mutable_attr())["N"].set_i(node->input_size() - 1);
+        properties->ClearInputProperties(node->name());
       }
+      return true;
     }
   }
-
-  return Status::OK();
+  return false;
 }
 
 Status ConstantFolding::RunOptimizationPass(Cluster* cluster,
diff --git a/tensorflow/core/grappler/optimizers/constant_folding.h b/tensorflow/core/grappler/optimizers/constant_folding.h
index a694f1721a..2096576538 100644
--- a/tensorflow/core/grappler/optimizers/constant_folding.h
+++ b/tensorflow/core/grappler/optimizers/constant_folding.h
@@ -101,6 +101,11 @@ class ConstantFolding : public GraphOptimizer {
   Status RunOptimizationPass(Cluster* cluster, const GrapplerItem& item,
                              GraphDef* output);
 
+  // Applies partial constant folding for Concat which is not commutative.
+  // Returns true if the transformation applied successfully.
+  bool PartialConcatConstFolding(GraphDef* optimized_graph,
+                                 GraphProperties* properties, NodeDef* node);
+
   // Points to an externally provided device or to owned_device_;
   RewriterConfig::Toggle opt_level_;
   DeviceBase* cpu_device_;
-- 
GitLab


From eb5ee79cb3108bb036fc4a6d465f6ef6e12f4a3a Mon Sep 17 00:00:00 2001
From: Yu-Cheng Ling <ycling@google.com>
Date: Mon, 7 May 2018 11:27:02 -0700
Subject: [PATCH 306/395] Release notes for TensorFlow Lite.

PiperOrigin-RevId: 195693362
---
 tensorflow/contrib/lite/RELEASE.md | 8 ++++++++
 1 file changed, 8 insertions(+)
 create mode 100644 tensorflow/contrib/lite/RELEASE.md

diff --git a/tensorflow/contrib/lite/RELEASE.md b/tensorflow/contrib/lite/RELEASE.md
new file mode 100644
index 0000000000..8fd63d5cee
--- /dev/null
+++ b/tensorflow/contrib/lite/RELEASE.md
@@ -0,0 +1,8 @@
+# Release 0.1.7
+
+* TensorFlow Lite 0.1.7 is based on tag `tflite-v0.1.7` (git commit
+  fa1db5eb0da85b5baccc2a46d534fdeb3bb473d0).
+* To reproduce the iOS library, it's required to cherry pick git commit
+  f1f1d5172fe5bfeaeb2cf657ffc43ba744187bee to fix a dependency issue.
+* The code is based on TensorFlow 1.8.0 release candidate and it's very close
+  to TensorFlow 1.8.0 release.
-- 
GitLab


From f14123dc19be468b6776f057d45ddd4d40fef9b2 Mon Sep 17 00:00:00 2001
From: Igor Saprykin <isaprykin@google.com>
Date: Mon, 7 May 2018 11:48:31 -0700
Subject: [PATCH 307/395] Generalize the input to TPU distribution strategy. 
 Add cross-shard-replica sum.

TPUStrategy passes tests in minimize_loss_test.  That caused me to add a capability to have `iterations x cores` inputs of any structure.  I also resolved a big number of small issues and uncovered more things to resolve that are documented as todos.

PiperOrigin-RevId: 195696833
---
 .../contrib/distribute/python/combinations.py |   4 +
 .../distribute/python/minimize_loss_test.py   | 115 ++++++++++++------
 .../distribute/python/single_loss_example.py  |  20 ++-
 .../contrib/distribute/python/tpu_strategy.py |  13 +-
 .../contrib/distribute/python/values.py       |  57 +++++++--
 tensorflow/python/training/distribute.py      |   1 +
 6 files changed, 158 insertions(+), 52 deletions(-)

diff --git a/tensorflow/contrib/distribute/python/combinations.py b/tensorflow/contrib/distribute/python/combinations.py
index 946310aa6f..45d191127e 100644
--- a/tensorflow/contrib/distribute/python/combinations.py
+++ b/tensorflow/contrib/distribute/python/combinations.py
@@ -265,6 +265,10 @@ class NamedDistribution(object):
 one_device_strategy = NamedDistribution(
     "OneDeviceCPU", one_device_strategy.OneDeviceStrategy("/cpu:0"),
     None)
+tpu_strategy_single_iteration = NamedDistribution(
+    "TPUSingleIteration",
+    tpu_strategy.TPUStrategy(iterations_per_step=1),
+    required_tpu=True)
 tpu_strategy = NamedDistribution(
     "TPU", tpu_strategy.TPUStrategy(), required_tpu=True)
 mirrored_strategy_with_gpu_and_cpu = NamedDistribution(
diff --git a/tensorflow/contrib/distribute/python/minimize_loss_test.py b/tensorflow/contrib/distribute/python/minimize_loss_test.py
index e134fe34e1..d2054715f1 100644
--- a/tensorflow/contrib/distribute/python/minimize_loss_test.py
+++ b/tensorflow/contrib/distribute/python/minimize_loss_test.py
@@ -44,13 +44,16 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase):
           combinations.distributions_and_v1_optimizers(),
           combinations.combine(mode=["graph"], use_callable_loss=[True, False])
           + combinations.combine(mode=["eager"], use_callable_loss=[True]),
-          combinations.combine(is_tpu=[False])) +
-      combinations.combine(
-          distribution=[combinations.tpu_strategy],
-          optimizer_fn=[combinations.adam_optimizer_v1_fn],
-          mode=["graph"],
-          use_callable_loss=[False],
-          is_tpu=[True]))
+          combinations.combine(is_tpu=[False])) + combinations.combine(
+              distribution=[combinations.tpu_strategy],
+              optimizer_fn=[
+                  combinations.adam_optimizer_v1_fn,
+                  # TODO(isaprykin):  Make Adam v2 work with while_loops
+                  # and TPUs.
+              ],
+              mode=["graph"],
+              use_callable_loss=[False],
+              is_tpu=[True]))
   def testTrainNetwork(self, distribution, optimizer_fn, use_callable_loss,
                        is_tpu):
     with distribution.scope():
@@ -101,7 +104,8 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase):
           distribution=[combinations.tpu_strategy],
           optimizer_fn=[
               combinations.adam_optimizer_v1_fn,
-              combinations.gradient_descent_optimizer_v1_fn
+              combinations.gradient_descent_optimizer_v1_fn,
+              combinations.gradient_descent_optimizer_v2_fn,
           ],
           mode=["graph"],
           is_tpu=[True]))
@@ -171,13 +175,28 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase):
           set(created_variables))
 
   @combinations.generate(
-      combinations.times(combinations.distributions_and_v1_optimizers(),
-                         combinations.combine(
-                             mode=["graph", "eager"],
-                             momentum=[0.8, 0.9, 0.99],
-                             renorm=[False, True])))
+      combinations.times(
+          combinations.combine(momentum=[0.8, 0.9, 0.99], renorm=[False, True]),
+          combinations.times(
+              combinations.distributions_and_v1_optimizers(),
+              combinations.combine(
+                  mode=["graph", "eager"],
+                  is_tpu=[False],
+                  # TODO(isaprykin):  Allow False here.  Currently subsequent
+                  # towers will re-execute UPDATE_OPS of previous towers.
+                  update_ops_in_cross_tower_mode=[True])) +
+          combinations.combine(
+              distribution=[combinations.tpu_strategy_single_iteration],
+              optimizer_fn=[
+                  combinations.gradient_descent_optimizer_v1_fn,
+                  combinations.gradient_descent_optimizer_v2_fn
+              ],
+              mode=["graph"],
+              is_tpu=[True],
+              update_ops_in_cross_tower_mode=[False])))
   def testTrainNetworkWithBatchNorm(self, distribution, optimizer_fn, momentum,
-                                    renorm):
+                                    renorm, is_tpu,
+                                    update_ops_in_cross_tower_mode):
     """Verifies that moving mean updates are reduced across towers."""
     with distribution.scope():
       num_towers = len(distribution.worker_devices)
@@ -185,7 +204,8 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase):
           optimizer_fn,
           batch_per_epoch=num_towers,
           momentum=momentum,
-          renorm=renorm)
+          renorm=renorm,
+          update_ops_in_tower_mode=not update_ops_in_cross_tower_mode)
 
       # Disable prefetching since that makes the specific input on each device
       # to be non deterministic, and this test relies on specific input being
@@ -196,16 +216,18 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase):
           dataset_fn).make_one_shot_iterator()
 
       def run_step():
-        return control_flow_ops.group(
-            distribution.unwrap(
-                distribution.call_for_each_tower(
-                    model_fn,
-                    iterator.get_next(),
-                    run_concurrently=batchnorm.built)) +
-            ops.get_collection(ops.GraphKeys.UPDATE_OPS))
+        fetches = distribution.unwrap(
+            distribution.call_for_each_tower(
+                model_fn, iterator.get_next(),
+                run_concurrently=batchnorm.built))
+        if update_ops_in_cross_tower_mode:
+          fetches += ops.get_collection(ops.GraphKeys.UPDATE_OPS)
+        return control_flow_ops.group(fetches)
 
       if not context.executing_eagerly():
         with self.test_session() as sess:
+          if is_tpu:
+            sess.run(tpu.initialize_system())
           run_step = sess.make_callable(run_step())
         self.evaluate(variables_lib.global_variables_initializer())
 
@@ -229,22 +251,40 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase):
               expected_moving_mean - averaged_batch_mean(i)) * (1.0 - momentum))
           self.assertNear(expected_moving_means[i], moving_means[i], 0.0001)
 
+      if is_tpu:
+        with self.test_session() as sess:
+          sess.run(tpu.shutdown_system())
+
   @combinations.generate(
       combinations.times(
           combinations.combine(
-              distribution=[combinations.one_device_strategy,
-                            combinations.mirrored_strategy_with_gpu_and_cpu,
-                            combinations.mirrored_strategy_with_two_gpus],
-              optimizer_fn=[combinations.gradient_descent_optimizer_v1_fn,
-                            combinations.gradient_descent_optimizer_v2_fn],
-              loss_reduction=[losses_impl.Reduction.SUM,
-                              losses_impl.Reduction.MEAN,
-                              losses_impl.Reduction.SUM_OVER_BATCH_SIZE,
-                              losses_impl.Reduction.SUM_OVER_NONZERO_WEIGHTS]),
-          combinations.combine(mode=["graph"], use_callable_loss=[True, False])
-          + combinations.combine(mode=["eager"], use_callable_loss=[True])))
+              optimizer_fn=[
+                  combinations.gradient_descent_optimizer_v1_fn,
+                  combinations.gradient_descent_optimizer_v2_fn
+              ],
+              loss_reduction=[
+                  losses_impl.Reduction.SUM, losses_impl.Reduction.MEAN,
+                  losses_impl.Reduction.SUM_OVER_BATCH_SIZE,
+                  losses_impl.Reduction.SUM_OVER_NONZERO_WEIGHTS
+              ]),
+          combinations.times(
+              combinations.combine(
+                  distribution=[
+                      combinations.one_device_strategy,
+                      combinations.mirrored_strategy_with_gpu_and_cpu,
+                      combinations.mirrored_strategy_with_two_gpus
+                  ],
+                  is_tpu=[False]),
+              combinations.combine(
+                  mode=["graph"], use_callable_loss=[True, False]) +
+              combinations.combine(mode=["eager"], use_callable_loss=[True])) +
+          combinations.combine(
+              distribution=[combinations.tpu_strategy_single_iteration],
+              is_tpu=[True],
+              mode=["graph"],
+              use_callable_loss=[True, False])))
   def testMeanVsSum(self, distribution, optimizer_fn, loss_reduction,
-                    use_callable_loss):
+                    use_callable_loss, is_tpu):
     with distribution.scope():
       all_vars = []
 
@@ -280,12 +320,13 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase):
 
       if not context.executing_eagerly():
         with self.test_session() as sess:
+          if is_tpu:
+            sess.run(tpu.initialize_system())
           run_step = sess.make_callable(run_step())
         self.evaluate(variables_lib.global_variables_initializer())
 
       run_step()
 
-      self.assertEqual(distribution.num_towers, len(all_vars))
       v = all_vars[0]
       self.assertTrue(all([v is vi for vi in all_vars[1:]]))
       weight = numpy.squeeze(self.evaluate(distribution.fetch(v)))
@@ -312,6 +353,10 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase):
         # One of the mean loss reductions.
         self.assertNear(weight, 2 + 10.6, 0.0001)
 
+      if is_tpu:
+        with self.test_session() as sess:
+          sess.run(tpu.shutdown_system())
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/contrib/distribute/python/single_loss_example.py b/tensorflow/contrib/distribute/python/single_loss_example.py
index 0db0b59fca..d1fdb3279c 100644
--- a/tensorflow/contrib/distribute/python/single_loss_example.py
+++ b/tensorflow/contrib/distribute/python/single_loss_example.py
@@ -22,6 +22,7 @@ from tensorflow.contrib.data.python.ops import batching
 from tensorflow.contrib.distribute.python import step_fn
 from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
 from tensorflow.python.layers import core
 from tensorflow.python.layers import normalization
 from tensorflow.python.ops import array_ops
@@ -59,7 +60,7 @@ def minimize_loss_example(optimizer_fn,
     # TODO(isaprykin): map_and_batch with drop_remainder causes shapes to be
     # fully defined for TPU.  Remove this when XLA supports dynamic shapes.
     return dataset.apply(
-        batching.map_and_batch(lambda x: x, batch_size=2, drop_remainder=True))
+        batching.map_and_batch(lambda x: x, batch_size=1, drop_remainder=True))
 
   # An Optimizer instance is created either outside or inside model_fn.
   outer_optimizer = None
@@ -68,11 +69,10 @@ def minimize_loss_example(optimizer_fn,
 
   layer = core.Dense(1, use_bias=use_bias)
 
-  def model_fn(xs):
+  def model_fn(x):
     """A very simple model written by the user."""
 
     def loss_fn():
-      x = math_ops.reduce_mean(xs, keepdims=True)
       y = array_ops.reshape(layer(x), []) - constant_op.constant(1.)
       return y * y
 
@@ -89,7 +89,8 @@ def minimize_loss_example(optimizer_fn,
 def batchnorm_example(optimizer_fn,
                       batch_per_epoch=1,
                       momentum=0.9,
-                      renorm=False):
+                      renorm=False,
+                      update_ops_in_tower_mode=False):
   """Example of non-distribution-aware legacy code with batch normalization."""
 
   def dataset_fn():
@@ -103,12 +104,19 @@ def batchnorm_example(optimizer_fn,
   optimizer = optimizer_fn()
   batchnorm = normalization.BatchNormalization(
       renorm=renorm, momentum=momentum, fused=False)
+  layer = core.Dense(1, use_bias=False)
 
   def model_fn(x):
+    """A model that uses batchnorm."""
 
     def loss_fn():
-      y = math_ops.reduce_sum(batchnorm(x, training=True), axis=1)
-      loss = math_ops.reduce_mean(y - constant_op.constant(1.))
+      y = batchnorm(x, training=True)
+      with ops.control_dependencies(
+          ops.get_collection(ops.GraphKeys.UPDATE_OPS)
+          if update_ops_in_tower_mode else []):
+        loss = math_ops.reduce_mean(
+            math_ops.reduce_sum(layer(y)) - constant_op.constant(1.))
+      # `x` and `y` will be fetched by the gradient computation, but not `loss`.
       return loss
 
     # Callable loss.
diff --git a/tensorflow/contrib/distribute/python/tpu_strategy.py b/tensorflow/contrib/distribute/python/tpu_strategy.py
index a7e4fe80f3..75441786a6 100644
--- a/tensorflow/contrib/distribute/python/tpu_strategy.py
+++ b/tensorflow/contrib/distribute/python/tpu_strategy.py
@@ -33,7 +33,6 @@ from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.util import nest
 
 
-# TODO(isaprykin):  Consider whether inheriting is really appropriate.
 class TPUStrategy(one_device_strategy.OneDeviceStrategy):
   """Experimental TPU distribution strategy implementation."""
 
@@ -73,7 +72,6 @@ class TPUStrategy(one_device_strategy.OneDeviceStrategy):
     def infeed_input(i):
       """Get input, split it and then enqueue."""
       iteration_inputs = [f.get(i) for f in feeds()]
-
       infeed_inputs = [[inputs_per_core[core_id]
                         for inputs_per_core in iteration_inputs]
                        for core_id in range(self._num_cores_per_host)]
@@ -117,3 +115,14 @@ class TPUStrategy(one_device_strategy.OneDeviceStrategy):
           iterate_on_tpu, [], num_shards=self._num_cores_per_host)
 
     return control_flow_ops.group(tpu_result, enqueue_ops)
+
+  def _reduce(self, method_string, value, destinations):
+    del destinations  # TPU is graph mode only.  Rely on implicit Send/Recv.
+    if method_string == 'mean':
+      # TODO(jhseu):  Revisit once we support model-parallelism.
+      value *= (1. / self._num_cores_per_host)
+    return tpu_ops.cross_replica_sum(value)
+
+  @property
+  def num_towers(self):
+    return self._num_cores_per_host
diff --git a/tensorflow/contrib/distribute/python/values.py b/tensorflow/contrib/distribute/python/values.py
index aaf177d07e..b04734f1a3 100644
--- a/tensorflow/contrib/distribute/python/values.py
+++ b/tensorflow/contrib/distribute/python/values.py
@@ -672,11 +672,12 @@ class MultiWorkerDataset(object):
     return MultiWorkerDataIterator(iterators, self._worker_device_map)
 
 
-class PerIteration(object):
-  """Holds input for multiple iterations at once."""
+class _PerKey(object):
+  """Holds data associated by keys."""
 
-  def __init__(self, index):
-    self._index = index
+  def __init__(self, *index):
+    # pylint: disable=protected-access
+    self._index = list(index)
 
   def get(self, iteration):
     return array_ops.gather(self._index, iteration)
@@ -687,6 +688,24 @@ class PerIteration(object):
   def get_dtype(self):
     return self._index[-1][-1].dtype
 
+  def __str__(self):
+    return "%s:%s" % (self.__class__.__name__, self._index)
+
+  def __repr__(self):
+    return "%s(%r)" % (self.__class__.__name__, self._index)
+
+
+class PerIteration(_PerKey):
+  """Holds input for multiple iterations at once."""
+
+  def __init__(self, *index):
+    # pylint: disable=protected-access
+    super(PerIteration, self).__init__(*[batch._index for batch in index])
+
+
+class Batches(_PerKey):
+  pass
+
 
 class MultiIterator(object):
   """Iterator that returns results of multiple get_next()s."""
@@ -697,11 +716,31 @@ class MultiIterator(object):
     self._batches_per_iteration = batches_per_iteration
 
   def get_next(self, name=None):
-    return PerIteration([[
-        self._dataset_iterator.get_next(name=name)
-        for _ in range(self._batches_per_iteration)
-    ]
-                         for _ in range(self._iterations)])
+    """Return PerIteration with `iterations x batches_per_iteration` inputs."""
+    data = []
+    for _ in range(self._batches_per_iteration):
+      batch = []
+      for _ in range(self._iterations):
+        batch.append(self._dataset_iterator.get_next(name=name))
+      data.append(batch)
+
+    # Here is an example.  Suppose each get_next returns a tuple of two tensors.
+    # For 3 `iterations` and 2 `batches_per_iteration`, the `data` is:
+    # [[(a,z), (b,y), (c,x)], [(A,Z), (B,Y), (C,X)]]
+    #
+    # After the first `map_structure` it gets transformed to:
+    #  [(Batches(a, A), Batches(z, Z)),
+    #   (Batches(b, B), Batches(y, Y)),
+    #   (Batches(c, C), Batches(x, X))]
+    #
+    # After the second `map_structure` it gets transformed to a tuple of:
+    # (PerIteration([Batches(a, A), Batches(b, B), Batches(c, C)]),
+    #  PerIteration([Batches(z, Z), Batches(y, Y), Batches(x, X)]))
+
+    data = nest.map_structure(Batches, *data)
+    data = nest.map_structure(PerIteration, *data)
+
+    return data
 
   @property
   def initializer(self):
diff --git a/tensorflow/python/training/distribute.py b/tensorflow/python/training/distribute.py
index 21f81ee187..b60f87c05f 100644
--- a/tensorflow/python/training/distribute.py
+++ b/tensorflow/python/training/distribute.py
@@ -816,6 +816,7 @@ class DistributionStrategy(object):
     # TODO(josh11b): Return an unwrapped value if colocate_with is a
     # single device.
     _require_cross_tower_context(self)
+    assert method_string in ("sum", "mean")
     return self._reduce(method_string, value, destinations)
 
   def _reduce(self, method_string, value, destinations):
-- 
GitLab


From aa57960b545ca25223568e366d99b0a4be7a03da Mon Sep 17 00:00:00 2001
From: Alexandre Passos <apassos@google.com>
Date: Mon, 7 May 2018 11:49:08 -0700
Subject: [PATCH 308/395] Register bool scatter_update for resource variables
 Fixes #17784

PiperOrigin-RevId: 195696915
---
 tensorflow/core/kernels/resource_variable_ops.cc         | 9 +++++++++
 tensorflow/core/kernels/scatter_functor_gpu.cu.cc        | 2 ++
 .../python/kernel_tests/resource_variable_ops_test.py    | 9 +++++++++
 3 files changed, 20 insertions(+)

diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc
index a8bcc7f7dc..03cc414905 100644
--- a/tensorflow/core/kernels/resource_variable_ops.cc
+++ b/tensorflow/core/kernels/resource_variable_ops.cc
@@ -703,6 +703,8 @@ TF_CALL_REAL_NUMBER_TYPES(REGISTER_SCATTER_MINMAX_CPU);
 
 REGISTER_SCATTER_KERNEL(string, CPU, "ResourceScatterUpdate",
                         scatter_op::UpdateOp::ASSIGN);
+REGISTER_SCATTER_KERNEL(bool, CPU, "ResourceScatterUpdate",
+                        scatter_op::UpdateOp::ASSIGN);
 REGISTER_SCATTER_KERNEL(Variant, CPU, "ResourceScatterUpdate",
                         scatter_op::UpdateOp::ASSIGN);
 
@@ -725,6 +727,13 @@ REGISTER_KERNEL_BUILDER(Name("ResourceScatterUpdate")
                             .TypeConstraint<int32>("Tindices"),
                         ResourceScatterUpdateOp<GPUDevice, Variant, int32,
                                                 scatter_op::UpdateOp::ASSIGN>)
+REGISTER_KERNEL_BUILDER(Name("ResourceScatterUpdate")
+                            .Device(DEVICE_GPU)
+                            .HostMemory("resource")
+                            .TypeConstraint<bool>("dtype")
+                            .TypeConstraint<int32>("Tindices"),
+                        ResourceScatterUpdateOp<GPUDevice, bool, int32,
+                                                scatter_op::UpdateOp::ASSIGN>)
 REGISTER_KERNEL_BUILDER(Name("ResourceScatterUpdate")
                             .Device(DEVICE_GPU)
                             .HostMemory("resource")
diff --git a/tensorflow/core/kernels/scatter_functor_gpu.cu.cc b/tensorflow/core/kernels/scatter_functor_gpu.cu.cc
index 59911bf0d2..bdc878594a 100644
--- a/tensorflow/core/kernels/scatter_functor_gpu.cu.cc
+++ b/tensorflow/core/kernels/scatter_functor_gpu.cu.cc
@@ -42,6 +42,8 @@ typedef Eigen::GpuDevice GPUDevice;
 
 DEFINE_GPU_SPECS(float);
 DEFINE_GPU_SPECS(double);
+DEFINE_GPU_SPECS_OP(bool, int32, scatter_op::UpdateOp::ASSIGN);
+DEFINE_GPU_SPECS_OP(bool, int64, scatter_op::UpdateOp::ASSIGN);
 // TODO(b/27222123): The following fails to compile due to lack of support for
 // fp16.
 // TF_CALL_GPU_NUMBER_TYPES(DEFINE_GPU_SPECS);
diff --git a/tensorflow/python/kernel_tests/resource_variable_ops_test.py b/tensorflow/python/kernel_tests/resource_variable_ops_test.py
index 984192258c..3daf07ea63 100644
--- a/tensorflow/python/kernel_tests/resource_variable_ops_test.py
+++ b/tensorflow/python/kernel_tests/resource_variable_ops_test.py
@@ -400,6 +400,15 @@ class ResourceVariableOpsTest(test_util.TensorFlowTestCase):
               resource_variable_ops.var_is_initialized_op(abc.handle)),
           True)
 
+  def testScatterBool(self):
+    with context.eager_mode():
+      ref = resource_variable_ops.ResourceVariable(
+          [False, True, False], trainable=False)
+      indices = math_ops.range(3)
+      updates = constant_op.constant([True, True, True])
+      state_ops.scatter_update(ref, indices, updates)
+      self.assertAllEqual(ref.read_value(), [True, True, True])
+
   @test_util.run_in_graph_and_eager_modes()
   def testConstraintArg(self):
     constraint = lambda x: x
-- 
GitLab


From abe83fe35ed3b4b245471d58811b03170fda857d Mon Sep 17 00:00:00 2001
From: Gunhan Gulsoy <gunan@google.com>
Date: Mon, 7 May 2018 11:52:46 -0700
Subject: [PATCH 309/395] Disable autograph cfg_test in windows.

PiperOrigin-RevId: 195697446
---
 tensorflow/contrib/autograph/pyct/static_analysis/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/contrib/autograph/pyct/static_analysis/BUILD b/tensorflow/contrib/autograph/pyct/static_analysis/BUILD
index 68fbdf6953..8064a967cd 100644
--- a/tensorflow/contrib/autograph/pyct/static_analysis/BUILD
+++ b/tensorflow/contrib/autograph/pyct/static_analysis/BUILD
@@ -48,6 +48,7 @@ py_test(
     name = "cfg_test",
     srcs = ["cfg_test.py"],
     srcs_version = "PY2AND3",
+    tags = ["no_windows"],
     deps = [
         ":static_analysis",
         "//tensorflow/contrib/autograph/pyct",
-- 
GitLab


From 0297d9c1a64270e266a7aeb48f81c78f0a31f63b Mon Sep 17 00:00:00 2001
From: Shivani Agrawal <shivaniagrawal@google.com>
Date: Mon, 7 May 2018 12:03:20 -0700
Subject: [PATCH 310/395] [tf.data] Patch to unref iterator_resource in
 DeserializeIteratorOp.

PiperOrigin-RevId: 195698980
---
 tensorflow/core/kernels/data/iterator_ops.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/data/iterator_ops.cc b/tensorflow/core/kernels/data/iterator_ops.cc
index a2f6c5fe2c..b6bf0ecd09 100644
--- a/tensorflow/core/kernels/data/iterator_ops.cc
+++ b/tensorflow/core/kernels/data/iterator_ops.cc
@@ -1051,7 +1051,7 @@ class DeserializeIteratorOp : public OpKernel {
     IteratorResource* iterator_resource;
     OP_REQUIRES_OK(
         ctx, LookupResource(ctx, HandleFromInput(ctx, 0), &iterator_resource));
-
+    core::ScopedUnref unref_iterator(iterator_resource);
     Variant variant = ctx->input(1).scalar<Variant>()();
     auto* wrapper = variant.get<IteratorStateVariant>();
     OP_REQUIRES(ctx, wrapper != nullptr,
-- 
GitLab


From fb6d927a06a1cff15a71f6b47c207fafbaad6a57 Mon Sep 17 00:00:00 2001
From: Justin Lebar <jlebar@google.com>
Date: Mon, 7 May 2018 12:10:15 -0700
Subject: [PATCH 311/395] [XLA] Add FusionKind matcher to pattern_matcher.h.

PiperOrigin-RevId: 195700319
---
 .../compiler/xla/service/pattern_matcher.h    | 34 +++++++++++++++++++
 .../xla/service/pattern_matcher_test.cc       | 23 +++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/tensorflow/compiler/xla/service/pattern_matcher.h b/tensorflow/compiler/xla/service/pattern_matcher.h
index 586f6ef7a9..d3bc47e61e 100644
--- a/tensorflow/compiler/xla/service/pattern_matcher.h
+++ b/tensorflow/compiler/xla/service/pattern_matcher.h
@@ -702,6 +702,30 @@ class HloInstructionPatternOperandImpl {
   HloInstructionPattern<OperandType, OperandImpl> operand_;
 };
 
+// An HloInstructionPattern implementation that matches only if the instruction
+// is a fusion node with a particular kind.
+template <typename Previous>
+class HloInstructionPatternFusionKindImpl {
+ public:
+  explicit constexpr HloInstructionPatternFusionKindImpl(
+      const Previous& previous, ::xla::HloInstruction::FusionKind kind)
+      : previous_(previous), kind_(kind) {}
+
+  bool Match(const ::xla::HloInstruction* inst) const {
+    return previous_.Match(inst) && inst->opcode() == HloOpcode::kFusion &&
+           inst->fusion_kind() == kind_;
+  }
+
+  bool Match(::xla::HloInstruction* inst) const {
+    return previous_.Match(inst) && inst->opcode() == HloOpcode::kFusion &&
+           inst->fusion_kind() == kind_;
+  }
+
+ private:
+  Previous previous_;
+  ::xla::HloInstruction::FusionKind kind_;
+};
+
 // A pattern that matches HloInstructions.
 template <typename HloInstructionType, typename Impl>
 class HloInstructionPattern {
@@ -807,6 +831,16 @@ class HloInstructionPattern {
         matched_inst_);
   }
 
+  // Modifies the pattern to match only if the instruction is a fusion node with
+  // the given kind.
+  constexpr HloInstructionPattern<HloInstructionType,
+                                  HloInstructionPatternFusionKindImpl<Impl>>
+  WithFusionKind(HloInstruction::FusionKind kind) const {
+    return HloInstructionPattern<HloInstructionType,
+                                 HloInstructionPatternFusionKindImpl<Impl>>(
+        HloInstructionPatternFusionKindImpl<Impl>(impl_, kind), matched_inst_);
+  }
+
  private:
   Impl impl_;
   HloInstructionType** matched_inst_;
diff --git a/tensorflow/compiler/xla/service/pattern_matcher_test.cc b/tensorflow/compiler/xla/service/pattern_matcher_test.cc
index c88157c312..204e8c9920 100644
--- a/tensorflow/compiler/xla/service/pattern_matcher_test.cc
+++ b/tensorflow/compiler/xla/service/pattern_matcher_test.cc
@@ -170,5 +170,28 @@ TEST(PatternMatcherTest, TupleShape) {
       Match(&tuple_shape, match::Shape().WithSubshape({0, 0}, match::Shape())));
 }
 
+TEST(PatternMatcherTest, FusionKind) {
+  constexpr char kModuleStr[] = R"(
+    HloModule test_module
+
+    fused_computation {
+      ROOT fp0 = f32[] parameter(0)
+    }
+
+    ENTRY while.v11 {
+      p0 = f32[] parameter(0)
+      ROOT fusion = f32[] fusion(p0), kind=kLoop, calls=fused_computation
+    })";
+  TF_ASSERT_OK_AND_ASSIGN(auto hlo_module, tools::Parse(kModuleStr));
+
+  auto* root = hlo_module->entry_computation()->root_instruction();
+  EXPECT_TRUE(Match(
+      root, match::Op().WithFusionKind(HloInstruction::FusionKind::kLoop)));
+  EXPECT_FALSE(Match(
+      root, match::Op().WithFusionKind(HloInstruction::FusionKind::kInput)));
+  EXPECT_FALSE(Match(root->operand(0), match::Op().WithFusionKind(
+                                           HloInstruction::FusionKind::kLoop)));
+}
+
 }  // namespace
 }  // namespace xla
-- 
GitLab


From a75f3f533f8e769af3cc11b5125ceb5db8c14479 Mon Sep 17 00:00:00 2001
From: Bixia Zheng <bixia@google.com>
Date: Mon, 7 May 2018 12:15:52 -0700
Subject: [PATCH 312/395] [TF:XLA:GPU] Allow the use of linear address when
 there are size one dimensions in a tensor.

The current implementation of EmitArrayElementAddress incorrectly concludes
that having a size one dimension in a tensor indicates broadcasting is needed
and the linear address can't be used to access the tensor. We fix this by
leaving LinearValidOnShape to decide whether the linear address can be used to
access the tensor. This enables the vectorization of loads/stores in unrolled
elementwise op kernels when other criteria are met.

Add a test case.

PiperOrigin-RevId: 195701194
---
 .../compiler/xla/service/llvm_ir/ir_array.cc  | 22 +++++++++----------
 1 file changed, 10 insertions(+), 12 deletions(-)

diff --git a/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc b/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc
index 3312a88844..7323abeb20 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc
+++ b/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc
@@ -333,18 +333,7 @@ llvm::Value* IrArray::EmitArrayElementAddress(
   }
   CHECK_EQ(index.size(), ShapeUtil::Rank(*shape_));
 
-  std::vector<llvm::Value*> actual_index;
-  bool is_implicit_broadcast = false;
-  // We perform broadcasting when the operand shape has dimension(s) of size
-  // 1. In this case we fix the index value for that dimension to zero. This
-  // effectively broadcasts along this dimension.
-  for (int64 i = 0; i < index.size(); ++i) {
-    auto dim = shape_->dimensions(i);
-    actual_index.push_back(dim == 1 ? ir_builder->getInt64(0) : index[i]);
-    is_implicit_broadcast |= dim == 1;
-  }
-
-  if (!is_implicit_broadcast && index.LinearValidOnShape(*shape_)) {
+  if (index.LinearValidOnShape(*shape_)) {
     llvm::Module* module =
         ir_builder->GetInsertBlock()->getParent()->getParent();
     return ir_builder->CreateInBoundsGEP(
@@ -354,6 +343,15 @@ llvm::Value* IrArray::EmitArrayElementAddress(
         {index.linear()}, llvm_ir::AsStringRef(name));
   }
 
+  std::vector<llvm::Value*> actual_index;
+  for (int64 i = 0; i < index.size(); ++i) {
+    // When dimension i is of size 1, LLVM optimization is able to replace
+    // index[i] with 0. However, setting index[i] to 0 here still allows LLVM to
+    // produce better code in some cases.
+    auto dim = shape_->dimensions(i);
+    actual_index.push_back(dim == 1 ? ir_builder->getInt64(0) : index[i]);
+  }
+
   // "base_ptr_" has the type of "<ir_type_for_its_shape>*"
   // (e.g. [3 x [2 x float]]*). Therefore, the address of the indexed element
   // should be computed by
-- 
GitLab


From c3fef21c4ddf34fd68ab2cd44b0be497b5303b4e Mon Sep 17 00:00:00 2001
From: Ian Langmore <langmore@google.com>
Date: Mon, 7 May 2018 12:17:02 -0700
Subject: [PATCH 313/395] Add 'optonly' directive to linear_operator_circulant
 tests.

PiperOrigin-RevId: 195701399
---
 tensorflow/python/kernel_tests/linalg/BUILD | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/kernel_tests/linalg/BUILD b/tensorflow/python/kernel_tests/linalg/BUILD
index 052f11f92e..91be80322c 100644
--- a/tensorflow/python/kernel_tests/linalg/BUILD
+++ b/tensorflow/python/kernel_tests/linalg/BUILD
@@ -85,7 +85,10 @@ cuda_py_test(
         "//tensorflow/python:platform_test",
     ],
     shard_count = 5,
-    tags = ["noasan"],  # times out b/63678675
+    tags = [
+        "noasan",  # times out, b/63678675
+        "optonly",  # times out, b/79171797
+    ],
 )
 
 cuda_py_test(
-- 
GitLab


From 6f3a890d91e6dbeb811aed23d0eb59abaa8c469f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 12:37:36 -0700
Subject: [PATCH 314/395] Adding Greater/GreaterEqual/LessEqual ops to
 complement Less.

PiperOrigin-RevId: 195704492
---
 tensorflow/contrib/lite/builtin_ops.h         |   3 +
 .../lite/g3doc/tf_ops_compatibility.md        |  39 ++
 .../contrib/lite/kernels/comparisons.cc       | 160 +++++---
 .../contrib/lite/kernels/comparisons_test.cc  | 207 ++++++++++-
 .../contrib/lite/kernels/internal/BUILD       |   1 +
 .../contrib/lite/kernels/internal/common.h    |  14 +
 .../internal/optimized/optimized_ops.h        |  11 +
 .../internal/reference/reference_ops.h        | 191 +++++++---
 tensorflow/contrib/lite/kernels/register.cc   |   6 +
 tensorflow/contrib/lite/model.cc              |   5 +-
 tensorflow/contrib/lite/nnapi_delegate.cc     |   3 +
 tensorflow/contrib/lite/schema/schema.fbs     |  15 +
 .../contrib/lite/schema/schema_generated.h    | 348 +++++++++++++++++-
 tensorflow/contrib/lite/testing/BUILD         |   3 +
 .../contrib/lite/testing/generate_examples.py | 102 +++++
 .../testing/generated_examples_zip_test.cc    |   3 +
 .../contrib/lite/toco/export_tensorflow.cc    |  21 ++
 .../propagate_fixed_sizes.cc                  |   6 +-
 .../toco/graph_transformations/quantize.cc    |   9 +-
 .../contrib/lite/toco/tflite/operator.cc      |   7 +
 20 files changed, 1051 insertions(+), 103 deletions(-)

diff --git a/tensorflow/contrib/lite/builtin_ops.h b/tensorflow/contrib/lite/builtin_ops.h
index d66b72843a..778933f569 100644
--- a/tensorflow/contrib/lite/builtin_ops.h
+++ b/tensorflow/contrib/lite/builtin_ops.h
@@ -86,6 +86,9 @@ typedef enum {
   kTfLiteBuiltinLess = 58,
   kTfLiteBuiltinNeg = 59,
   kTfLiteBuiltinPadv2 = 60,
+  kTfLiteBuiltinGreater = 61,
+  kTfLiteBuiltinGreaterEqual = 62,
+  kTfLiteBuiltinLessEqual = 63,
 } TfLiteBuiltinOperator;
 
 #ifdef __cplusplus
diff --git a/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md b/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md
index 0051ee84ec..fc57b8f28b 100644
--- a/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md
+++ b/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md
@@ -281,6 +281,32 @@ Options {
 }
 ```
 
+**GREATER**
+
+```
+Inputs {
+  0: a tensor
+  1: a tensor
+}
+Outputs {
+  0: a tensor of type bool, true whenever an element of the first tensor is
+  greater than the corresponding element of the second tensor.
+}
+```
+
+**GREATER_EQUAL**
+
+```
+Inputs {
+  0: a tensor
+  1: a tensor
+}
+Outputs {
+  0: a tensor of type bool, true whenever an element of the first tensor is
+  greater than or equal to the corresponding element of the second tensor.
+}
+```
+
 **L2_NORMALIZATION**
 
 ```
@@ -325,6 +351,19 @@ Outputs {
 }
 ```
 
+**LESS_EQUAL**
+
+```
+Inputs {
+  0: a tensor
+  1: a tensor
+}
+Outputs {
+  0: a tensor of type bool, true whenever an element of the first tensor is less
+  than or equal to the corresponding element of the second tensor.
+}
+```
+
 **LOCAL_RESPONSE_NORMALIZATION**
 
 ```
diff --git a/tensorflow/contrib/lite/kernels/comparisons.cc b/tensorflow/contrib/lite/kernels/comparisons.cc
index 87c413cb98..2885ce032b 100644
--- a/tensorflow/contrib/lite/kernels/comparisons.cc
+++ b/tensorflow/contrib/lite/kernels/comparisons.cc
@@ -28,7 +28,7 @@ constexpr int kInputTensor1 = 0;
 constexpr int kInputTensor2 = 1;
 constexpr int kOutputTensor = 0;
 
-TfLiteStatus LessPrepare(TfLiteContext* context, TfLiteNode* node) {
+TfLiteStatus ComparisonPrepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
 
@@ -56,61 +56,139 @@ TfLiteStatus LessPrepare(TfLiteContext* context, TfLiteNode* node) {
   return context->ResizeTensor(context, output, output_size);
 }
 
-TfLiteStatus LessEval(TfLiteContext* context, TfLiteNode* node) {
+#define TF_LITE_COMPARISON(type, opname, requires_broadcast)    \
+  requires_broadcast                                            \
+      ? reference_ops::Broadcast##opname(                       \
+            GetTensorData<type>(input1), GetTensorDims(input1), \
+            GetTensorData<type>(input2), GetTensorDims(input2), \
+            GetTensorData<bool>(output), GetTensorDims(output)) \
+      : reference_ops::opname(                                  \
+            GetTensorData<type>(input1), GetTensorDims(input1), \
+            GetTensorData<type>(input2), GetTensorDims(input2), \
+            GetTensorData<bool>(output), GetTensorDims(output));
+
+TfLiteStatus GreaterEval(TfLiteContext* context, TfLiteNode* node) {
   TfLiteTensor* input1 = GetInput(context, node, kInputTensor1);
   TfLiteTensor* input2 = GetInput(context, node, kInputTensor2);
   TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+  bool requires_broadcast = !HaveSameShapes(input1, input2);
+  // TODO(renjieliu): Support quantized data.
+  switch (input1->type) {
+    case kTfLiteFloat32:
+      TF_LITE_COMPARISON(float, Greater, requires_broadcast);
+      break;
+    case kTfLiteInt32:
+      TF_LITE_COMPARISON(int32_t, Greater, requires_broadcast);
+      break;
+    case kTfLiteInt64:
+      TF_LITE_COMPARISON(int64_t, Greater, requires_broadcast);
+      break;
+    default:
+      context->ReportError(context,
+                           "Does not support type other than float|int");
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
 
+TfLiteStatus GreaterEqualEval(TfLiteContext* context, TfLiteNode* node) {
+  TfLiteTensor* input1 = GetInput(context, node, kInputTensor1);
+  TfLiteTensor* input2 = GetInput(context, node, kInputTensor2);
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
   bool requires_broadcast = !HaveSameShapes(input1, input2);
+  // TODO(renjieliu): Support quantized data.
+  switch (input1->type) {
+    case kTfLiteFloat32:
+      TF_LITE_COMPARISON(float, GreaterEqual, requires_broadcast);
+      break;
+    case kTfLiteInt32:
+      TF_LITE_COMPARISON(int32_t, GreaterEqual, requires_broadcast);
+      break;
+    case kTfLiteInt64:
+      TF_LITE_COMPARISON(int64_t, GreaterEqual, requires_broadcast);
+      break;
+    default:
+      context->ReportError(context,
+                           "Does not support type other than float|int");
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
 
-#define TF_LITE_LESS(type, opname)                                          \
-  reference_ops::opname(GetTensorData<type>(input1), GetTensorDims(input1), \
-                        GetTensorData<type>(input2), GetTensorDims(input2), \
-                        GetTensorData<bool>(output), GetTensorDims(output));
+TfLiteStatus LessEval(TfLiteContext* context, TfLiteNode* node) {
+  TfLiteTensor* input1 = GetInput(context, node, kInputTensor1);
+  TfLiteTensor* input2 = GetInput(context, node, kInputTensor2);
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+  bool requires_broadcast = !HaveSameShapes(input1, input2);
+  // TODO(renjieliu): Support quantized data.
+  switch (input1->type) {
+    case kTfLiteFloat32:
+      TF_LITE_COMPARISON(float, Less, requires_broadcast);
+      break;
+    case kTfLiteInt32:
+      TF_LITE_COMPARISON(int32_t, Less, requires_broadcast);
+      break;
+    case kTfLiteInt64:
+      TF_LITE_COMPARISON(int64_t, Less, requires_broadcast);
+      break;
+    default:
+      context->ReportError(context,
+                           "Does not support type other than float|int");
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
 
+TfLiteStatus LessEqualEval(TfLiteContext* context, TfLiteNode* node) {
+  TfLiteTensor* input1 = GetInput(context, node, kInputTensor1);
+  TfLiteTensor* input2 = GetInput(context, node, kInputTensor2);
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+  bool requires_broadcast = !HaveSameShapes(input1, input2);
   // TODO(renjieliu): Support quantized data.
-  if (requires_broadcast) {
-    switch (input1->type) {
-      case kTfLiteFloat32:
-        TF_LITE_LESS(float, BroadcastLess);
-        break;
-      case kTfLiteInt32:
-        TF_LITE_LESS(int32_t, BroadcastLess);
-        break;
-      case kTfLiteInt64:
-        TF_LITE_LESS(int64_t, BroadcastLess);
-        break;
-      default:
-        context->ReportError(context,
-                             "Does not support type other than float|int");
-        return kTfLiteError;
-    }
-  } else {
-    switch (input1->type) {
-      case kTfLiteFloat32:
-        TF_LITE_LESS(float, Less);
-        break;
-      case kTfLiteInt32:
-        TF_LITE_LESS(int32_t, Less);
-        break;
-      case kTfLiteInt64:
-        TF_LITE_LESS(int64_t, Less);
-        break;
-      default:
-        context->ReportError(context,
-                             "Does not support type other than float|int");
-        return kTfLiteError;
-    }
+  switch (input1->type) {
+    case kTfLiteFloat32:
+      TF_LITE_COMPARISON(float, LessEqual, requires_broadcast);
+      break;
+    case kTfLiteInt32:
+      TF_LITE_COMPARISON(int32_t, LessEqual, requires_broadcast);
+      break;
+    case kTfLiteInt64:
+      TF_LITE_COMPARISON(int64_t, LessEqual, requires_broadcast);
+      break;
+    default:
+      context->ReportError(context,
+                           "Does not support type other than float|int");
+      return kTfLiteError;
   }
-#undef TF_LITE_LESS
   return kTfLiteOk;
 }
 
 }  // namespace comparisons
 
+TfLiteRegistration* Register_GREATER() {
+  static TfLiteRegistration r = {nullptr, nullptr,
+                                 comparisons::ComparisonPrepare,
+                                 comparisons::GreaterEval};
+  return &r;
+}
+
+TfLiteRegistration* Register_GREATER_EQUAL() {
+  static TfLiteRegistration r = {nullptr, nullptr,
+                                 comparisons::ComparisonPrepare,
+                                 comparisons::GreaterEqualEval};
+  return &r;
+}
+
 TfLiteRegistration* Register_LESS() {
-  static TfLiteRegistration r = {nullptr, nullptr, comparisons::LessPrepare,
-                                 comparisons::LessEval};
+  static TfLiteRegistration r = {
+      nullptr, nullptr, comparisons::ComparisonPrepare, comparisons::LessEval};
+  return &r;
+}
+
+TfLiteRegistration* Register_LESS_EQUAL() {
+  static TfLiteRegistration r = {nullptr, nullptr,
+                                 comparisons::ComparisonPrepare,
+                                 comparisons::LessEqualEval};
   return &r;
 }
 
diff --git a/tensorflow/contrib/lite/kernels/comparisons_test.cc b/tensorflow/contrib/lite/kernels/comparisons_test.cc
index da2d7f8589..835d238d36 100644
--- a/tensorflow/contrib/lite/kernels/comparisons_test.cc
+++ b/tensorflow/contrib/lite/kernels/comparisons_test.cc
@@ -23,6 +23,139 @@ namespace {
 
 using ::testing::ElementsAreArray;
 
+class GreaterOpModel : public SingleOpModel {
+ public:
+  GreaterOpModel(std::initializer_list<int> input1_shape,
+                 std::initializer_list<int> input2_shape,
+                 TensorType input_type) {
+    input1_ = AddInput(input_type);
+    input2_ = AddInput(input_type);
+    output_ = AddOutput(TensorType_BOOL);
+    SetBuiltinOp(BuiltinOperator_GREATER, BuiltinOptions_GreaterOptions,
+                 CreateGreaterOptions(builder_).Union());
+    BuildInterpreter({input1_shape, input2_shape});
+  }
+
+  int input1() { return input1_; }
+  int input2() { return input2_; }
+
+  std::vector<bool> GetOutput() { return ExtractVector<bool>(output_); }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ private:
+  int input1_;
+  int input2_;
+  int output_;
+};
+
+TEST(ComparisonsTest, GreaterFloat) {
+  GreaterOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_FLOAT32);
+  model.PopulateTensor<float>(model.input1(), {0.1, 0.9, 0.7, 0.3});
+  model.PopulateTensor<float>(model.input2(), {0.1, 0.2, 0.6, 0.5});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({false, true, true, false}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(ComparisonsTest, GreaterInt) {
+  GreaterOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT32);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {1, 2, 7, 5});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({false, true, false, false}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(ComparisonsTest, GreaterBroadcast) {
+  GreaterOpModel model({1, 1, 1, 4}, {1, 1, 1, 1}, TensorType_INT32);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {7});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({false, true, false, false}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(ComparisonsTest, GreaterBroadcastTwoD) {
+  GreaterOpModel model({1, 1, 2, 4}, {1, 1, 1, 4}, TensorType_INT32);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3, 2, 4, 2, 8});
+  model.PopulateTensor<int>(model.input2(), {7, 1, 2, 4});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({false, true, true, false,
+                                                   false, true, false, true}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 2, 4}));
+}
+
+class GreaterEqualOpModel : public SingleOpModel {
+ public:
+  GreaterEqualOpModel(std::initializer_list<int> input1_shape,
+                      std::initializer_list<int> input2_shape,
+                      TensorType input_type) {
+    input1_ = AddInput(input_type);
+    input2_ = AddInput(input_type);
+    output_ = AddOutput(TensorType_BOOL);
+    SetBuiltinOp(BuiltinOperator_GREATER_EQUAL,
+                 BuiltinOptions_GreaterEqualOptions,
+                 CreateGreaterEqualOptions(builder_).Union());
+    BuildInterpreter({input1_shape, input2_shape});
+  }
+
+  int input1() { return input1_; }
+  int input2() { return input2_; }
+
+  std::vector<bool> GetOutput() { return ExtractVector<bool>(output_); }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ private:
+  int input1_;
+  int input2_;
+  int output_;
+};
+
+TEST(ComparisonsTest, GreaterEqualFloat) {
+  GreaterEqualOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_FLOAT32);
+  model.PopulateTensor<float>(model.input1(), {0.1, 0.9, 0.7, 0.3});
+  model.PopulateTensor<float>(model.input2(), {0.1, 0.2, 0.6, 0.5});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({true, true, true, false}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(ComparisonsTest, GreaterEqualInt) {
+  GreaterEqualOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT32);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {1, 2, 7, 5});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({false, true, true, false}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(ComparisonsTest, GreaterEqualBroadcast) {
+  GreaterEqualOpModel model({1, 1, 1, 4}, {1, 1, 1, 1}, TensorType_INT32);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {7});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({false, true, true, false}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(ComparisonsTest, GreaterEqualBroadcastTwoD) {
+  GreaterEqualOpModel model({1, 1, 2, 4}, {1, 1, 1, 4}, TensorType_INT32);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3, 2, 4, 2, 8});
+  model.PopulateTensor<int>(model.input2(), {7, 1, 2, 4});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({false, true, true, false,
+                                                   false, true, true, true}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 2, 4}));
+}
+
 class LessOpModel : public SingleOpModel {
  public:
   LessOpModel(std::initializer_list<int> input1_shape,
@@ -47,7 +180,7 @@ class LessOpModel : public SingleOpModel {
   int output_;
 };
 
-TEST(ArgMaxOpTest, LessFloat) {
+TEST(ComparisonsTest, LessFloat) {
   LessOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_FLOAT32);
   model.PopulateTensor<float>(model.input1(), {0.1, 0.9, 0.7, 0.3});
   model.PopulateTensor<float>(model.input2(), {0.1, 0.2, 0.6, 0.5});
@@ -57,7 +190,7 @@ TEST(ArgMaxOpTest, LessFloat) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
 }
 
-TEST(ArgMaxOpTest, LessInt) {
+TEST(ComparisonsTest, LessInt) {
   LessOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT32);
   model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
   model.PopulateTensor<int>(model.input2(), {1, 2, 6, 5});
@@ -67,7 +200,7 @@ TEST(ArgMaxOpTest, LessInt) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
 }
 
-TEST(ArgMaxOpTest, LessBroadcast) {
+TEST(ComparisonsTest, LessBroadcast) {
   LessOpModel model({1, 1, 1, 4}, {1, 1, 1, 1}, TensorType_INT32);
   model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
   model.PopulateTensor<int>(model.input2(), {7});
@@ -77,7 +210,7 @@ TEST(ArgMaxOpTest, LessBroadcast) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
 }
 
-TEST(ArgMaxOpTest, LessBroadcastTwoD) {
+TEST(ComparisonsTest, LessBroadcastTwoD) {
   LessOpModel model({1, 1, 2, 4}, {1, 1, 1, 4}, TensorType_INT32);
   model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3, 2, 4, 6, 8});
   model.PopulateTensor<int>(model.input2(), {7, 1, 2, 4});
@@ -88,6 +221,72 @@ TEST(ArgMaxOpTest, LessBroadcastTwoD) {
   EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 2, 4}));
 }
 
+class LessEqualOpModel : public SingleOpModel {
+ public:
+  LessEqualOpModel(std::initializer_list<int> input1_shape,
+                   std::initializer_list<int> input2_shape,
+                   TensorType input_type) {
+    input1_ = AddInput(input_type);
+    input2_ = AddInput(input_type);
+    output_ = AddOutput(TensorType_BOOL);
+    SetBuiltinOp(BuiltinOperator_LESS_EQUAL, BuiltinOptions_LessEqualOptions,
+                 CreateLessEqualOptions(builder_).Union());
+    BuildInterpreter({input1_shape, input2_shape});
+  }
+
+  int input1() { return input1_; }
+  int input2() { return input2_; }
+
+  std::vector<bool> GetOutput() { return ExtractVector<bool>(output_); }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ private:
+  int input1_;
+  int input2_;
+  int output_;
+};
+
+TEST(ComparisonsTest, LessEqualFloat) {
+  LessEqualOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_FLOAT32);
+  model.PopulateTensor<float>(model.input1(), {0.1, 0.9, 0.7, 0.3});
+  model.PopulateTensor<float>(model.input2(), {0.1, 0.2, 0.6, 0.5});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({true, false, false, true}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(ComparisonsTest, LessEqualInt) {
+  LessEqualOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT32);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {1, 2, 7, 5});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({true, false, true, true}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(ComparisonsTest, LessEqualBroadcast) {
+  LessEqualOpModel model({1, 1, 1, 4}, {1, 1, 1, 1}, TensorType_INT32);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {7});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({true, false, true, true}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(ComparisonsTest, LessEqualBroadcastTwoD) {
+  LessEqualOpModel model({1, 1, 2, 4}, {1, 1, 1, 4}, TensorType_INT32);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3, 2, 4, 2, 8});
+  model.PopulateTensor<int>(model.input2(), {7, 1, 2, 4});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({true, false, false, true,
+                                                   true, false, true, false}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 2, 4}));
+}
+
 }  // namespace
 }  // namespace tflite
 
diff --git a/tensorflow/contrib/lite/kernels/internal/BUILD b/tensorflow/contrib/lite/kernels/internal/BUILD
index df29172f83..7ec4782f96 100644
--- a/tensorflow/contrib/lite/kernels/internal/BUILD
+++ b/tensorflow/contrib/lite/kernels/internal/BUILD
@@ -157,6 +157,7 @@ cc_library(
         ":quantization_util",
         ":strided_slice_logic",
         ":types",
+        ":reference_base",
         ":round",
         "//third_party/eigen3",
         "@gemmlowp",
diff --git a/tensorflow/contrib/lite/kernels/internal/common.h b/tensorflow/contrib/lite/kernels/internal/common.h
index 18601df22c..ede95dfee0 100644
--- a/tensorflow/contrib/lite/kernels/internal/common.h
+++ b/tensorflow/contrib/lite/kernels/internal/common.h
@@ -113,6 +113,20 @@ inline int32 MultiplyByQuantizedMultiplier(int32 x, int32 quantized_multiplier,
                              right_shift);
 }
 
+template <typename T>
+int CountLeadingZeros(T integer_input) {
+  static_assert(std::is_unsigned<T>::value,
+                "Only unsigned integer types handled.");
+  const T one_in_leading_positive = static_cast<T>(1)
+                                    << (std::numeric_limits<T>::digits - 1);
+  int leading_zeros = 0;
+  while (integer_input < one_in_leading_positive) {
+    integer_input <<= 1;
+    ++leading_zeros;
+  }
+  return leading_zeros;
+}
+
 }  // namespace tflite
 
 #endif  // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMMON_H_
diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
index e2a1a6996d..c506c5636c 100644
--- a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
@@ -31,6 +31,7 @@ limitations under the License.
 #include "public/gemmlowp.h"
 #include "tensorflow/contrib/lite/kernels/internal/common.h"
 #include "tensorflow/contrib/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h"
 #include "tensorflow/contrib/lite/kernels/internal/round.h"
 #include "tensorflow/contrib/lite/kernels/internal/strided_slice_logic.h"
 #include "tensorflow/contrib/lite/kernels/internal/types.h"
@@ -38,6 +39,16 @@ limitations under the License.
 namespace tflite {
 namespace optimized_ops {
 
+// Unoptimized reference ops:
+using reference_ops::BroadcastGreater;
+using reference_ops::BroadcastGreaterEqual;
+using reference_ops::BroadcastLess;
+using reference_ops::BroadcastLessEqual;
+using reference_ops::Greater;
+using reference_ops::GreaterEqual;
+using reference_ops::Less;
+using reference_ops::LessEqual;
+
 // Make a local VectorMap typedef allowing to map a float array
 // as a Eigen vector expression. The std::conditional here is to
 // construct the suitable Eigen type for the constness of the
diff --git a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
index 05e6ca8e7e..93dba1cc8e 100644
--- a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
@@ -35,35 +35,6 @@ limitations under the License.
 namespace tflite {
 namespace reference_ops {
 
-inline int32 MultiplyByQuantizedMultiplierSmallerThanOne(
-    int32 x, int32 quantized_multiplier, int right_shift) {
-  using gemmlowp::RoundingDivideByPOT;
-  using gemmlowp::SaturatingRoundingDoublingHighMul;
-  return RoundingDivideByPOT(
-      SaturatingRoundingDoublingHighMul(x, quantized_multiplier), right_shift);
-}
-
-inline int32 MultiplyByQuantizedMultiplierGreaterThanOne(
-    int32 x, int32 quantized_multiplier, int left_shift) {
-  using gemmlowp::SaturatingRoundingDoublingHighMul;
-  return SaturatingRoundingDoublingHighMul(x * (1 << left_shift),
-                                           quantized_multiplier);
-}
-
-template <typename T>
-int CountLeadingZeros(T integer_input) {
-  static_assert(std::is_unsigned<T>::value,
-                "Only unsigned integer types handled.");
-  const T one_in_leading_positive = static_cast<T>(1)
-                                    << (std::numeric_limits<T>::digits - 1);
-  int leading_zeros = 0;
-  while (integer_input < one_in_leading_positive) {
-    integer_input <<= 1;
-    ++leading_zeros;
-  }
-  return leading_zeros;
-}
-
 // DO NOT USE THIS STRUCT FOR NEW FUNCTIONALITY BEYOND IMPLEMENTING ELEMENT-WISE
 // BROADCASTING.
 //
@@ -3614,17 +3585,29 @@ inline void TransposeConv(const float* input_data, const Dims<4>& input_dims,
 }
 
 template <typename T>
-inline void Less(int64_t num_elements, const T* input1, const T* input2,
-                 bool* output) {
-  for (int64_t i = 0; i < num_elements; ++i) {
-    output[i] = input1[i] < input2[i];
-  }
+inline bool GreaterFn(T lhs, T rhs) {
+  return lhs > rhs;
+}
+template <typename T>
+inline bool GreaterEqualFn(T lhs, T rhs) {
+  return lhs >= rhs;
+}
+template <typename T>
+inline bool LessFn(T lhs, T rhs) {
+  return lhs < rhs;
+}
+template <typename T>
+inline bool LessEqualFn(T lhs, T rhs) {
+  return lhs <= rhs;
 }
 
 template <typename T>
-inline void Less(const T* input1_data, const Dims<4>& input1_dims,
-                 const T* input2_data, const Dims<4>& input2_dims,
-                 bool* output_data, const Dims<4>& output_dims) {
+using ComparisonFn = bool (*)(T, T);
+
+template <typename T, ComparisonFn<T> F>
+inline void Comparison(const T* input1_data, const Dims<4>& input1_dims,
+                       const T* input2_data, const Dims<4>& input2_dims,
+                       bool* output_data, const Dims<4>& output_dims) {
   const int64_t batches =
       MatchingArraySize(input1_dims, 3, input2_dims, 3, output_dims, 3);
   const int64_t height =
@@ -3633,31 +3616,149 @@ inline void Less(const T* input1_data, const Dims<4>& input1_dims,
       MatchingArraySize(input1_dims, 1, input2_dims, 1, output_dims, 1);
   const int64_t depth =
       MatchingArraySize(input1_dims, 0, input2_dims, 0, output_dims, 0);
-  Less(batches * height * width * depth, input1_data, input2_data, output_data);
+  for (int64_t i = 0; i < batches * height * width * depth; ++i) {
+    output_data[i] = F(input1_data[i], input2_data[i]);
+  }
 }
 
-template <typename T1, typename T2>
-inline void BroadcastLess(T1* input1_data, const Dims<4>& input1_dims,
-                          T2* input2_data, const Dims<4>& input2_dims,
-                          bool* output_data, const Dims<4>& output_dims) {
-  gemmlowp::ScopedProfilingLabel label("BroadcastLess");
+template <typename T, ComparisonFn<T> F>
+inline void Comparison(int left_shift, const T* input1_data,
+                       const Dims<4>& input1_dims, int32 input1_offset,
+                       int32 input1_multiplier, int input1_shift,
+                       const T* input2_data, const Dims<4>& input2_dims,
+                       int32 input2_offset, int32 input2_multiplier,
+                       int input2_shift, bool* output_data,
+                       const Dims<4>& output_dims) {
+  const int64_t batches =
+      MatchingArraySize(input1_dims, 3, input2_dims, 3, output_dims, 3);
+  const int64_t height =
+      MatchingArraySize(input1_dims, 2, input2_dims, 2, output_dims, 2);
+  const int64_t width =
+      MatchingArraySize(input1_dims, 1, input2_dims, 1, output_dims, 1);
+  const int64_t depth =
+      MatchingArraySize(input1_dims, 0, input2_dims, 0, output_dims, 0);
+  for (int64_t i = 0; i < batches * height * width * depth; ++i) {
+    const int32 input1_val = input1_offset + input1_data[i];
+    const int32 input2_val = input2_offset + input2_data[i];
+    const int32 shifted_input1_val = input1_val * (1 << left_shift);
+    const int32 shifted_input2_val = input2_val * (1 << left_shift);
+    const int32 scaled_input1_val = MultiplyByQuantizedMultiplierSmallerThanOne(
+        shifted_input1_val, input1_multiplier, input1_shift);
+    const int32 scaled_input2_val = MultiplyByQuantizedMultiplierSmallerThanOne(
+        shifted_input2_val, input2_multiplier, input2_shift);
+    output_data[i] = F(scaled_input1_val, scaled_input2_val);
+  }
+}
+
+template <typename T, ComparisonFn<T> F>
+inline void BroadcastComparison(const T* input1_data,
+                                const Dims<4>& input1_dims,
+                                const T* input2_data,
+                                const Dims<4>& input2_dims, bool* output_data,
+                                const Dims<4>& output_dims) {
   NdArrayDesc<4> desc1;
   NdArrayDesc<4> desc2;
   NdArrayDescsForElementwiseBroadcast(input1_dims, input2_dims, &desc1, &desc2);
+  for (int b = 0; b < ArraySize(output_dims, 3); ++b) {
+    for (int y = 0; y < ArraySize(output_dims, 2); ++y) {
+      for (int x = 0; x < ArraySize(output_dims, 1); ++x) {
+        for (int c = 0; c < ArraySize(output_dims, 0); ++c) {
+          output_data[Offset(output_dims, c, x, y, b)] =
+              F(input1_data[SubscriptToIndex(desc1, c, x, y, b)],
+                input2_data[SubscriptToIndex(desc2, c, x, y, b)]);
+        }
+      }
+    }
+  }
+}
 
+template <typename T, ComparisonFn<T> F>
+inline void BroadcastComparison(int left_shift, const T* input1_data,
+                                const Dims<4>& input1_dims, int32 input1_offset,
+                                int32 input1_multiplier, int input1_shift,
+                                const T* input2_data,
+                                const Dims<4>& input2_dims, int32 input2_offset,
+                                int32 input2_multiplier, int input2_shift,
+                                bool* output_data, const Dims<4>& output_dims) {
+  NdArrayDesc<4> desc1;
+  NdArrayDesc<4> desc2;
+  NdArrayDescsForElementwiseBroadcast(input1_dims, input2_dims, &desc1, &desc2);
   for (int b = 0; b < ArraySize(output_dims, 3); ++b) {
     for (int y = 0; y < ArraySize(output_dims, 2); ++y) {
       for (int x = 0; x < ArraySize(output_dims, 1); ++x) {
         for (int c = 0; c < ArraySize(output_dims, 0); ++c) {
+          const int32 input1_val =
+              input1_offset + input1_data[SubscriptToIndex(desc1, c, x, y, b)];
+          const int32 input2_val =
+              input2_offset + input2_data[SubscriptToIndex(desc2, c, x, y, b)];
+          const int32 shifted_input1_val = input1_val * (1 << left_shift);
+          const int32 shifted_input2_val = input2_val * (1 << left_shift);
+          const int32 scaled_input1_val =
+              MultiplyByQuantizedMultiplierSmallerThanOne(
+                  shifted_input1_val, input1_multiplier, input1_shift);
+          const int32 scaled_input2_val =
+              MultiplyByQuantizedMultiplierSmallerThanOne(
+                  shifted_input2_val, input2_multiplier, input2_shift);
           output_data[Offset(output_dims, c, x, y, b)] =
-              input1_data[SubscriptToIndex(desc1, c, x, y, b)] <
-              input2_data[SubscriptToIndex(desc2, c, x, y, b)];
+              F(scaled_input1_val, scaled_input2_val);
         }
       }
     }
   }
 }
 
+#define TFLITE_COMPARISON_OP(name)                                            \
+  template <typename T>                                                       \
+  inline void name(const T* input1_data, const Dims<4>& input1_dims,          \
+                   const T* input2_data, const Dims<4>& input2_dims,          \
+                   bool* output_data, const Dims<4>& output_dims) {           \
+    gemmlowp::ScopedProfilingLabel label(#name);                              \
+    Comparison<T, name##Fn>(input1_data, input1_dims, input2_data,            \
+                            input2_dims, output_data, output_dims);           \
+  }                                                                           \
+  template <typename T>                                                       \
+  inline void name(                                                           \
+      int left_shift, const T* input1_data, const Dims<4>& input1_dims,       \
+      int32 input1_offset, int32 input1_multiplier, int input1_shift,         \
+      const T* input2_data, const Dims<4>& input2_dims, int32 input2_offset,  \
+      int32 input2_multiplier, int input2_shift, bool* output_data,           \
+      const Dims<4>& output_dims) {                                           \
+    gemmlowp::ScopedProfilingLabel label(#name "/8bit");                      \
+    BroadcastComparison<T, name##Fn>(left_shift, input1_data, input1_dims,    \
+                                     input1_offset, input1_multiplier,        \
+                                     input1_shift, input2_data, input2_dims,  \
+                                     input2_offset, input2_multiplier,        \
+                                     input2_shift, output_data, output_dims); \
+  }                                                                           \
+  template <typename T>                                                       \
+  inline void Broadcast##name(                                                \
+      const T* input1_data, const Dims<4>& input1_dims, const T* input2_data, \
+      const Dims<4>& input2_dims, bool* output_data,                          \
+      const Dims<4>& output_dims) {                                           \
+    gemmlowp::ScopedProfilingLabel label("Broadcast" #name);                  \
+    BroadcastComparison<T, name##Fn>(input1_data, input1_dims, input2_data,   \
+                                     input2_dims, output_data, output_dims);  \
+  }                                                                           \
+  template <typename T>                                                       \
+  inline void Broadcast##name(                                                \
+      int left_shift, const T* input1_data, const Dims<4>& input1_dims,       \
+      int32 input1_offset, int32 input1_multiplier, int input1_shift,         \
+      const T* input2_data, const Dims<4>& input2_dims, int32 input2_offset,  \
+      int32 input2_multiplier, int input2_shift, bool* output_data,           \
+      const Dims<4>& output_dims) {                                           \
+    gemmlowp::ScopedProfilingLabel label("Broadcast" #name "/8bit");          \
+    BroadcastComparison<T, name##Fn>(left_shift, input1_data, input1_dims,    \
+                                     input1_offset, input1_multiplier,        \
+                                     input1_shift, input2_data, input2_dims,  \
+                                     input2_offset, input2_multiplier,        \
+                                     input2_shift, output_data, output_dims); \
+  }
+TFLITE_COMPARISON_OP(Greater);
+TFLITE_COMPARISON_OP(GreaterEqual);
+TFLITE_COMPARISON_OP(Less);
+TFLITE_COMPARISON_OP(LessEqual);
+#undef TFLITE_COMPARISON_OP
+
 }  // namespace reference_ops
 }  // namespace tflite
 
diff --git a/tensorflow/contrib/lite/kernels/register.cc b/tensorflow/contrib/lite/kernels/register.cc
index a6ea874546..40855891a6 100644
--- a/tensorflow/contrib/lite/kernels/register.cc
+++ b/tensorflow/contrib/lite/kernels/register.cc
@@ -80,7 +80,10 @@ TfLiteRegistration* Register_PRELU();
 TfLiteRegistration* Register_MAXIMUM();
 TfLiteRegistration* Register_MINIMUM();
 TfLiteRegistration* Register_ARG_MAX();
+TfLiteRegistration* Register_GREATER();
+TfLiteRegistration* Register_GREATER_EQUAL();
 TfLiteRegistration* Register_LESS();
+TfLiteRegistration* Register_LESS_EQUAL();
 TfLiteRegistration* Register_FLOOR();
 TfLiteRegistration* Register_NEG();
 
@@ -144,7 +147,10 @@ BuiltinOpResolver::BuiltinOpResolver() {
   AddBuiltin(BuiltinOperator_MAXIMUM, Register_MAXIMUM());
   AddBuiltin(BuiltinOperator_MINIMUM, Register_MINIMUM());
   AddBuiltin(BuiltinOperator_ARG_MAX, Register_ARG_MAX());
+  AddBuiltin(BuiltinOperator_GREATER, Register_GREATER());
+  AddBuiltin(BuiltinOperator_GREATER_EQUAL, Register_GREATER_EQUAL());
   AddBuiltin(BuiltinOperator_LESS, Register_LESS());
+  AddBuiltin(BuiltinOperator_LESS_EQUAL, Register_LESS_EQUAL());
   AddBuiltin(BuiltinOperator_FLOOR, Register_FLOOR());
   AddBuiltin(BuiltinOperator_NEG, Register_NEG());
 
diff --git a/tensorflow/contrib/lite/model.cc b/tensorflow/contrib/lite/model.cc
index 6253570fa2..21c2181377 100644
--- a/tensorflow/contrib/lite/model.cc
+++ b/tensorflow/contrib/lite/model.cc
@@ -672,7 +672,10 @@ TfLiteStatus ParseOpData(const Operator* op, BuiltinOperator op_type,
       *builtin_data = reinterpret_cast<void*>(params);
       break;
     }
-    case BuiltinOperator_LESS: {
+    case BuiltinOperator_GREATER:
+    case BuiltinOperator_GREATER_EQUAL:
+    case BuiltinOperator_LESS:
+    case BuiltinOperator_LESS_EQUAL: {
       break;
     }
     case BuiltinOperator_DELEGATE: {
diff --git a/tensorflow/contrib/lite/nnapi_delegate.cc b/tensorflow/contrib/lite/nnapi_delegate.cc
index b4c46917bf..e903af87b7 100644
--- a/tensorflow/contrib/lite/nnapi_delegate.cc
+++ b/tensorflow/contrib/lite/nnapi_delegate.cc
@@ -372,7 +372,10 @@ void AddOpsAndParams(tflite::Interpreter* interpreter,
       case tflite::BuiltinOperator_MAXIMUM:
       case tflite::BuiltinOperator_MINIMUM:
       case tflite::BuiltinOperator_ARG_MAX:
+      case tflite::BuiltinOperator_GREATER:
+      case tflite::BuiltinOperator_GREATER_EQUAL:
       case tflite::BuiltinOperator_LESS:
+      case tflite::BuiltinOperator_LESS_EQUAL:
       case tflite::BuiltinOperator_NEG:
         FATAL("Op code %d is currently not delegated to NNAPI", builtin);
         nn_op_type = -1;  // set to invalid
diff --git a/tensorflow/contrib/lite/schema/schema.fbs b/tensorflow/contrib/lite/schema/schema.fbs
index 84ff3b16bd..9409e76233 100644
--- a/tensorflow/contrib/lite/schema/schema.fbs
+++ b/tensorflow/contrib/lite/schema/schema.fbs
@@ -138,6 +138,9 @@ enum BuiltinOperator : byte {
   LESS = 58,
   NEG = 59,
   PADV2 = 60,
+  GREATER = 61,
+  GREATER_EQUAL = 62,
+  LESS_EQUAL = 63,
 }
 
 // Options for the builtin operators.
@@ -183,7 +186,10 @@ union BuiltinOptions {
   DequantizeOptions,
   MaximumMinimumOptions,
   ArgMaxOptions,
+  GreaterOptions,
+  GreaterEqualOptions,
   LessOptions,
+  LessEqualOptions,
   NegOptions,
 }
 
@@ -410,9 +416,18 @@ table ArgMaxOptions {
   output_type : TensorType;
 }
 
+table GreaterOptions {
+}
+
+table GreaterEqualOptions {
+}
+
 table LessOptions {
 }
 
+table LessEqualOptions {
+}
+
 table NegOptions {
 }
 
diff --git a/tensorflow/contrib/lite/schema/schema_generated.h b/tensorflow/contrib/lite/schema/schema_generated.h
index 8855e4ad58..ae3b33063e 100755
--- a/tensorflow/contrib/lite/schema/schema_generated.h
+++ b/tensorflow/contrib/lite/schema/schema_generated.h
@@ -154,9 +154,18 @@ struct MaximumMinimumOptionsT;
 struct ArgMaxOptions;
 struct ArgMaxOptionsT;
 
+struct GreaterOptions;
+struct GreaterOptionsT;
+
+struct GreaterEqualOptions;
+struct GreaterEqualOptionsT;
+
 struct LessOptions;
 struct LessOptionsT;
 
+struct LessEqualOptions;
+struct LessEqualOptionsT;
+
 struct NegOptions;
 struct NegOptionsT;
 
@@ -280,11 +289,14 @@ enum BuiltinOperator {
   BuiltinOperator_LESS = 58,
   BuiltinOperator_NEG = 59,
   BuiltinOperator_PADV2 = 60,
+  BuiltinOperator_GREATER = 61,
+  BuiltinOperator_GREATER_EQUAL = 62,
+  BuiltinOperator_LESS_EQUAL = 63,
   BuiltinOperator_MIN = BuiltinOperator_ADD,
-  BuiltinOperator_MAX = BuiltinOperator_PADV2
+  BuiltinOperator_MAX = BuiltinOperator_LESS_EQUAL
 };
 
-inline BuiltinOperator (&EnumValuesBuiltinOperator())[60] {
+inline BuiltinOperator (&EnumValuesBuiltinOperator())[63] {
   static BuiltinOperator values[] = {
     BuiltinOperator_ADD,
     BuiltinOperator_AVERAGE_POOL_2D,
@@ -345,7 +357,10 @@ inline BuiltinOperator (&EnumValuesBuiltinOperator())[60] {
     BuiltinOperator_MINIMUM,
     BuiltinOperator_LESS,
     BuiltinOperator_NEG,
-    BuiltinOperator_PADV2
+    BuiltinOperator_PADV2,
+    BuiltinOperator_GREATER,
+    BuiltinOperator_GREATER_EQUAL,
+    BuiltinOperator_LESS_EQUAL
   };
   return values;
 }
@@ -413,6 +428,9 @@ inline const char **EnumNamesBuiltinOperator() {
     "LESS",
     "NEG",
     "PADV2",
+    "GREATER",
+    "GREATER_EQUAL",
+    "LESS_EQUAL",
     nullptr
   };
   return names;
@@ -466,13 +484,16 @@ enum BuiltinOptions {
   BuiltinOptions_DequantizeOptions = 39,
   BuiltinOptions_MaximumMinimumOptions = 40,
   BuiltinOptions_ArgMaxOptions = 41,
-  BuiltinOptions_LessOptions = 42,
-  BuiltinOptions_NegOptions = 43,
+  BuiltinOptions_GreaterOptions = 42,
+  BuiltinOptions_GreaterEqualOptions = 43,
+  BuiltinOptions_LessOptions = 44,
+  BuiltinOptions_LessEqualOptions = 45,
+  BuiltinOptions_NegOptions = 46,
   BuiltinOptions_MIN = BuiltinOptions_NONE,
   BuiltinOptions_MAX = BuiltinOptions_NegOptions
 };
 
-inline BuiltinOptions (&EnumValuesBuiltinOptions())[44] {
+inline BuiltinOptions (&EnumValuesBuiltinOptions())[47] {
   static BuiltinOptions values[] = {
     BuiltinOptions_NONE,
     BuiltinOptions_Conv2DOptions,
@@ -516,7 +537,10 @@ inline BuiltinOptions (&EnumValuesBuiltinOptions())[44] {
     BuiltinOptions_DequantizeOptions,
     BuiltinOptions_MaximumMinimumOptions,
     BuiltinOptions_ArgMaxOptions,
+    BuiltinOptions_GreaterOptions,
+    BuiltinOptions_GreaterEqualOptions,
     BuiltinOptions_LessOptions,
+    BuiltinOptions_LessEqualOptions,
     BuiltinOptions_NegOptions
   };
   return values;
@@ -566,7 +590,10 @@ inline const char **EnumNamesBuiltinOptions() {
     "DequantizeOptions",
     "MaximumMinimumOptions",
     "ArgMaxOptions",
+    "GreaterOptions",
+    "GreaterEqualOptions",
     "LessOptions",
+    "LessEqualOptions",
     "NegOptions",
     nullptr
   };
@@ -746,10 +773,22 @@ template<> struct BuiltinOptionsTraits<ArgMaxOptions> {
   static const BuiltinOptions enum_value = BuiltinOptions_ArgMaxOptions;
 };
 
+template<> struct BuiltinOptionsTraits<GreaterOptions> {
+  static const BuiltinOptions enum_value = BuiltinOptions_GreaterOptions;
+};
+
+template<> struct BuiltinOptionsTraits<GreaterEqualOptions> {
+  static const BuiltinOptions enum_value = BuiltinOptions_GreaterEqualOptions;
+};
+
 template<> struct BuiltinOptionsTraits<LessOptions> {
   static const BuiltinOptions enum_value = BuiltinOptions_LessOptions;
 };
 
+template<> struct BuiltinOptionsTraits<LessEqualOptions> {
+  static const BuiltinOptions enum_value = BuiltinOptions_LessEqualOptions;
+};
+
 template<> struct BuiltinOptionsTraits<NegOptions> {
   static const BuiltinOptions enum_value = BuiltinOptions_NegOptions;
 };
@@ -1113,6 +1152,22 @@ struct BuiltinOptionsUnion {
     return type == BuiltinOptions_ArgMaxOptions ?
       reinterpret_cast<const ArgMaxOptionsT *>(value) : nullptr;
   }
+  GreaterOptionsT *AsGreaterOptions() {
+    return type == BuiltinOptions_GreaterOptions ?
+      reinterpret_cast<GreaterOptionsT *>(value) : nullptr;
+  }
+  const GreaterOptionsT *AsGreaterOptions() const {
+    return type == BuiltinOptions_GreaterOptions ?
+      reinterpret_cast<const GreaterOptionsT *>(value) : nullptr;
+  }
+  GreaterEqualOptionsT *AsGreaterEqualOptions() {
+    return type == BuiltinOptions_GreaterEqualOptions ?
+      reinterpret_cast<GreaterEqualOptionsT *>(value) : nullptr;
+  }
+  const GreaterEqualOptionsT *AsGreaterEqualOptions() const {
+    return type == BuiltinOptions_GreaterEqualOptions ?
+      reinterpret_cast<const GreaterEqualOptionsT *>(value) : nullptr;
+  }
   LessOptionsT *AsLessOptions() {
     return type == BuiltinOptions_LessOptions ?
       reinterpret_cast<LessOptionsT *>(value) : nullptr;
@@ -1121,6 +1176,14 @@ struct BuiltinOptionsUnion {
     return type == BuiltinOptions_LessOptions ?
       reinterpret_cast<const LessOptionsT *>(value) : nullptr;
   }
+  LessEqualOptionsT *AsLessEqualOptions() {
+    return type == BuiltinOptions_LessEqualOptions ?
+      reinterpret_cast<LessEqualOptionsT *>(value) : nullptr;
+  }
+  const LessEqualOptionsT *AsLessEqualOptions() const {
+    return type == BuiltinOptions_LessEqualOptions ?
+      reinterpret_cast<const LessEqualOptionsT *>(value) : nullptr;
+  }
   NegOptionsT *AsNegOptions() {
     return type == BuiltinOptions_NegOptions ?
       reinterpret_cast<NegOptionsT *>(value) : nullptr;
@@ -4056,6 +4119,86 @@ inline flatbuffers::Offset<ArgMaxOptions> CreateArgMaxOptions(
 
 flatbuffers::Offset<ArgMaxOptions> CreateArgMaxOptions(flatbuffers::FlatBufferBuilder &_fbb, const ArgMaxOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
 
+struct GreaterOptionsT : public flatbuffers::NativeTable {
+  typedef GreaterOptions TableType;
+  GreaterOptionsT() {
+  }
+};
+
+struct GreaterOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
+  typedef GreaterOptionsT NativeTableType;
+  bool Verify(flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           verifier.EndTable();
+  }
+  GreaterOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  void UnPackTo(GreaterOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  static flatbuffers::Offset<GreaterOptions> Pack(flatbuffers::FlatBufferBuilder &_fbb, const GreaterOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
+};
+
+struct GreaterOptionsBuilder {
+  flatbuffers::FlatBufferBuilder &fbb_;
+  flatbuffers::uoffset_t start_;
+  explicit GreaterOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  GreaterOptionsBuilder &operator=(const GreaterOptionsBuilder &);
+  flatbuffers::Offset<GreaterOptions> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = flatbuffers::Offset<GreaterOptions>(end);
+    return o;
+  }
+};
+
+inline flatbuffers::Offset<GreaterOptions> CreateGreaterOptions(
+    flatbuffers::FlatBufferBuilder &_fbb) {
+  GreaterOptionsBuilder builder_(_fbb);
+  return builder_.Finish();
+}
+
+flatbuffers::Offset<GreaterOptions> CreateGreaterOptions(flatbuffers::FlatBufferBuilder &_fbb, const GreaterOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
+
+struct GreaterEqualOptionsT : public flatbuffers::NativeTable {
+  typedef GreaterEqualOptions TableType;
+  GreaterEqualOptionsT() {
+  }
+};
+
+struct GreaterEqualOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
+  typedef GreaterEqualOptionsT NativeTableType;
+  bool Verify(flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           verifier.EndTable();
+  }
+  GreaterEqualOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  void UnPackTo(GreaterEqualOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  static flatbuffers::Offset<GreaterEqualOptions> Pack(flatbuffers::FlatBufferBuilder &_fbb, const GreaterEqualOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
+};
+
+struct GreaterEqualOptionsBuilder {
+  flatbuffers::FlatBufferBuilder &fbb_;
+  flatbuffers::uoffset_t start_;
+  explicit GreaterEqualOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  GreaterEqualOptionsBuilder &operator=(const GreaterEqualOptionsBuilder &);
+  flatbuffers::Offset<GreaterEqualOptions> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = flatbuffers::Offset<GreaterEqualOptions>(end);
+    return o;
+  }
+};
+
+inline flatbuffers::Offset<GreaterEqualOptions> CreateGreaterEqualOptions(
+    flatbuffers::FlatBufferBuilder &_fbb) {
+  GreaterEqualOptionsBuilder builder_(_fbb);
+  return builder_.Finish();
+}
+
+flatbuffers::Offset<GreaterEqualOptions> CreateGreaterEqualOptions(flatbuffers::FlatBufferBuilder &_fbb, const GreaterEqualOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
+
 struct LessOptionsT : public flatbuffers::NativeTable {
   typedef LessOptions TableType;
   LessOptionsT() {
@@ -4096,6 +4239,46 @@ inline flatbuffers::Offset<LessOptions> CreateLessOptions(
 
 flatbuffers::Offset<LessOptions> CreateLessOptions(flatbuffers::FlatBufferBuilder &_fbb, const LessOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
 
+struct LessEqualOptionsT : public flatbuffers::NativeTable {
+  typedef LessEqualOptions TableType;
+  LessEqualOptionsT() {
+  }
+};
+
+struct LessEqualOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
+  typedef LessEqualOptionsT NativeTableType;
+  bool Verify(flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           verifier.EndTable();
+  }
+  LessEqualOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  void UnPackTo(LessEqualOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  static flatbuffers::Offset<LessEqualOptions> Pack(flatbuffers::FlatBufferBuilder &_fbb, const LessEqualOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
+};
+
+struct LessEqualOptionsBuilder {
+  flatbuffers::FlatBufferBuilder &fbb_;
+  flatbuffers::uoffset_t start_;
+  explicit LessEqualOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  LessEqualOptionsBuilder &operator=(const LessEqualOptionsBuilder &);
+  flatbuffers::Offset<LessEqualOptions> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = flatbuffers::Offset<LessEqualOptions>(end);
+    return o;
+  }
+};
+
+inline flatbuffers::Offset<LessEqualOptions> CreateLessEqualOptions(
+    flatbuffers::FlatBufferBuilder &_fbb) {
+  LessEqualOptionsBuilder builder_(_fbb);
+  return builder_.Finish();
+}
+
+flatbuffers::Offset<LessEqualOptions> CreateLessEqualOptions(flatbuffers::FlatBufferBuilder &_fbb, const LessEqualOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
+
 struct NegOptionsT : public flatbuffers::NativeTable {
   typedef NegOptions TableType;
   NegOptionsT() {
@@ -4376,9 +4559,18 @@ struct Operator FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
   const ArgMaxOptions *builtin_options_as_ArgMaxOptions() const {
     return builtin_options_type() == BuiltinOptions_ArgMaxOptions ? static_cast<const ArgMaxOptions *>(builtin_options()) : nullptr;
   }
+  const GreaterOptions *builtin_options_as_GreaterOptions() const {
+    return builtin_options_type() == BuiltinOptions_GreaterOptions ? static_cast<const GreaterOptions *>(builtin_options()) : nullptr;
+  }
+  const GreaterEqualOptions *builtin_options_as_GreaterEqualOptions() const {
+    return builtin_options_type() == BuiltinOptions_GreaterEqualOptions ? static_cast<const GreaterEqualOptions *>(builtin_options()) : nullptr;
+  }
   const LessOptions *builtin_options_as_LessOptions() const {
     return builtin_options_type() == BuiltinOptions_LessOptions ? static_cast<const LessOptions *>(builtin_options()) : nullptr;
   }
+  const LessEqualOptions *builtin_options_as_LessEqualOptions() const {
+    return builtin_options_type() == BuiltinOptions_LessEqualOptions ? static_cast<const LessEqualOptions *>(builtin_options()) : nullptr;
+  }
   const NegOptions *builtin_options_as_NegOptions() const {
     return builtin_options_type() == BuiltinOptions_NegOptions ? static_cast<const NegOptions *>(builtin_options()) : nullptr;
   }
@@ -4572,10 +4764,22 @@ template<> inline const ArgMaxOptions *Operator::builtin_options_as<ArgMaxOption
   return builtin_options_as_ArgMaxOptions();
 }
 
+template<> inline const GreaterOptions *Operator::builtin_options_as<GreaterOptions>() const {
+  return builtin_options_as_GreaterOptions();
+}
+
+template<> inline const GreaterEqualOptions *Operator::builtin_options_as<GreaterEqualOptions>() const {
+  return builtin_options_as_GreaterEqualOptions();
+}
+
 template<> inline const LessOptions *Operator::builtin_options_as<LessOptions>() const {
   return builtin_options_as_LessOptions();
 }
 
+template<> inline const LessEqualOptions *Operator::builtin_options_as<LessEqualOptions>() const {
+  return builtin_options_as_LessEqualOptions();
+}
+
 template<> inline const NegOptions *Operator::builtin_options_as<NegOptions>() const {
   return builtin_options_as_NegOptions();
 }
@@ -6206,6 +6410,52 @@ inline flatbuffers::Offset<ArgMaxOptions> CreateArgMaxOptions(flatbuffers::FlatB
       _output_type);
 }
 
+inline GreaterOptionsT *GreaterOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const {
+  auto _o = new GreaterOptionsT();
+  UnPackTo(_o, _resolver);
+  return _o;
+}
+
+inline void GreaterOptions::UnPackTo(GreaterOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const {
+  (void)_o;
+  (void)_resolver;
+}
+
+inline flatbuffers::Offset<GreaterOptions> GreaterOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const GreaterOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) {
+  return CreateGreaterOptions(_fbb, _o, _rehasher);
+}
+
+inline flatbuffers::Offset<GreaterOptions> CreateGreaterOptions(flatbuffers::FlatBufferBuilder &_fbb, const GreaterOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) {
+  (void)_rehasher;
+  (void)_o;
+  struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const GreaterOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va;
+  return tflite::CreateGreaterOptions(
+      _fbb);
+}
+
+inline GreaterEqualOptionsT *GreaterEqualOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const {
+  auto _o = new GreaterEqualOptionsT();
+  UnPackTo(_o, _resolver);
+  return _o;
+}
+
+inline void GreaterEqualOptions::UnPackTo(GreaterEqualOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const {
+  (void)_o;
+  (void)_resolver;
+}
+
+inline flatbuffers::Offset<GreaterEqualOptions> GreaterEqualOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const GreaterEqualOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) {
+  return CreateGreaterEqualOptions(_fbb, _o, _rehasher);
+}
+
+inline flatbuffers::Offset<GreaterEqualOptions> CreateGreaterEqualOptions(flatbuffers::FlatBufferBuilder &_fbb, const GreaterEqualOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) {
+  (void)_rehasher;
+  (void)_o;
+  struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const GreaterEqualOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va;
+  return tflite::CreateGreaterEqualOptions(
+      _fbb);
+}
+
 inline LessOptionsT *LessOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const {
   auto _o = new LessOptionsT();
   UnPackTo(_o, _resolver);
@@ -6229,6 +6479,29 @@ inline flatbuffers::Offset<LessOptions> CreateLessOptions(flatbuffers::FlatBuffe
       _fbb);
 }
 
+inline LessEqualOptionsT *LessEqualOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const {
+  auto _o = new LessEqualOptionsT();
+  UnPackTo(_o, _resolver);
+  return _o;
+}
+
+inline void LessEqualOptions::UnPackTo(LessEqualOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const {
+  (void)_o;
+  (void)_resolver;
+}
+
+inline flatbuffers::Offset<LessEqualOptions> LessEqualOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const LessEqualOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) {
+  return CreateLessEqualOptions(_fbb, _o, _rehasher);
+}
+
+inline flatbuffers::Offset<LessEqualOptions> CreateLessEqualOptions(flatbuffers::FlatBufferBuilder &_fbb, const LessEqualOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) {
+  (void)_rehasher;
+  (void)_o;
+  struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const LessEqualOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va;
+  return tflite::CreateLessEqualOptions(
+      _fbb);
+}
+
 inline NegOptionsT *NegOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const {
   auto _o = new NegOptionsT();
   UnPackTo(_o, _resolver);
@@ -6599,10 +6872,22 @@ inline bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *ob
       auto ptr = reinterpret_cast<const ArgMaxOptions *>(obj);
       return verifier.VerifyTable(ptr);
     }
+    case BuiltinOptions_GreaterOptions: {
+      auto ptr = reinterpret_cast<const GreaterOptions *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    case BuiltinOptions_GreaterEqualOptions: {
+      auto ptr = reinterpret_cast<const GreaterEqualOptions *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
     case BuiltinOptions_LessOptions: {
       auto ptr = reinterpret_cast<const LessOptions *>(obj);
       return verifier.VerifyTable(ptr);
     }
+    case BuiltinOptions_LessEqualOptions: {
+      auto ptr = reinterpret_cast<const LessEqualOptions *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
     case BuiltinOptions_NegOptions: {
       auto ptr = reinterpret_cast<const NegOptions *>(obj);
       return verifier.VerifyTable(ptr);
@@ -6789,10 +7074,22 @@ inline void *BuiltinOptionsUnion::UnPack(const void *obj, BuiltinOptions type, c
       auto ptr = reinterpret_cast<const ArgMaxOptions *>(obj);
       return ptr->UnPack(resolver);
     }
+    case BuiltinOptions_GreaterOptions: {
+      auto ptr = reinterpret_cast<const GreaterOptions *>(obj);
+      return ptr->UnPack(resolver);
+    }
+    case BuiltinOptions_GreaterEqualOptions: {
+      auto ptr = reinterpret_cast<const GreaterEqualOptions *>(obj);
+      return ptr->UnPack(resolver);
+    }
     case BuiltinOptions_LessOptions: {
       auto ptr = reinterpret_cast<const LessOptions *>(obj);
       return ptr->UnPack(resolver);
     }
+    case BuiltinOptions_LessEqualOptions: {
+      auto ptr = reinterpret_cast<const LessEqualOptions *>(obj);
+      return ptr->UnPack(resolver);
+    }
     case BuiltinOptions_NegOptions: {
       auto ptr = reinterpret_cast<const NegOptions *>(obj);
       return ptr->UnPack(resolver);
@@ -6967,10 +7264,22 @@ inline flatbuffers::Offset<void> BuiltinOptionsUnion::Pack(flatbuffers::FlatBuff
       auto ptr = reinterpret_cast<const ArgMaxOptionsT *>(value);
       return CreateArgMaxOptions(_fbb, ptr, _rehasher).Union();
     }
+    case BuiltinOptions_GreaterOptions: {
+      auto ptr = reinterpret_cast<const GreaterOptionsT *>(value);
+      return CreateGreaterOptions(_fbb, ptr, _rehasher).Union();
+    }
+    case BuiltinOptions_GreaterEqualOptions: {
+      auto ptr = reinterpret_cast<const GreaterEqualOptionsT *>(value);
+      return CreateGreaterEqualOptions(_fbb, ptr, _rehasher).Union();
+    }
     case BuiltinOptions_LessOptions: {
       auto ptr = reinterpret_cast<const LessOptionsT *>(value);
       return CreateLessOptions(_fbb, ptr, _rehasher).Union();
     }
+    case BuiltinOptions_LessEqualOptions: {
+      auto ptr = reinterpret_cast<const LessEqualOptionsT *>(value);
+      return CreateLessEqualOptions(_fbb, ptr, _rehasher).Union();
+    }
     case BuiltinOptions_NegOptions: {
       auto ptr = reinterpret_cast<const NegOptionsT *>(value);
       return CreateNegOptions(_fbb, ptr, _rehasher).Union();
@@ -7145,10 +7454,22 @@ inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) FL
       value = new ArgMaxOptionsT(*reinterpret_cast<ArgMaxOptionsT *>(u.value));
       break;
     }
+    case BuiltinOptions_GreaterOptions: {
+      value = new GreaterOptionsT(*reinterpret_cast<GreaterOptionsT *>(u.value));
+      break;
+    }
+    case BuiltinOptions_GreaterEqualOptions: {
+      value = new GreaterEqualOptionsT(*reinterpret_cast<GreaterEqualOptionsT *>(u.value));
+      break;
+    }
     case BuiltinOptions_LessOptions: {
       value = new LessOptionsT(*reinterpret_cast<LessOptionsT *>(u.value));
       break;
     }
+    case BuiltinOptions_LessEqualOptions: {
+      value = new LessEqualOptionsT(*reinterpret_cast<LessEqualOptionsT *>(u.value));
+      break;
+    }
     case BuiltinOptions_NegOptions: {
       value = new NegOptionsT(*reinterpret_cast<NegOptionsT *>(u.value));
       break;
@@ -7365,11 +7686,26 @@ inline void BuiltinOptionsUnion::Reset() {
       delete ptr;
       break;
     }
+    case BuiltinOptions_GreaterOptions: {
+      auto ptr = reinterpret_cast<GreaterOptionsT *>(value);
+      delete ptr;
+      break;
+    }
+    case BuiltinOptions_GreaterEqualOptions: {
+      auto ptr = reinterpret_cast<GreaterEqualOptionsT *>(value);
+      delete ptr;
+      break;
+    }
     case BuiltinOptions_LessOptions: {
       auto ptr = reinterpret_cast<LessOptionsT *>(value);
       delete ptr;
       break;
     }
+    case BuiltinOptions_LessEqualOptions: {
+      auto ptr = reinterpret_cast<LessEqualOptionsT *>(value);
+      delete ptr;
+      break;
+    }
     case BuiltinOptions_NegOptions: {
       auto ptr = reinterpret_cast<NegOptionsT *>(value);
       delete ptr;
diff --git a/tensorflow/contrib/lite/testing/BUILD b/tensorflow/contrib/lite/testing/BUILD
index ca1390fdeb..6749e63552 100644
--- a/tensorflow/contrib/lite/testing/BUILD
+++ b/tensorflow/contrib/lite/testing/BUILD
@@ -33,9 +33,12 @@ gen_zipped_test_files(
         "fused_batch_norm.zip",
         "gather.zip",
         "global_batch_norm.zip",
+        "greater.zip",
+        "greater_equal.zip",
         "l2_pool.zip",
         "l2norm.zip",
         "less.zip",
+        "less_equal.zip",
         "local_response_norm.zip",
         "log_softmax.zip",
         "max_pool.zip",
diff --git a/tensorflow/contrib/lite/testing/generate_examples.py b/tensorflow/contrib/lite/testing/generate_examples.py
index 6fe0f491d0..7a658d43d3 100644
--- a/tensorflow/contrib/lite/testing/generate_examples.py
+++ b/tensorflow/contrib/lite/testing/generate_examples.py
@@ -2055,6 +2055,74 @@ def make_arg_max_tests(zip_path):
   make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs)
 
 
+def make_greater_tests(zip_path):
+  """Make a set of tests to do greater."""
+
+  test_parameters = [{
+      "input_dtype": [tf.float32, tf.int32, tf.int64],
+      "input_shape_pair": [([1, 1, 1, 3], [1, 1, 1, 3]),
+                           ([2, 3, 4, 5], [2, 3, 4, 5]), ([2, 3, 3], [2, 3]),
+                           ([5, 5], [1]), ([10], [2, 4, 10])],
+  }]
+
+  def build_graph(parameters):
+    """Build the greater op testing graph."""
+    input_value1 = tf.placeholder(
+        dtype=parameters["input_dtype"],
+        name="input1",
+        shape=parameters["input_shape_pair"][0])
+    input_value2 = tf.placeholder(
+        dtype=parameters["input_dtype"],
+        name="input2",
+        shape=parameters["input_shape_pair"][1])
+    out = tf.greater(input_value1, input_value2)
+    return [input_value1, input_value2], [out]
+
+  def build_inputs(parameters, sess, inputs, outputs):
+    input_value1 = create_tensor_data(parameters["input_dtype"],
+                                      parameters["input_shape_pair"][0])
+    input_value2 = create_tensor_data(parameters["input_dtype"],
+                                      parameters["input_shape_pair"][1])
+    return [input_value1, input_value2], sess.run(
+        outputs, feed_dict=dict(zip(inputs, [input_value1, input_value2])))
+
+  make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs)
+
+
+def make_greater_equal_tests(zip_path):
+  """Make a set of tests to do greater_equal."""
+
+  test_parameters = [{
+      "input_dtype": [tf.float32, tf.int32, tf.int64],
+      "input_shape_pair": [([1, 1, 1, 3], [1, 1, 1, 3]),
+                           ([2, 3, 4, 5], [2, 3, 4, 5]), ([2, 3, 3], [2, 3]),
+                           ([5, 5], [1]), ([10], [2, 4, 10])],
+  }]
+
+  def build_graph(parameters):
+    """Build the greater_equal op testing graph."""
+    input_value1 = tf.placeholder(
+        dtype=parameters["input_dtype"],
+        name="input1",
+        shape=parameters["input_shape_pair"][0])
+    input_value2 = tf.placeholder(
+        dtype=parameters["input_dtype"],
+        name="input2",
+        shape=parameters["input_shape_pair"][1])
+    out = tf.greater_equal(input_value1, input_value2)
+    return [input_value1, input_value2], [out]
+
+  def build_inputs(parameters, sess, inputs, outputs):
+    input_value1 = create_tensor_data(parameters["input_dtype"],
+                                      parameters["input_shape_pair"][0])
+    input_value2 = create_tensor_data(parameters["input_dtype"],
+                                      parameters["input_shape_pair"][1])
+    return [input_value1, input_value2], sess.run(
+        outputs, feed_dict=dict(zip(inputs, [input_value1, input_value2])))
+
+  make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs)
+
+
 def make_less_tests(zip_path):
   """Make a set of tests to do less."""
 
@@ -2089,6 +2157,40 @@ def make_less_tests(zip_path):
   make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs)
 
 
+def make_less_equal_tests(zip_path):
+  """Make a set of tests to do less_equal."""
+
+  test_parameters = [{
+      "input_dtype": [tf.float32, tf.int32, tf.int64],
+      "input_shape_pair": [([1, 1, 1, 3], [1, 1, 1, 3]),
+                           ([2, 3, 4, 5], [2, 3, 4, 5]), ([2, 3, 3], [2, 3]),
+                           ([5, 5], [1]), ([10], [2, 4, 10])],
+  }]
+
+  def build_graph(parameters):
+    """Build the less_equal op testing graph."""
+    input_value1 = tf.placeholder(
+        dtype=parameters["input_dtype"],
+        name="input1",
+        shape=parameters["input_shape_pair"][0])
+    input_value2 = tf.placeholder(
+        dtype=parameters["input_dtype"],
+        name="input2",
+        shape=parameters["input_shape_pair"][1])
+    out = tf.less_equal(input_value1, input_value2)
+    return [input_value1, input_value2], [out]
+
+  def build_inputs(parameters, sess, inputs, outputs):
+    input_value1 = create_tensor_data(parameters["input_dtype"],
+                                      parameters["input_shape_pair"][0])
+    input_value2 = create_tensor_data(parameters["input_dtype"],
+                                      parameters["input_shape_pair"][1])
+    return [input_value1, input_value2], sess.run(
+        outputs, feed_dict=dict(zip(inputs, [input_value1, input_value2])))
+
+  make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs)
+
+
 def make_floor_tests(zip_path):
   """Make a set of tests to do floor."""
 
diff --git a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
index 96681952c9..2ce14f3b38 100644
--- a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
+++ b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
@@ -258,9 +258,12 @@ INSTANTIATE_TESTS(fully_connected)
 INSTANTIATE_TESTS(fused_batch_norm)
 INSTANTIATE_TESTS(gather)
 INSTANTIATE_TESTS(global_batch_norm)
+INSTANTIATE_TESTS(greater)
+INSTANTIATE_TESTS(greater_equal)
 INSTANTIATE_TESTS(l2_pool)
 INSTANTIATE_TESTS(l2norm)
 INSTANTIATE_TESTS(less)
+INSTANTIATE_TESTS(less_equal)
 INSTANTIATE_TESTS(local_response_norm)
 INSTANTIATE_TESTS(log_softmax)
 INSTANTIATE_TESTS(max_pool)
diff --git a/tensorflow/contrib/lite/toco/export_tensorflow.cc b/tensorflow/contrib/lite/toco/export_tensorflow.cc
index 9e899cf977..53df1987b3 100644
--- a/tensorflow/contrib/lite/toco/export_tensorflow.cc
+++ b/tensorflow/contrib/lite/toco/export_tensorflow.cc
@@ -1702,6 +1702,19 @@ void ConvertRandomUniformOperator(const Model& model,
   (*new_op->mutable_attr())["seed2"].set_i(src_op.seed2);
 }
 
+void ConvertComparisonOperator(const Model& model, const Operator& src_op,
+                               const char* op_name,
+                               GraphDef* tensorflow_graph) {
+  auto* comparison_op = tensorflow_graph->add_node();
+  comparison_op->set_op(op_name);
+  comparison_op->set_name(src_op.outputs[0]);
+  CHECK_EQ(src_op.inputs.size(), 2);
+  *comparison_op->add_input() = src_op.inputs[0];
+  *comparison_op->add_input() = src_op.inputs[1];
+  const auto data_type = GetTensorFlowDataType(model, src_op.inputs[0]);
+  (*comparison_op->mutable_attr())["T"].set_type(data_type);
+}
+
 void ConvertOperator(const Model& model, const Operator& src_op,
                      GraphDef* tensorflow_graph) {
   if (src_op.fused_activation_function != FusedActivationFunctionType::kNone) {
@@ -1893,6 +1906,14 @@ void ConvertOperator(const Model& model, const Operator& src_op,
     ConvertRandomUniformOperator(
         model, static_cast<const RandomUniformOperator&>(src_op),
         tensorflow_graph);
+  } else if (src_op.type == OperatorType::kTensorFlowGreater) {
+    ConvertComparisonOperator(model, src_op, "Greater", tensorflow_graph);
+  } else if (src_op.type == OperatorType::kTensorFlowGreaterEqual) {
+    ConvertComparisonOperator(model, src_op, "GreaterEqual", tensorflow_graph);
+  } else if (src_op.type == OperatorType::kTensorFlowLess) {
+    ConvertComparisonOperator(model, src_op, "Less", tensorflow_graph);
+  } else if (src_op.type == OperatorType::kTensorFlowLessEqual) {
+    ConvertComparisonOperator(model, src_op, "LessEqual", tensorflow_graph);
   } else {
     LOG(FATAL) << "Unhandled operator type " << OperatorTypeName(src_op.type);
   }
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
index 9b0e232132..a081abea55 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
@@ -499,8 +499,8 @@ void ProcessTensorFlowReshapeOperator(Model* model,
       << op->outputs[0] << "\". Are your input shapes correct?";
 }
 
-void ProcessSimpleOperator(Model* model, Operator* op) {
-  const auto& input_array = model->GetArray(op->inputs[0]);
+void ProcessSimpleOperator(Model* model, Operator* op, int input_index) {
+  const auto& input_array = model->GetArray(op->inputs[input_index]);
   // Yield until input dims have been resolved.
   if (!input_array.has_shape()) {
     return;
@@ -1499,7 +1499,7 @@ bool PropagateFixedSizes::Run(Model* model, std::size_t op_index) {
     case OperatorType::kCast:
     case OperatorType::kFloor:
     case OperatorType::kExp:
-      ProcessSimpleOperator(model, op);
+      ProcessSimpleOperator(model, op, 0);
       break;
     case OperatorType::kGather:
       ProcessGatherOperator(model, static_cast<GatherOperator*>(op));
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc b/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc
index 58e214b76b..a1ca7371c8 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/quantize.cc
@@ -55,7 +55,11 @@ bool SupportsQuantization(const Operator& op) {
          type == OperatorType::kStridedSlice ||
          type == OperatorType::kDepthToSpace ||
          type == OperatorType::kLstmCell || type == OperatorType::kGather ||
-         type == OperatorType::kTranspose || type == OperatorType::kMean;
+         type == OperatorType::kTranspose || type == OperatorType::kMean ||
+         type == OperatorType::kTensorFlowGreater ||
+         type == OperatorType::kTensorFlowGreaterEqual ||
+         type == OperatorType::kTensorFlowLess ||
+         type == OperatorType::kTensorFlowLessEqual;
 }
 
 const MinMax& GetOrComputeMinMax(Model* model, const string& array_name) {
@@ -257,8 +261,7 @@ bool ChooseHardcodedQuantizationForOperatorOutput(
         IsExactlyRepresentable(0., *quantized_data_type, *quantization_params));
     return true;
   }
-  if ((op.type == OperatorType::kLogistic) ||
-      (op.type == OperatorType::kSoftmax)) {
+  if (op.type == OperatorType::kLogistic || op.type == OperatorType::kSoftmax) {
     // Logistic and Softmax have range: [0, 1].
     //
     // For Logistic, 0.5 should be exactly representable, as implementations
diff --git a/tensorflow/contrib/lite/toco/tflite/operator.cc b/tensorflow/contrib/lite/toco/tflite/operator.cc
index df784a2a76..a008e63351 100644
--- a/tensorflow/contrib/lite/toco/tflite/operator.cc
+++ b/tensorflow/contrib/lite/toco/tflite/operator.cc
@@ -915,9 +915,16 @@ std::vector<std::unique_ptr<BaseOperator>> BuildOperatorList() {
       "MAXIMUM", OperatorType::kTensorFlowMaximum));
   ops.emplace_back(new SimpleOperator<TensorFlowMinimumOperator>(
       "MINIMUM", OperatorType::kTensorFlowMinimum));
+  ops.emplace_back(new SimpleOperator<TensorFlowGreaterOperator>(
+      "GREATER", OperatorType::kTensorFlowGreater));
+  ops.emplace_back(new SimpleOperator<TensorFlowGreaterEqualOperator>(
+      "GREATER_EQUAL", OperatorType::kTensorFlowGreaterEqual));
   ops.emplace_back(new SimpleOperator<TensorFlowLessOperator>(
       "LESS", OperatorType::kTensorFlowLess));
+  ops.emplace_back(new SimpleOperator<TensorFlowLessEqualOperator>(
+      "LESS_EQUAL", OperatorType::kTensorFlowLessEqual));
   ops.emplace_back(new SimpleOperator<NegOperator>("NEG", OperatorType::kNeg));
+
   return ops;
 }
 }  // namespace
-- 
GitLab


From a0496a1646fd13188cc985889ca325c004674a17 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 12:41:07 -0700
Subject: [PATCH 315/395] Make test tensorflow/python/keras:resnet50_test be
 size "medium"

This test sometimes runs longer than 60s, and has been getting flaky timeouts
as a result.  With a longer timeout, it succeeds reliably.

PiperOrigin-RevId: 195704998
---
 tensorflow/python/keras/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/keras/BUILD b/tensorflow/python/keras/BUILD
index 1b66f58939..37b24841bd 100755
--- a/tensorflow/python/keras/BUILD
+++ b/tensorflow/python/keras/BUILD
@@ -395,7 +395,7 @@ py_test(
 
 py_test(
     name = "resnet50_test",
-    size = "small",
+    size = "medium",
     srcs = ["_impl/keras/applications/resnet50_test.py"],
     srcs_version = "PY2AND3",
     deps = [
-- 
GitLab


From 91fb950c266345ca5b689038adad5e1c31d36b57 Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Mon, 7 May 2018 12:48:38 -0700
Subject: [PATCH 316/395] Rename HloDotWithContractDimsMatcher to
 HloDotWithContractingDimsMatcher

This is a typo I introduced in cr/195514907.

PiperOrigin-RevId: 195706006
---
 tensorflow/compiler/xla/service/hlo_matchers.cc | 4 ++--
 tensorflow/compiler/xla/service/hlo_matchers.h  | 6 +++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/compiler/xla/service/hlo_matchers.cc b/tensorflow/compiler/xla/service/hlo_matchers.cc
index 41ce9c1762..7e4b883435 100644
--- a/tensorflow/compiler/xla/service/hlo_matchers.cc
+++ b/tensorflow/compiler/xla/service/hlo_matchers.cc
@@ -198,7 +198,7 @@ void HloShardingMatcher::DescribeTo(std::ostream* os) const {
   }
 }
 
-bool HloDotWithContractDimsMatcher::MatchAndExplain(
+bool HloDotWithContractingDimsMatcher::MatchAndExplain(
     const HloInstruction* instruction,
     ::testing::MatchResultListener* listener) const {
   if (!HloMatcher::MatchAndExplain(instruction, listener)) {
@@ -227,7 +227,7 @@ bool HloDotWithContractDimsMatcher::MatchAndExplain(
   return true;
 }
 
-void HloDotWithContractDimsMatcher::DescribeTo(std::ostream* os) const {
+void HloDotWithContractingDimsMatcher::DescribeTo(std::ostream* os) const {
   HloMatcher::DescribeTo(os);
   *os << " with lhs_contracting_dims={" << lhs_contracting_dim_
       << "} and rhs_contracting_dims={" << rhs_contracting_dim_ << "}";
diff --git a/tensorflow/compiler/xla/service/hlo_matchers.h b/tensorflow/compiler/xla/service/hlo_matchers.h
index 75231beac7..c33bdadf1c 100644
--- a/tensorflow/compiler/xla/service/hlo_matchers.h
+++ b/tensorflow/compiler/xla/service/hlo_matchers.h
@@ -133,9 +133,9 @@ class HloShardingMatcher
 
 // Matches a Dot HLO instruction with specific LHS and RHS contracting
 // dimensions.
-class HloDotWithContractDimsMatcher : public HloMatcher {
+class HloDotWithContractingDimsMatcher : public HloMatcher {
  public:
-  explicit HloDotWithContractDimsMatcher(
+  explicit HloDotWithContractingDimsMatcher(
       ::testing::Matcher<const HloInstruction*> lhs,
       ::testing::Matcher<const HloInstruction*> rhs, int64 lhs_contracting_dim,
       int64 rhs_contracting_dim)
@@ -350,7 +350,7 @@ inline ::testing::Matcher<const ::xla::HloInstruction*> Dot(
     ::testing::Matcher<const HloInstruction*> rhs_matcher,
     int64 lhs_contracting_dim, int64 rhs_contracting_dim) {
   return ::testing::MakeMatcher(
-      new ::xla::testing::HloDotWithContractDimsMatcher(
+      new ::xla::testing::HloDotWithContractingDimsMatcher(
           lhs_matcher, rhs_matcher, lhs_contracting_dim, rhs_contracting_dim));
 }
 
-- 
GitLab


From 914c971c7b690661754e83549325c5deadd9e62d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 13:18:33 -0700
Subject: [PATCH 317/395] Specialize functions only once per unique context.

PiperOrigin-RevId: 195710562
---
 .../grappler/optimizers/function_optimizer.cc | 140 +++++++++++++++++-
 .../optimizers/function_optimizer_test.cc     | 117 +++++++++++++++
 .../optimizers/meta_optimizer_test.cc         |  19 ++-
 tensorflow/core/grappler/utils/functions.cc   |  57 +++++++
 tensorflow/core/grappler/utils/functions.h    |  21 ++-
 .../core/grappler/utils/functions_test.cc     |  38 +++++
 6 files changed, 376 insertions(+), 16 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/function_optimizer.cc b/tensorflow/core/grappler/optimizers/function_optimizer.cc
index 1bec9086f7..a44e1ee7f9 100644
--- a/tensorflow/core/grappler/optimizers/function_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/function_optimizer.cc
@@ -14,10 +14,13 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/core/grappler/optimizers/function_optimizer.h"
+
 #include <unordered_map>
+
 #include "tensorflow/core/common_runtime/device_mgr.h"
 #include "tensorflow/core/common_runtime/function.h"
 #include "tensorflow/core/common_runtime/process_function_library_runtime.h"
+#include "tensorflow/core/framework/attr_value_util.h"
 #include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/function.pb.h"
 #include "tensorflow/core/framework/graph_def_util.h"
@@ -74,6 +77,73 @@ string UniqueSpecializedFunctionName(const FunctionDef& func,
   return unique_name;
 }
 
+// Specialized function instantiation type parameters, body parameters, and
+// const inputs.
+struct FunctionSpecializationSignature {
+  string func_name;
+  std::unordered_map<string, DataType> type_parameters;
+  std::unordered_map<string, AttrValue> body_parameters;
+  std::unordered_map<int, string> const_inputs;
+
+  bool operator==(const FunctionSpecializationSignature& other) const {
+    bool equals = func_name == other.func_name &&
+                  type_parameters == other.type_parameters &&
+                  const_inputs == other.const_inputs;
+
+    if (!equals) return false;
+
+    // Equality is not defined for AttrValue.
+    if (body_parameters.size() != other.body_parameters.size()) return false;
+
+    for (const auto& lhs : body_parameters) {
+      auto it = other.body_parameters.find(lhs.first);
+      if (it == other.body_parameters.end()) return false;
+      if (!AreAttrValuesEqual(lhs.second, (*it).second)) return false;
+    }
+
+    return true;
+  }
+
+  struct Hash {
+    uint64 operator()(FunctionSpecializationSignature const& s) const {
+      uint64 h = Hash64(s.func_name);
+
+      // Use std::map for deterministic iteration order.
+
+      std::map<string, DataType> types(s.type_parameters.begin(),
+                                       s.type_parameters.end());
+      for (const auto& pair : types) {
+        AttrValue attr_value;
+        attr_value.set_type(pair.second);
+        h = Hash64Combine(Hash64(pair.first), h);
+        h = Hash64Combine(AttrValueHash(attr_value), h);
+      }
+
+      std::map<string, AttrValue> body(s.body_parameters.begin(),
+                                       s.body_parameters.end());
+      for (const auto& pair : body) {
+        h = Hash64Combine(Hash64(pair.first), h);
+        h = Hash64Combine(AttrValueHash(pair.second), h);
+      }
+
+      std::map<int, string> inputs(s.const_inputs.begin(),
+                                   s.const_inputs.end());
+      for (const auto& pair : inputs) {
+        h = Hash64Combine(std::hash<int>()(pair.first), h);
+        h = Hash64Combine(Hash64(pair.second), h);
+      }
+
+      return h;
+    }
+  };
+};
+
+struct FunctionSpecialization {
+  string specialized_func_name;
+  std::unordered_set<string> const_inputs;
+  std::unordered_set<string> control_deps;
+};
+
 class FunctionOptimizerContext {
  public:
   explicit FunctionOptimizerContext(RewriterConfig::Toggle opt_level,
@@ -108,6 +178,16 @@ class FunctionOptimizerContext {
     return gtl::FindWithDefault(inlined_functions_, name, nullptr);
   }
 
+  const FunctionSpecialization* FindFunctionSpecialization(
+      const FunctionSpecializationSignature& sig) const {
+    return gtl::FindOrNull(specialized_functions_, sig);
+  }
+
+  void AddSpecializedFunction(const FunctionSpecializationSignature& sig,
+                              const FunctionSpecialization& specialized_func) {
+    specialized_functions_.emplace(sig, specialized_func);
+  }
+
  private:
   void InitializeTrulyConstNodes(const GrapplerItem& item) {
     std::unordered_set<string> feed_nodes;
@@ -148,6 +228,12 @@ class FunctionOptimizerContext {
   // Nodes that are Const and not in feed.
   std::unordered_map<string, const NodeDef*> truly_const_nodes_;
 
+  // Specialized functions.
+  std::unordered_map<FunctionSpecializationSignature,
+                     const FunctionSpecialization,
+                     FunctionSpecializationSignature::Hash>
+      specialized_functions_;
+
   TF_DISALLOW_COPY_AND_ASSIGN(FunctionOptimizerContext);
 };
 
@@ -303,14 +389,34 @@ void RemovePushedDownConstInputs(const std::unordered_set<string>& const_inputs,
 
     for (const string& ctrl : control_deps) {
       if (existing_control_deps.find(ctrl) == existing_control_deps.end()) {
-        VLOG(3) << "Forward control dependency to function caller node: input="
-                << ctrl;
+        VLOG(3) << "Forward control dependency: input=" << ctrl;
         specialized_func_node->add_input(ctrl);
       }
     }
   }
 }
 
+Status InitializeFunctionSpecializationSignature(
+    const NodeDef& func_node, const FunctionDef& func,
+    const AttrValueMap& func_attr, const FunctionOptimizerContext& ctx,
+    FunctionSpecializationSignature* sig) {
+  sig->func_name = func.signature().name();
+
+  TF_RETURN_IF_ERROR(
+      InstantiationTypeParameters(func, func_attr, &sig->type_parameters));
+  TF_RETURN_IF_ERROR(
+      InstantiationBodyParameters(func, func_attr, &sig->body_parameters));
+
+  for (int i = 0; i < func_node.input_size(); ++i) {
+    const string& input = func_node.input(i);
+    if (ctx.IsTrulyConst(input)) {
+      sig->const_inputs.emplace(i, input);
+    }
+  }
+
+  return Status::OK();
+}
+
 Status SpecializeFunction(const NodeDef& func_node, const FunctionDef& func,
                           FunctionOptimizerContext* ctx,
                           GraphDef* optimized_graph) {
@@ -320,6 +426,32 @@ Status SpecializeFunction(const NodeDef& func_node, const FunctionDef& func,
   const std::unordered_map<string, AttrValue> func_attr(
       func_node.attr().begin(), func_node.attr().end());
 
+  FunctionSpecializationSignature signature;
+  TF_RETURN_IF_ERROR(InitializeFunctionSpecializationSignature(
+      func_node, func, func_attr, *ctx, &signature));
+
+  // Check if function was already specialized for identical context.
+  const FunctionSpecialization* already_specialized =
+      ctx->FindFunctionSpecialization(signature);
+
+  if (already_specialized) {
+    VLOG(2) << "Function was already specialized in identical context: "
+               "specialized_name="
+            << already_specialized->specialized_func_name;
+
+    // Add a function call node for the specialized function.
+    NodeDef* specialized_func_node = optimized_graph->add_node();
+    *specialized_func_node = func_node;
+    specialized_func_node->set_op(already_specialized->specialized_func_name);
+
+    RemovePushedDownConstInputs(already_specialized->const_inputs,
+                                already_specialized->control_deps,
+                                specialized_func_node);
+
+    return Status::OK();
+  }
+
+  // Add a new specialized function definition to the library.
   const auto& flib = ctx->function_library();
 
   // Make a GrapplerFunctionItem and convert it back to FunctionDef after
@@ -358,6 +490,10 @@ Status SpecializeFunction(const NodeDef& func_node, const FunctionDef& func,
   // Update specialized node to remove inputs for pushed down consts.
   RemovePushedDownConstInputs(const_inputs, control_deps,
                               specialized_func_node);
+
+  ctx->AddSpecializedFunction(
+      signature, {specialized_func_name, const_inputs, control_deps});
+
   return Status::OK();
 }
 
diff --git a/tensorflow/core/grappler/optimizers/function_optimizer_test.cc b/tensorflow/core/grappler/optimizers/function_optimizer_test.cc
index 147a264421..a2dbab3ded 100644
--- a/tensorflow/core/grappler/optimizers/function_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/function_optimizer_test.cc
@@ -718,5 +718,122 @@ TEST_F(FunctionOptimizerTest, SpecializeFunction_PushDownConstInput) {
   test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
 }
 
+TEST_F(FunctionOptimizerTest, SpecializeFunction_OncePerUniqueContext) {
+  using test::function::NDef;
+
+  FunctionOptimizer optimizer(RewriterConfig::DEFAULT);
+
+  // Mark MyMul as noinline.
+  FunctionDef mul_func = FunctionDefHelper::Create(
+      "MyMul", {"x:T", "y:T"}, {"z:T"}, {"T: {float, int32}"},
+      {{{"output"}, "Mul", {"x", "y"}, {{"T", "$T"}}}},
+      /* Mapping between function returns and function node outputs. */
+      {{"z", "output:z:0"}});
+  (*mul_func.mutable_attr())["_noinline"].set_b(true);
+  std::vector<FunctionDef> function_library = {mul_func};
+
+  const Tensor kTwo = test::AsScalar<float>(2.0);
+  const Tensor kThree = test::AsScalar<float>(3.0);
+
+  GrapplerItem item;
+  item.graph = test::function::GDef(
+      {NDef("init", "NoOp", {}, {}, kDevice),
+
+       // Float placeholders.
+       NDef("xf", "Placeholder", {}, {{"dtype", DT_FLOAT}}, kDevice),
+       NDef("yf", "Placeholder", {}, {{"dtype", DT_FLOAT}}, kDevice),
+
+       // Int32 placeholders.
+       NDef("xi", "Placeholder", {}, {{"dtype", DT_INT32}}, kDevice),
+       NDef("yi", "Placeholder", {}, {{"dtype", DT_INT32}}, kDevice),
+
+       // Consts. Control inputs has to be attached to specialized func calls.
+       NDef("two", "Const", {"^init", "^xf"},
+            {{"dtype", DT_FLOAT}, {"value", kTwo}}, kDevice),
+       NDef("three", "Const", {"^init", "^xf"},
+            {{"dtype", DT_FLOAT}, {"value", kThree}}, kDevice),
+
+       // Specialization #1: DT_FLOAT type parameter.
+       NDef("mul_1", "MyMul", {"xf", "yf"}, {{"T", DT_FLOAT}}, kDevice),
+       NDef("mul_2", "MyMul", {"yf", "xf"}, {{"T", DT_FLOAT}}, kDevice),
+
+       // Specialization #2: DT_INT32 type parameter.
+       NDef("mul_3", "MyMul", {"xi", "yi"}, {{"T", DT_INT32}}, kDevice),
+
+       // Specialization #3: DT_FLOAT type parameter + const input kTwo.
+       NDef("mul_4", "MyMul", {"xf", "two"}, {{"T", DT_FLOAT}}, kDevice),
+       NDef("mul_5", "MyMul", {"yf", "two"}, {{"T", DT_FLOAT}}, kDevice),
+
+       // Specialization #4: DT_FLOAT type parameter + const input kThree.
+       NDef("mul_6", "MyMul", {"three", "xf"}, {{"T", DT_FLOAT}}, kDevice)},
+      function_library);
+
+  GraphDef output;
+  TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
+
+  // Make sure that MyMul was specialized once per unique context.
+  EXPECT_EQ(4, output.library().function_size());
+
+  // And graph nodes calling specialized functions.
+  int count = 0;
+  for (const NodeDef& node : output.node()) {
+    if (node.name() == "mul_1" && count++) {
+      EXPECT_EQ("MyMul_specialized_for_mul_1", node.op());
+      ASSERT_EQ(2, node.input_size());
+      EXPECT_EQ("xf", node.input(0));
+      EXPECT_EQ("yf", node.input(1));
+
+    } else if (node.name() == "mul_2" && count++) {
+      EXPECT_EQ("MyMul_specialized_for_mul_1", node.op());
+      ASSERT_EQ(2, node.input_size());
+      EXPECT_EQ("yf", node.input(0));
+      EXPECT_EQ("xf", node.input(1));
+
+    } else if (node.name() == "mul_3" && count++) {
+      EXPECT_EQ("MyMul_specialized_for_mul_3", node.op());
+      ASSERT_EQ(2, node.input_size());
+      EXPECT_EQ("xi", node.input(0));
+      EXPECT_EQ("yi", node.input(1));
+
+    } else if (node.name() == "mul_4" && count++) {
+      EXPECT_EQ("MyMul_specialized_for_mul_4", node.op());
+      ASSERT_EQ(2, node.input_size());
+      EXPECT_EQ("xf", node.input(0));
+      EXPECT_EQ("^init", node.input(1));
+
+    } else if (node.name() == "mul_5" && count++) {
+      EXPECT_EQ("MyMul_specialized_for_mul_4", node.op());
+      ASSERT_EQ(3, node.input_size());
+      EXPECT_EQ("yf", node.input(0));
+      EXPECT_EQ("^init", node.input(1));
+      EXPECT_EQ("^xf", node.input(2));
+
+    } else if (node.name() == "mul_6" && count++) {
+      EXPECT_EQ("MyMul_specialized_for_mul_6", node.op());
+      ASSERT_EQ(2, node.input_size());
+      EXPECT_EQ("xf", node.input(0));
+      EXPECT_EQ("^init", node.input(1));
+    }
+  }
+  EXPECT_EQ(6, count);
+
+  // And that graph evaluation yields the same result.
+  Tensor pi = test::AsScalar<float>(3.14f);
+  Tensor four = test::AsScalar<int32>(4);
+  item.fetch = {"mul_1", "mul_2", "mul_3", "mul_4", "mul_5", "mul_6"};
+  item.feed = {{"xf", pi}, {"yf", pi}, {"xi", four}, {"yi", four}};
+
+  auto tensors_expected = EvaluateFetchNodes(item);
+  GrapplerItem optimized(item, std::move(output));
+  auto tensors = EvaluateFetchNodes(optimized);
+
+  test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
+  test::ExpectTensorEqual<float>(tensors_expected[1], tensors[1]);
+  test::ExpectTensorEqual<int32>(tensors_expected[2], tensors[2]);
+  test::ExpectTensorEqual<float>(tensors_expected[3], tensors[3]);
+  test::ExpectTensorEqual<float>(tensors_expected[4], tensors[4]);
+  test::ExpectTensorEqual<float>(tensors_expected[5], tensors[5]);
+}
+
 }  // namespace grappler
 }  // namespace tensorflow
diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc b/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc
index 887a988af9..8247cce339 100644
--- a/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc
@@ -163,30 +163,28 @@ TEST_F(MetaOptimizerTest, OptimizeFunctionLibrary) {
                                            output.library());
 
   // Specialized and optimized functions should be added to the graph.
-  EXPECT_EQ(6, optimized_flib.num_functions());
+  EXPECT_EQ(5, optimized_flib.num_functions());
 
   // MyQuadratic should be specialized once:
   //   0. 'quadratic' node in the main graph
   const string optimized_0 = "MyQuadratic_specialized_for_quadratic";
 
   // MySquare should be specialized and optimized for 3 instantiations:
-  //   1. 'square' node in the main graph
-  //   2. 'square' node in the MyQuadratic specialization
-  //   3. 'quadratic' node in the MyQuadratic specialization
+  //   1.  'square' node in the main graph
+  //   2.  'square' node in the MyQuadratic specialization
+  //   3*. 'quadratic' node in the MyQuadratic specialization
+  //        has identical instantiation context to #2
 
   const string optimized_1 = "MySquare_specialized_for_square";
   const string optimized_2 = "MySquare_specialized_for_square_1";
-  const string optimized_3 = "MySquare_specialized_for_quadratic";
 
   const FunctionDef* optimized_func_0 = optimized_flib.Find(optimized_0);
   const FunctionDef* optimized_func_1 = optimized_flib.Find(optimized_1);
   const FunctionDef* optimized_func_2 = optimized_flib.Find(optimized_2);
-  const FunctionDef* optimized_func_3 = optimized_flib.Find(optimized_3);
 
   ASSERT_NE(optimized_func_0, nullptr);
   ASSERT_NE(optimized_func_1, nullptr);
   ASSERT_NE(optimized_func_2, nullptr);
-  ASSERT_NE(optimized_func_3, nullptr);
 
   // Graph should call optimized function.
   int count = 0;
@@ -205,13 +203,14 @@ TEST_F(MetaOptimizerTest, OptimizeFunctionLibrary) {
     if (node.name() == "square" && count++) {
       EXPECT_EQ(optimized_2, node.op());
     } else if (node.name() == "quadratic" && count++) {
-      EXPECT_EQ(optimized_3, node.op());
+      // Share specialized function with the 'square' node.
+      EXPECT_EQ(optimized_2, node.op());
     }
   }
   EXPECT_EQ(2, count);
 
-  const std::vector<const FunctionDef*> optimized_funcs = {
-      optimized_func_1, optimized_func_1, optimized_func_3};
+  const std::vector<const FunctionDef*> optimized_funcs = {optimized_func_1,
+                                                           optimized_func_2};
 
   // MyMul should be inlined into all optimized versions of MySquare.
   for (const FunctionDef* optimized_func : optimized_funcs) {
diff --git a/tensorflow/core/grappler/utils/functions.cc b/tensorflow/core/grappler/utils/functions.cc
index 79b823fa2d..34603f9869 100644
--- a/tensorflow/core/grappler/utils/functions.cc
+++ b/tensorflow/core/grappler/utils/functions.cc
@@ -417,6 +417,63 @@ bool IsParametrized(const FunctionDef& func) {
   return HasParametrizedType(func) || HasParametrizedBody(func);
 }
 
+Status InstantiationTypeParameters(
+    const FunctionDef& func, const AttrValueMap& func_instantiation_attr,
+    std::unordered_map<string, DataType>* type_parameters) {
+  if (!type_parameters->empty()) {
+    return errors::InvalidArgument("Type parameters output map must be empty");
+  }
+
+  GrapplerFunctionItemInstantiation instantiation(&func_instantiation_attr);
+
+  const auto resolve_type_attr = [&](const OpDef::ArgDef& arg) {
+    // Check if it's unknown and unresolved type.
+    if (arg.type() == DT_INVALID &&
+        type_parameters->find(arg.type_attr()) == type_parameters->end()) {
+      DataType data_type;
+      TF_RETURN_IF_ERROR(instantiation.GetArgType(arg, &data_type));
+      type_parameters->insert({arg.type_attr(), data_type});
+    }
+    return Status::OK();
+  };
+
+  for (const auto& input : func.signature().input_arg())
+    TF_RETURN_IF_ERROR(resolve_type_attr(input));
+  for (const auto& output : func.signature().output_arg())
+    TF_RETURN_IF_ERROR(resolve_type_attr(output));
+
+  return Status::OK();
+}
+
+Status InstantiationBodyParameters(
+    const FunctionDef& func, const AttrValueMap& func_instantiation_attr,
+    std::unordered_map<string, AttrValue>* body_parameters) {
+  if (!body_parameters->empty()) {
+    return errors::InvalidArgument("Body parameters output map must be empty");
+  }
+
+  for (const NodeDef& func_body_node : func.node_def()) {
+    for (auto& attr : func_body_node.attr()) {
+      const string& placeholder = attr.second.placeholder();
+
+      if (placeholder.empty() ||
+          body_parameters->find(placeholder) != body_parameters->end()) {
+        continue;
+      }
+
+      auto it = func_instantiation_attr.find(placeholder);
+      if (it != func_instantiation_attr.end()) {
+        body_parameters->emplace(placeholder, it->second);
+      } else {
+        return errors::InvalidArgument("Can't resolve placeholder: ",
+                                       placeholder);
+      }
+    }
+  }
+
+  return Status::OK();
+}
+
 Status MakeGrapplerFunctionItem(const FunctionDef& func,
                                 const AttrValueMap& func_instantiation_attr,
                                 const FunctionLibraryDefinition& flib,
diff --git a/tensorflow/core/grappler/utils/functions.h b/tensorflow/core/grappler/utils/functions.h
index d9d71b80eb..4641bf5252 100644
--- a/tensorflow/core/grappler/utils/functions.h
+++ b/tensorflow/core/grappler/utils/functions.h
@@ -191,6 +191,19 @@ bool HasParametrizedBody(const FunctionDef& func);
 // Check if function has parametrized type or body.
 bool IsParametrized(const FunctionDef& func);
 
+// Resolve function instantiation type parameters from the attributes of the
+// caller node. Return error if type can't be resolved.
+Status InstantiationTypeParameters(
+    const FunctionDef& func, const AttrValueMap& func_instantiation_attr,
+    std::unordered_map<string, DataType>* type_parameters);
+
+// Resolve function instantiation body parameters (values for the function body
+// attr placeholders) from the attributes of the caller node. Return error if
+// type can't be resolved.
+Status InstantiationBodyParameters(
+    const FunctionDef& func, const AttrValueMap& func_instantiation_attr,
+    std::unordered_map<string, AttrValue>* body_parameters);
+
 // Register GrapplerFunctionItem input arg expansion and function body outputs
 // in the GrapplerFunctionConnectivity. Use function library definition to
 // lookup function body nodes output names and ranges.
@@ -205,10 +218,10 @@ Status ReplaceInputWithConst(const NodeDef& input_const, int input_position,
 // Make a GrapplerFunctionItem from the function definition and function
 // instantiation attributes (caller node attributes). Returns error if the given
 // function def cannot be converted (e.g. not all attributes are defined).
-Status MakeGrapplerFunctionItem(
-    const FunctionDef& func,
-    const std::unordered_map<string, AttrValue>& func_instantiation_attr,
-    const FunctionLibraryDefinition& flib, GrapplerFunctionItem* item);
+Status MakeGrapplerFunctionItem(const FunctionDef& func,
+                                const AttrValueMap& func_instantiation_attr,
+                                const FunctionLibraryDefinition& flib,
+                                GrapplerFunctionItem* item);
 
 // Make a GrapplerFunction item from the function definition. Function must be
 // fully defined (no type or body parametrization).
diff --git a/tensorflow/core/grappler/utils/functions_test.cc b/tensorflow/core/grappler/utils/functions_test.cc
index fa6fec70ff..15d8437438 100644
--- a/tensorflow/core/grappler/utils/functions_test.cc
+++ b/tensorflow/core/grappler/utils/functions_test.cc
@@ -54,6 +54,44 @@ TEST_F(FunctionsTest, IsParametrized) {
   EXPECT_FALSE(IsParametrized(non_parametrized_func));
 }
 
+TEST_F(FunctionsTest, InstantiationParameters) {
+  // Function definition is invalid, only type/body parameters are important.
+  FunctionDef func = FunctionDefHelper::Create(
+      "ParametrizedFunc",
+      /* inputs */
+      {"input1:A", "input2:B", "input3:float"},
+      /* outputs */
+      {"output1: A", "output2:C"},
+      /* type parameters */
+      {"A: {float, double}", "B: {float, int32}", "C: {float, double}"},
+      /* function body*/
+      {{{"output"}, "FakeOp", {"input1", "input2"}, {{"key", "$key"}}}},
+      /* Mapping between function returns and function node outputs. */
+      {{"x", "cx:output:0"}, {"y", "cy:output:0"}});
+
+  std::unordered_map<string, AttrValue> func_instantiation_attr;
+  func_instantiation_attr["key"].set_s("key-value");
+  func_instantiation_attr["A"].set_type(DT_FLOAT);
+  func_instantiation_attr["B"].set_type(DT_INT32);
+  func_instantiation_attr["C"].set_type(DT_DOUBLE);
+
+  std::unordered_map<string, DataType> type_parameters;
+  TF_EXPECT_OK(InstantiationTypeParameters(func, func_instantiation_attr,
+                                           &type_parameters));
+
+  ASSERT_EQ(3, type_parameters.size());
+  EXPECT_EQ(DT_FLOAT, type_parameters["A"]);
+  EXPECT_EQ(DT_INT32, type_parameters["B"]);
+  EXPECT_EQ(DT_DOUBLE, type_parameters["C"]);
+
+  std::unordered_map<string, AttrValue> body_parameters;
+  TF_EXPECT_OK(InstantiationBodyParameters(func, func_instantiation_attr,
+                                           &body_parameters));
+
+  ASSERT_EQ(1, body_parameters.size());
+  EXPECT_EQ("key-value", body_parameters["key"].s());
+}
+
 TEST_F(FunctionsTest, GrapplerFunctionConnectivity_ExpandFunctionDefInput) {
   GrapplerFunctionConnectivity connectivity;
 
-- 
GitLab


From bd8d7440d7121dc1e92c4794ca1d18d0e9eb0a17 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 13:24:12 -0700
Subject: [PATCH 318/395] Fixes for accessing variables with a MirroredStrategy
 in a cross-tower context:

* only provide read-only access to variables via get()

* don't fail if use the variable isn't copied to the current device in
  get()

* make _as_graph_element() return the aggregate value for tower-local
  variables (instead of the incorrect previous behavior of returning
  the primary)

PiperOrigin-RevId: 195711474
---
 .../python/mirrored_strategy_multigpu_test.py | 33 ++++++++------
 .../contrib/distribute/python/values.py       | 44 ++++++++++++++++---
 2 files changed, 56 insertions(+), 21 deletions(-)

diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py b/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py
index 6c5c055070..3635bd2e34 100644
--- a/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py
+++ b/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py
@@ -370,22 +370,27 @@ class MirroredStrategyVariableCreationTest(test.TestCase):
       expected_sum = 0.0
       expected_mean = 0.0
       for i, d in enumerate(dist.worker_devices):
-        # Test access within a device scope, should see different values.
-        with ops.device(d):
-          v_sum_value = self.evaluate(ret_v_sum.read_value())
-          v_mean_value = self.evaluate(ret_v_mean.read_value())
-          expected = i + 3.0
-          self.assertEqual(expected, v_sum_value)
-          expected_sum += expected
-          expected = i * 6.0
-          self.assertEqual(expected, v_mean_value)
-          expected_mean += expected
-
-      # fetch() should return the value you get by applying the
-      # reduction across all towers.
-      self.assertEqual(expected_sum, self.evaluate(dist.fetch(ret_v_sum)))
+        # Should see different values on different devices.
+        v_sum_value = self.evaluate(ret_v_sum.get(d).read_value())
+        v_mean_value = self.evaluate(ret_v_mean.get(d).read_value())
+        expected = i + 3.0
+        self.assertEqual(expected, v_sum_value)
+        expected_sum += expected
+        expected = i * 6.0
+        self.assertEqual(expected, v_mean_value)
+        expected_mean += expected
       expected_mean /= len(dist.worker_devices)
+
+      # Without get(device), should return the value you get by
+      # applying the reduction across all towers (whether you use
+      # fetch(), get(), or nothing).
+      self.assertEqual(expected_sum, self.evaluate(dist.fetch(ret_v_sum)))
       self.assertEqual(expected_mean, self.evaluate(dist.fetch(ret_v_mean)))
+      self.assertEqual(expected_sum, self.evaluate(ret_v_sum.get()))
+      self.assertEqual(expected_mean, self.evaluate(ret_v_mean.get()))
+      if not context.executing_eagerly():
+        self.assertEqual(expected_sum, self.evaluate(ret_v_sum))
+        self.assertEqual(expected_mean, self.evaluate(ret_v_mean))
 
   # NOTE(priyag): Names and name scopes are ignored in eager, hence we are not
   # testing this in eager mode.
diff --git a/tensorflow/contrib/distribute/python/values.py b/tensorflow/contrib/distribute/python/values.py
index b04734f1a3..759f3c3599 100644
--- a/tensorflow/contrib/distribute/python/values.py
+++ b/tensorflow/contrib/distribute/python/values.py
@@ -34,6 +34,7 @@ from tensorflow.python.framework import device as tf_device
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
 from tensorflow.python.training import checkpointable
 from tensorflow.python.training import device_util
 from tensorflow.python.training import distribute as distribute_lib
@@ -60,7 +61,7 @@ class DistributedValues(object):
       else:
         device = distribute_lib.get_update_device()
         if device is None:
-          device = device_util.current()
+          return self._get_cross_tower()
     device = device_util.canonicalize(device)
     try:
       return self._index[device]
@@ -231,12 +232,6 @@ class DistributedVariable(DistributedDelegate):
                               self._primary_var.op.type)
     return self.get().op
 
-  def _as_graph_element(self):
-    # pylint: disable=protected-access
-    if distribute_lib.get_cross_tower_context():
-      return self._primary_var._as_graph_element()
-    return self.get()._as_graph_element()
-
   def _should_act_as_resource_variable(self):
     """Pass resource_variable_ops.is_resource_variable check."""
     pass
@@ -320,6 +315,18 @@ class MirroredVariable(DistributedVariable, Mirrored,
   def assign(self, *args, **kwargs):
     return self.get(device=_get_update_device()).assign(*args, **kwargs)
 
+  def _get_cross_tower(self):
+    device = device_util.canonicalize(device_util.current())
+    if device in self._index:
+      return array_ops.identity(self._index[device])
+    return array_ops.identity(self._primary_var)
+
+  def _as_graph_element(self):
+    # pylint: disable=protected-access
+    if distribute_lib.get_cross_tower_context():
+      return self._primary_var._as_graph_element()
+    return self.get()._as_graph_element()
+
   def _gather_saveables_for_checkpoint(self):
     """Overrides CheckpointableBase method.
 
@@ -364,6 +371,12 @@ class _TowerLocalSaveable(saver.BaseSaverBuilder.SaveableObject):
         for d, v in six.iteritems(self._tower_local_variable._index)])  # pylint: disable=protected-access
 
 
+def _assert_tower_context():
+  if not distribute_lib.get_tower_context():
+    raise RuntimeError(
+        "Tower-local variables may only be assigned in a tower context.")
+
+
 class TowerLocalVariable(DistributedVariable, PerDevice,
                          checkpointable.CheckpointableBase):
   """Holds a map from device to variables whose values are reduced on save."""
@@ -374,18 +387,35 @@ class TowerLocalVariable(DistributedVariable, PerDevice,
     super(TowerLocalVariable, self).__init__(index)
 
   def assign_sub(self, *args, **kwargs):
+    _assert_tower_context()
     return self.get().assign_sub(*args, **kwargs)
 
   def assign_add(self, *args, **kwargs):
+    _assert_tower_context()
     return self.get().assign_add(*args, **kwargs)
 
   def assign(self, *args, **kwargs):
+    _assert_tower_context()
     return self.get().assign(*args, **kwargs)
 
   @property
   def reduce_method(self):
     return self._reduce_method
 
+  def _get_cross_tower(self):
+    all_components = tuple(self._index.values())
+    # TODO(josh11b): Use a strategy-specific method.
+    total = math_ops.add_n(all_components)
+    if self._reduce_method == "mean":
+      return total * (1./ len(all_components))
+    return total
+
+  def _as_graph_element(self):
+    # pylint: disable=protected-access
+    if distribute_lib.get_cross_tower_context():
+      return self._get_cross_tower()
+    return self.get()._as_graph_element()
+
   def _gather_saveables_for_checkpoint(self):
     """Overrides CheckpointableBase method.
 
-- 
GitLab


From 9a0c8453e4d12cb8cca6e72a8e6a19a4a3ba21b2 Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Mon, 7 May 2018 14:00:09 -0700
Subject: [PATCH 319/395] [TF:XLA] Bump open source llvm revision to r331624

PiperOrigin-RevId: 195717497
---
 tensorflow/workspace.bzl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index b67c4bf2ac..8f499976de 100644
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -452,11 +452,11 @@ def tf_workspace(path_prefix="", tf_repo_name=""):
   tf_http_archive(
       name = "llvm",
       urls = [
-          "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/b3f6a6a61625296bb532a65c0bf51b91b05b3361.tar.gz",
-          "https://github.com/llvm-mirror/llvm/archive/b3f6a6a61625296bb532a65c0bf51b91b05b3361.tar.gz",
+          "https://mirror.bazel.build/github.com/llvm-mirror/llvm/archive/7b8a8728fbd27086efbf3c57cf2bb35a557108c9.tar.gz",
+          "https://github.com/llvm-mirror/llvm/archive/7b8a8728fbd27086efbf3c57cf2bb35a557108c9.tar.gz",
       ],
-      sha256 = "93895b289a78a47a1e75652e12a1b9a6c119f086a509b00e0084cf2bb944b709",
-      strip_prefix = "llvm-b3f6a6a61625296bb532a65c0bf51b91b05b3361",
+      sha256 = "c620859c3ae5818f316de4837f340b3bba1646f8add0a28e6d4da34ce47e3969",
+      strip_prefix = "llvm-7b8a8728fbd27086efbf3c57cf2bb35a557108c9",
       build_file = clean_dep("//third_party/llvm:llvm.BUILD"),
   )
 
-- 
GitLab


From 92acc302beefd2b596b505d9278f8b0b46239cea Mon Sep 17 00:00:00 2001
From: Anna R <annarev@google.com>
Date: Mon, 7 May 2018 14:03:15 -0700
Subject: [PATCH 320/395] Change deprecation_version field in api_def.proto to
 a string.

PiperOrigin-RevId: 195718061
---
 tensorflow/core/framework/api_def.proto      | 6 ++++--
 tensorflow/core/framework/op_gen_lib.cc      | 3 ---
 tensorflow/core/framework/op_gen_lib_test.cc | 1 -
 3 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/framework/api_def.proto b/tensorflow/core/framework/api_def.proto
index 98c38efc0e..e878ab620b 100644
--- a/tensorflow/core/framework/api_def.proto
+++ b/tensorflow/core/framework/api_def.proto
@@ -55,8 +55,10 @@ message ApiDef {
     // use a snake_case convention instead of CamelCase.
     string name = 1;
 
-    // First GraphDef version at which the op is disallowed.
-    int32 deprecation_version = 2;
+    // If this endpoint is deprecated, set deprecation_message to a
+    // message that should be logged when the endpoint is used.
+    // The message should indicate alternative endpoint to use, if any.
+    string deprecation_message = 2;
   }
   repeated Endpoint endpoint = 3;
 
diff --git a/tensorflow/core/framework/op_gen_lib.cc b/tensorflow/core/framework/op_gen_lib.cc
index 5e14043625..3d7920a6e2 100644
--- a/tensorflow/core/framework/op_gen_lib.cc
+++ b/tensorflow/core/framework/op_gen_lib.cc
@@ -306,9 +306,6 @@ void InitApiDefFromOpDef(const OpDef& op_def, ApiDef* api_def) {
 
   auto* endpoint = api_def->add_endpoint();
   endpoint->set_name(op_def.name());
-  if (op_def.has_deprecation()) {
-    endpoint->set_deprecation_version(op_def.deprecation().version());
-  }
 
   for (const auto& op_in_arg : op_def.input_arg()) {
     auto* api_in_arg = api_def->add_in_arg();
diff --git a/tensorflow/core/framework/op_gen_lib_test.cc b/tensorflow/core/framework/op_gen_lib_test.cc
index 857b1c8dbc..e0e77c7449 100644
--- a/tensorflow/core/framework/op_gen_lib_test.cc
+++ b/tensorflow/core/framework/op_gen_lib_test.cc
@@ -189,7 +189,6 @@ TEST(OpGenLibTest, ApiDefInitializedFromOpDef) {
 visibility: VISIBLE
 endpoint {
   name: "testop"
-  deprecation_version: 123
 }
 in_arg {
   name: "arg_a"
-- 
GitLab


From fa2132ab65f92ea40c94152dba105a9f86a0a555 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 14:13:23 -0700
Subject: [PATCH 321/395] Use 64bit aggregation for gradients and hessians
 since the 32 bit version is numerically unstable for large minibatches.

PiperOrigin-RevId: 195719795
---
 .../lib/learner/batch/ordinal_split_handler.py      | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/tensorflow/contrib/boosted_trees/lib/learner/batch/ordinal_split_handler.py b/tensorflow/contrib/boosted_trees/lib/learner/batch/ordinal_split_handler.py
index 9d6cc9245a..f06b73c00d 100644
--- a/tensorflow/contrib/boosted_trees/lib/learner/batch/ordinal_split_handler.py
+++ b/tensorflow/contrib/boosted_trees/lib/learner/batch/ordinal_split_handler.py
@@ -501,11 +501,18 @@ def sparse_make_stats_update(
         example_partition_ids)
 
     # Compute aggregate stats for each partition.
+    # Since unsorted_segment_sum can be numerically unstable, use 64bit
+    # operation.
+    gradients64 = math_ops.cast(gradients, dtypes.float64)
+    hessians64 = math_ops.cast(hessians, dtypes.float64)
     per_partition_gradients = math_ops.unsorted_segment_sum(
-        gradients, mapped_partitions, array_ops.size(unique_partitions))
+        gradients64, mapped_partitions, array_ops.size(unique_partitions))
     per_partition_hessians = math_ops.unsorted_segment_sum(
-        hessians, mapped_partitions, array_ops.size(unique_partitions))
-
+        hessians64, mapped_partitions, array_ops.size(unique_partitions))
+    per_partition_gradients = math_ops.cast(per_partition_gradients,
+                                            dtypes.float32)
+    per_partition_hessians = math_ops.cast(per_partition_hessians,
+                                           dtypes.float32)
     # Prepend a bias feature per partition that accumulates the stats for all
     # examples in that partition.
     bias_feature_ids = array_ops.fill(
-- 
GitLab


From 544dcc5092a7bf49a5d3a43e25c0f29f087062dd Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 14:15:37 -0700
Subject: [PATCH 322/395] Move PadV2Options to the end, in order to maintain
 schema compatibility.

PiperOrigin-RevId: 195720133
---
 tensorflow/contrib/lite/schema/schema.fbs     |   6 +-
 .../contrib/lite/schema/schema_generated.h    | 304 +++++++++---------
 2 files changed, 155 insertions(+), 155 deletions(-)

diff --git a/tensorflow/contrib/lite/schema/schema.fbs b/tensorflow/contrib/lite/schema/schema.fbs
index 9409e76233..3ec91e505d 100644
--- a/tensorflow/contrib/lite/schema/schema.fbs
+++ b/tensorflow/contrib/lite/schema/schema.fbs
@@ -167,7 +167,6 @@ union BuiltinOptions {
   EmbeddingLookupSparseOptions,
   MulOptions,
   PadOptions,
-  PadV2Options,
   GatherOptions,
   BatchToSpaceNDOptions,
   SpaceToBatchNDOptions,
@@ -186,11 +185,12 @@ union BuiltinOptions {
   DequantizeOptions,
   MaximumMinimumOptions,
   ArgMaxOptions,
+  LessOptions,
+  NegOptions,
+  PadV2Options,
   GreaterOptions,
   GreaterEqualOptions,
-  LessOptions,
   LessEqualOptions,
-  NegOptions,
 }
 
 enum Padding : byte { SAME, VALID }
diff --git a/tensorflow/contrib/lite/schema/schema_generated.h b/tensorflow/contrib/lite/schema/schema_generated.h
index ae3b33063e..c6e4dab454 100755
--- a/tensorflow/contrib/lite/schema/schema_generated.h
+++ b/tensorflow/contrib/lite/schema/schema_generated.h
@@ -465,32 +465,32 @@ enum BuiltinOptions {
   BuiltinOptions_EmbeddingLookupSparseOptions = 20,
   BuiltinOptions_MulOptions = 21,
   BuiltinOptions_PadOptions = 22,
-  BuiltinOptions_PadV2Options = 23,
-  BuiltinOptions_GatherOptions = 24,
-  BuiltinOptions_BatchToSpaceNDOptions = 25,
-  BuiltinOptions_SpaceToBatchNDOptions = 26,
-  BuiltinOptions_TransposeOptions = 27,
-  BuiltinOptions_MeanOptions = 28,
-  BuiltinOptions_SubOptions = 29,
-  BuiltinOptions_DivOptions = 30,
-  BuiltinOptions_SqueezeOptions = 31,
-  BuiltinOptions_SequenceRNNOptions = 32,
-  BuiltinOptions_StridedSliceOptions = 33,
-  BuiltinOptions_ExpOptions = 34,
-  BuiltinOptions_TopKV2Options = 35,
-  BuiltinOptions_SplitOptions = 36,
-  BuiltinOptions_LogSoftmaxOptions = 37,
-  BuiltinOptions_CastOptions = 38,
-  BuiltinOptions_DequantizeOptions = 39,
-  BuiltinOptions_MaximumMinimumOptions = 40,
-  BuiltinOptions_ArgMaxOptions = 41,
-  BuiltinOptions_GreaterOptions = 42,
-  BuiltinOptions_GreaterEqualOptions = 43,
-  BuiltinOptions_LessOptions = 44,
-  BuiltinOptions_LessEqualOptions = 45,
-  BuiltinOptions_NegOptions = 46,
+  BuiltinOptions_GatherOptions = 23,
+  BuiltinOptions_BatchToSpaceNDOptions = 24,
+  BuiltinOptions_SpaceToBatchNDOptions = 25,
+  BuiltinOptions_TransposeOptions = 26,
+  BuiltinOptions_MeanOptions = 27,
+  BuiltinOptions_SubOptions = 28,
+  BuiltinOptions_DivOptions = 29,
+  BuiltinOptions_SqueezeOptions = 30,
+  BuiltinOptions_SequenceRNNOptions = 31,
+  BuiltinOptions_StridedSliceOptions = 32,
+  BuiltinOptions_ExpOptions = 33,
+  BuiltinOptions_TopKV2Options = 34,
+  BuiltinOptions_SplitOptions = 35,
+  BuiltinOptions_LogSoftmaxOptions = 36,
+  BuiltinOptions_CastOptions = 37,
+  BuiltinOptions_DequantizeOptions = 38,
+  BuiltinOptions_MaximumMinimumOptions = 39,
+  BuiltinOptions_ArgMaxOptions = 40,
+  BuiltinOptions_LessOptions = 41,
+  BuiltinOptions_NegOptions = 42,
+  BuiltinOptions_PadV2Options = 43,
+  BuiltinOptions_GreaterOptions = 44,
+  BuiltinOptions_GreaterEqualOptions = 45,
+  BuiltinOptions_LessEqualOptions = 46,
   BuiltinOptions_MIN = BuiltinOptions_NONE,
-  BuiltinOptions_MAX = BuiltinOptions_NegOptions
+  BuiltinOptions_MAX = BuiltinOptions_LessEqualOptions
 };
 
 inline BuiltinOptions (&EnumValuesBuiltinOptions())[47] {
@@ -518,7 +518,6 @@ inline BuiltinOptions (&EnumValuesBuiltinOptions())[47] {
     BuiltinOptions_EmbeddingLookupSparseOptions,
     BuiltinOptions_MulOptions,
     BuiltinOptions_PadOptions,
-    BuiltinOptions_PadV2Options,
     BuiltinOptions_GatherOptions,
     BuiltinOptions_BatchToSpaceNDOptions,
     BuiltinOptions_SpaceToBatchNDOptions,
@@ -537,11 +536,12 @@ inline BuiltinOptions (&EnumValuesBuiltinOptions())[47] {
     BuiltinOptions_DequantizeOptions,
     BuiltinOptions_MaximumMinimumOptions,
     BuiltinOptions_ArgMaxOptions,
+    BuiltinOptions_LessOptions,
+    BuiltinOptions_NegOptions,
+    BuiltinOptions_PadV2Options,
     BuiltinOptions_GreaterOptions,
     BuiltinOptions_GreaterEqualOptions,
-    BuiltinOptions_LessOptions,
-    BuiltinOptions_LessEqualOptions,
-    BuiltinOptions_NegOptions
+    BuiltinOptions_LessEqualOptions
   };
   return values;
 }
@@ -571,7 +571,6 @@ inline const char **EnumNamesBuiltinOptions() {
     "EmbeddingLookupSparseOptions",
     "MulOptions",
     "PadOptions",
-    "PadV2Options",
     "GatherOptions",
     "BatchToSpaceNDOptions",
     "SpaceToBatchNDOptions",
@@ -590,11 +589,12 @@ inline const char **EnumNamesBuiltinOptions() {
     "DequantizeOptions",
     "MaximumMinimumOptions",
     "ArgMaxOptions",
+    "LessOptions",
+    "NegOptions",
+    "PadV2Options",
     "GreaterOptions",
     "GreaterEqualOptions",
-    "LessOptions",
     "LessEqualOptions",
-    "NegOptions",
     nullptr
   };
   return names;
@@ -697,10 +697,6 @@ template<> struct BuiltinOptionsTraits<PadOptions> {
   static const BuiltinOptions enum_value = BuiltinOptions_PadOptions;
 };
 
-template<> struct BuiltinOptionsTraits<PadV2Options> {
-  static const BuiltinOptions enum_value = BuiltinOptions_PadV2Options;
-};
-
 template<> struct BuiltinOptionsTraits<GatherOptions> {
   static const BuiltinOptions enum_value = BuiltinOptions_GatherOptions;
 };
@@ -773,6 +769,18 @@ template<> struct BuiltinOptionsTraits<ArgMaxOptions> {
   static const BuiltinOptions enum_value = BuiltinOptions_ArgMaxOptions;
 };
 
+template<> struct BuiltinOptionsTraits<LessOptions> {
+  static const BuiltinOptions enum_value = BuiltinOptions_LessOptions;
+};
+
+template<> struct BuiltinOptionsTraits<NegOptions> {
+  static const BuiltinOptions enum_value = BuiltinOptions_NegOptions;
+};
+
+template<> struct BuiltinOptionsTraits<PadV2Options> {
+  static const BuiltinOptions enum_value = BuiltinOptions_PadV2Options;
+};
+
 template<> struct BuiltinOptionsTraits<GreaterOptions> {
   static const BuiltinOptions enum_value = BuiltinOptions_GreaterOptions;
 };
@@ -781,18 +789,10 @@ template<> struct BuiltinOptionsTraits<GreaterEqualOptions> {
   static const BuiltinOptions enum_value = BuiltinOptions_GreaterEqualOptions;
 };
 
-template<> struct BuiltinOptionsTraits<LessOptions> {
-  static const BuiltinOptions enum_value = BuiltinOptions_LessOptions;
-};
-
 template<> struct BuiltinOptionsTraits<LessEqualOptions> {
   static const BuiltinOptions enum_value = BuiltinOptions_LessEqualOptions;
 };
 
-template<> struct BuiltinOptionsTraits<NegOptions> {
-  static const BuiltinOptions enum_value = BuiltinOptions_NegOptions;
-};
-
 struct BuiltinOptionsUnion {
   BuiltinOptions type;
   void *value;
@@ -1000,14 +1000,6 @@ struct BuiltinOptionsUnion {
     return type == BuiltinOptions_PadOptions ?
       reinterpret_cast<const PadOptionsT *>(value) : nullptr;
   }
-  PadV2OptionsT *AsPadV2Options() {
-    return type == BuiltinOptions_PadV2Options ?
-      reinterpret_cast<PadV2OptionsT *>(value) : nullptr;
-  }
-  const PadV2OptionsT *AsPadV2Options() const {
-    return type == BuiltinOptions_PadV2Options ?
-      reinterpret_cast<const PadV2OptionsT *>(value) : nullptr;
-  }
   GatherOptionsT *AsGatherOptions() {
     return type == BuiltinOptions_GatherOptions ?
       reinterpret_cast<GatherOptionsT *>(value) : nullptr;
@@ -1152,6 +1144,30 @@ struct BuiltinOptionsUnion {
     return type == BuiltinOptions_ArgMaxOptions ?
       reinterpret_cast<const ArgMaxOptionsT *>(value) : nullptr;
   }
+  LessOptionsT *AsLessOptions() {
+    return type == BuiltinOptions_LessOptions ?
+      reinterpret_cast<LessOptionsT *>(value) : nullptr;
+  }
+  const LessOptionsT *AsLessOptions() const {
+    return type == BuiltinOptions_LessOptions ?
+      reinterpret_cast<const LessOptionsT *>(value) : nullptr;
+  }
+  NegOptionsT *AsNegOptions() {
+    return type == BuiltinOptions_NegOptions ?
+      reinterpret_cast<NegOptionsT *>(value) : nullptr;
+  }
+  const NegOptionsT *AsNegOptions() const {
+    return type == BuiltinOptions_NegOptions ?
+      reinterpret_cast<const NegOptionsT *>(value) : nullptr;
+  }
+  PadV2OptionsT *AsPadV2Options() {
+    return type == BuiltinOptions_PadV2Options ?
+      reinterpret_cast<PadV2OptionsT *>(value) : nullptr;
+  }
+  const PadV2OptionsT *AsPadV2Options() const {
+    return type == BuiltinOptions_PadV2Options ?
+      reinterpret_cast<const PadV2OptionsT *>(value) : nullptr;
+  }
   GreaterOptionsT *AsGreaterOptions() {
     return type == BuiltinOptions_GreaterOptions ?
       reinterpret_cast<GreaterOptionsT *>(value) : nullptr;
@@ -1168,14 +1184,6 @@ struct BuiltinOptionsUnion {
     return type == BuiltinOptions_GreaterEqualOptions ?
       reinterpret_cast<const GreaterEqualOptionsT *>(value) : nullptr;
   }
-  LessOptionsT *AsLessOptions() {
-    return type == BuiltinOptions_LessOptions ?
-      reinterpret_cast<LessOptionsT *>(value) : nullptr;
-  }
-  const LessOptionsT *AsLessOptions() const {
-    return type == BuiltinOptions_LessOptions ?
-      reinterpret_cast<const LessOptionsT *>(value) : nullptr;
-  }
   LessEqualOptionsT *AsLessEqualOptions() {
     return type == BuiltinOptions_LessEqualOptions ?
       reinterpret_cast<LessEqualOptionsT *>(value) : nullptr;
@@ -1184,14 +1192,6 @@ struct BuiltinOptionsUnion {
     return type == BuiltinOptions_LessEqualOptions ?
       reinterpret_cast<const LessEqualOptionsT *>(value) : nullptr;
   }
-  NegOptionsT *AsNegOptions() {
-    return type == BuiltinOptions_NegOptions ?
-      reinterpret_cast<NegOptionsT *>(value) : nullptr;
-  }
-  const NegOptionsT *AsNegOptions() const {
-    return type == BuiltinOptions_NegOptions ?
-      reinterpret_cast<const NegOptionsT *>(value) : nullptr;
-  }
 };
 
 bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *obj, BuiltinOptions type);
@@ -4502,9 +4502,6 @@ struct Operator FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
   const PadOptions *builtin_options_as_PadOptions() const {
     return builtin_options_type() == BuiltinOptions_PadOptions ? static_cast<const PadOptions *>(builtin_options()) : nullptr;
   }
-  const PadV2Options *builtin_options_as_PadV2Options() const {
-    return builtin_options_type() == BuiltinOptions_PadV2Options ? static_cast<const PadV2Options *>(builtin_options()) : nullptr;
-  }
   const GatherOptions *builtin_options_as_GatherOptions() const {
     return builtin_options_type() == BuiltinOptions_GatherOptions ? static_cast<const GatherOptions *>(builtin_options()) : nullptr;
   }
@@ -4559,21 +4556,24 @@ struct Operator FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
   const ArgMaxOptions *builtin_options_as_ArgMaxOptions() const {
     return builtin_options_type() == BuiltinOptions_ArgMaxOptions ? static_cast<const ArgMaxOptions *>(builtin_options()) : nullptr;
   }
+  const LessOptions *builtin_options_as_LessOptions() const {
+    return builtin_options_type() == BuiltinOptions_LessOptions ? static_cast<const LessOptions *>(builtin_options()) : nullptr;
+  }
+  const NegOptions *builtin_options_as_NegOptions() const {
+    return builtin_options_type() == BuiltinOptions_NegOptions ? static_cast<const NegOptions *>(builtin_options()) : nullptr;
+  }
+  const PadV2Options *builtin_options_as_PadV2Options() const {
+    return builtin_options_type() == BuiltinOptions_PadV2Options ? static_cast<const PadV2Options *>(builtin_options()) : nullptr;
+  }
   const GreaterOptions *builtin_options_as_GreaterOptions() const {
     return builtin_options_type() == BuiltinOptions_GreaterOptions ? static_cast<const GreaterOptions *>(builtin_options()) : nullptr;
   }
   const GreaterEqualOptions *builtin_options_as_GreaterEqualOptions() const {
     return builtin_options_type() == BuiltinOptions_GreaterEqualOptions ? static_cast<const GreaterEqualOptions *>(builtin_options()) : nullptr;
   }
-  const LessOptions *builtin_options_as_LessOptions() const {
-    return builtin_options_type() == BuiltinOptions_LessOptions ? static_cast<const LessOptions *>(builtin_options()) : nullptr;
-  }
   const LessEqualOptions *builtin_options_as_LessEqualOptions() const {
     return builtin_options_type() == BuiltinOptions_LessEqualOptions ? static_cast<const LessEqualOptions *>(builtin_options()) : nullptr;
   }
-  const NegOptions *builtin_options_as_NegOptions() const {
-    return builtin_options_type() == BuiltinOptions_NegOptions ? static_cast<const NegOptions *>(builtin_options()) : nullptr;
-  }
   const flatbuffers::Vector<uint8_t> *custom_options() const {
     return GetPointer<const flatbuffers::Vector<uint8_t> *>(VT_CUSTOM_OPTIONS);
   }
@@ -4688,10 +4688,6 @@ template<> inline const PadOptions *Operator::builtin_options_as<PadOptions>() c
   return builtin_options_as_PadOptions();
 }
 
-template<> inline const PadV2Options *Operator::builtin_options_as<PadV2Options>() const {
-  return builtin_options_as_PadV2Options();
-}
-
 template<> inline const GatherOptions *Operator::builtin_options_as<GatherOptions>() const {
   return builtin_options_as_GatherOptions();
 }
@@ -4764,6 +4760,18 @@ template<> inline const ArgMaxOptions *Operator::builtin_options_as<ArgMaxOption
   return builtin_options_as_ArgMaxOptions();
 }
 
+template<> inline const LessOptions *Operator::builtin_options_as<LessOptions>() const {
+  return builtin_options_as_LessOptions();
+}
+
+template<> inline const NegOptions *Operator::builtin_options_as<NegOptions>() const {
+  return builtin_options_as_NegOptions();
+}
+
+template<> inline const PadV2Options *Operator::builtin_options_as<PadV2Options>() const {
+  return builtin_options_as_PadV2Options();
+}
+
 template<> inline const GreaterOptions *Operator::builtin_options_as<GreaterOptions>() const {
   return builtin_options_as_GreaterOptions();
 }
@@ -4772,18 +4780,10 @@ template<> inline const GreaterEqualOptions *Operator::builtin_options_as<Greate
   return builtin_options_as_GreaterEqualOptions();
 }
 
-template<> inline const LessOptions *Operator::builtin_options_as<LessOptions>() const {
-  return builtin_options_as_LessOptions();
-}
-
 template<> inline const LessEqualOptions *Operator::builtin_options_as<LessEqualOptions>() const {
   return builtin_options_as_LessEqualOptions();
 }
 
-template<> inline const NegOptions *Operator::builtin_options_as<NegOptions>() const {
-  return builtin_options_as_NegOptions();
-}
-
 struct OperatorBuilder {
   flatbuffers::FlatBufferBuilder &fbb_;
   flatbuffers::uoffset_t start_;
@@ -6796,10 +6796,6 @@ inline bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *ob
       auto ptr = reinterpret_cast<const PadOptions *>(obj);
       return verifier.VerifyTable(ptr);
     }
-    case BuiltinOptions_PadV2Options: {
-      auto ptr = reinterpret_cast<const PadV2Options *>(obj);
-      return verifier.VerifyTable(ptr);
-    }
     case BuiltinOptions_GatherOptions: {
       auto ptr = reinterpret_cast<const GatherOptions *>(obj);
       return verifier.VerifyTable(ptr);
@@ -6872,6 +6868,18 @@ inline bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *ob
       auto ptr = reinterpret_cast<const ArgMaxOptions *>(obj);
       return verifier.VerifyTable(ptr);
     }
+    case BuiltinOptions_LessOptions: {
+      auto ptr = reinterpret_cast<const LessOptions *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    case BuiltinOptions_NegOptions: {
+      auto ptr = reinterpret_cast<const NegOptions *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
+    case BuiltinOptions_PadV2Options: {
+      auto ptr = reinterpret_cast<const PadV2Options *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
     case BuiltinOptions_GreaterOptions: {
       auto ptr = reinterpret_cast<const GreaterOptions *>(obj);
       return verifier.VerifyTable(ptr);
@@ -6880,18 +6888,10 @@ inline bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *ob
       auto ptr = reinterpret_cast<const GreaterEqualOptions *>(obj);
       return verifier.VerifyTable(ptr);
     }
-    case BuiltinOptions_LessOptions: {
-      auto ptr = reinterpret_cast<const LessOptions *>(obj);
-      return verifier.VerifyTable(ptr);
-    }
     case BuiltinOptions_LessEqualOptions: {
       auto ptr = reinterpret_cast<const LessEqualOptions *>(obj);
       return verifier.VerifyTable(ptr);
     }
-    case BuiltinOptions_NegOptions: {
-      auto ptr = reinterpret_cast<const NegOptions *>(obj);
-      return verifier.VerifyTable(ptr);
-    }
     default: return false;
   }
 }
@@ -6998,10 +6998,6 @@ inline void *BuiltinOptionsUnion::UnPack(const void *obj, BuiltinOptions type, c
       auto ptr = reinterpret_cast<const PadOptions *>(obj);
       return ptr->UnPack(resolver);
     }
-    case BuiltinOptions_PadV2Options: {
-      auto ptr = reinterpret_cast<const PadV2Options *>(obj);
-      return ptr->UnPack(resolver);
-    }
     case BuiltinOptions_GatherOptions: {
       auto ptr = reinterpret_cast<const GatherOptions *>(obj);
       return ptr->UnPack(resolver);
@@ -7074,6 +7070,18 @@ inline void *BuiltinOptionsUnion::UnPack(const void *obj, BuiltinOptions type, c
       auto ptr = reinterpret_cast<const ArgMaxOptions *>(obj);
       return ptr->UnPack(resolver);
     }
+    case BuiltinOptions_LessOptions: {
+      auto ptr = reinterpret_cast<const LessOptions *>(obj);
+      return ptr->UnPack(resolver);
+    }
+    case BuiltinOptions_NegOptions: {
+      auto ptr = reinterpret_cast<const NegOptions *>(obj);
+      return ptr->UnPack(resolver);
+    }
+    case BuiltinOptions_PadV2Options: {
+      auto ptr = reinterpret_cast<const PadV2Options *>(obj);
+      return ptr->UnPack(resolver);
+    }
     case BuiltinOptions_GreaterOptions: {
       auto ptr = reinterpret_cast<const GreaterOptions *>(obj);
       return ptr->UnPack(resolver);
@@ -7082,18 +7090,10 @@ inline void *BuiltinOptionsUnion::UnPack(const void *obj, BuiltinOptions type, c
       auto ptr = reinterpret_cast<const GreaterEqualOptions *>(obj);
       return ptr->UnPack(resolver);
     }
-    case BuiltinOptions_LessOptions: {
-      auto ptr = reinterpret_cast<const LessOptions *>(obj);
-      return ptr->UnPack(resolver);
-    }
     case BuiltinOptions_LessEqualOptions: {
       auto ptr = reinterpret_cast<const LessEqualOptions *>(obj);
       return ptr->UnPack(resolver);
     }
-    case BuiltinOptions_NegOptions: {
-      auto ptr = reinterpret_cast<const NegOptions *>(obj);
-      return ptr->UnPack(resolver);
-    }
     default: return nullptr;
   }
 }
@@ -7188,10 +7188,6 @@ inline flatbuffers::Offset<void> BuiltinOptionsUnion::Pack(flatbuffers::FlatBuff
       auto ptr = reinterpret_cast<const PadOptionsT *>(value);
       return CreatePadOptions(_fbb, ptr, _rehasher).Union();
     }
-    case BuiltinOptions_PadV2Options: {
-      auto ptr = reinterpret_cast<const PadV2OptionsT *>(value);
-      return CreatePadV2Options(_fbb, ptr, _rehasher).Union();
-    }
     case BuiltinOptions_GatherOptions: {
       auto ptr = reinterpret_cast<const GatherOptionsT *>(value);
       return CreateGatherOptions(_fbb, ptr, _rehasher).Union();
@@ -7264,6 +7260,18 @@ inline flatbuffers::Offset<void> BuiltinOptionsUnion::Pack(flatbuffers::FlatBuff
       auto ptr = reinterpret_cast<const ArgMaxOptionsT *>(value);
       return CreateArgMaxOptions(_fbb, ptr, _rehasher).Union();
     }
+    case BuiltinOptions_LessOptions: {
+      auto ptr = reinterpret_cast<const LessOptionsT *>(value);
+      return CreateLessOptions(_fbb, ptr, _rehasher).Union();
+    }
+    case BuiltinOptions_NegOptions: {
+      auto ptr = reinterpret_cast<const NegOptionsT *>(value);
+      return CreateNegOptions(_fbb, ptr, _rehasher).Union();
+    }
+    case BuiltinOptions_PadV2Options: {
+      auto ptr = reinterpret_cast<const PadV2OptionsT *>(value);
+      return CreatePadV2Options(_fbb, ptr, _rehasher).Union();
+    }
     case BuiltinOptions_GreaterOptions: {
       auto ptr = reinterpret_cast<const GreaterOptionsT *>(value);
       return CreateGreaterOptions(_fbb, ptr, _rehasher).Union();
@@ -7272,18 +7280,10 @@ inline flatbuffers::Offset<void> BuiltinOptionsUnion::Pack(flatbuffers::FlatBuff
       auto ptr = reinterpret_cast<const GreaterEqualOptionsT *>(value);
       return CreateGreaterEqualOptions(_fbb, ptr, _rehasher).Union();
     }
-    case BuiltinOptions_LessOptions: {
-      auto ptr = reinterpret_cast<const LessOptionsT *>(value);
-      return CreateLessOptions(_fbb, ptr, _rehasher).Union();
-    }
     case BuiltinOptions_LessEqualOptions: {
       auto ptr = reinterpret_cast<const LessEqualOptionsT *>(value);
       return CreateLessEqualOptions(_fbb, ptr, _rehasher).Union();
     }
-    case BuiltinOptions_NegOptions: {
-      auto ptr = reinterpret_cast<const NegOptionsT *>(value);
-      return CreateNegOptions(_fbb, ptr, _rehasher).Union();
-    }
     default: return 0;
   }
 }
@@ -7378,10 +7378,6 @@ inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) FL
       value = new PadOptionsT(*reinterpret_cast<PadOptionsT *>(u.value));
       break;
     }
-    case BuiltinOptions_PadV2Options: {
-      value = new PadV2OptionsT(*reinterpret_cast<PadV2OptionsT *>(u.value));
-      break;
-    }
     case BuiltinOptions_GatherOptions: {
       value = new GatherOptionsT(*reinterpret_cast<GatherOptionsT *>(u.value));
       break;
@@ -7454,6 +7450,18 @@ inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) FL
       value = new ArgMaxOptionsT(*reinterpret_cast<ArgMaxOptionsT *>(u.value));
       break;
     }
+    case BuiltinOptions_LessOptions: {
+      value = new LessOptionsT(*reinterpret_cast<LessOptionsT *>(u.value));
+      break;
+    }
+    case BuiltinOptions_NegOptions: {
+      value = new NegOptionsT(*reinterpret_cast<NegOptionsT *>(u.value));
+      break;
+    }
+    case BuiltinOptions_PadV2Options: {
+      value = new PadV2OptionsT(*reinterpret_cast<PadV2OptionsT *>(u.value));
+      break;
+    }
     case BuiltinOptions_GreaterOptions: {
       value = new GreaterOptionsT(*reinterpret_cast<GreaterOptionsT *>(u.value));
       break;
@@ -7462,18 +7470,10 @@ inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) FL
       value = new GreaterEqualOptionsT(*reinterpret_cast<GreaterEqualOptionsT *>(u.value));
       break;
     }
-    case BuiltinOptions_LessOptions: {
-      value = new LessOptionsT(*reinterpret_cast<LessOptionsT *>(u.value));
-      break;
-    }
     case BuiltinOptions_LessEqualOptions: {
       value = new LessEqualOptionsT(*reinterpret_cast<LessEqualOptionsT *>(u.value));
       break;
     }
-    case BuiltinOptions_NegOptions: {
-      value = new NegOptionsT(*reinterpret_cast<NegOptionsT *>(u.value));
-      break;
-    }
     default:
       break;
   }
@@ -7591,11 +7591,6 @@ inline void BuiltinOptionsUnion::Reset() {
       delete ptr;
       break;
     }
-    case BuiltinOptions_PadV2Options: {
-      auto ptr = reinterpret_cast<PadV2OptionsT *>(value);
-      delete ptr;
-      break;
-    }
     case BuiltinOptions_GatherOptions: {
       auto ptr = reinterpret_cast<GatherOptionsT *>(value);
       delete ptr;
@@ -7686,28 +7681,33 @@ inline void BuiltinOptionsUnion::Reset() {
       delete ptr;
       break;
     }
-    case BuiltinOptions_GreaterOptions: {
-      auto ptr = reinterpret_cast<GreaterOptionsT *>(value);
+    case BuiltinOptions_LessOptions: {
+      auto ptr = reinterpret_cast<LessOptionsT *>(value);
       delete ptr;
       break;
     }
-    case BuiltinOptions_GreaterEqualOptions: {
-      auto ptr = reinterpret_cast<GreaterEqualOptionsT *>(value);
+    case BuiltinOptions_NegOptions: {
+      auto ptr = reinterpret_cast<NegOptionsT *>(value);
       delete ptr;
       break;
     }
-    case BuiltinOptions_LessOptions: {
-      auto ptr = reinterpret_cast<LessOptionsT *>(value);
+    case BuiltinOptions_PadV2Options: {
+      auto ptr = reinterpret_cast<PadV2OptionsT *>(value);
       delete ptr;
       break;
     }
-    case BuiltinOptions_LessEqualOptions: {
-      auto ptr = reinterpret_cast<LessEqualOptionsT *>(value);
+    case BuiltinOptions_GreaterOptions: {
+      auto ptr = reinterpret_cast<GreaterOptionsT *>(value);
       delete ptr;
       break;
     }
-    case BuiltinOptions_NegOptions: {
-      auto ptr = reinterpret_cast<NegOptionsT *>(value);
+    case BuiltinOptions_GreaterEqualOptions: {
+      auto ptr = reinterpret_cast<GreaterEqualOptionsT *>(value);
+      delete ptr;
+      break;
+    }
+    case BuiltinOptions_LessEqualOptions: {
+      auto ptr = reinterpret_cast<LessEqualOptionsT *>(value);
       delete ptr;
       break;
     }
-- 
GitLab


From cd065ca7be11a4c87c9a5e68271cbc2d9aaaa260 Mon Sep 17 00:00:00 2001
From: Justin Lebar <jlebar@google.com>
Date: Mon, 7 May 2018 14:23:19 -0700
Subject: [PATCH 323/395] [XLA] Shard compilation of HloEvaluator.

PiperOrigin-RevId: 195721404
---
 tensorflow/compiler/xla/service/BUILD         |   17 +-
 .../compiler/xla/service/hlo_evaluator.cc     | 2102 +----------------
 .../compiler/xla/service/hlo_evaluator.h      |   48 +-
 .../xla/service/hlo_evaluator_typed_visitor.h | 2102 +++++++++++++++++
 .../hlo_evaluator_typed_visitor_bfloat16.cc   |   22 +
 .../hlo_evaluator_typed_visitor_bool.cc       |   22 +
 .../hlo_evaluator_typed_visitor_complex64.cc  |   22 +
 .../hlo_evaluator_typed_visitor_double.cc     |   22 +
 .../hlo_evaluator_typed_visitor_float.cc      |   22 +
 .../hlo_evaluator_typed_visitor_half.cc       |   22 +
 .../hlo_evaluator_typed_visitor_int32.cc      |   22 +
 .../hlo_evaluator_typed_visitor_int64.cc      |   22 +
 .../hlo_evaluator_typed_visitor_int8.cc       |   22 +
 .../hlo_evaluator_typed_visitor_uint32.cc     |   22 +
 .../hlo_evaluator_typed_visitor_uint64.cc     |   22 +
 .../hlo_evaluator_typed_visitor_uint8.cc      |   22 +
 16 files changed, 2440 insertions(+), 2093 deletions(-)
 create mode 100644 tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h
 create mode 100644 tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_bfloat16.cc
 create mode 100644 tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_bool.cc
 create mode 100644 tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_complex64.cc
 create mode 100644 tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_double.cc
 create mode 100644 tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_float.cc
 create mode 100644 tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_half.cc
 create mode 100644 tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int32.cc
 create mode 100644 tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int64.cc
 create mode 100644 tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int8.cc
 create mode 100644 tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint32.cc
 create mode 100644 tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint64.cc
 create mode 100644 tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint8.cc

diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD
index 714c1e8754..ec67e19b23 100644
--- a/tensorflow/compiler/xla/service/BUILD
+++ b/tensorflow/compiler/xla/service/BUILD
@@ -200,7 +200,22 @@ tf_cc_test(
 
 cc_library(
     name = "hlo_evaluator",
-    srcs = ["hlo_evaluator.cc"],
+    srcs = [
+        "hlo_evaluator.cc",
+        "hlo_evaluator_typed_visitor.h",
+        "hlo_evaluator_typed_visitor_bfloat16.cc",
+        "hlo_evaluator_typed_visitor_bool.cc",
+        "hlo_evaluator_typed_visitor_complex64.cc",
+        "hlo_evaluator_typed_visitor_double.cc",
+        "hlo_evaluator_typed_visitor_float.cc",
+        "hlo_evaluator_typed_visitor_half.cc",
+        "hlo_evaluator_typed_visitor_int32.cc",
+        "hlo_evaluator_typed_visitor_int64.cc",
+        "hlo_evaluator_typed_visitor_int8.cc",
+        "hlo_evaluator_typed_visitor_uint32.cc",
+        "hlo_evaluator_typed_visitor_uint64.cc",
+        "hlo_evaluator_typed_visitor_uint8.cc",
+    ],
     hdrs = ["hlo_evaluator.h"],
     deps = [
         ":hlo",
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.cc b/tensorflow/compiler/xla/service/hlo_evaluator.cc
index 8cf94123b7..fffe1923ba 100644
--- a/tensorflow/compiler/xla/service/hlo_evaluator.cc
+++ b/tensorflow/compiler/xla/service/hlo_evaluator.cc
@@ -29,6 +29,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/map_util.h"
 #include "tensorflow/compiler/xla/primitive_util.h"
 #include "tensorflow/compiler/xla/ptr_util.h"
+#include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
 #include "tensorflow/compiler/xla/service/hlo_opcode.h"
 #include "tensorflow/compiler/xla/service/hlo_query.h"
@@ -42,7 +43,6 @@ limitations under the License.
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/core/stringpiece.h"
-#include "tensorflow/core/lib/gtl/optional.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/protobuf.h"
 #include "tensorflow/core/platform/types.h"
@@ -53,19 +53,6 @@ namespace {
 
 using tensorflow::gtl::ArraySlice;
 using tensorflow::gtl::FlatSet;
-using tensorflow::gtl::optional;
-
-template <typename T>
-struct is_complex_t : public std::false_type {};
-
-template <>
-struct is_complex_t<complex64> : public std::true_type {};
-
-template <typename T>
-struct is_complex64_t : public std::false_type {};
-
-template <>
-struct is_complex64_t<complex64> : public std::true_type {};
 
 template <typename OperandT>
 StatusOr<std::unique_ptr<Literal>> Compare(const Shape& shape, HloOpcode opcode,
@@ -147,2092 +134,47 @@ StatusOr<std::unique_ptr<Literal>> Compare<complex64>(
   return std::move(result);
 }
 
-template <typename ReturnT, typename NativeT>
-StatusOr<std::unique_ptr<Literal>> ElementWiseUnaryOpImpl(
-    HloInstruction* instruction,
-    const std::function<ReturnT(NativeT)>& unary_op,
-    const Literal& operand_literal) {
-  const auto shape = instruction->shape();
-  const auto* operand = instruction->operand(0);
-
-  // TODO(b/35950897, b/27796129): add DCHECK back once implicit broadcast is
-  // removed.
-  if (!ShapeUtil::SameDimensions(shape, operand->shape())) {
-    return Unimplemented(
-        "Implicit broadcasting is currently unsupported in HLO evaluator "
-        "Shape Mismatch: %s vs %s",
-        ShapeUtil::HumanString(shape).c_str(),
-        ShapeUtil::HumanString(operand->shape()).c_str());
-  }
-
-  auto result = Literal::CreateFromShape(shape);
-
-  TF_RETURN_IF_ERROR(
-      result->Populate<ReturnT>([&](ArraySlice<int64> multi_index) {
-        return unary_op(operand_literal.Get<NativeT>(multi_index));
-      }));
-  return std::move(result);
-}
-
-// For one particular placement of a window in a base shape (the placement is
-// represented as `window_count_index`), iterates inside the window. Translates
-// the window index into base index. If the base index is within bound, call `f`
-// with the base index.
-void IterateThroughWindow(
-    const Shape& window_shape, const Window& window, const Shape& base_shape,
-    const ArraySlice<int64>& window_count_index,
-    const std::function<void(const std::vector<int64>&)>& f) {
-  const int64 rank = ShapeUtil::Rank(base_shape);
-  DimensionVector window_index(rank);
-  std::fill(window_index.begin(), window_index.end(), 0);
-  do {
-    std::vector<int64> base_index(rank);
-    bool out_of_bound = false;
-    for (int64 i = 0; i < rank; ++i) {
-      base_index[i] = window_count_index[i] * window.dimensions(i).stride() +
-                      window_index[i] - window.dimensions(i).padding_low();
-      if (base_index[i] < 0 || base_index[i] >= base_shape.dimensions(i)) {
-        out_of_bound = true;
-        break;
-      }
-    }
-    if (!out_of_bound) {
-      f(base_index);
-    }
-  } while (IndexUtil::BumpIndices(window_shape, &window_index));
-}
-
-// Creates a vector of multipliers which can be used to create a linear index
-// into shape.
-//
-// Given the multidimensional index {i1, ..., iN} and
-// M = MakeDimMultipliers(shape), the corresponding linear index LI is simply
-//
-//   LI = i1 * M[1] + i2 * M[2] + ... + iN * M[N].
-//
-// This lets you calculate LI given the multidimensional indices in any order.
-DimensionVector MakeDimMultipliers(const Shape& shape) {
-  DimensionVector v(ShapeUtil::Rank(shape));
-  int64 scale = 1;
-  for (auto dim : LayoutUtil::MinorToMajor(shape)) {
-    v[dim] = scale;
-    scale *= shape.dimensions(dim);
-  }
-  return v;
-}
-
 }  // namespace
 
-template <typename ReturnT, typename ElementwiseT>
-class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault {
- public:
-  explicit TypedVisitor(HloEvaluator* p) : parent_(p) {}
-
-  // The following higher-order functions convert a function with ElementwiseT
-  // to a function with ReturnT.
-  std::function<ReturnT(ReturnT)> ConvertUnaryFunction(
-      const std::function<ElementwiseT(ElementwiseT)>& unary_op) {
-    return [&unary_op](ReturnT arg) {
-      return static_cast<ReturnT>(unary_op(static_cast<ElementwiseT>(arg)));
-    };
-  }
-  std::function<ReturnT(ReturnT, ReturnT)> ConvertBinaryFunction(
-      const std::function<ElementwiseT(ElementwiseT, ElementwiseT)>&
-          binary_op) {
-    return [&binary_op](ReturnT arg1, ReturnT arg2) {
-      return static_cast<ReturnT>(binary_op(static_cast<ElementwiseT>(arg1),
-                                            static_cast<ElementwiseT>(arg2)));
-    };
-  }
-  std::function<ReturnT(ReturnT, ReturnT, ReturnT)> ConvertTernaryFunction(
-      const std::function<ElementwiseT(ElementwiseT, ElementwiseT,
-                                       ElementwiseT)>& ternary_op) {
-    return [&ternary_op](ReturnT arg1, ReturnT arg2, ReturnT arg3) {
-      return static_cast<ReturnT>(ternary_op(static_cast<ElementwiseT>(arg1),
-                                             static_cast<ElementwiseT>(arg2),
-                                             static_cast<ElementwiseT>(arg3)));
-    };
-  }
-
-  Status DefaultAction(HloInstruction* hlo_instruction) override {
-    return Unimplemented("unhandled HLO ops for HloEvaluator: %s.",
-                         HloOpcodeString(hlo_instruction->opcode()).c_str());
-  }
-
-  // TODO(b/35950897): many of the stl functions used in the handlers are not
-  // overloaded for every XLA primitive types.
-
-  template <typename NativeT,
-            typename std::enable_if<std::is_unsigned<NativeT>::value>::type* =
-                nullptr>
-  Status HandleAbs(HloInstruction* abs) {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[abs],
-                        ElementWiseUnaryOp(abs, [](NativeT elem_operand) {
-                          return elem_operand;
-                        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<std::is_signed<NativeT>::value>::type* = nullptr>
-  Status HandleAbs(HloInstruction* abs) {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[abs],
-                        ElementWiseUnaryOp(abs, [](NativeT elem_operand) {
-                          return std::abs(elem_operand);
-                        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<is_complex64_t<NativeT>::value>::type* = nullptr>
-  Status HandleAbs(HloInstruction* abs) {
-    const Literal& operand_literal =
-        parent_->GetEvaluatedLiteralFor(abs->operand(0));
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[abs],
-        (ElementWiseUnaryOpImpl<float, NativeT>(
-            abs, [](NativeT elem_operand) { return std::abs(elem_operand); },
-            operand_literal)));
-
-    return Status::OK();
-  }
-
-  Status HandleAbs(HloInstruction* abs) override {
-    // If the operand is of C64 type, the return type of abs will be F32.
-    // However, ElementwiseT would still be the return type, F32, and thus
-    // specifying the ElementwiseT explicitly as C64 is needed below.
-    if (abs->operand(0)->shape().element_type() == C64) {
-      return HandleAbs<complex64>(abs);
-    }
-    return HandleAbs<ElementwiseT>(abs);
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<!is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleRound(HloInstruction* round) {
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[round],
-        ElementWiseUnaryOp(round, [](ElementwiseT elem_operand) {
-          return std::round(elem_operand);
-        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleRound(HloInstruction* round) {
-    return InvalidArgument("Unsupported type for Round");
-  }
-
-  Status HandleRound(HloInstruction* round) override {
-    return HandleRound<ReturnT>(round);
-  }
-
-  Status HandleBroadcast(HloInstruction* broadcast) override {
-    parent_->evaluated_[broadcast] =
-        Literal::CreateFromShape(broadcast->shape());
-    auto output = parent_->evaluated_[broadcast].get();
-    const Literal& operand_to_broadcast =
-        parent_->GetEvaluatedLiteralFor(broadcast->operand(0));
-    std::vector<int64> broadcast_indices(
-        ShapeUtil::Rank(broadcast->operand(0)->shape()), 0);
-
-    TF_RET_CHECK(broadcast->dimensions().size() ==
-                 ShapeUtil::Rank(operand_to_broadcast.shape()))
-        << "broadcast dimensions is of size: " << broadcast->dimensions().size()
-        << " and rank of operand_to_broadcast is: "
-        << ShapeUtil::Rank(operand_to_broadcast.shape());
-    // Checks that operand's dimensions are the same as the broadcast's
-    // dimensions along the dimensions to be broadcasted.
-    for (int64 i = 0; i < broadcast->dimensions().size(); ++i) {
-      TF_RET_CHECK(broadcast->shape().dimensions(broadcast->dimensions(i)) ==
-                   operand_to_broadcast.shape().dimensions(i));
-    }
-
-    return output->Populate<ReturnT>([&](ArraySlice<int64> multi_index) {
-      for (int64 i = 0; i < broadcast->dimensions().size(); ++i) {
-        broadcast_indices[i] = multi_index[broadcast->dimensions(i)];
-      }
-      return operand_to_broadcast.Get<ReturnT>(broadcast_indices);
-    });
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<!is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleCeil(HloInstruction* ceil) {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[ceil],
-                        ElementWiseUnaryOp(ceil, [](ElementwiseT elem_operand) {
-                          return std::ceil(elem_operand);
-                        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleCeil(HloInstruction* ceil) {
-    return InvalidArgument("Unsupported type for Ceil");
-  }
-
-  Status HandleCeil(HloInstruction* ceil) override {
-    return HandleCeil<ReturnT>(ceil);
-  }
-
-  Status HandleConvert(HloInstruction* convert) override {
-    const HloInstruction* operand = convert->operand(0);
-    TF_RET_CHECK(ShapeUtil::SameDimensions(operand->shape(), convert->shape()));
-    TF_ASSIGN_OR_RETURN(std::unique_ptr<Literal> result,
-                        parent_->GetEvaluatedLiteralFor(operand).Convert(
-                            convert->shape().element_type()));
-
-    if (LayoutUtil::LayoutsInShapesEqual(result->shape(), convert->shape())) {
-      parent_->evaluated_[convert] = std::move(result);
-    } else {
-      parent_->evaluated_[convert] =
-          result->Relayout(convert->shape().layout());
-    }
-    return Status::OK();
-  }
-
-  Status HandleBitcastConvert(HloInstruction* convert) override {
-    const HloInstruction* operand = convert->operand(0);
-    TF_RET_CHECK(ShapeUtil::SameDimensions(operand->shape(), convert->shape()));
-    TF_ASSIGN_OR_RETURN(std::unique_ptr<Literal> result,
-                        parent_->GetEvaluatedLiteralFor(operand).BitcastConvert(
-                            convert->shape().element_type()));
-
-    if (LayoutUtil::LayoutsInShapesEqual(result->shape(), convert->shape())) {
-      parent_->evaluated_[convert] = std::move(result);
-    } else {
-      parent_->evaluated_[convert] =
-          result->Relayout(convert->shape().layout());
-    }
-    return Status::OK();
-  }
-
-  Status HandleExp(HloInstruction* exp) override {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[exp],
-                        ElementWiseUnaryOp(exp, [](ElementwiseT elem_operand) {
-                          return std::exp(elem_operand);
-                        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<!is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleFloor(HloInstruction* floor) {
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[floor],
-        ElementWiseUnaryOp(floor, [](ElementwiseT elem_operand) {
-          return std::floor(elem_operand);
-        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleFloor(HloInstruction* floor) {
-    return InvalidArgument("Unsupported type for Floor");
-  }
-
-  Status HandleFloor(HloInstruction* floor) override {
-    return HandleFloor<ReturnT>(floor);
-  }
-
-  Status HandleLog(HloInstruction* log) override {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[log],
-                        ElementWiseUnaryOp(log, [](ElementwiseT elem_operand) {
-                          return std::log(elem_operand);
-                        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<
-                std::is_integral<NativeT>::value &&
-                !std::is_same<NativeT, bool>::value>::type* = nullptr>
-  Status HandleNot(HloInstruction* not_) {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[not_],
-                        ElementWiseUnaryOp(not_, [](ElementwiseT elem_operand) {
-                          return ~elem_operand;
-                        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT, typename std::enable_if<std::is_floating_point<
-                                  NativeT>::value>::type* = nullptr>
-  Status HandleNot(HloInstruction* not_) {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[not_],
-                        ElementWiseUnaryOp(not_, [](ElementwiseT elem_operand) {
-                          return !elem_operand;
-                        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<std::is_same<NativeT, bool>::value>::type* =
-                nullptr>
-  Status HandleNot(HloInstruction* not_) {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[not_],
-                        ElementWiseUnaryOp(not_, [](ElementwiseT elem_operand) {
-                          return !elem_operand;
-                        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleNot(HloInstruction* not_) {
-    return InvalidArgument("Unsupported type for Not");
-  }
-
-  Status HandleNot(HloInstruction* not_) override {
-    return HandleNot<ElementwiseT>(not_);
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<
-                std::is_signed<NativeT>::value &&
-                !std::is_floating_point<NativeT>::value>::type* = nullptr>
-  Status HandleNegate(HloInstruction* negate) {
-    using type = typename std::make_unsigned<NativeT>::type;
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[negate],
-        ElementWiseUnaryOp(negate, [](ElementwiseT elem_operand) {
-          return NativeT(-type(elem_operand));
-        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<
-                !std::is_signed<NativeT>::value ||
-                std::is_floating_point<NativeT>::value>::type* = nullptr>
-  Status HandleNegate(HloInstruction* negate) {
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[negate],
-        ElementWiseUnaryOp(
-            negate, [](ElementwiseT elem_operand) { return -elem_operand; }));
-    return Status::OK();
-  }
-
-  Status HandleNegate(HloInstruction* negate) override {
-    return HandleNegate<ReturnT>(negate);
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<!is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleSign(HloInstruction* sign) {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[sign],
-                        ElementWiseUnaryOp(sign, [](ElementwiseT elem_operand) {
-                          return (ElementwiseT(0) < elem_operand) -
-                                 (elem_operand < ElementwiseT(0));
-                        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleSign(HloInstruction* sign) {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[sign],
-                        ElementWiseUnaryOp(sign, [](ElementwiseT elem_operand) {
-                          auto abs_val = std::abs(elem_operand);
-                          return 0 == abs_val ? ElementwiseT(0)
-                                              : elem_operand / abs_val;
-                        }));
-    return Status::OK();
-  }
-
-  Status HandleSign(HloInstruction* sign) override {
-    return HandleSign<ReturnT>(sign);
-  }
-
-  template <typename NativeT, typename std::enable_if<std::is_floating_point<
-                                  NativeT>::value>::type* = nullptr>
-  Status HandleAtan2(HloInstruction* atan2) {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[atan2],
-                        ElementWiseBinaryOp(atan2, [](ElementwiseT lhs_elem,
-                                                      ElementwiseT rhs_elem) {
-                          return std::atan2(lhs_elem, rhs_elem);
-                        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT, typename std::enable_if<!std::is_floating_point<
-                                  NativeT>::value>::type* = nullptr>
-  Status HandleAtan2(HloInstruction* atan2) {
-    return InvalidArgument("Unsupported type for Atan2");
-  }
-
-  Status HandleAtan2(HloInstruction* atan2) override {
-    return HandleAtan2<ElementwiseT>(atan2);
-  }
-
-  Status HandleTanh(HloInstruction* tanh) override {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[tanh],
-                        ElementWiseUnaryOp(tanh, [](ElementwiseT elem_operand) {
-                          return std::tanh(elem_operand);
-                        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<
-                std::is_signed<NativeT>::value &&
-                !std::is_floating_point<NativeT>::value>::type* = nullptr>
-  Status HandleMultiply(HloInstruction* multiply) {
-    using type = typename std::make_unsigned<NativeT>::type;
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[multiply],
-        ElementWiseBinaryOp(multiply,
-                            [](ElementwiseT lhs_elem, ElementwiseT rhs_elem) {
-                              return NativeT(type(lhs_elem) * type(rhs_elem));
-                            }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<std::is_unsigned<NativeT>::value ||
-                              std::is_floating_point<NativeT>::value ||
-                              is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleMultiply(HloInstruction* multiply) {
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[multiply],
-        ElementWiseBinaryOp(multiply,
-                            [](ElementwiseT lhs_elem, ElementwiseT rhs_elem) {
-                              return lhs_elem * rhs_elem;
-                            }));
-    return Status::OK();
-  }
-
-  Status HandleMultiply(HloInstruction* multiply) override {
-    return HandleMultiply<ElementwiseT>(multiply);
-  }
-
-  Status HandleSubtract(HloInstruction* subtract) override {
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[subtract],
-        ElementWiseBinaryOp(subtract,
-                            [](ElementwiseT lhs_elem, ElementwiseT rhs_elem) {
-                              return lhs_elem - rhs_elem;
-                            }));
-    return Status::OK();
-  }
-
-  Status HandleAdd(HloInstruction* add) override {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[add],
-                        ElementWiseBinaryOp(add, [](ElementwiseT lhs_elem,
-                                                    ElementwiseT rhs_elem) {
-                          return lhs_elem + rhs_elem;
-                        }));
-    return Status::OK();
-  }
-
-  Status HandleDivide(HloInstruction* divide) override {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[divide],
-                        ElementWiseBinaryOp(divide, [](ElementwiseT lhs_elem,
-                                                       ElementwiseT rhs_elem) {
-                          return lhs_elem / rhs_elem;
-                        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<std::is_integral<NativeT>::value>::type* =
-                nullptr>
-  Status HandleMaximum(HloInstruction* maximum) {
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[maximum],
-        ElementWiseBinaryOp(maximum, [](ElementwiseT lhs, ElementwiseT rhs) {
-          return std::max(lhs, rhs);
-        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT, typename std::enable_if<std::is_floating_point<
-                                  NativeT>::value>::type* = nullptr>
-  Status HandleMaximum(HloInstruction* maximum) {
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[maximum],
-        ElementWiseBinaryOp(maximum, [](ElementwiseT lhs, ElementwiseT rhs) {
-          return ((lhs >= rhs) || std::isnan(lhs)) ? lhs : rhs;
-        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleMaximum(HloInstruction* maximum) {
-    return InvalidArgument("Unsupported type for Maximum");
-  }
-
-  Status HandleMaximum(HloInstruction* maximum) override {
-    return HandleMaximum<ElementwiseT>(maximum);
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<std::is_integral<NativeT>::value>::type* =
-                nullptr>
-  Status HandleMinimum(HloInstruction* minimum) {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[minimum],
-                        ElementWiseBinaryOp(minimum, [](ElementwiseT lhs_el,
-                                                        ElementwiseT rhs_el) {
-                          return std::min(lhs_el, rhs_el);
-                        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT, typename std::enable_if<std::is_floating_point<
-                                  NativeT>::value>::type* = nullptr>
-  Status HandleMinimum(HloInstruction* minimum) {
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[minimum],
-        ElementWiseBinaryOp(minimum, [](ElementwiseT lhs_el,
-                                        ElementwiseT rhs_el) {
-          return ((lhs_el <= rhs_el) || std::isnan(lhs_el)) ? lhs_el : rhs_el;
-        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleMinimum(HloInstruction* minimum) {
-    return InvalidArgument("Unsupported type for Minimum");
-  }
-
-  Status HandleMinimum(HloInstruction* minimum) override {
-    return HandleMinimum<ElementwiseT>(minimum);
-  }
-
-  Status HandlePower(HloInstruction* power) override {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[power],
-                        ElementWiseBinaryOp(power, [](ElementwiseT lhs_el,
-                                                      ElementwiseT rhs_el) {
-                          return std::pow(lhs_el, rhs_el);
-                        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<!is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleRemainder(HloInstruction* remainder) {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[remainder],
-                        ElementWiseBinaryOp(remainder, [](ElementwiseT lhs_el,
-                                                          ElementwiseT rhs_el) {
-                          return std::fmod(lhs_el, rhs_el);
-                        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleRemainder(HloInstruction* remainder) {
-    return InvalidArgument("Unsupported type for Remainder");
-  }
-
-  Status HandleRemainder(HloInstruction* remainder) override {
-    return HandleRemainder<ElementwiseT>(remainder);
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<std::is_integral<NativeT>::value>::type* =
-                nullptr>
-  Status HandleAnd(HloInstruction* and_) {
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[and_],
-        ElementWiseBinaryOp(and_, [](ElementwiseT lhs_el, ElementwiseT rhs_el) {
-          return lhs_el & rhs_el;
-        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT, typename std::enable_if<std::is_floating_point<
-                                  NativeT>::value>::type* = nullptr>
-  Status HandleAnd(HloInstruction* and_) {
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[and_],
-        ElementWiseBinaryOp(and_, [](ElementwiseT lhs_el, ElementwiseT rhs_el) {
-          return lhs_el && rhs_el;
-        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleAnd(HloInstruction* and_) {
-    return InvalidArgument("Unsupported type for And");
-  }
-
-  Status HandleAnd(HloInstruction* and_) override {
-    return HandleAnd<ElementwiseT>(and_);
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<std::is_integral<NativeT>::value>::type* =
-                nullptr>
-  Status HandleOr(HloInstruction* or_) {
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[or_],
-        ElementWiseBinaryOp(or_, [](ElementwiseT lhs_el, ElementwiseT rhs_el) {
-          return lhs_el | rhs_el;
-        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT, typename std::enable_if<std::is_floating_point<
-                                  NativeT>::value>::type* = nullptr>
-  Status HandleOr(HloInstruction* or_) {
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[or_],
-        ElementWiseBinaryOp(or_, [](ElementwiseT lhs_el, ElementwiseT rhs_el) {
-          return lhs_el || rhs_el;
-        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleOr(HloInstruction* or_) {
-    return InvalidArgument("Unsupported type for Or");
-  }
-
-  Status HandleOr(HloInstruction* or_) override {
-    return HandleOr<ElementwiseT>(or_);
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<
-                std::is_integral<NativeT>::value &&
-                !std::is_same<NativeT, bool>::value>::type* = nullptr>
-  Status HandleShiftLeft(HloInstruction* shl) {
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[shl],
-        ElementWiseBinaryOp(shl, [](NativeT lhs_elem, NativeT rhs_elem) {
-          return IsShiftOutOfBounds<NativeT>(rhs_elem) ? 0
-                                                       : (lhs_elem << rhs_elem);
-        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<!std::is_integral<NativeT>::value ||
-                                    std::is_same<NativeT, bool>::value>::type* =
-                nullptr>
-  Status HandleShiftLeft(HloInstruction*) {
-    return InvalidArgument("Unsupported type for ShiftLeft");
-  }
-
-  Status HandleShiftLeft(HloInstruction* shl) override {
-    return HandleShiftLeft<ElementwiseT>(shl);
-  }
-  template <typename NativeT,
-            typename std::enable_if<
-                std::is_integral<NativeT>::value &&
-                !std::is_same<NativeT, bool>::value>::type* = nullptr>
-  Status HandleShiftRightArithmetic(HloInstruction* shr) {
-    typedef typename std::make_signed<NativeT>::type SignedT;
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[shr],
-        ElementWiseBinaryOp(shr, [](NativeT lhs_elem, NativeT rhs_elem) {
-          SignedT lhs_signed = static_cast<SignedT>(lhs_elem);
-          if (IsShiftOutOfBounds<NativeT>(rhs_elem)) {
-            return lhs_signed < 0 ? static_cast<SignedT>(-1) : 0;
-          } else {
-            return lhs_signed >> rhs_elem;
-          }
-        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<!std::is_integral<NativeT>::value ||
-                                    std::is_same<NativeT, bool>::value>::type* =
-                nullptr>
-  Status HandleShiftRightArithmetic(HloInstruction*) {
-    return InvalidArgument("Unsupported type for ShiftRightArithmetic");
-  }
-
-  Status HandleShiftRightArithmetic(HloInstruction* shra) override {
-    return HandleShiftRightArithmetic<ElementwiseT>(shra);
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<
-                std::is_integral<NativeT>::value &&
-                !std::is_same<NativeT, bool>::value>::type* = nullptr>
-  Status HandleShiftRightLogical(HloInstruction* shr) {
-    typedef typename std::make_unsigned<NativeT>::type UnsignedT;
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[shr],
-        ElementWiseBinaryOp(shr, [](NativeT lhs_elem, NativeT rhs_elem) {
-          // If shift amount is greater than the number of bits, then return 0.
-          if (IsShiftOutOfBounds<NativeT>(rhs_elem)) {
-            return static_cast<NativeT>(0);
-          }
-          return static_cast<NativeT>(static_cast<UnsignedT>(lhs_elem) >>
-                                      rhs_elem);
-        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<!std::is_integral<NativeT>::value ||
-                                    std::is_same<NativeT, bool>::value>::type* =
-                nullptr>
-  Status HandleShiftRightLogical(HloInstruction*) {
-    return InvalidArgument("Unsupported type for ShiftRightLogical");
-  }
-
-  Status HandleShiftRightLogical(HloInstruction* shrl) override {
-    return HandleShiftRightLogical<ElementwiseT>(shrl);
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<!is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleClamp(HloInstruction* clamp) {
-    std::function<ElementwiseT(ElementwiseT, ElementwiseT, ElementwiseT)>
-        clamp_op = [](ElementwiseT low, ElementwiseT value, ElementwiseT high) {
-          return std::fmin(high, std::fmax(value, low));
-        };
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[clamp],
-        ElementwiseTernaryOp(clamp,
-                             std::move(ConvertTernaryFunction(clamp_op))));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleClamp(HloInstruction*) {
-    return InvalidArgument("Unsupported type for Clamp");
-  }
-
-  Status HandleClamp(HloInstruction* clamp) override {
-    return HandleClamp<ElementwiseT>(clamp);
-  }
-
-  Status HandleSelect(HloInstruction* select) override {
-    CHECK(!ShapeUtil::IsScalar(select->operand(0)->shape()));
-    CHECK(!ShapeUtil::IsTuple(select->shape()));
-    std::function<ReturnT(bool, ReturnT, ReturnT)> select_op =
-        [](bool pred, ReturnT on_true, ReturnT on_false) {
-          if (pred) {
-            return on_true;
-          }
-          return on_false;
-        };
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[select],
-                        ElementwiseTernaryOp(select, std::move(select_op)));
-    return Status::OK();
-  }
-
-  Status HandleReverse(HloInstruction* reverse) override {
-    const auto result_shape = reverse->shape();
-    const auto reverse_dimensions = reverse->dimensions();
-
-    auto operand = reverse->operand(0);
-    TF_ASSIGN_OR_RETURN(auto inferred_return_shape,
-                        ShapeInference::InferReverseShape(operand->shape(),
-                                                          reverse_dimensions));
-
-    TF_RET_CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape))
-        << "return shape set to: " << ShapeUtil::HumanString(result_shape)
-        << " but is inferred to be: "
-        << ShapeUtil::HumanString(inferred_return_shape);
-
-    const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand);
-    auto result = Literal::CreateFromShape(result_shape);
-
-    TF_RETURN_IF_ERROR(
-        result->Populate<ReturnT>([&](ArraySlice<int64> out_index) {
-          std::vector<int64> from_index(out_index.begin(), out_index.end());
-          for (const int64 dim : reverse_dimensions) {
-            from_index[dim] = result_shape.dimensions(dim) - 1 - out_index[dim];
-          }
-          return operand_literal.Get<ReturnT>(from_index);
-        }));
-
-    parent_->evaluated_[reverse] = std::move(result);
-    return Status::OK();
-  }
-
-  Status HandleConvolution(HloInstruction* conv) override {
-    auto lhs = conv->operand(0);
-    auto rhs = conv->operand(1);
-    const auto& window = conv->window();
-    const Shape& result_shape = conv->shape();
-    const Shape& lhs_shape = lhs->shape();
-    const Shape& rhs_shape = rhs->shape();
-
-    TF_CHECK_OK(ShapeUtil::ValidateShape(lhs_shape));
-    TF_CHECK_OK(ShapeUtil::ValidateShape(rhs_shape));
-    CHECK(ShapeUtil::IsArray(lhs_shape));
-    CHECK(ShapeUtil::IsArray(rhs_shape));
-    CHECK(ShapeUtil::SameElementType(lhs_shape, rhs_shape));
-    CHECK(ShapeUtil::SameElementType(lhs_shape, result_shape));
-
-    const auto& dnums = conv->convolution_dimension_numbers();
-    const int64 num_spatial_dims = dnums.output_spatial_dimensions_size();
-    CHECK_EQ(num_spatial_dims, dnums.input_spatial_dimensions_size());
-    CHECK_EQ(num_spatial_dims, dnums.kernel_spatial_dimensions_size());
-    CHECK_GE(num_spatial_dims, 0);
-    CHECK_EQ(window.dimensions_size(), num_spatial_dims);
-
-    const auto lhs_rank = ShapeUtil::Rank(lhs_shape);
-    const auto rhs_rank = ShapeUtil::Rank(rhs_shape);
-
-    CHECK_EQ(num_spatial_dims + 2, lhs_rank);
-    CHECK_EQ(num_spatial_dims + 2, rhs_rank);
-
-    TF_ASSIGN_OR_RETURN(auto inferred_return_shape,
-                        ShapeInference::InferConvolveShape(lhs_shape, rhs_shape,
-                                                           window, dnums));
-    CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape))
-        << "return shape set to: " << ShapeUtil::HumanString(result_shape)
-        << " but is inferred to be: "
-        << ShapeUtil::HumanString(inferred_return_shape);
-
-    const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs);
-    const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs);
-
-    std::vector<int64> window_dimension_sizes;
-    for (auto i : dnums.kernel_spatial_dimensions()) {
-      window_dimension_sizes.push_back(ShapeUtil::GetDimension(rhs_shape, i));
-    }
-
-    const Shape& window_shape =
-        ShapeUtil::MakeShape(rhs_shape.element_type(), window_dimension_sizes);
-
-    DimensionVector lhs_dim_multipliers = MakeDimMultipliers(lhs_shape);
-    DimensionVector rhs_dim_multipliers = MakeDimMultipliers(rhs_shape);
-
-    auto lhs_literal_data = lhs_literal.data<ReturnT>();
-    auto rhs_literal_data = rhs_literal.data<ReturnT>();
-
-    auto func = [&window_shape, &dnums, &lhs_shape, &rhs_shape, &window,
-                 &lhs_dim_multipliers, &rhs_dim_multipliers, lhs_literal_data,
-                 rhs_literal_data](ArraySlice<int64> out_index) {
-      // Dimension number applicable for input (lhs).
-      const int64 input_batch_dim = dnums.input_batch_dimension();
-      const int64 input_z_dim = dnums.input_feature_dimension();
-      // Dimension number applicable for kernel (rhs).
-      const int64 kernel_input_z_dim = dnums.kernel_input_feature_dimension();
-      const int64 kernel_output_z_dim = dnums.kernel_output_feature_dimension();
-      // Dimension number applicable for output.
-      const int64 output_batch_dim = dnums.output_batch_dimension();
-      const int64 output_z_dim = dnums.output_feature_dimension();
-
-      const int64 z_size = ShapeUtil::GetDimension(lhs_shape, input_z_dim);
-
-      ElementwiseT result_val = static_cast<ElementwiseT>(0);
-      DimensionVector rhs_spatial_index(dnums.kernel_spatial_dimensions_size(),
-                                        0);
-
-      // Convolve input feature with kernel.
-      do {
-        for (int64 iz = 0; iz < z_size; ++iz) {
-          int64 lhs_linear_index = 0;
-          lhs_linear_index += out_index[output_batch_dim] *
-                              lhs_dim_multipliers[input_batch_dim];
-          lhs_linear_index += iz * lhs_dim_multipliers[input_z_dim];
-
-          int64 rhs_linear_index = 0;
-          rhs_linear_index += out_index[output_z_dim] *
-                              rhs_dim_multipliers[kernel_output_z_dim];
-          rhs_linear_index += iz * rhs_dim_multipliers[kernel_input_z_dim];
-
-          // Find corresponding spatial dimension index for input (lhs).
-          for (int64 ki = 0; ki < rhs_spatial_index.size(); ++ki) {
-            // Spatial dimension number for input (lhs) and output.
-            const int64 input_spatial_dim = dnums.input_spatial_dimensions(ki);
-            const int64 output_spatial_dim =
-                dnums.output_spatial_dimensions(ki);
-
-            // Calculate lhs (input) index without taking base dilation into
-            // account.
-            const auto& window_dim = window.dimensions(ki);
-            const int64 undilated_index =
-                out_index[output_spatial_dim] * window_dim.stride() -
-                window_dim.padding_low() +
-                rhs_spatial_index[ki] * window_dim.window_dilation();
-            // Skip if the lhs (input) index is to be dilated.  As an
-            // optimization, skip this mod if there's no dilation.
-            if (window_dim.base_dilation() > 1 &&
-                undilated_index % window_dim.base_dilation() != 0) {
-              goto cnt;
-            }
-
-            // Calculate the actual lhs (input) index after dilation.  As an
-            // optimization, skip this integer divide if there's no dilation.
-            int64 lhs_spatial_index;
-            if (window_dim.base_dilation() > 1) {
-              lhs_spatial_index = undilated_index / window_dim.base_dilation();
-            } else {
-              lhs_spatial_index = undilated_index;
-            }
-            lhs_linear_index +=
-                lhs_spatial_index * lhs_dim_multipliers[input_spatial_dim];
-
-            // Skip if input index is not in bounds.
-            if (!(lhs_spatial_index >= 0 &&
-                  lhs_spatial_index <
-                      lhs_shape.dimensions(input_spatial_dim))) {
-              goto cnt;
-            }
-
-            rhs_linear_index +=
-                (window_dim.window_reversal()
-                     ? ((window_dim.size() - 1) - rhs_spatial_index[ki])
-                     : rhs_spatial_index[ki]) *
-                rhs_dim_multipliers[dnums.kernel_spatial_dimensions(ki)];
-          }
-
-          result_val +=
-              static_cast<ElementwiseT>(lhs_literal_data[lhs_linear_index]) *
-              static_cast<ElementwiseT>(rhs_literal_data[rhs_linear_index]);
-        }
-      cnt : {}
-      } while (IndexUtil::BumpIndices(window_shape, &rhs_spatial_index));
-
-      return static_cast<ReturnT>(result_val);
-    };
-
-    auto result = Literal::CreateFromShape(result_shape);
-    TF_RETURN_IF_ERROR(result->PopulateParallel<ReturnT>(func));
-
-    parent_->evaluated_[conv] = std::move(result);
-    return Status::OK();
-  }
-
-  Status HandleDot(HloInstruction* dot) override {
-    auto lhs = dot->operand(0);
-    auto rhs = dot->operand(1);
-    CHECK(ShapeUtil::IsArray(dot->shape()));
-    CHECK(ShapeUtil::IsArray(lhs->shape()));
-    CHECK(ShapeUtil::IsArray(rhs->shape()));
-
-    const auto& dnums = dot->dot_dimension_numbers();
-
-    const auto lhs_rank = ShapeUtil::Rank(lhs->shape());
-    const auto rhs_rank = ShapeUtil::Rank(rhs->shape());
-
-    CHECK(ShapeUtil::SameElementType(lhs->shape(), rhs->shape()));
-    CHECK(ShapeUtil::SameElementType(lhs->shape(), dot->shape()));
-
-    // There must be 1 and only 1 Contracting dimension for lhs and rhs.
-    CHECK_EQ(dnums.lhs_contracting_dimensions_size(), 1);
-    CHECK_EQ(dnums.rhs_contracting_dimensions_size(), 1);
-    const int64 lhs_contracting_dimension = dnums.lhs_contracting_dimensions(0);
-    const int64 rhs_contracting_dimension = dnums.rhs_contracting_dimensions(0);
-    // Contracted dimension sizes must be the same.
-    CHECK_EQ(lhs->shape().dimensions(lhs_contracting_dimension),
-             rhs->shape().dimensions(rhs_contracting_dimension))
-        << "lhs contracted dimension: "
-        << lhs->shape().dimensions(lhs_contracting_dimension)
-        << " rhs contracted dimension: "
-        << rhs->shape().dimensions(rhs_contracting_dimension);
-    const int64 contracted_dimension_size =
-        lhs->shape().dimensions(lhs_contracting_dimension);
-
-    const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs);
-    const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs);
-
-    auto result = Literal::CreateFromShape(dot->shape());
-
-    CHECK_EQ(dnums.lhs_batch_dimensions_size(),
-             dnums.rhs_batch_dimensions_size());
-
-    std::vector<int64> lhs_non_contracting_dims;
-    for (int64 i = 0; i < lhs_rank; i++) {
-      if (i != lhs_contracting_dimension) {
-        lhs_non_contracting_dims.push_back(i);
-      }
-    }
-
-    std::vector<int64> rhs_non_batch_non_contracting_dims;
-    FlatSet<int64> batch_dims_set(dnums.rhs_batch_dimensions().begin(),
-                                  dnums.rhs_batch_dimensions().end());
-    for (int64 i = 0; i < rhs_rank; i++) {
-      if (i != rhs_contracting_dimension && batch_dims_set.count(i) == 0) {
-        rhs_non_batch_non_contracting_dims.push_back(i);
-      }
-    }
-
-    const int64 batch_dim_size = dnums.lhs_batch_dimensions_size();
-    const int64 lhs_non_contracting_size = lhs_non_contracting_dims.size();
-
-    DimensionVector lhs_index(lhs_rank);
-    DimensionVector rhs_index(rhs_rank);
-    TF_RETURN_IF_ERROR(
-        result->Populate<ReturnT>([&](ArraySlice<int64> result_index) {
-          ElementwiseT result_val = static_cast<ElementwiseT>(0);
-
-          // Find the corresponding non-contracting indices for lhs and rhs.
-          //
-          // For `result_index`, its batch dimension, if exists, will be at the
-          // same dimension as the batch dimension of lhs and rhs. More
-          // specifically:
-          // - For lhs, the non-contracting dimensions, including the batch
-          // dimension have the same index as the `result_index`.
-          // - For rhs, the batch dimension is set seperately from other
-          // non-contracting dimensions, since these other non-contracting
-          // dimensions in rhs follow the non-contracting dimensions of lhs in
-          // the resulting index.
-          //
-          // As an example, for a resulting index:
-          //  result_index [result_batch, result_x, result_y]
-          // the effecting lhs and rhs indices are:
-          //  lhs [result_batch, lhs_non_contracting_dim, contracting_dim
-          //  rhs [result_batch, contracting_dim, rhs_non_contracting_dim]
-          // `result_x` is only affected by the lhs_non_contracting_dim and
-          // likewise `result_y` only depends on rhs_non_contracting_dim.
-          //
-          // so we can look up the lhs and rhs indices by:
-          //
-          // lhs:
-          //  batch index is the same as `result_batch`.
-          //    non-contracting dimension is the same as
-          //    result_index[lhs_non_contracting_dim]
-          // rhs:
-          //  batch index: the same as `result_batch`.
-          //  non-contracting dimension index: *not* the same as
-          //    result_index[rhs_non_contractng_dim], since the
-          //    non-contracting dimensions of lhs are included in the
-          //    result_index first. Instead, the non_contracting_dim of rhs must
-          //    be calculated as following:
-          //      lhs_non_contracting_dimensions_size +
-          //      (rhs_non_batch_non_contracting_dim - batch_dim_size) - 1
-          //
-          //    Note that (rhs_non_batch_contracting_dim - batch_dim_size) is
-          //    the index offset to the result_index that only depends on
-          //    the non_batch and non-contracting dimensions of rhs. -1 at the
-          //    end translates size to index.
-          for (auto i : lhs_non_contracting_dims) {
-            lhs_index[i] = result_index[i];
-          }
-          for (auto i : dnums.rhs_batch_dimensions()) {
-            rhs_index[i] = result_index[i];
-          }
-          for (auto i : rhs_non_batch_non_contracting_dims) {
-            const int64 rhs_non_batch_non_contracting_dim =
-                lhs_non_contracting_size + (i - batch_dim_size) - 1;
-            rhs_index[i] = result_index[rhs_non_batch_non_contracting_dim];
-          }
-
-          // Accumulates resulting product along the contracted dimension.
-          for (int64 i = 0; i < contracted_dimension_size; ++i) {
-            lhs_index[lhs_contracting_dimension] = i;
-            rhs_index[rhs_contracting_dimension] = i;
-
-            result_val +=
-                static_cast<ElementwiseT>(lhs_literal.Get<ReturnT>(lhs_index)) *
-                static_cast<ElementwiseT>(rhs_literal.Get<ReturnT>(rhs_index));
-          }
-
-          return static_cast<ReturnT>(result_val);
-        }));
-
-    parent_->evaluated_[dot] = std::move(result);
-    return Status::OK();
-  }
-
-  Status HandlePad(HloInstruction* pad) override {
-    CHECK(!ShapeUtil::IsTuple(pad->operand(0)->shape()));
-    // Padding value must be scalar.
-    CHECK(ShapeUtil::IsScalar(pad->operand(1)->shape()));
-    CHECK_EQ(ShapeUtil::Rank(pad->operand(0)->shape()),
-             pad->padding_config().dimensions_size());
-
-    TF_ASSIGN_OR_RETURN(auto inferred_return_shape,
-                        ShapeInference::InferPadShape(
-                            /*operand_shape=*/pad->operand(0)->shape(),
-                            /*padding_value_shape=*/pad->operand(1)->shape(),
-                            /*padding_config=*/pad->padding_config()));
-    CHECK(ShapeUtil::Compatible(pad->shape(), inferred_return_shape))
-        << "return shape is set to: " << ShapeUtil::HumanString(pad->shape())
-        << "but is inferred to be: "
-        << ShapeUtil::HumanString(inferred_return_shape);
-
-    // Create new HLO of padded shape with padding value.
-    ReturnT scalar =
-        parent_->GetEvaluatedLiteralFor(pad->operand(1)).Get<ReturnT>({});
-    auto result = Literal::CreateFromShape(pad->shape());
-    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(
-        [&scalar](ArraySlice<int64> multi_index) { return scalar; }));
-
-    const Literal& evaluated_operand =
-        parent_->GetEvaluatedLiteralFor(pad->operand(0));
-
-    std::vector<int64> input_index(ShapeUtil::Rank(evaluated_operand.shape()),
-                                   0);
-    std::vector<int64> target_index(ShapeUtil::Rank(result->shape()), 0);
-
-    // Loop through each element of the operand, assign them to the
-    // corresponding index of the resulting padded literal.
-    const PaddingConfig& pad_config = pad->padding_config();
-
-    auto func = [&](ArraySlice<int64> input_index) {
-      for (auto i = 0; i < input_index.size(); ++i) {
-        // Interior padding occurs logically before edge padding, so in the case
-        // of negative edge padding elements are removed from the
-        // interior-padded operand.
-        target_index[i] =
-            pad_config.dimensions(i).edge_padding_low() +
-            input_index[i] * (pad_config.dimensions(i).interior_padding() + 1);
-
-        // Account for negative low and high padding: skip assignment if the
-        // any target index is out of range.
-        if (!(target_index[i] >= 0 &&
-              target_index[i] < pad->shape().dimensions(i))) {
-          return true;
-        }
-      }
-      result->Set<ReturnT>(target_index,
-                           evaluated_operand.Get<ReturnT>(input_index));
-      return true;
-    };
-
-    std::vector<int64> zero_base(evaluated_operand.shape().dimensions_size(),
-                                 0);
-    std::vector<int64> step(evaluated_operand.shape().dimensions_size(), 1);
-
-    ShapeUtil::ForEachIndex(
-        evaluated_operand.shape(), zero_base,
-        AsInt64Slice(evaluated_operand.shape().dimensions()), step, func);
-
-    parent_->evaluated_[pad] = std::move(result);
-    return Status::OK();
-  }
-
-  Status HandleDynamicSlice(HloInstruction* dynamic_slice) override {
-    auto operand = dynamic_slice->operand(0);
-    auto start_indices = dynamic_slice->operand(1);
-    auto result_shape = dynamic_slice->shape();
-    TF_ASSIGN_OR_RETURN(auto inferred_return_shape,
-                        ShapeInference::InferDynamicSliceShape(
-                            operand->shape(), start_indices->shape(),
-                            dynamic_slice->dynamic_slice_sizes()));
-    TF_RET_CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape))
-        << "return shape is set to: " << ShapeUtil::HumanString(result_shape)
-        << "but is inferred to be: "
-        << ShapeUtil::HumanString(inferred_return_shape);
-    TF_RET_CHECK(
-        primitive_util::IsIntegralType(start_indices->shape().element_type()));
-
-    const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand);
-    const Literal& start_indices_literal =
-        parent_->GetEvaluatedLiteralFor(start_indices);
-
-    switch (start_indices->shape().element_type()) {
-      case S32: {
-        TF_ASSIGN_OR_RETURN(
-            parent_->evaluated_[dynamic_slice],
-            DynamicSlice<int32>(operand_literal, start_indices_literal,
-                                result_shape));
-      } break;
-      case S64: {
-        TF_ASSIGN_OR_RETURN(
-            parent_->evaluated_[dynamic_slice],
-            DynamicSlice<int64>(operand_literal, start_indices_literal,
-                                result_shape));
-      } break;
-      case U32: {
-        TF_ASSIGN_OR_RETURN(
-            parent_->evaluated_[dynamic_slice],
-            DynamicSlice<uint32>(operand_literal, start_indices_literal,
-                                 result_shape));
-      } break;
-      case U64: {
-        TF_ASSIGN_OR_RETURN(
-            parent_->evaluated_[dynamic_slice],
-            DynamicSlice<uint64>(operand_literal, start_indices_literal,
-                                 result_shape));
-      } break;
-      default:
-        LOG(FATAL) << "HandleDynamicSlice: unhandled primitive type for "
-                      "start_indices: "
-                   << PrimitiveType_Name(start_indices->shape().element_type());
-    }
-
-    return Status::OK();
-  }
-
-  Status HandleDynamicUpdateSlice(
-      HloInstruction* dynamic_update_slice) override {
-    auto operand = dynamic_update_slice->operand(0);
-    auto update = dynamic_update_slice->operand(1);
-    auto start_indices = dynamic_update_slice->operand(2);
-    auto result_shape = dynamic_update_slice->shape();
-    TF_ASSIGN_OR_RETURN(
-        auto inferred_return_shape,
-        ShapeInference::InferDynamicUpdateSliceShape(
-            operand->shape(), update->shape(), start_indices->shape()));
-    TF_RET_CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape))
-        << "return shape is set to: " << ShapeUtil::HumanString(result_shape)
-        << "but is inferred to be: "
-        << ShapeUtil::HumanString(inferred_return_shape);
-    TF_RET_CHECK(
-        primitive_util::IsIntegralType(start_indices->shape().element_type()));
-    TF_RET_CHECK(ShapeUtil::Compatible(result_shape, operand->shape()));
-
-    const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand);
-    const Literal& update_literal = parent_->GetEvaluatedLiteralFor(update);
-    const Literal& start_indices_literal =
-        parent_->GetEvaluatedLiteralFor(start_indices);
-
-    switch (start_indices->shape().element_type()) {
-      case S32: {
-        TF_ASSIGN_OR_RETURN(
-            parent_->evaluated_[dynamic_update_slice],
-            DynamicUpdateSlice<int32>(operand_literal, update_literal,
-                                      start_indices_literal));
-      } break;
-      case S64: {
-        TF_ASSIGN_OR_RETURN(
-            parent_->evaluated_[dynamic_update_slice],
-            DynamicUpdateSlice<int64>(operand_literal, update_literal,
-                                      start_indices_literal));
-      } break;
-      case U32: {
-        TF_ASSIGN_OR_RETURN(
-            parent_->evaluated_[dynamic_update_slice],
-            DynamicUpdateSlice<uint32>(operand_literal, update_literal,
-                                       start_indices_literal));
-      } break;
-      case U64: {
-        TF_ASSIGN_OR_RETURN(
-            parent_->evaluated_[dynamic_update_slice],
-            DynamicUpdateSlice<uint64>(operand_literal, update_literal,
-                                       start_indices_literal));
-      } break;
-      default:
-        LOG(FATAL) << "HandleDynamicUpdateSlice: unhandled primitive type for "
-                      "start_indices: "
-                   << PrimitiveType_Name(start_indices->shape().element_type());
-    }
-
-    return Status::OK();
-  }
-
-  template <typename NativeT>
-  StatusOr<std::unique_ptr<Literal>> MapImpl(HloInstruction* map) {
-    auto operands = map->operands();
-    HloComputation* computation = map->to_apply();
-
-    auto result = Literal::CreateFromShape(map->shape());
-
-    HloEvaluator embedded_evaluator(parent_->max_loop_iterations_);
-    TF_RETURN_IF_ERROR(
-        result->Populate<ReturnT>([&](ArraySlice<int64> multi_index) {
-          std::vector<std::unique_ptr<Literal>> arg_literals;
-          arg_literals.reserve(operands.size());
-
-          // Construct scalar literal parameters to be passed to the map
-          // computation.
-          for (auto operand : operands) {
-            const Literal& arg_literal =
-                parent_->GetEvaluatedLiteralFor(operand);
-
-            auto curr_val = arg_literal.Get<NativeT>(multi_index);
-            auto curr_val_literal = Literal::CreateR0<NativeT>(curr_val);
-
-            arg_literals.push_back(std::move(curr_val_literal));
-          }
-
-          std::unique_ptr<Literal> computed_result =
-              embedded_evaluator
-                  .Evaluate<std::unique_ptr<Literal>>(*computation,
-                                                      arg_literals)
-                  .ConsumeValueOrDie();
-          // Clear visit states so that the we can use the evaluate again on
-          // the same computation.
-          embedded_evaluator.ResetVisitStates();
-
-          return computed_result->Get<ReturnT>({});
-        }));
-    return std::move(result);
-  }
-
-  Status HandleMap(HloInstruction* map) override {
-    switch (map->operand(0)->shape().element_type()) {
-      case PRED: {
-        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<bool>(map));
-        break;
-      }
-      case U8: {
-        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<uint8>(map));
-        break;
-      }
-      case U32: {
-        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<uint32>(map));
-        break;
-      }
-      case U64: {
-        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<uint64>(map));
-        break;
-      }
-      case S8: {
-        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<int8>(map));
-        break;
-      }
-      case S32: {
-        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<int32>(map));
-        break;
-      }
-      case S64: {
-        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<int64>(map));
-        break;
-      }
-      case F16: {
-        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map],
-                            MapImpl<Eigen::half>(map));
-        break;
-      }
-      case F32: {
-        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<float>(map));
-        break;
-      }
-      case F64: {
-        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<double>(map));
-        break;
-      }
-      case C64: {
-        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<complex64>(map));
-        break;
-      }
-      default:
-        LOG(FATAL) << "HandleMap: unhandled primitive type for "
-                      "input operand: "
-                   << PrimitiveType_Name(
-                          map->operand(0)->shape().element_type());
-    }
-
-    return Status::OK();
-  }
-
-  Status HandleReduce(HloInstruction* reduce) override {
-    auto arg = reduce->operand(0);
-    auto init_value = reduce->operand(1);
-    ArraySlice<int64> dimensions(reduce->dimensions());
-    HloComputation* function = reduce->to_apply();
-    TF_RET_CHECK(ShapeUtil::Rank(reduce->shape()) ==
-                 ShapeUtil::Rank(arg->shape()) - dimensions.size());
-    TF_ASSIGN_OR_RETURN(auto inferred_return_shape,
-                        ShapeInference::InferReduceShape(
-                            /*arg=*/arg->shape(),
-                            /*init_value=*/init_value->shape(),
-                            /*dimensions_to_reduce=*/dimensions,
-                            /*to_apply=*/function->ComputeProgramShape()));
-    TF_RET_CHECK(ShapeUtil::Compatible(reduce->shape(), inferred_return_shape))
-        << "return shape is set to: " << ShapeUtil::HumanString(reduce->shape())
-        << "but is inferred to be: "
-        << ShapeUtil::HumanString(inferred_return_shape);
-
-    const Literal& arg_literal = parent_->GetEvaluatedLiteralFor(arg);
-    VLOG(3) << "HandleReduce arg_literal: " << arg_literal.ToString();
-    const Literal& init_literal = parent_->GetEvaluatedLiteralFor(init_value);
-    VLOG(3) << "HandleReduce init_literal: " << init_literal.ToString();
-    TF_RET_CHECK(ShapeUtil::IsScalar(init_literal.shape()));
-    auto init_scalar = init_literal.Get<ReturnT>({});
-
-    auto result = Literal::CreateFromShape(reduce->shape());
-
-    const auto arg_dimensions = AsInt64Slice(arg_literal.shape().dimensions());
-    std::vector<int64> arg_dim_steps(arg_dimensions.size());
-    std::vector<int64> arg_dim_counts(arg_dimensions.size());
-    for (const int64 dim : dimensions) {
-      arg_dim_steps[dim] = 1;
-      arg_dim_counts[dim] = arg_dimensions[dim];
-    }
-
-    // Map each dimension in the result to a dimension in arg that isn't
-    // being reduced.
-    std::vector<int64> result_to_arg_index;
-    for (int64 i = 0; i < arg_dimensions.size(); ++i) {
-      if (arg_dim_steps[i] == 0) {
-        result_to_arg_index.push_back(i);
-      }
-    }
-
-    HloEvaluator embedded_evaluator(parent_->max_loop_iterations_);
-    // For each resulting dimension, calculate and assign computed value.
-    TF_RETURN_IF_ERROR(
-        result->Populate<ReturnT>([&](ArraySlice<int64> multi_index) {
-          ReturnT result_val = init_scalar;
-
-          std::vector<int64> base(arg_dimensions.size());
-          for (int64 i = 0; i < multi_index.size(); ++i) {
-            base[result_to_arg_index[i]] = multi_index[i];
-          }
-
-          // When the reduction is addition of floats, accumulate in a double
-          // for better precision. Also, avoid creating Literals for the
-          // intermediate results; it's much faster.
-          if (ShapeUtil::ElementIsFloating(init_literal.shape()) &&
-              IsScalarAdd(function)) {
-            double computed_result = 0;
-            auto func = [&](ArraySlice<int64> input_index) {
-              computed_result += arg_literal.Get<float>(input_index);
-              return true;
-            };
-            ShapeUtil::ForEachIndex(arg_literal.shape(), base, arg_dim_counts,
-                                    arg_dim_steps, func);
-            return static_cast<ReturnT>(computed_result);
-          }
-          auto func = [&](ArraySlice<int64> input_index) {
-            auto curr_val = arg_literal.Get<ReturnT>(input_index);
-
-            // Evaluate computation with specified literal operands.
-            auto curr_val_literal = Literal::CreateR0<ReturnT>(curr_val);
-            auto result_val_literal = Literal::CreateR0<ReturnT>(result_val);
-            std::vector<const Literal*> args = {result_val_literal.get(),
-                                                curr_val_literal.get()};
-
-            std::unique_ptr<Literal> computed_result =
-                embedded_evaluator.Evaluate<const Literal*>(*function, args)
-                    .ConsumeValueOrDie();
-            // Clear visit states so that we can use the evaluator again on
-            // the same computation.
-            embedded_evaluator.ResetVisitStates();
-            // Assign computed result to result_val.
-            result_val = computed_result->Get<ReturnT>({});
-            return true;
-          };
-          // Computes one element of the result, reducing all dimensions that
-          // contribute to that element.
-          ShapeUtil::ForEachIndex(arg_literal.shape(), base, arg_dim_counts,
-                                  arg_dim_steps, func);
-          return result_val;
-        }));
-
-    parent_->evaluated_[reduce] = std::move(result);
-    return Status::OK();
-  }
-
-  bool IsScalarAdd(HloComputation* computation) {
-    HloInstruction* instruction = computation->root_instruction();
-    if (instruction->opcode() == HloOpcode::kAdd &&
-        computation->num_parameters() == 2) {
-      const HloInstruction* lhs = instruction->operand(0);
-      const HloInstruction* rhs = instruction->operand(1);
-      return lhs->opcode() == HloOpcode::kParameter &&
-             ShapeUtil::IsScalar(lhs->shape()) &&
-             rhs->opcode() == HloOpcode::kParameter &&
-             ShapeUtil::IsScalar(rhs->shape()) && lhs != rhs;
-    }
-    return false;
-  }
-
-  Status HandleSelectAndScatter(HloInstruction* select_and_scatter) override {
-    auto operand = select_and_scatter->operand(0);
-    auto source = select_and_scatter->operand(1);
-    const Window& window = select_and_scatter->window();
-
-    const Literal& init_literal =
-        parent_->GetEvaluatedLiteralFor(select_and_scatter->operand(2));
-    TF_RET_CHECK(ShapeUtil::IsScalar(init_literal.shape()));
-    auto init_scalar = init_literal.Get<ReturnT>({});
-
-    auto result = Literal::CreateFromShape(select_and_scatter->shape());
-
-    // Initialize result array with the init value.
-    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(
-        [&](ArraySlice<int64> output_index) { return init_scalar; }));
-
-    std::vector<int64> window_dimension_sizes;
-    for (const auto& window_dimension : window.dimensions()) {
-      window_dimension_sizes.push_back(window_dimension.size());
-    }
-    const Shape window_shape = ShapeUtil::MakeShape(
-        operand->shape().element_type(), window_dimension_sizes);
-
-    HloComputation* select = select_and_scatter->select();
-    HloComputation* scatter = select_and_scatter->scatter();
-
-    const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand);
-    const Literal& source_literal = parent_->GetEvaluatedLiteralFor(source);
-
-    int64 rank = ShapeUtil::Rank(operand_literal.shape());
-
-    HloEvaluator embedded_evaluator(parent_->max_loop_iterations_);
-    DimensionVector source_index(rank);
-
-    std::fill(source_index.begin(), source_index.end(), 0);
-    do {
-      // For each element in `source`, we place a window in `operand`. For each
-      // window placement, we iterate inside the window twice:
-      //
-      // 1. Find the selected index by applying `select` function to all
-      // elements. E.g., If the `select` function is GreaterEqual, the first
-      // iteration through the window finds the biggest value and returns its
-      // index.
-      //
-      // 2. Using the selected index, scatter value from `source` to result. We
-      // do this by iterating through the window, and compare each index with
-      // the selected index.
-      optional<ReturnT> selected_val;
-      optional<std::vector<int64>> selected_index;
-
-      IterateThroughWindow(
-          window_shape, window, operand_literal.shape(), source_index,
-          [&](const std::vector<int64>& operand_index) {
-            auto curr_val = operand_literal.Get<ReturnT>(operand_index);
-            if (!selected_val) {
-              selected_val = curr_val;
-              selected_index = operand_index;
-            }
-            const auto curr_val_literal = Literal::CreateR0<ReturnT>(curr_val);
-            const auto selected_val_literal =
-                Literal::CreateR0<ReturnT>(*selected_val);
-
-            const std::vector<const Literal*> args = {
-                selected_val_literal.get(), curr_val_literal.get()};
-            std::unique_ptr<Literal> computed_result =
-                embedded_evaluator.Evaluate<const Literal*>(*select, args)
-                    .ConsumeValueOrDie();
-            bool selected = !computed_result->Get<bool>({});
-            if (selected) {
-              selected_val = curr_val;
-              selected_index = operand_index;
-            }
-            embedded_evaluator.ResetVisitStates();
-          });
-
-      IterateThroughWindow(
-          window_shape, window, operand_literal.shape(), source_index,
-          [&](const std::vector<int64>& operand_index) {
-            if (std::equal(operand_index.begin(), operand_index.end(),
-                           selected_index->begin())) {
-              auto source = source_literal.Get<ReturnT>(source_index);
-              auto scattered = result->Get<ReturnT>(operand_index);
-              const auto source_literal = Literal::CreateR0<ReturnT>(source);
-              const auto scattered_literal =
-                  Literal::CreateR0<ReturnT>(scattered);
-
-              const std::vector<const Literal*> args = {
-                  source_literal.get(), scattered_literal.get()};
-              std::unique_ptr<Literal> computed_result =
-                  embedded_evaluator.Evaluate<const Literal*>(*scatter, args)
-                      .ConsumeValueOrDie();
-              result->Set(operand_index, computed_result->Get<ReturnT>({}));
-              // Clear visit states so that the we can use the evaluator again
-              // on the same computation.
-              embedded_evaluator.ResetVisitStates();
-            }
-          });
-    } while (IndexUtil::BumpIndices(source->shape(), &source_index));
-
-    parent_->evaluated_[select_and_scatter] = std::move(result);
-    return Status::OK();
-  }
-
-  Status HandleReduceWindow(HloInstruction* reduce_window) override {
-    auto operand = reduce_window->operand(0);
-    const Window& window = reduce_window->window();
-    HloComputation* function = reduce_window->to_apply();
-    TF_ASSIGN_OR_RETURN(
-        auto inferred_return_shape,
-        ShapeInference::InferReduceWindowShape(
-            /*operand_shape=*/reduce_window->operand(0)->shape(),
-            /*init_value=*/reduce_window->operand(1)->shape(), window,
-            /*to_apply_shape=*/function->ComputeProgramShape()));
-    TF_RET_CHECK(
-        ShapeUtil::Compatible(reduce_window->shape(), inferred_return_shape))
-        << "return shape is set to: "
-        << ShapeUtil::HumanStringWithLayout(reduce_window->shape())
-        << "but is inferred to be: "
-        << ShapeUtil::HumanStringWithLayout(inferred_return_shape);
-
-    const Literal& operand_literal =
-        parent_->GetEvaluatedLiteralFor(reduce_window->operand(0));
-    VLOG(3) << "HandleReduceWindow arg_literal: " << operand_literal.ToString();
-    const Literal& init_literal =
-        parent_->GetEvaluatedLiteralFor(reduce_window->operand(1));
-    VLOG(3) << "HandleReduceWindow init_literal: " << init_literal.ToString();
-    TF_RET_CHECK(ShapeUtil::IsScalar(init_literal.shape()));
-    auto init_scalar = init_literal.Get<ReturnT>({});
-
-    auto result = Literal::CreateFromShape(reduce_window->shape());
-
-    // Creates a Shape object from window, for iteration below.
-    std::vector<int64> window_dimension_sizes;
-    for (const auto& window_dimension : window.dimensions()) {
-      window_dimension_sizes.push_back(window_dimension.size());
-    }
-    const Shape window_shape = ShapeUtil::MakeShape(
-        operand->shape().element_type(), window_dimension_sizes);
-
-    DimensionVector window_index(window.dimensions_size());
-    DimensionVector operand_index(ShapeUtil::Rank(operand_literal.shape()));
-
-    HloEvaluator embedded_evaluator(parent_->max_loop_iterations_);
-    // For each resulting dimension, calculate and assign computed value.
-    TF_RETURN_IF_ERROR(
-        result->Populate<ReturnT>([&](ArraySlice<int64> output_index) {
-          ReturnT result_val = init_scalar;
-
-          std::fill(window_index.begin(), window_index.end(), 0);
-          std::fill(operand_index.begin(), operand_index.end(), 0);
-
-          IterateThroughWindow(
-              window_shape, window, operand_literal.shape(), output_index,
-              [&](const std::vector<int64>& operand_index) {
-                auto curr_val = operand_literal.Get<ReturnT>(operand_index);
-
-                // Evaluate computation with specified literal operands.
-                const auto curr_val_literal =
-                    Literal::CreateR0<ReturnT>(curr_val);
-                const auto result_val_literal =
-                    Literal::CreateR0<ReturnT>(result_val);
-                const std::vector<const Literal*> args = {
-                    result_val_literal.get(), curr_val_literal.get()};
-                std::unique_ptr<Literal> computed_result =
-                    embedded_evaluator.Evaluate<const Literal*>(*function, args)
-                        .ConsumeValueOrDie();
-
-                // Clear visit states so that the we can use the evaluate again
-                // on the same computation.
-                embedded_evaluator.ResetVisitStates();
-
-                result_val = computed_result->Get<ReturnT>({});
-              });
-
-          return result_val;
-        }));
-
-    parent_->evaluated_[reduce_window] = std::move(result);
-    return Status::OK();
-  }
-
-  Status HandleSlice(HloInstruction* slice) override {
-    auto operand = slice->operand(0);
-    const Shape& shape = slice->shape();
-    TF_ASSIGN_OR_RETURN(auto inferred_return_shape,
-                        ShapeInference::InferSliceShape(
-                            operand->shape(), slice->slice_starts(),
-                            slice->slice_limits(), slice->slice_strides()));
-    TF_RET_CHECK(ShapeUtil::Compatible(shape, inferred_return_shape))
-        << "return shape set to: " << ShapeUtil::HumanString(shape)
-        << " but is inferred to be: "
-        << ShapeUtil::HumanString(inferred_return_shape);
-
-    const int64 rank = ShapeUtil::Rank(operand->shape());
-    const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand);
-    auto func = [&](ArraySlice<int64> out_index) {
-      DimensionVector operand_index(rank);
-      for (int64 i = 0; i < rank; ++i) {
-        operand_index[i] =
-            slice->slice_starts(i) + out_index[i] * slice->slice_strides(i);
-      }
-      return operand_literal.Get<ReturnT>(operand_index);
-    };
-
-    auto result = Literal::CreateFromDimensions(
-        shape.element_type(), AsInt64Slice(shape.dimensions()));
-    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(func));
-    parent_->evaluated_[slice] = std::move(result);
-    return Status::OK();
-  }
-
-  // Enable CLZ only for int32 and uint32.
-  template <
-      typename NativeT,
-      typename std::enable_if<
-          (std::is_floating_point<NativeT>::value ||
-           std::is_integral<NativeT>::value || is_complex_t<NativeT>::value) &&
-          !(std::is_same<NativeT, uint32>::value ||
-            std::is_same<NativeT, int32>::value)>::type* = nullptr>
-  Status HandleClz(HloInstruction* clz) {
-    return InvalidArgument("Unsupported type for Clz");
-  }
-
-  template <typename NativeT,
-            typename std::enable_if<
-                std::is_same<NativeT, uint32>::value ||
-                std::is_same<NativeT, int32>::value>::type* = nullptr>
-  Status HandleClz(HloInstruction* clz) {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[clz],
-                        ElementWiseUnaryOp(clz, [](ElementwiseT elem_operand) {
-                          return 31 - tensorflow::Log2Floor(elem_operand);
-                        }));
-    return Status::OK();
-  }
-
-  Status HandleClz(HloInstruction* clz) override {
-    return HandleClz<ElementwiseT>(clz);
-  }
-
-  template <typename NativeT, typename std::enable_if<std::is_floating_point<
-                                  NativeT>::value>::type* = nullptr>
-  Status HandleSin(HloInstruction* sin) {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[sin],
-                        ElementWiseUnaryOp(sin, [](ElementwiseT elem_operand) {
-                          return std::sin(elem_operand);
-                        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<std::is_integral<NativeT>::value ||
-                              is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleSin(HloInstruction* sin) {
-    return InvalidArgument("Unsupported type for Sin");
-  }
-
-  Status HandleSin(HloInstruction* sin) override {
-    return HandleSin<ElementwiseT>(sin);
-  }
-
-  template <typename NativeT, typename std::enable_if<std::is_floating_point<
-                                  NativeT>::value>::type* = nullptr>
-  Status HandleCos(HloInstruction* cos) {
-    TF_ASSIGN_OR_RETURN(parent_->evaluated_[cos],
-                        ElementWiseUnaryOp(cos, [](ElementwiseT elem_operand) {
-                          return std::cos(elem_operand);
-                        }));
-    return Status::OK();
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<std::is_integral<NativeT>::value ||
-                              is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleCos(HloInstruction* cos) {
-    return InvalidArgument("Unsupported type for Cos");
-  }
-
-  Status HandleCos(HloInstruction* cos) override {
-    return HandleCos<ElementwiseT>(cos);
-  }
-
-  template <typename NativeT, typename std::enable_if<std::is_same<
-                                  float, NativeT>::value>::type* = nullptr>
-  Status HandleReducePrecision(HloInstruction* reduce_precision) {
-    TF_ASSIGN_OR_RETURN(
-        parent_->evaluated_[reduce_precision],
-        ElementWiseUnaryOp(reduce_precision, [reduce_precision](
-                                                 ElementwiseT elem) {
-          uint32_t value_as_int = tensorflow::bit_cast<uint32_t>(elem);
-          const uint32_t mantissa_bits = reduce_precision->mantissa_bits();
-          const uint32_t exponent_bits = reduce_precision->exponent_bits();
-
-          // Code is based on the CPU/GPU implementation in LLVM-emitting code.
-          //
-          // Bits in float type:
-          //   mantissa : bits [0:22]
-          //   exponent : bits [23:30]
-          //   sign     : bits [31]
-          if (mantissa_bits < 23) {
-            const uint32_t last_mantissa_bit_mask = 1u << (23 - mantissa_bits);
-
-            // Compute rounding bias for round-to-nearest with ties to even.
-            // This is equal to a base value of 0111... plus one bit if the last
-            // remaining mantissa bit is 1.
-            const uint32_t base_rounding_bias =
-                (last_mantissa_bit_mask >> 1) - 1;
-            const uint32_t x_last_mantissa_bit =
-                (value_as_int & last_mantissa_bit_mask) >> (23 - mantissa_bits);
-            const uint32_t x_rounding_bias =
-                x_last_mantissa_bit + base_rounding_bias;
-
-            // Add rounding bias, and mask out truncated bits.  Note that the
-            // case where adding the rounding bias overflows into the exponent
-            // bits is correct; the non-masked mantissa bits will all be zero,
-            // and the exponent will be incremented by one.
-            const uint32_t truncation_mask = ~(last_mantissa_bit_mask - 1);
-            value_as_int = value_as_int + x_rounding_bias;
-            value_as_int = value_as_int & truncation_mask;
-          }
-          if (exponent_bits < 8) {
-            // Masks for f32 values.
-            const uint32_t f32_sign_bit_mask = 1u << 31;
-            const uint32_t f32_exp_bits_mask = 0xffu << 23;
-
-            // An exponent of 2^(n-1)-1 -- that is, 0111... with the zero in the
-            // most- significant bit -- is equal to 1.0f for all exponent sizes.
-            // Adding 2^(n-1)-1 to this gives us the highest non-infinite
-            // exponent for a bit- size of n, and subtracting 2^(n-1)-1 from
-            // this gives us the lowest' exponent (corresponding to 0.0f).
-            //
-            // Thus, the f32 exponent corresponding to the highest non-infinite
-            // exponent for a bit size of n is (2^7-1) + 2^(n-1)-1, and the f32
-            // exponent corresponding to the lowest exponent for a bit size of n
-            // is (2^7-1) - 2^(n-1)-1.
-            //
-            // Note that we have already checked that exponents_bits >= 1.
-            const uint32_t f32_exponent_bias = (1 << 7) - 1;
-            const uint32_t reduced_exponent_bias =
-                (1 << (exponent_bits - 1)) - 1;
-            const uint32_t reduced_max_exponent =
-                f32_exponent_bias + reduced_exponent_bias;
-            const uint32_t reduced_min_exponent =
-                f32_exponent_bias - reduced_exponent_bias;
-
-            // Do we overflow or underflow?
-            const uint32_t x_exponent = value_as_int & f32_exp_bits_mask;
-            const bool x_overflows = x_exponent > (reduced_max_exponent << 23);
-            const bool x_underflows =
-                x_exponent <= (reduced_min_exponent << 23);
-
-            // Compute appropriately-signed values of zero and infinity.
-            const uint32_t x_signed_zero = value_as_int & f32_sign_bit_mask;
-            const uint32_t x_signed_inf = x_signed_zero | f32_exp_bits_mask;
-
-            // Force to zero or infinity if overflow or underflow.  (Note that
-            // this truncates all denormal values to zero, rather than rounding
-            // them.)
-            value_as_int = x_overflows ? x_signed_inf : value_as_int;
-            value_as_int = x_underflows ? x_signed_zero : value_as_int;
-          }
-
-          float reduced_result = tensorflow::bit_cast<float>(value_as_int);
-          if (std::isnan(elem)) {
-            reduced_result = mantissa_bits > 0
-                                 ? elem
-                                 : std::numeric_limits<float>::infinity();
-          }
-          return reduced_result;
-        }));
-    return Status::OK();
-  }
-
-  template <typename NativeT, typename std::enable_if<std::is_same<
-                                  double, NativeT>::value>::type* = nullptr>
-  Status HandleReducePrecision(HloInstruction* reduce_precision) {
-    return InvalidArgument("Double not supported for reduce precision");
-  }
-
-  template <
-      typename NativeT,
-      typename std::enable_if<std::is_integral<NativeT>::value ||
-                              is_complex_t<NativeT>::value>::type* = nullptr>
-  Status HandleReducePrecision(HloInstruction* reduce_precision) {
-    return InvalidArgument("Unsupported type for reduce precision");
-  }
-
-  Status HandleReducePrecision(HloInstruction* reduce_precision) override {
-    return HandleReducePrecision<ElementwiseT>(reduce_precision);
-  }
-
- private:
-  template <typename IndexT>
-  StatusOr<std::unique_ptr<Literal>> DynamicSlice(
-      const Literal& operand_literal, const Literal& start_indices_literal,
-      const Shape& result_shape) {
-    auto start_indices_typed = start_indices_literal.data<IndexT>();
-    std::vector<int64> start(start_indices_typed.begin(),
-                             start_indices_typed.end());
-
-    std::vector<int64> operand_indices(start.size());
-
-    auto result = Literal::CreateFromShape(result_shape);
-    TF_RETURN_IF_ERROR(
-        result->Populate<ReturnT>([&](ArraySlice<int64> multi_index) {
-          for (int64 i = 0; i < operand_indices.size(); ++i) {
-            CHECK_GE(multi_index[i] + start[i], 0);
-            // Mod is only used here to be consistent with the existing
-            // backends' behavior.
-            operand_indices[i] = (multi_index[i] + start[i]) %
-                                 operand_literal.shape().dimensions(i);
-          }
-
-          auto result = operand_literal.Get<ReturnT>(operand_indices);
-          return result;
-        }));
-
-    return std::move(result);
-  }
-
-  template <typename IndexT>
-  StatusOr<std::unique_ptr<Literal>> DynamicUpdateSlice(
-      const Literal& operand_literal, const Literal& update_literal,
-      const Literal& start_indices_literal) {
-    auto result = operand_literal.CloneToUnique();
-    auto start_indices_typed = start_indices_literal.data<IndexT>();
-    const auto rank = ShapeUtil::Rank(result->shape());
-    std::vector<int64> start(rank, 0);
-    for (int64 i = 0; i < rank; ++i) {
-      // All other implementations currently wrap-around the index, so this
-      // should do so as well.
-      start[i] = (start_indices_typed[i] % result->shape().dimensions(i));
-      start[i] += (start[i] < 0) * result->shape().dimensions(i);
-    }
-    std::vector<int64> result_index(rank, 0);
-
-    auto func = [&](ArraySlice<int64> update_index) {
-      std::transform(update_index.begin(), update_index.end(), start.begin(),
-                     result_index.begin(), std::plus<int64>());
-      // Same as above, wrap-around only to match other implementations'
-      // semantics.
-      std::transform(result_index.begin(), result_index.end(),
-                     result->shape().dimensions().begin(), result_index.begin(),
-                     std::modulus<int64>());
-      result->Set<ReturnT>(result_index,
-                           update_literal.Get<ReturnT>(update_index));
-      return true;
-    };
-
-    std::vector<int64> base(update_literal.shape().dimensions_size(), 0);
-    std::vector<int64> step(update_literal.shape().dimensions_size(), 1);
-    ShapeUtil::ForEachIndex(update_literal.shape(), base,
-                            AsInt64Slice(update_literal.shape().dimensions()),
-                            step, func);
-
-    return std::move(result);
-  }
-
-  StatusOr<std::unique_ptr<Literal>> ElementWiseUnaryOp(
-      HloInstruction* instruction,
-      const std::function<ElementwiseT(ElementwiseT)>& unary_op) {
-    const Literal& operand_literal =
-        parent_->GetEvaluatedLiteralFor(instruction->operand(0));
-    TF_ASSIGN_OR_RETURN(
-        auto result_literal,
-        (ElementWiseUnaryOpImpl<ReturnT, ReturnT>(
-            instruction, ConvertUnaryFunction(unary_op), operand_literal)));
-
-    return std::move(result_literal);
-  }
-
-  StatusOr<std::unique_ptr<Literal>> ElementWiseBinaryOp(
-      HloInstruction* instruction,
-      const std::function<ElementwiseT(ElementwiseT, ElementwiseT)>&
-          binary_op) {
-    const auto shape = instruction->shape();
-    const auto* lhs = instruction->operand(0);
-    const auto* rhs = instruction->operand(1);
-
-    // TODO(b/35950897, b/27796129): add DCHECK back once implicit broadcast
-    // is removed.
-    if (!(ShapeUtil::SameDimensions(shape, rhs->shape()) &&
-          ShapeUtil::SameDimensions(lhs->shape(), rhs->shape()))) {
-      return Unimplemented(
-          "Implicit broadcasting is currently unsupported in HLO evaluator "
-          "Shape Mismatch: %s vs %s vs %s: ",
-          ShapeUtil::HumanString(shape).c_str(),
-          ShapeUtil::HumanString(lhs->shape()).c_str(),
-          ShapeUtil::HumanString(rhs->shape()).c_str());
-    }
-
-    const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs);
-    const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs);
-
-    auto result = Literal::CreateFromShape(shape);
-
-    TF_RETURN_IF_ERROR(
-        result->Populate<ReturnT>([&](ArraySlice<int64> multi_index) {
-          return ConvertBinaryFunction(binary_op)(
-              lhs_literal.Get<ReturnT>(multi_index),
-              rhs_literal.Get<ReturnT>(multi_index));
-        }));
-    return std::move(result);
-  }
-
-  template <typename LhsType, typename RhsType, typename EhsType>
-  StatusOr<std::unique_ptr<Literal>> ElementwiseTernaryOp(
-      HloInstruction* instruction,
-      const std::function<ReturnT(LhsType, RhsType, EhsType)>& ternary_op) {
-    const auto shape = instruction->shape();
-    const auto* lhs = instruction->operand(0);
-    const auto* rhs = instruction->operand(1);
-    const auto* ehs = instruction->operand(2);
-
-    // TODO(b/35950897, b/27796129): add DCHECK back once implicit
-    // broadcast is removed.
-    if (!(ShapeUtil::SameDimensions(shape, lhs->shape()) &&
-          ShapeUtil::SameDimensions(lhs->shape(), rhs->shape()) &&
-          ShapeUtil::SameDimensions(rhs->shape(), ehs->shape()))) {
-      return Unimplemented(
-          "Implicit broadcasting is currently unsupported in HLO evaluator "
-          "Shape Mismatch: %s vs %s vs %s vs %s: ",
-          ShapeUtil::HumanString(shape).c_str(),
-          ShapeUtil::HumanString(lhs->shape()).c_str(),
-          ShapeUtil::HumanString(rhs->shape()).c_str(),
-          ShapeUtil::HumanString(ehs->shape()).c_str());
-    }
-
-    const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs);
-    const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs);
-    const Literal& ehs_literal = parent_->GetEvaluatedLiteralFor(ehs);
-
-    auto result = Literal::CreateFromShape(shape);
-
-    TF_RETURN_IF_ERROR(
-        result->Populate<ReturnT>([&](ArraySlice<int64> multi_index) {
-          return ternary_op(lhs_literal.Get<LhsType>(multi_index),
-                            rhs_literal.Get<RhsType>(multi_index),
-                            ehs_literal.Get<EhsType>(multi_index));
-        }));
-
-    return std::move(result);
-  }
-
-  template <typename NativeT>
-  static bool IsShiftOutOfBounds(NativeT rhs) {
-    typedef typename std::make_unsigned<NativeT>::type UnsignedT;
-    UnsignedT lhs_size_unsigned = sizeof(NativeT) * CHAR_BIT;
-    UnsignedT rhs_unsigned = static_cast<UnsignedT>(rhs);
-    return rhs_unsigned >= lhs_size_unsigned;
-  }
-
-  HloEvaluator* parent_;
-};  // class HloEvaluator::TypedVisitor
-
 HloEvaluator::HloEvaluator(int64 max_loop_iterations)
     : max_loop_iterations_(max_loop_iterations) {
-  typed_visitors_[PRED] = MakeUnique<TypedVisitor<bool>>(this);
-  typed_visitors_[U8] = MakeUnique<TypedVisitor<uint8>>(this);
+  typed_visitors_[PRED] = MakeUnique<HloEvaluatorTypedVisitor<bool>>(this);
+  typed_visitors_[U8] = MakeUnique<HloEvaluatorTypedVisitor<uint8>>(this);
   typed_visitors_[U16] = MakeUnique<FunctionVisitor>([](HloInstruction*) {
     return Unimplemented(
-        "HloEvaluator::TypedVisitor: unhandled primitive type: U16.");
+        "HloEvaluator::HloEvaluatorTypedVisitor: unhandled primitive type: "
+        "U16.");
   });
-  typed_visitors_[U32] = MakeUnique<TypedVisitor<uint32>>(this);
-  typed_visitors_[U64] = MakeUnique<TypedVisitor<uint64>>(this);
-  typed_visitors_[S8] = MakeUnique<TypedVisitor<int8>>(this);
+  typed_visitors_[U32] = MakeUnique<HloEvaluatorTypedVisitor<uint32>>(this);
+  typed_visitors_[U64] = MakeUnique<HloEvaluatorTypedVisitor<uint64>>(this);
+  typed_visitors_[S8] = MakeUnique<HloEvaluatorTypedVisitor<int8>>(this);
   typed_visitors_[S16] = MakeUnique<FunctionVisitor>([](HloInstruction*) {
     return Unimplemented(
-        "HloEvaluator::TypedVisitor: unhandled primitive type: S16.");
+        "HloEvaluator::HloEvaluatorTypedVisitor: unhandled primitive type: "
+        "S16.");
   });
-  typed_visitors_[S32] = MakeUnique<TypedVisitor<int32>>(this);
-  typed_visitors_[S64] = MakeUnique<TypedVisitor<int64>>(this);
-  typed_visitors_[F16] = MakeUnique<TypedVisitor<Eigen::half, float>>(this);
-  typed_visitors_[F32] = MakeUnique<TypedVisitor<float>>(this);
-  typed_visitors_[F64] = MakeUnique<TypedVisitor<double>>(this);
-  typed_visitors_[C64] = MakeUnique<TypedVisitor<complex64>>(this);
+  typed_visitors_[S32] = MakeUnique<HloEvaluatorTypedVisitor<int32>>(this);
+  typed_visitors_[S64] = MakeUnique<HloEvaluatorTypedVisitor<int64>>(this);
+  typed_visitors_[F16] =
+      MakeUnique<HloEvaluatorTypedVisitor<Eigen::half, float>>(this);
+  typed_visitors_[F32] = MakeUnique<HloEvaluatorTypedVisitor<float>>(this);
+  typed_visitors_[F64] = MakeUnique<HloEvaluatorTypedVisitor<double>>(this);
+  typed_visitors_[C64] = MakeUnique<HloEvaluatorTypedVisitor<complex64>>(this);
 
   // Most of the evaluator computations we use don't support BF16 (e.g.,
   // std::ceil, std::tanh). To make evaluator work with BF16, we set all
   // elementwise computations to be done in F32 and do BF16<->F32 conversion
   // around the input and the output of the computations.
-  typed_visitors_[BF16] = MakeUnique<TypedVisitor<bfloat16, float>>(this);
+  typed_visitors_[BF16] =
+      MakeUnique<HloEvaluatorTypedVisitor<bfloat16, float>>(this);
 
   typed_visitors_[TUPLE] = MakeUnique<FunctionVisitor>([](HloInstruction*) {
     return Unimplemented(
-        "HloEvaluator::TypedVistor: unhandled primitive type: TUPLE.");
+        "HloEvaluatorTypedVisitor: unhandled primitive type: TUPLE.");
   });
   typed_visitors_[OPAQUE] = MakeUnique<FunctionVisitor>([](HloInstruction*) {
     return Unimplemented(
-        "HloEvaluator::TypedVisitor: unhandled primitive type: OPAQUE.");
+        "HloEvaluatorTypedVisitor: unhandled primitive type: OPAQUE.");
   });
 }
 
@@ -3034,7 +976,7 @@ Status HloEvaluator::HandleSelect(HloInstruction* select) {
 
   // If predicate is of scalar type, no element-wise selection would be needed.
   // This would also handle output array of tuple types as the DefaultAction
-  // would go through the TypedVisitor which doesn't handle tuples.
+  // would go through the HloEvaluatorTypedVisitor which doesn't handle tuples.
   if (ShapeUtil::IsScalar(pred.shape())) {
     if (pred.Get<bool>({})) {
       evaluated_[select] = on_true.CloneToUnique();
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.h b/tensorflow/compiler/xla/service/hlo_evaluator.h
index c0dcee0c3e..cc5676ea7b 100644
--- a/tensorflow/compiler/xla/service/hlo_evaluator.h
+++ b/tensorflow/compiler/xla/service/hlo_evaluator.h
@@ -109,19 +109,16 @@ class HloEvaluator : public DfsHloVisitorWithDefault {
           substitutions);
 
  protected:
-  // Templated DfsHloVisitor. Typically ReturnT here indicates the resulting
-  // literal type of each evaluated Handle* method of a TypedVisitor.
-  // There are however a few notable exceptions to this rule, notably:
-  // - HandleCompare and HandleIsFinite: where the resulting literal type is
-  // always boolean.
-  // These operations are handled outside of the parent HloEvaluator handlers
-  // instead of from within TypedVisitor.
+  // Make HloEvaluatorTypedVisitor a friend because it is logically part of this
+  // class.
   //
-  // Type params:
-  //   - ReturnT: The type of input and output of each operation.
-  //   - ElementwiseT: The type in which internal computation are done.
-  template <typename ReturnT, typename ElementwiseT = ReturnT>
-  class TypedVisitor;
+  // A straightforward implementation would be to make it a nested class
+  // declared and defined in hlo_evaluator.cc.  Instead HloEvaluatorTypedVisitor
+  // lives as a separate class with its own header because its template gets
+  // instantiated many times and we want to use extern templates to shard out
+  // the compilation of those instantiations across multiple cc files.
+  template <typename ReturnT, typename ElementwiseT>
+  friend class HloEvaluatorTypedVisitor;
 
   // Wraps around instruction handling to infer types before dispatching to
   // the corresponding typed Visitor.
@@ -169,6 +166,33 @@ class HloEvaluator : public DfsHloVisitorWithDefault {
   Status HandleSelect(HloInstruction* select) override;
 
  private:
+  template <typename ReturnT, typename NativeT>
+  static StatusOr<std::unique_ptr<Literal>> ElementWiseUnaryOpImpl(
+      HloInstruction* instruction,
+      const std::function<ReturnT(NativeT)>& unary_op,
+      const Literal& operand_literal) {
+    const auto shape = instruction->shape();
+    const auto* operand = instruction->operand(0);
+
+    // TODO(b/35950897, b/27796129): add DCHECK back once implicit broadcast is
+    // removed.
+    if (!ShapeUtil::SameDimensions(shape, operand->shape())) {
+      return Unimplemented(
+          "Implicit broadcasting is currently unsupported in HLO evaluator "
+          "Shape Mismatch: %s vs %s",
+          ShapeUtil::HumanString(shape).c_str(),
+          ShapeUtil::HumanString(operand->shape()).c_str());
+    }
+
+    auto result = Literal::CreateFromShape(shape);
+
+    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(
+        [&](tensorflow::gtl::ArraySlice<int64> multi_index) {
+          return unary_op(operand_literal.Get<NativeT>(multi_index));
+        }));
+    return std::move(result);
+  }
+
   // Returns the already-evaluated literal result for the instruction.
   // A Constant instruction is considered evaluated and its literal will be
   // returned directly without looking up the cache.
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h
new file mode 100644
index 0000000000..f1cb363478
--- /dev/null
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h
@@ -0,0 +1,2102 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_TYPED_VISITOR_H_
+#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_TYPED_VISITOR_H_
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
+#include "tensorflow/compiler/xla/service/shape_inference.h"
+#include "tensorflow/core/lib/core/casts.h"
+#include "tensorflow/core/lib/gtl/optional.h"
+
+namespace xla {
+
+// TODO(b/79274244): We'd like these type traits to live inside of
+// HloEvaluatorTypedVisitor so they don't pollute namespace xla, but that
+// crashes clang in the frontend.
+//
+// Anyway this is relatively safe as-is because hlo_evaluator_typed_visitor.h is
+// a "private" header that's not exposed outside of hlo_evaluator.cc.
+template <typename T>
+using is_complex_t = std::is_same<T, complex64>;
+template <typename T>
+using is_complex64_t = std::is_same<T, complex64>;
+
+// Templated DfsHloVisitor for use by HloEvaluator.
+//
+// Typically ReturnT here indicates the resulting literal type of each evaluated
+// Handle* method of a TypedVisitor.  There are however a few notable exceptions
+// to this rule, notably:
+// - HandleCompare and HandleIsFinite: where the resulting literal type is
+//   always boolean.
+// These operations are handled outside of the parent HloEvaluator handlers
+// instead of from within TypedVisitor.
+//
+// Type params:
+//   - ReturnT: The type of input and output of each operation.
+//   - ElementwiseT: The type in which internal computation are done.
+//
+// This a logically a private part of HloEvaluator.  It lives in this header
+// file rather than in hlo_evaluator.cc because we use extern templates and a
+// bunch of independent cc files to speed up compiling the many instantiations
+// of this class.
+template <typename ReturnT, typename ElementwiseT = ReturnT>
+class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault {
+ public:
+  explicit HloEvaluatorTypedVisitor(HloEvaluator* p) : parent_(p) {}
+
+  // The following higher-order functions convert a function with ElementwiseT
+  // to a function with ReturnT.
+  std::function<ReturnT(ReturnT)> ConvertUnaryFunction(
+      const std::function<ElementwiseT(ElementwiseT)>& unary_op) {
+    return [&unary_op](ReturnT arg) {
+      return static_cast<ReturnT>(unary_op(static_cast<ElementwiseT>(arg)));
+    };
+  }
+  std::function<ReturnT(ReturnT, ReturnT)> ConvertBinaryFunction(
+      const std::function<ElementwiseT(ElementwiseT, ElementwiseT)>&
+          binary_op) {
+    return [&binary_op](ReturnT arg1, ReturnT arg2) {
+      return static_cast<ReturnT>(binary_op(static_cast<ElementwiseT>(arg1),
+                                            static_cast<ElementwiseT>(arg2)));
+    };
+  }
+  std::function<ReturnT(ReturnT, ReturnT, ReturnT)> ConvertTernaryFunction(
+      const std::function<ElementwiseT(ElementwiseT, ElementwiseT,
+                                       ElementwiseT)>& ternary_op) {
+    return [&ternary_op](ReturnT arg1, ReturnT arg2, ReturnT arg3) {
+      return static_cast<ReturnT>(ternary_op(static_cast<ElementwiseT>(arg1),
+                                             static_cast<ElementwiseT>(arg2),
+                                             static_cast<ElementwiseT>(arg3)));
+    };
+  }
+
+  Status DefaultAction(HloInstruction* hlo_instruction) override {
+    return Unimplemented("unhandled HLO ops for HloEvaluator: %s.",
+                         HloOpcodeString(hlo_instruction->opcode()).c_str());
+  }
+
+  // TODO(b/35950897): many of the stl functions used in the handlers are not
+  // overloaded for every XLA primitive type.
+
+  template <typename NativeT,
+            typename std::enable_if<std::is_unsigned<NativeT>::value>::type* =
+                nullptr>
+  Status HandleAbs(HloInstruction* abs) {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[abs],
+                        ElementWiseUnaryOp(abs, [](NativeT elem_operand) {
+                          return elem_operand;
+                        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<std::is_signed<NativeT>::value>::type* = nullptr>
+  Status HandleAbs(HloInstruction* abs) {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[abs],
+                        ElementWiseUnaryOp(abs, [](NativeT elem_operand) {
+                          return std::abs(elem_operand);
+                        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<is_complex64_t<NativeT>::value>::type* = nullptr>
+  Status HandleAbs(HloInstruction* abs) {
+    const Literal& operand_literal =
+        parent_->GetEvaluatedLiteralFor(abs->operand(0));
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[abs],
+        (HloEvaluator::ElementWiseUnaryOpImpl<float, NativeT>(
+            abs, [](NativeT elem_operand) { return std::abs(elem_operand); },
+            operand_literal)));
+
+    return Status::OK();
+  }
+
+  Status HandleAbs(HloInstruction* abs) override {
+    // If the operand is of C64 type, the return type of abs will be F32.
+    // However, ElementwiseT would still be the return type, F32, and thus
+    // specifying the ElementwiseT explicitly as C64 is needed below.
+    if (abs->operand(0)->shape().element_type() == C64) {
+      return HandleAbs<complex64>(abs);
+    }
+    return HandleAbs<ElementwiseT>(abs);
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<!is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleRound(HloInstruction* round) {
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[round],
+        ElementWiseUnaryOp(round, [](ElementwiseT elem_operand) {
+          return std::round(elem_operand);
+        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleRound(HloInstruction* round) {
+    return InvalidArgument("Unsupported type for Round");
+  }
+
+  Status HandleRound(HloInstruction* round) override {
+    return HandleRound<ReturnT>(round);
+  }
+
+  Status HandleBroadcast(HloInstruction* broadcast) override {
+    parent_->evaluated_[broadcast] =
+        Literal::CreateFromShape(broadcast->shape());
+    auto output = parent_->evaluated_[broadcast].get();
+    const Literal& operand_to_broadcast =
+        parent_->GetEvaluatedLiteralFor(broadcast->operand(0));
+    std::vector<int64> broadcast_indices(
+        ShapeUtil::Rank(broadcast->operand(0)->shape()), 0);
+
+    TF_RET_CHECK(broadcast->dimensions().size() ==
+                 ShapeUtil::Rank(operand_to_broadcast.shape()))
+        << "broadcast dimensions is of size: " << broadcast->dimensions().size()
+        << " and rank of operand_to_broadcast is: "
+        << ShapeUtil::Rank(operand_to_broadcast.shape());
+    // Checks that operand's dimensions are the same as the broadcast's
+    // dimensions along the dimensions to be broadcasted.
+    for (int64 i = 0; i < broadcast->dimensions().size(); ++i) {
+      TF_RET_CHECK(broadcast->shape().dimensions(broadcast->dimensions(i)) ==
+                   operand_to_broadcast.shape().dimensions(i));
+    }
+
+    return output->Populate<ReturnT>(
+        [&](tensorflow::gtl::ArraySlice<int64> multi_index) {
+          for (int64 i = 0; i < broadcast->dimensions().size(); ++i) {
+            broadcast_indices[i] = multi_index[broadcast->dimensions(i)];
+          }
+          return operand_to_broadcast.Get<ReturnT>(broadcast_indices);
+        });
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<!is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleCeil(HloInstruction* ceil) {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[ceil],
+                        ElementWiseUnaryOp(ceil, [](ElementwiseT elem_operand) {
+                          return std::ceil(elem_operand);
+                        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleCeil(HloInstruction* ceil) {
+    return InvalidArgument("Unsupported type for Ceil");
+  }
+
+  Status HandleCeil(HloInstruction* ceil) override {
+    return HandleCeil<ReturnT>(ceil);
+  }
+
+  Status HandleConvert(HloInstruction* convert) override {
+    const HloInstruction* operand = convert->operand(0);
+    TF_RET_CHECK(ShapeUtil::SameDimensions(operand->shape(), convert->shape()));
+    TF_ASSIGN_OR_RETURN(std::unique_ptr<Literal> result,
+                        parent_->GetEvaluatedLiteralFor(operand).Convert(
+                            convert->shape().element_type()));
+
+    if (LayoutUtil::LayoutsInShapesEqual(result->shape(), convert->shape())) {
+      parent_->evaluated_[convert] = std::move(result);
+    } else {
+      parent_->evaluated_[convert] =
+          result->Relayout(convert->shape().layout());
+    }
+    return Status::OK();
+  }
+
+  Status HandleBitcastConvert(HloInstruction* convert) override {
+    const HloInstruction* operand = convert->operand(0);
+    TF_RET_CHECK(ShapeUtil::SameDimensions(operand->shape(), convert->shape()));
+    TF_ASSIGN_OR_RETURN(std::unique_ptr<Literal> result,
+                        parent_->GetEvaluatedLiteralFor(operand).BitcastConvert(
+                            convert->shape().element_type()));
+
+    if (LayoutUtil::LayoutsInShapesEqual(result->shape(), convert->shape())) {
+      parent_->evaluated_[convert] = std::move(result);
+    } else {
+      parent_->evaluated_[convert] =
+          result->Relayout(convert->shape().layout());
+    }
+    return Status::OK();
+  }
+
+  Status HandleExp(HloInstruction* exp) override {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[exp],
+                        ElementWiseUnaryOp(exp, [](ElementwiseT elem_operand) {
+                          return std::exp(elem_operand);
+                        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<!is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleFloor(HloInstruction* floor) {
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[floor],
+        ElementWiseUnaryOp(floor, [](ElementwiseT elem_operand) {
+          return std::floor(elem_operand);
+        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleFloor(HloInstruction* floor) {
+    return InvalidArgument("Unsupported type for Floor");
+  }
+
+  Status HandleFloor(HloInstruction* floor) override {
+    return HandleFloor<ReturnT>(floor);
+  }
+
+  Status HandleLog(HloInstruction* log) override {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[log],
+                        ElementWiseUnaryOp(log, [](ElementwiseT elem_operand) {
+                          return std::log(elem_operand);
+                        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<
+                std::is_integral<NativeT>::value &&
+                !std::is_same<NativeT, bool>::value>::type* = nullptr>
+  Status HandleNot(HloInstruction* not_) {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[not_],
+                        ElementWiseUnaryOp(not_, [](ElementwiseT elem_operand) {
+                          return ~elem_operand;
+                        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT, typename std::enable_if<std::is_floating_point<
+                                  NativeT>::value>::type* = nullptr>
+  Status HandleNot(HloInstruction* not_) {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[not_],
+                        ElementWiseUnaryOp(not_, [](ElementwiseT elem_operand) {
+                          return !elem_operand;
+                        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<std::is_same<NativeT, bool>::value>::type* =
+                nullptr>
+  Status HandleNot(HloInstruction* not_) {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[not_],
+                        ElementWiseUnaryOp(not_, [](ElementwiseT elem_operand) {
+                          return !elem_operand;
+                        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleNot(HloInstruction* not_) {
+    return InvalidArgument("Unsupported type for Not");
+  }
+
+  Status HandleNot(HloInstruction* not_) override {
+    return HandleNot<ElementwiseT>(not_);
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<
+                std::is_signed<NativeT>::value &&
+                !std::is_floating_point<NativeT>::value>::type* = nullptr>
+  Status HandleNegate(HloInstruction* negate) {
+    using type = typename std::make_unsigned<NativeT>::type;
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[negate],
+        ElementWiseUnaryOp(negate, [](ElementwiseT elem_operand) {
+          return NativeT(-type(elem_operand));
+        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<
+                !std::is_signed<NativeT>::value ||
+                std::is_floating_point<NativeT>::value>::type* = nullptr>
+  Status HandleNegate(HloInstruction* negate) {
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[negate],
+        ElementWiseUnaryOp(
+            negate, [](ElementwiseT elem_operand) { return -elem_operand; }));
+    return Status::OK();
+  }
+
+  Status HandleNegate(HloInstruction* negate) override {
+    return HandleNegate<ReturnT>(negate);
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<!is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleSign(HloInstruction* sign) {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[sign],
+                        ElementWiseUnaryOp(sign, [](ElementwiseT elem_operand) {
+                          return (ElementwiseT(0) < elem_operand) -
+                                 (elem_operand < ElementwiseT(0));
+                        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleSign(HloInstruction* sign) {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[sign],
+                        ElementWiseUnaryOp(sign, [](ElementwiseT elem_operand) {
+                          auto abs_val = std::abs(elem_operand);
+                          return 0 == abs_val ? ElementwiseT(0)
+                                              : elem_operand / abs_val;
+                        }));
+    return Status::OK();
+  }
+
+  Status HandleSign(HloInstruction* sign) override {
+    return HandleSign<ReturnT>(sign);
+  }
+
+  template <typename NativeT, typename std::enable_if<std::is_floating_point<
+                                  NativeT>::value>::type* = nullptr>
+  Status HandleAtan2(HloInstruction* atan2) {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[atan2],
+                        ElementWiseBinaryOp(atan2, [](ElementwiseT lhs_elem,
+                                                      ElementwiseT rhs_elem) {
+                          return std::atan2(lhs_elem, rhs_elem);
+                        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT, typename std::enable_if<!std::is_floating_point<
+                                  NativeT>::value>::type* = nullptr>
+  Status HandleAtan2(HloInstruction* atan2) {
+    return InvalidArgument("Unsupported type for Atan2");
+  }
+
+  Status HandleAtan2(HloInstruction* atan2) override {
+    return HandleAtan2<ElementwiseT>(atan2);
+  }
+
+  Status HandleTanh(HloInstruction* tanh) override {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[tanh],
+                        ElementWiseUnaryOp(tanh, [](ElementwiseT elem_operand) {
+                          return std::tanh(elem_operand);
+                        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<
+                std::is_signed<NativeT>::value &&
+                !std::is_floating_point<NativeT>::value>::type* = nullptr>
+  Status HandleMultiply(HloInstruction* multiply) {
+    using type = typename std::make_unsigned<NativeT>::type;
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[multiply],
+        ElementWiseBinaryOp(multiply,
+                            [](ElementwiseT lhs_elem, ElementwiseT rhs_elem) {
+                              return NativeT(type(lhs_elem) * type(rhs_elem));
+                            }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<std::is_unsigned<NativeT>::value ||
+                              std::is_floating_point<NativeT>::value ||
+                              is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleMultiply(HloInstruction* multiply) {
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[multiply],
+        ElementWiseBinaryOp(multiply,
+                            [](ElementwiseT lhs_elem, ElementwiseT rhs_elem) {
+                              return lhs_elem * rhs_elem;
+                            }));
+    return Status::OK();
+  }
+
+  Status HandleMultiply(HloInstruction* multiply) override {
+    return HandleMultiply<ElementwiseT>(multiply);
+  }
+
+  Status HandleSubtract(HloInstruction* subtract) override {
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[subtract],
+        ElementWiseBinaryOp(subtract,
+                            [](ElementwiseT lhs_elem, ElementwiseT rhs_elem) {
+                              return lhs_elem - rhs_elem;
+                            }));
+    return Status::OK();
+  }
+
+  Status HandleAdd(HloInstruction* add) override {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[add],
+                        ElementWiseBinaryOp(add, [](ElementwiseT lhs_elem,
+                                                    ElementwiseT rhs_elem) {
+                          return lhs_elem + rhs_elem;
+                        }));
+    return Status::OK();
+  }
+
+  Status HandleDivide(HloInstruction* divide) override {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[divide],
+                        ElementWiseBinaryOp(divide, [](ElementwiseT lhs_elem,
+                                                       ElementwiseT rhs_elem) {
+                          return lhs_elem / rhs_elem;
+                        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<std::is_integral<NativeT>::value>::type* =
+                nullptr>
+  Status HandleMaximum(HloInstruction* maximum) {
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[maximum],
+        ElementWiseBinaryOp(maximum, [](ElementwiseT lhs, ElementwiseT rhs) {
+          return std::max(lhs, rhs);
+        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT, typename std::enable_if<std::is_floating_point<
+                                  NativeT>::value>::type* = nullptr>
+  Status HandleMaximum(HloInstruction* maximum) {
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[maximum],
+        ElementWiseBinaryOp(maximum, [](ElementwiseT lhs, ElementwiseT rhs) {
+          return ((lhs >= rhs) || std::isnan(lhs)) ? lhs : rhs;
+        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleMaximum(HloInstruction* maximum) {
+    return InvalidArgument("Unsupported type for Maximum");
+  }
+
+  Status HandleMaximum(HloInstruction* maximum) override {
+    return HandleMaximum<ElementwiseT>(maximum);
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<std::is_integral<NativeT>::value>::type* =
+                nullptr>
+  Status HandleMinimum(HloInstruction* minimum) {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[minimum],
+                        ElementWiseBinaryOp(minimum, [](ElementwiseT lhs_el,
+                                                        ElementwiseT rhs_el) {
+                          return std::min(lhs_el, rhs_el);
+                        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT, typename std::enable_if<std::is_floating_point<
+                                  NativeT>::value>::type* = nullptr>
+  Status HandleMinimum(HloInstruction* minimum) {
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[minimum],
+        ElementWiseBinaryOp(minimum, [](ElementwiseT lhs_el,
+                                        ElementwiseT rhs_el) {
+          return ((lhs_el <= rhs_el) || std::isnan(lhs_el)) ? lhs_el : rhs_el;
+        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleMinimum(HloInstruction* minimum) {
+    return InvalidArgument("Unsupported type for Minimum");
+  }
+
+  Status HandleMinimum(HloInstruction* minimum) override {
+    return HandleMinimum<ElementwiseT>(minimum);
+  }
+
+  Status HandlePower(HloInstruction* power) override {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[power],
+                        ElementWiseBinaryOp(power, [](ElementwiseT lhs_el,
+                                                      ElementwiseT rhs_el) {
+                          return std::pow(lhs_el, rhs_el);
+                        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<!is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleRemainder(HloInstruction* remainder) {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[remainder],
+                        ElementWiseBinaryOp(remainder, [](ElementwiseT lhs_el,
+                                                          ElementwiseT rhs_el) {
+                          return std::fmod(lhs_el, rhs_el);
+                        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleRemainder(HloInstruction* remainder) {
+    return InvalidArgument("Unsupported type for Remainder");
+  }
+
+  Status HandleRemainder(HloInstruction* remainder) override {
+    return HandleRemainder<ElementwiseT>(remainder);
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<std::is_integral<NativeT>::value>::type* =
+                nullptr>
+  Status HandleAnd(HloInstruction* and_) {
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[and_],
+        ElementWiseBinaryOp(and_, [](ElementwiseT lhs_el, ElementwiseT rhs_el) {
+          return lhs_el & rhs_el;
+        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT, typename std::enable_if<std::is_floating_point<
+                                  NativeT>::value>::type* = nullptr>
+  Status HandleAnd(HloInstruction* and_) {
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[and_],
+        ElementWiseBinaryOp(and_, [](ElementwiseT lhs_el, ElementwiseT rhs_el) {
+          return lhs_el && rhs_el;
+        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleAnd(HloInstruction* and_) {
+    return InvalidArgument("Unsupported type for And");
+  }
+
+  Status HandleAnd(HloInstruction* and_) override {
+    return HandleAnd<ElementwiseT>(and_);
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<std::is_integral<NativeT>::value>::type* =
+                nullptr>
+  Status HandleOr(HloInstruction* or_) {
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[or_],
+        ElementWiseBinaryOp(or_, [](ElementwiseT lhs_el, ElementwiseT rhs_el) {
+          return lhs_el | rhs_el;
+        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT, typename std::enable_if<std::is_floating_point<
+                                  NativeT>::value>::type* = nullptr>
+  Status HandleOr(HloInstruction* or_) {
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[or_],
+        ElementWiseBinaryOp(or_, [](ElementwiseT lhs_el, ElementwiseT rhs_el) {
+          return lhs_el || rhs_el;
+        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleOr(HloInstruction* or_) {
+    return InvalidArgument("Unsupported type for Or");
+  }
+
+  Status HandleOr(HloInstruction* or_) override {
+    return HandleOr<ElementwiseT>(or_);
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<
+                std::is_integral<NativeT>::value &&
+                !std::is_same<NativeT, bool>::value>::type* = nullptr>
+  Status HandleShiftLeft(HloInstruction* shl) {
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[shl],
+        ElementWiseBinaryOp(shl, [](NativeT lhs_elem, NativeT rhs_elem) {
+          return IsShiftOutOfBounds<NativeT>(rhs_elem) ? 0
+                                                       : (lhs_elem << rhs_elem);
+        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<!std::is_integral<NativeT>::value ||
+                                    std::is_same<NativeT, bool>::value>::type* =
+                nullptr>
+  Status HandleShiftLeft(HloInstruction*) {
+    return InvalidArgument("Unsupported type for ShiftLeft");
+  }
+
+  Status HandleShiftLeft(HloInstruction* shl) override {
+    return HandleShiftLeft<ElementwiseT>(shl);
+  }
+  template <typename NativeT,
+            typename std::enable_if<
+                std::is_integral<NativeT>::value &&
+                !std::is_same<NativeT, bool>::value>::type* = nullptr>
+  Status HandleShiftRightArithmetic(HloInstruction* shr) {
+    typedef typename std::make_signed<NativeT>::type SignedT;
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[shr],
+        ElementWiseBinaryOp(shr, [](NativeT lhs_elem, NativeT rhs_elem) {
+          SignedT lhs_signed = static_cast<SignedT>(lhs_elem);
+          if (IsShiftOutOfBounds<NativeT>(rhs_elem)) {
+            return lhs_signed < 0 ? static_cast<SignedT>(-1) : 0;
+          } else {
+            return lhs_signed >> rhs_elem;
+          }
+        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<!std::is_integral<NativeT>::value ||
+                                    std::is_same<NativeT, bool>::value>::type* =
+                nullptr>
+  Status HandleShiftRightArithmetic(HloInstruction*) {
+    return InvalidArgument("Unsupported type for ShiftRightArithmetic");
+  }
+
+  Status HandleShiftRightArithmetic(HloInstruction* shra) override {
+    return HandleShiftRightArithmetic<ElementwiseT>(shra);
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<
+                std::is_integral<NativeT>::value &&
+                !std::is_same<NativeT, bool>::value>::type* = nullptr>
+  Status HandleShiftRightLogical(HloInstruction* shr) {
+    typedef typename std::make_unsigned<NativeT>::type UnsignedT;
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[shr],
+        ElementWiseBinaryOp(shr, [](NativeT lhs_elem, NativeT rhs_elem) {
+          // If shift amount is greater than the number of bits, then return 0.
+          if (IsShiftOutOfBounds<NativeT>(rhs_elem)) {
+            return static_cast<NativeT>(0);
+          }
+          return static_cast<NativeT>(static_cast<UnsignedT>(lhs_elem) >>
+                                      rhs_elem);
+        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<!std::is_integral<NativeT>::value ||
+                                    std::is_same<NativeT, bool>::value>::type* =
+                nullptr>
+  Status HandleShiftRightLogical(HloInstruction*) {
+    return InvalidArgument("Unsupported type for ShiftRightLogical");
+  }
+
+  Status HandleShiftRightLogical(HloInstruction* shrl) override {
+    return HandleShiftRightLogical<ElementwiseT>(shrl);
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<!is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleClamp(HloInstruction* clamp) {
+    std::function<ElementwiseT(ElementwiseT, ElementwiseT, ElementwiseT)>
+        clamp_op = [](ElementwiseT low, ElementwiseT value, ElementwiseT high) {
+          return std::fmin(high, std::fmax(value, low));
+        };
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[clamp],
+        ElementwiseTernaryOp(clamp,
+                             std::move(ConvertTernaryFunction(clamp_op))));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleClamp(HloInstruction*) {
+    return InvalidArgument("Unsupported type for Clamp");
+  }
+
+  Status HandleClamp(HloInstruction* clamp) override {
+    return HandleClamp<ElementwiseT>(clamp);
+  }
+
+  Status HandleSelect(HloInstruction* select) override {
+    CHECK(!ShapeUtil::IsScalar(select->operand(0)->shape()));
+    CHECK(!ShapeUtil::IsTuple(select->shape()));
+    std::function<ReturnT(bool, ReturnT, ReturnT)> select_op =
+        [](bool pred, ReturnT on_true, ReturnT on_false) {
+          if (pred) {
+            return on_true;
+          }
+          return on_false;
+        };
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[select],
+                        ElementwiseTernaryOp(select, std::move(select_op)));
+    return Status::OK();
+  }
+
+  Status HandleReverse(HloInstruction* reverse) override {
+    const auto result_shape = reverse->shape();
+    const auto reverse_dimensions = reverse->dimensions();
+
+    auto operand = reverse->operand(0);
+    TF_ASSIGN_OR_RETURN(auto inferred_return_shape,
+                        ShapeInference::InferReverseShape(operand->shape(),
+                                                          reverse_dimensions));
+
+    TF_RET_CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape))
+        << "return shape set to: " << ShapeUtil::HumanString(result_shape)
+        << " but is inferred to be: "
+        << ShapeUtil::HumanString(inferred_return_shape);
+
+    const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand);
+    auto result = Literal::CreateFromShape(result_shape);
+
+    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(
+        [&](tensorflow::gtl::ArraySlice<int64> out_index) {
+          std::vector<int64> from_index(out_index.begin(), out_index.end());
+          for (const int64 dim : reverse_dimensions) {
+            from_index[dim] = result_shape.dimensions(dim) - 1 - out_index[dim];
+          }
+          return operand_literal.Get<ReturnT>(from_index);
+        }));
+
+    parent_->evaluated_[reverse] = std::move(result);
+    return Status::OK();
+  }
+
+  Status HandleConvolution(HloInstruction* conv) override {
+    auto lhs = conv->operand(0);
+    auto rhs = conv->operand(1);
+    const auto& window = conv->window();
+    const Shape& result_shape = conv->shape();
+    const Shape& lhs_shape = lhs->shape();
+    const Shape& rhs_shape = rhs->shape();
+
+    TF_CHECK_OK(ShapeUtil::ValidateShape(lhs_shape));
+    TF_CHECK_OK(ShapeUtil::ValidateShape(rhs_shape));
+    CHECK(ShapeUtil::IsArray(lhs_shape));
+    CHECK(ShapeUtil::IsArray(rhs_shape));
+    CHECK(ShapeUtil::SameElementType(lhs_shape, rhs_shape));
+    CHECK(ShapeUtil::SameElementType(lhs_shape, result_shape));
+
+    const auto& dnums = conv->convolution_dimension_numbers();
+    const int64 num_spatial_dims = dnums.output_spatial_dimensions_size();
+    CHECK_EQ(num_spatial_dims, dnums.input_spatial_dimensions_size());
+    CHECK_EQ(num_spatial_dims, dnums.kernel_spatial_dimensions_size());
+    CHECK_GE(num_spatial_dims, 0);
+    CHECK_EQ(window.dimensions_size(), num_spatial_dims);
+
+    const auto lhs_rank = ShapeUtil::Rank(lhs_shape);
+    const auto rhs_rank = ShapeUtil::Rank(rhs_shape);
+
+    CHECK_EQ(num_spatial_dims + 2, lhs_rank);
+    CHECK_EQ(num_spatial_dims + 2, rhs_rank);
+
+    TF_ASSIGN_OR_RETURN(auto inferred_return_shape,
+                        ShapeInference::InferConvolveShape(lhs_shape, rhs_shape,
+                                                           window, dnums));
+    CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape))
+        << "return shape set to: " << ShapeUtil::HumanString(result_shape)
+        << " but is inferred to be: "
+        << ShapeUtil::HumanString(inferred_return_shape);
+
+    const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs);
+    const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs);
+
+    std::vector<int64> window_dimension_sizes;
+    for (auto i : dnums.kernel_spatial_dimensions()) {
+      window_dimension_sizes.push_back(ShapeUtil::GetDimension(rhs_shape, i));
+    }
+
+    const Shape& window_shape =
+        ShapeUtil::MakeShape(rhs_shape.element_type(), window_dimension_sizes);
+
+    DimensionVector lhs_dim_multipliers = MakeDimMultipliers(lhs_shape);
+    DimensionVector rhs_dim_multipliers = MakeDimMultipliers(rhs_shape);
+
+    auto lhs_literal_data = lhs_literal.data<ReturnT>();
+    auto rhs_literal_data = rhs_literal.data<ReturnT>();
+
+    auto func = [&window_shape, &dnums, &lhs_shape, &rhs_shape, &window,
+                 &lhs_dim_multipliers, &rhs_dim_multipliers, lhs_literal_data,
+                 rhs_literal_data](
+                    tensorflow::gtl::ArraySlice<int64> out_index) {
+      // Dimension number applicable for input (lhs).
+      const int64 input_batch_dim = dnums.input_batch_dimension();
+      const int64 input_z_dim = dnums.input_feature_dimension();
+      // Dimension number applicable for kernel (rhs).
+      const int64 kernel_input_z_dim = dnums.kernel_input_feature_dimension();
+      const int64 kernel_output_z_dim = dnums.kernel_output_feature_dimension();
+      // Dimension number applicable for output.
+      const int64 output_batch_dim = dnums.output_batch_dimension();
+      const int64 output_z_dim = dnums.output_feature_dimension();
+
+      const int64 z_size = ShapeUtil::GetDimension(lhs_shape, input_z_dim);
+
+      ElementwiseT result_val = static_cast<ElementwiseT>(0);
+      DimensionVector rhs_spatial_index(dnums.kernel_spatial_dimensions_size(),
+                                        0);
+
+      // Convolve input feature with kernel.
+      do {
+        for (int64 iz = 0; iz < z_size; ++iz) {
+          int64 lhs_linear_index = 0;
+          lhs_linear_index += out_index[output_batch_dim] *
+                              lhs_dim_multipliers[input_batch_dim];
+          lhs_linear_index += iz * lhs_dim_multipliers[input_z_dim];
+
+          int64 rhs_linear_index = 0;
+          rhs_linear_index += out_index[output_z_dim] *
+                              rhs_dim_multipliers[kernel_output_z_dim];
+          rhs_linear_index += iz * rhs_dim_multipliers[kernel_input_z_dim];
+
+          // Find corresponding spatial dimension index for input (lhs).
+          for (int64 ki = 0; ki < rhs_spatial_index.size(); ++ki) {
+            // Spatial dimension number for input (lhs) and output.
+            const int64 input_spatial_dim = dnums.input_spatial_dimensions(ki);
+            const int64 output_spatial_dim =
+                dnums.output_spatial_dimensions(ki);
+
+            // Calculate lhs (input) index without taking base dilation into
+            // account.
+            const auto& window_dim = window.dimensions(ki);
+            const int64 undilated_index =
+                out_index[output_spatial_dim] * window_dim.stride() -
+                window_dim.padding_low() +
+                rhs_spatial_index[ki] * window_dim.window_dilation();
+            // Skip if the lhs (input) index is to be dilated.  As an
+            // optimization, skip this mod if there's no dilation.
+            if (window_dim.base_dilation() > 1 &&
+                undilated_index % window_dim.base_dilation() != 0) {
+              goto cnt;
+            }
+
+            // Calculate the actual lhs (input) index after dilation.  As an
+            // optimization, skip this integer divide if there's no dilation.
+            int64 lhs_spatial_index;
+            if (window_dim.base_dilation() > 1) {
+              lhs_spatial_index = undilated_index / window_dim.base_dilation();
+            } else {
+              lhs_spatial_index = undilated_index;
+            }
+            lhs_linear_index +=
+                lhs_spatial_index * lhs_dim_multipliers[input_spatial_dim];
+
+            // Skip if input index is not in bounds.
+            if (!(lhs_spatial_index >= 0 &&
+                  lhs_spatial_index <
+                      lhs_shape.dimensions(input_spatial_dim))) {
+              goto cnt;
+            }
+
+            rhs_linear_index +=
+                (window_dim.window_reversal()
+                     ? ((window_dim.size() - 1) - rhs_spatial_index[ki])
+                     : rhs_spatial_index[ki]) *
+                rhs_dim_multipliers[dnums.kernel_spatial_dimensions(ki)];
+          }
+
+          result_val +=
+              static_cast<ElementwiseT>(lhs_literal_data[lhs_linear_index]) *
+              static_cast<ElementwiseT>(rhs_literal_data[rhs_linear_index]);
+        }
+      cnt : {}
+      } while (IndexUtil::BumpIndices(window_shape, &rhs_spatial_index));
+
+      return static_cast<ReturnT>(result_val);
+    };
+
+    auto result = Literal::CreateFromShape(result_shape);
+    TF_RETURN_IF_ERROR(result->PopulateParallel<ReturnT>(func));
+
+    parent_->evaluated_[conv] = std::move(result);
+    return Status::OK();
+  }
+
+  Status HandleDot(HloInstruction* dot) override {
+    auto lhs = dot->operand(0);
+    auto rhs = dot->operand(1);
+    CHECK(ShapeUtil::IsArray(dot->shape()));
+    CHECK(ShapeUtil::IsArray(lhs->shape()));
+    CHECK(ShapeUtil::IsArray(rhs->shape()));
+
+    const auto& dnums = dot->dot_dimension_numbers();
+
+    const auto lhs_rank = ShapeUtil::Rank(lhs->shape());
+    const auto rhs_rank = ShapeUtil::Rank(rhs->shape());
+
+    CHECK(ShapeUtil::SameElementType(lhs->shape(), rhs->shape()));
+    CHECK(ShapeUtil::SameElementType(lhs->shape(), dot->shape()));
+
+    // There must be 1 and only 1 Contracting dimension for lhs and rhs.
+    CHECK_EQ(dnums.lhs_contracting_dimensions_size(), 1);
+    CHECK_EQ(dnums.rhs_contracting_dimensions_size(), 1);
+    const int64 lhs_contracting_dimension = dnums.lhs_contracting_dimensions(0);
+    const int64 rhs_contracting_dimension = dnums.rhs_contracting_dimensions(0);
+    // Contracted dimension sizes must be the same.
+    CHECK_EQ(lhs->shape().dimensions(lhs_contracting_dimension),
+             rhs->shape().dimensions(rhs_contracting_dimension))
+        << "lhs contracted dimension: "
+        << lhs->shape().dimensions(lhs_contracting_dimension)
+        << " rhs contracted dimension: "
+        << rhs->shape().dimensions(rhs_contracting_dimension);
+    const int64 contracted_dimension_size =
+        lhs->shape().dimensions(lhs_contracting_dimension);
+
+    const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs);
+    const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs);
+
+    auto result = Literal::CreateFromShape(dot->shape());
+
+    CHECK_EQ(dnums.lhs_batch_dimensions_size(),
+             dnums.rhs_batch_dimensions_size());
+
+    std::vector<int64> lhs_non_contracting_dims;
+    for (int64 i = 0; i < lhs_rank; i++) {
+      if (i != lhs_contracting_dimension) {
+        lhs_non_contracting_dims.push_back(i);
+      }
+    }
+
+    std::vector<int64> rhs_non_batch_non_contracting_dims;
+    tensorflow::gtl::FlatSet<int64> batch_dims_set(
+        dnums.rhs_batch_dimensions().begin(),
+        dnums.rhs_batch_dimensions().end());
+    for (int64 i = 0; i < rhs_rank; i++) {
+      if (i != rhs_contracting_dimension && batch_dims_set.count(i) == 0) {
+        rhs_non_batch_non_contracting_dims.push_back(i);
+      }
+    }
+
+    const int64 batch_dim_size = dnums.lhs_batch_dimensions_size();
+    const int64 lhs_non_contracting_size = lhs_non_contracting_dims.size();
+
+    DimensionVector lhs_index(lhs_rank);
+    DimensionVector rhs_index(rhs_rank);
+    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(
+        [&](tensorflow::gtl::ArraySlice<int64> result_index) {
+          ElementwiseT result_val = static_cast<ElementwiseT>(0);
+
+          // Find the corresponding non-contracting indices for lhs and rhs.
+          //
+          // For `result_index`, its batch dimension, if exists, will be at the
+          // same dimension as the batch dimension of lhs and rhs. More
+          // specifically:
+          // - For lhs, the non-contracting dimensions, including the batch
+          // dimension have the same index as the `result_index`.
+          // - For rhs, the batch dimension is set seperately from other
+          // non-contracting dimensions, since these other non-contracting
+          // dimensions in rhs follow the non-contracting dimensions of lhs in
+          // the resulting index.
+          //
+          // As an example, for a resulting index:
+          //  result_index [result_batch, result_x, result_y]
+          // the effecting lhs and rhs indices are:
+          //  lhs [result_batch, lhs_non_contracting_dim, contracting_dim
+          //  rhs [result_batch, contracting_dim, rhs_non_contracting_dim]
+          // `result_x` is only affected by the lhs_non_contracting_dim and
+          // likewise `result_y` only depends on rhs_non_contracting_dim.
+          //
+          // so we can look up the lhs and rhs indices by:
+          //
+          // lhs:
+          //  batch index is the same as `result_batch`.
+          //    non-contracting dimension is the same as
+          //    result_index[lhs_non_contracting_dim]
+          // rhs:
+          //  batch index: the same as `result_batch`.
+          //  non-contracting dimension index: *not* the same as
+          //    result_index[rhs_non_contractng_dim], since the
+          //    non-contracting dimensions of lhs are included in the
+          //    result_index first. Instead, the non_contracting_dim of rhs must
+          //    be calculated as following:
+          //      lhs_non_contracting_dimensions_size +
+          //      (rhs_non_batch_non_contracting_dim - batch_dim_size) - 1
+          //
+          //    Note that (rhs_non_batch_contracting_dim - batch_dim_size) is
+          //    the index offset to the result_index that only depends on
+          //    the non_batch and non-contracting dimensions of rhs. -1 at the
+          //    end translates size to index.
+          for (auto i : lhs_non_contracting_dims) {
+            lhs_index[i] = result_index[i];
+          }
+          for (auto i : dnums.rhs_batch_dimensions()) {
+            rhs_index[i] = result_index[i];
+          }
+          for (auto i : rhs_non_batch_non_contracting_dims) {
+            const int64 rhs_non_batch_non_contracting_dim =
+                lhs_non_contracting_size + (i - batch_dim_size) - 1;
+            rhs_index[i] = result_index[rhs_non_batch_non_contracting_dim];
+          }
+
+          // Accumulates resulting product along the contracted dimension.
+          for (int64 i = 0; i < contracted_dimension_size; ++i) {
+            lhs_index[lhs_contracting_dimension] = i;
+            rhs_index[rhs_contracting_dimension] = i;
+
+            result_val +=
+                static_cast<ElementwiseT>(lhs_literal.Get<ReturnT>(lhs_index)) *
+                static_cast<ElementwiseT>(rhs_literal.Get<ReturnT>(rhs_index));
+          }
+
+          return static_cast<ReturnT>(result_val);
+        }));
+
+    parent_->evaluated_[dot] = std::move(result);
+    return Status::OK();
+  }
+
+  Status HandlePad(HloInstruction* pad) override {
+    CHECK(!ShapeUtil::IsTuple(pad->operand(0)->shape()));
+    // Padding value must be scalar.
+    CHECK(ShapeUtil::IsScalar(pad->operand(1)->shape()));
+    CHECK_EQ(ShapeUtil::Rank(pad->operand(0)->shape()),
+             pad->padding_config().dimensions_size());
+
+    TF_ASSIGN_OR_RETURN(auto inferred_return_shape,
+                        ShapeInference::InferPadShape(
+                            /*operand_shape=*/pad->operand(0)->shape(),
+                            /*padding_value_shape=*/pad->operand(1)->shape(),
+                            /*padding_config=*/pad->padding_config()));
+    CHECK(ShapeUtil::Compatible(pad->shape(), inferred_return_shape))
+        << "return shape is set to: " << ShapeUtil::HumanString(pad->shape())
+        << "but is inferred to be: "
+        << ShapeUtil::HumanString(inferred_return_shape);
+
+    // Create new HLO of padded shape with padding value.
+    ReturnT scalar =
+        parent_->GetEvaluatedLiteralFor(pad->operand(1)).Get<ReturnT>({});
+    auto result = Literal::CreateFromShape(pad->shape());
+    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(
+        [&scalar](tensorflow::gtl::ArraySlice<int64> multi_index) {
+          return scalar;
+        }));
+
+    const Literal& evaluated_operand =
+        parent_->GetEvaluatedLiteralFor(pad->operand(0));
+
+    std::vector<int64> input_index(ShapeUtil::Rank(evaluated_operand.shape()),
+                                   0);
+    std::vector<int64> target_index(ShapeUtil::Rank(result->shape()), 0);
+
+    // Loop through each element of the operand, assign them to the
+    // corresponding index of the resulting padded literal.
+    const PaddingConfig& pad_config = pad->padding_config();
+
+    auto func = [&](tensorflow::gtl::ArraySlice<int64> input_index) {
+      for (auto i = 0; i < input_index.size(); ++i) {
+        // Interior padding occurs logically before edge padding, so in the case
+        // of negative edge padding elements are removed from the
+        // interior-padded operand.
+        target_index[i] =
+            pad_config.dimensions(i).edge_padding_low() +
+            input_index[i] * (pad_config.dimensions(i).interior_padding() + 1);
+
+        // Account for negative low and high padding: skip assignment if the
+        // any target index is out of range.
+        if (!(target_index[i] >= 0 &&
+              target_index[i] < pad->shape().dimensions(i))) {
+          return true;
+        }
+      }
+      result->Set<ReturnT>(target_index,
+                           evaluated_operand.Get<ReturnT>(input_index));
+      return true;
+    };
+
+    std::vector<int64> zero_base(evaluated_operand.shape().dimensions_size(),
+                                 0);
+    std::vector<int64> step(evaluated_operand.shape().dimensions_size(), 1);
+
+    ShapeUtil::ForEachIndex(
+        evaluated_operand.shape(), zero_base,
+        AsInt64Slice(evaluated_operand.shape().dimensions()), step, func);
+
+    parent_->evaluated_[pad] = std::move(result);
+    return Status::OK();
+  }
+
+  Status HandleDynamicSlice(HloInstruction* dynamic_slice) override {
+    auto operand = dynamic_slice->operand(0);
+    auto start_indices = dynamic_slice->operand(1);
+    auto result_shape = dynamic_slice->shape();
+    TF_ASSIGN_OR_RETURN(auto inferred_return_shape,
+                        ShapeInference::InferDynamicSliceShape(
+                            operand->shape(), start_indices->shape(),
+                            dynamic_slice->dynamic_slice_sizes()));
+    TF_RET_CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape))
+        << "return shape is set to: " << ShapeUtil::HumanString(result_shape)
+        << "but is inferred to be: "
+        << ShapeUtil::HumanString(inferred_return_shape);
+    TF_RET_CHECK(
+        primitive_util::IsIntegralType(start_indices->shape().element_type()));
+
+    const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand);
+    const Literal& start_indices_literal =
+        parent_->GetEvaluatedLiteralFor(start_indices);
+
+    switch (start_indices->shape().element_type()) {
+      case S32: {
+        TF_ASSIGN_OR_RETURN(
+            parent_->evaluated_[dynamic_slice],
+            DynamicSlice<int32>(operand_literal, start_indices_literal,
+                                result_shape));
+      } break;
+      case S64: {
+        TF_ASSIGN_OR_RETURN(
+            parent_->evaluated_[dynamic_slice],
+            DynamicSlice<int64>(operand_literal, start_indices_literal,
+                                result_shape));
+      } break;
+      case U32: {
+        TF_ASSIGN_OR_RETURN(
+            parent_->evaluated_[dynamic_slice],
+            DynamicSlice<uint32>(operand_literal, start_indices_literal,
+                                 result_shape));
+      } break;
+      case U64: {
+        TF_ASSIGN_OR_RETURN(
+            parent_->evaluated_[dynamic_slice],
+            DynamicSlice<uint64>(operand_literal, start_indices_literal,
+                                 result_shape));
+      } break;
+      default:
+        LOG(FATAL) << "HandleDynamicSlice: unhandled primitive type for "
+                      "start_indices: "
+                   << PrimitiveType_Name(start_indices->shape().element_type());
+    }
+
+    return Status::OK();
+  }
+
+  Status HandleDynamicUpdateSlice(
+      HloInstruction* dynamic_update_slice) override {
+    auto operand = dynamic_update_slice->operand(0);
+    auto update = dynamic_update_slice->operand(1);
+    auto start_indices = dynamic_update_slice->operand(2);
+    auto result_shape = dynamic_update_slice->shape();
+    TF_ASSIGN_OR_RETURN(
+        auto inferred_return_shape,
+        ShapeInference::InferDynamicUpdateSliceShape(
+            operand->shape(), update->shape(), start_indices->shape()));
+    TF_RET_CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape))
+        << "return shape is set to: " << ShapeUtil::HumanString(result_shape)
+        << "but is inferred to be: "
+        << ShapeUtil::HumanString(inferred_return_shape);
+    TF_RET_CHECK(
+        primitive_util::IsIntegralType(start_indices->shape().element_type()));
+    TF_RET_CHECK(ShapeUtil::Compatible(result_shape, operand->shape()));
+
+    const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand);
+    const Literal& update_literal = parent_->GetEvaluatedLiteralFor(update);
+    const Literal& start_indices_literal =
+        parent_->GetEvaluatedLiteralFor(start_indices);
+
+    switch (start_indices->shape().element_type()) {
+      case S32: {
+        TF_ASSIGN_OR_RETURN(
+            parent_->evaluated_[dynamic_update_slice],
+            DynamicUpdateSlice<int32>(operand_literal, update_literal,
+                                      start_indices_literal));
+      } break;
+      case S64: {
+        TF_ASSIGN_OR_RETURN(
+            parent_->evaluated_[dynamic_update_slice],
+            DynamicUpdateSlice<int64>(operand_literal, update_literal,
+                                      start_indices_literal));
+      } break;
+      case U32: {
+        TF_ASSIGN_OR_RETURN(
+            parent_->evaluated_[dynamic_update_slice],
+            DynamicUpdateSlice<uint32>(operand_literal, update_literal,
+                                       start_indices_literal));
+      } break;
+      case U64: {
+        TF_ASSIGN_OR_RETURN(
+            parent_->evaluated_[dynamic_update_slice],
+            DynamicUpdateSlice<uint64>(operand_literal, update_literal,
+                                       start_indices_literal));
+      } break;
+      default:
+        LOG(FATAL) << "HandleDynamicUpdateSlice: unhandled primitive type for "
+                      "start_indices: "
+                   << PrimitiveType_Name(start_indices->shape().element_type());
+    }
+
+    return Status::OK();
+  }
+
+  template <typename NativeT>
+  StatusOr<std::unique_ptr<Literal>> MapImpl(HloInstruction* map) {
+    auto operands = map->operands();
+    HloComputation* computation = map->to_apply();
+
+    auto result = Literal::CreateFromShape(map->shape());
+
+    HloEvaluator embedded_evaluator(parent_->max_loop_iterations_);
+    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(
+        [&](tensorflow::gtl::ArraySlice<int64> multi_index) {
+          std::vector<std::unique_ptr<Literal>> arg_literals;
+          arg_literals.reserve(operands.size());
+
+          // Construct scalar literal parameters to be passed to the map
+          // computation.
+          for (auto operand : operands) {
+            const Literal& arg_literal =
+                parent_->GetEvaluatedLiteralFor(operand);
+
+            auto curr_val = arg_literal.Get<NativeT>(multi_index);
+            auto curr_val_literal = Literal::CreateR0<NativeT>(curr_val);
+
+            arg_literals.push_back(std::move(curr_val_literal));
+          }
+
+          std::unique_ptr<Literal> computed_result =
+              embedded_evaluator
+                  .Evaluate<std::unique_ptr<Literal>>(*computation,
+                                                      arg_literals)
+                  .ConsumeValueOrDie();
+          // Clear visit states so that the we can use the evaluate again on
+          // the same computation.
+          embedded_evaluator.ResetVisitStates();
+
+          return computed_result->Get<ReturnT>({});
+        }));
+    return std::move(result);
+  }
+
+  Status HandleMap(HloInstruction* map) override {
+    switch (map->operand(0)->shape().element_type()) {
+      case PRED: {
+        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<bool>(map));
+        break;
+      }
+      case U8: {
+        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<uint8>(map));
+        break;
+      }
+      case U32: {
+        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<uint32>(map));
+        break;
+      }
+      case U64: {
+        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<uint64>(map));
+        break;
+      }
+      case S8: {
+        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<int8>(map));
+        break;
+      }
+      case S32: {
+        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<int32>(map));
+        break;
+      }
+      case S64: {
+        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<int64>(map));
+        break;
+      }
+      case F16: {
+        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map],
+                            MapImpl<Eigen::half>(map));
+        break;
+      }
+      case F32: {
+        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<float>(map));
+        break;
+      }
+      case F64: {
+        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<double>(map));
+        break;
+      }
+      case C64: {
+        TF_ASSIGN_OR_RETURN(parent_->evaluated_[map], MapImpl<complex64>(map));
+        break;
+      }
+      default:
+        LOG(FATAL) << "HandleMap: unhandled primitive type for "
+                      "input operand: "
+                   << PrimitiveType_Name(
+                          map->operand(0)->shape().element_type());
+    }
+
+    return Status::OK();
+  }
+
+  Status HandleReduce(HloInstruction* reduce) override {
+    auto arg = reduce->operand(0);
+    auto init_value = reduce->operand(1);
+    tensorflow::gtl::ArraySlice<int64> dimensions(reduce->dimensions());
+    HloComputation* function = reduce->to_apply();
+    TF_RET_CHECK(ShapeUtil::Rank(reduce->shape()) ==
+                 ShapeUtil::Rank(arg->shape()) - dimensions.size());
+    TF_ASSIGN_OR_RETURN(auto inferred_return_shape,
+                        ShapeInference::InferReduceShape(
+                            /*arg=*/arg->shape(),
+                            /*init_value=*/init_value->shape(),
+                            /*dimensions_to_reduce=*/dimensions,
+                            /*to_apply=*/function->ComputeProgramShape()));
+    TF_RET_CHECK(ShapeUtil::Compatible(reduce->shape(), inferred_return_shape))
+        << "return shape is set to: " << ShapeUtil::HumanString(reduce->shape())
+        << "but is inferred to be: "
+        << ShapeUtil::HumanString(inferred_return_shape);
+
+    const Literal& arg_literal = parent_->GetEvaluatedLiteralFor(arg);
+    VLOG(3) << "HandleReduce arg_literal: " << arg_literal.ToString();
+    const Literal& init_literal = parent_->GetEvaluatedLiteralFor(init_value);
+    VLOG(3) << "HandleReduce init_literal: " << init_literal.ToString();
+    TF_RET_CHECK(ShapeUtil::IsScalar(init_literal.shape()));
+    auto init_scalar = init_literal.Get<ReturnT>({});
+
+    auto result = Literal::CreateFromShape(reduce->shape());
+
+    const auto arg_dimensions = AsInt64Slice(arg_literal.shape().dimensions());
+    std::vector<int64> arg_dim_steps(arg_dimensions.size());
+    std::vector<int64> arg_dim_counts(arg_dimensions.size());
+    for (const int64 dim : dimensions) {
+      arg_dim_steps[dim] = 1;
+      arg_dim_counts[dim] = arg_dimensions[dim];
+    }
+
+    // Map each dimension in the result to a dimension in arg that isn't
+    // being reduced.
+    std::vector<int64> result_to_arg_index;
+    for (int64 i = 0; i < arg_dimensions.size(); ++i) {
+      if (arg_dim_steps[i] == 0) {
+        result_to_arg_index.push_back(i);
+      }
+    }
+
+    HloEvaluator embedded_evaluator(parent_->max_loop_iterations_);
+    // For each resulting dimension, calculate and assign computed value.
+    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(
+        [&](tensorflow::gtl::ArraySlice<int64> multi_index) {
+          ReturnT result_val = init_scalar;
+
+          std::vector<int64> base(arg_dimensions.size());
+          for (int64 i = 0; i < multi_index.size(); ++i) {
+            base[result_to_arg_index[i]] = multi_index[i];
+          }
+
+          // When the reduction is addition of floats, accumulate in a double
+          // for better precision. Also, avoid creating Literals for the
+          // intermediate results; it's much faster.
+          if (ShapeUtil::ElementIsFloating(init_literal.shape()) &&
+              IsScalarAdd(function)) {
+            double computed_result = 0;
+            auto func = [&](tensorflow::gtl::ArraySlice<int64> input_index) {
+              computed_result += arg_literal.Get<float>(input_index);
+              return true;
+            };
+            ShapeUtil::ForEachIndex(arg_literal.shape(), base, arg_dim_counts,
+                                    arg_dim_steps, func);
+            return static_cast<ReturnT>(computed_result);
+          }
+          auto func = [&](tensorflow::gtl::ArraySlice<int64> input_index) {
+            auto curr_val = arg_literal.Get<ReturnT>(input_index);
+
+            // Evaluate computation with specified literal operands.
+            auto curr_val_literal = Literal::CreateR0<ReturnT>(curr_val);
+            auto result_val_literal = Literal::CreateR0<ReturnT>(result_val);
+            std::vector<const Literal*> args = {result_val_literal.get(),
+                                                curr_val_literal.get()};
+
+            std::unique_ptr<Literal> computed_result =
+                embedded_evaluator.Evaluate<const Literal*>(*function, args)
+                    .ConsumeValueOrDie();
+            // Clear visit states so that we can use the evaluator again on
+            // the same computation.
+            embedded_evaluator.ResetVisitStates();
+            // Assign computed result to result_val.
+            result_val = computed_result->Get<ReturnT>({});
+            return true;
+          };
+          // Computes one element of the result, reducing all dimensions that
+          // contribute to that element.
+          ShapeUtil::ForEachIndex(arg_literal.shape(), base, arg_dim_counts,
+                                  arg_dim_steps, func);
+          return result_val;
+        }));
+
+    parent_->evaluated_[reduce] = std::move(result);
+    return Status::OK();
+  }
+
+  bool IsScalarAdd(HloComputation* computation) {
+    HloInstruction* instruction = computation->root_instruction();
+    if (instruction->opcode() == HloOpcode::kAdd &&
+        computation->num_parameters() == 2) {
+      const HloInstruction* lhs = instruction->operand(0);
+      const HloInstruction* rhs = instruction->operand(1);
+      return lhs->opcode() == HloOpcode::kParameter &&
+             ShapeUtil::IsScalar(lhs->shape()) &&
+             rhs->opcode() == HloOpcode::kParameter &&
+             ShapeUtil::IsScalar(rhs->shape()) && lhs != rhs;
+    }
+    return false;
+  }
+
+  Status HandleSelectAndScatter(HloInstruction* select_and_scatter) override {
+    auto operand = select_and_scatter->operand(0);
+    auto source = select_and_scatter->operand(1);
+    const Window& window = select_and_scatter->window();
+
+    const Literal& init_literal =
+        parent_->GetEvaluatedLiteralFor(select_and_scatter->operand(2));
+    TF_RET_CHECK(ShapeUtil::IsScalar(init_literal.shape()));
+    auto init_scalar = init_literal.Get<ReturnT>({});
+
+    auto result = Literal::CreateFromShape(select_and_scatter->shape());
+
+    // Initialize result array with the init value.
+    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(
+        [&](tensorflow::gtl::ArraySlice<int64> output_index) {
+          return init_scalar;
+        }));
+
+    std::vector<int64> window_dimension_sizes;
+    for (const auto& window_dimension : window.dimensions()) {
+      window_dimension_sizes.push_back(window_dimension.size());
+    }
+    const Shape window_shape = ShapeUtil::MakeShape(
+        operand->shape().element_type(), window_dimension_sizes);
+
+    HloComputation* select = select_and_scatter->select();
+    HloComputation* scatter = select_and_scatter->scatter();
+
+    const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand);
+    const Literal& source_literal = parent_->GetEvaluatedLiteralFor(source);
+
+    int64 rank = ShapeUtil::Rank(operand_literal.shape());
+
+    HloEvaluator embedded_evaluator(parent_->max_loop_iterations_);
+    DimensionVector source_index(rank);
+
+    std::fill(source_index.begin(), source_index.end(), 0);
+    do {
+      // For each element in `source`, we place a window in `operand`. For each
+      // window placement, we iterate inside the window twice:
+      //
+      // 1. Find the selected index by applying `select` function to all
+      // elements. E.g., If the `select` function is GreaterEqual, the first
+      // iteration through the window finds the biggest value and returns its
+      // index.
+      //
+      // 2. Using the selected index, scatter value from `source` to result. We
+      // do this by iterating through the window, and compare each index with
+      // the selected index.
+      tensorflow::gtl::optional<ReturnT> selected_val;
+      tensorflow::gtl::optional<std::vector<int64>> selected_index;
+
+      IterateThroughWindow(
+          window_shape, window, operand_literal.shape(), source_index,
+          [&](const std::vector<int64>& operand_index) {
+            auto curr_val = operand_literal.Get<ReturnT>(operand_index);
+            if (!selected_val) {
+              selected_val = curr_val;
+              selected_index = operand_index;
+            }
+            const auto curr_val_literal = Literal::CreateR0<ReturnT>(curr_val);
+            const auto selected_val_literal =
+                Literal::CreateR0<ReturnT>(*selected_val);
+
+            const std::vector<const Literal*> args = {
+                selected_val_literal.get(), curr_val_literal.get()};
+            std::unique_ptr<Literal> computed_result =
+                embedded_evaluator.Evaluate<const Literal*>(*select, args)
+                    .ConsumeValueOrDie();
+            bool selected = !computed_result->Get<bool>({});
+            if (selected) {
+              selected_val = curr_val;
+              selected_index = operand_index;
+            }
+            embedded_evaluator.ResetVisitStates();
+          });
+
+      IterateThroughWindow(
+          window_shape, window, operand_literal.shape(), source_index,
+          [&](const std::vector<int64>& operand_index) {
+            if (std::equal(operand_index.begin(), operand_index.end(),
+                           selected_index->begin())) {
+              auto source = source_literal.Get<ReturnT>(source_index);
+              auto scattered = result->Get<ReturnT>(operand_index);
+              const auto source_literal = Literal::CreateR0<ReturnT>(source);
+              const auto scattered_literal =
+                  Literal::CreateR0<ReturnT>(scattered);
+
+              const std::vector<const Literal*> args = {
+                  source_literal.get(), scattered_literal.get()};
+              std::unique_ptr<Literal> computed_result =
+                  embedded_evaluator.Evaluate<const Literal*>(*scatter, args)
+                      .ConsumeValueOrDie();
+              result->Set(operand_index, computed_result->Get<ReturnT>({}));
+              // Clear visit states so that the we can use the evaluator again
+              // on the same computation.
+              embedded_evaluator.ResetVisitStates();
+            }
+          });
+    } while (IndexUtil::BumpIndices(source->shape(), &source_index));
+
+    parent_->evaluated_[select_and_scatter] = std::move(result);
+    return Status::OK();
+  }
+
+  Status HandleReduceWindow(HloInstruction* reduce_window) override {
+    auto operand = reduce_window->operand(0);
+    const Window& window = reduce_window->window();
+    HloComputation* function = reduce_window->to_apply();
+    TF_ASSIGN_OR_RETURN(
+        auto inferred_return_shape,
+        ShapeInference::InferReduceWindowShape(
+            /*operand_shape=*/reduce_window->operand(0)->shape(),
+            /*init_value=*/reduce_window->operand(1)->shape(), window,
+            /*to_apply_shape=*/function->ComputeProgramShape()));
+    TF_RET_CHECK(
+        ShapeUtil::Compatible(reduce_window->shape(), inferred_return_shape))
+        << "return shape is set to: "
+        << ShapeUtil::HumanStringWithLayout(reduce_window->shape())
+        << "but is inferred to be: "
+        << ShapeUtil::HumanStringWithLayout(inferred_return_shape);
+
+    const Literal& operand_literal =
+        parent_->GetEvaluatedLiteralFor(reduce_window->operand(0));
+    VLOG(3) << "HandleReduceWindow arg_literal: " << operand_literal.ToString();
+    const Literal& init_literal =
+        parent_->GetEvaluatedLiteralFor(reduce_window->operand(1));
+    VLOG(3) << "HandleReduceWindow init_literal: " << init_literal.ToString();
+    TF_RET_CHECK(ShapeUtil::IsScalar(init_literal.shape()));
+    auto init_scalar = init_literal.Get<ReturnT>({});
+
+    auto result = Literal::CreateFromShape(reduce_window->shape());
+
+    // Creates a Shape object from window, for iteration below.
+    std::vector<int64> window_dimension_sizes;
+    for (const auto& window_dimension : window.dimensions()) {
+      window_dimension_sizes.push_back(window_dimension.size());
+    }
+    const Shape window_shape = ShapeUtil::MakeShape(
+        operand->shape().element_type(), window_dimension_sizes);
+
+    DimensionVector window_index(window.dimensions_size());
+    DimensionVector operand_index(ShapeUtil::Rank(operand_literal.shape()));
+
+    HloEvaluator embedded_evaluator(parent_->max_loop_iterations_);
+    // For each resulting dimension, calculate and assign computed value.
+    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(
+        [&](tensorflow::gtl::ArraySlice<int64> output_index) {
+          ReturnT result_val = init_scalar;
+
+          std::fill(window_index.begin(), window_index.end(), 0);
+          std::fill(operand_index.begin(), operand_index.end(), 0);
+
+          IterateThroughWindow(
+              window_shape, window, operand_literal.shape(), output_index,
+              [&](const std::vector<int64>& operand_index) {
+                auto curr_val = operand_literal.Get<ReturnT>(operand_index);
+
+                // Evaluate computation with specified literal operands.
+                const auto curr_val_literal =
+                    Literal::CreateR0<ReturnT>(curr_val);
+                const auto result_val_literal =
+                    Literal::CreateR0<ReturnT>(result_val);
+                const std::vector<const Literal*> args = {
+                    result_val_literal.get(), curr_val_literal.get()};
+                std::unique_ptr<Literal> computed_result =
+                    embedded_evaluator.Evaluate<const Literal*>(*function, args)
+                        .ConsumeValueOrDie();
+
+                // Clear visit states so that the we can use the evaluate again
+                // on the same computation.
+                embedded_evaluator.ResetVisitStates();
+
+                result_val = computed_result->Get<ReturnT>({});
+              });
+
+          return result_val;
+        }));
+
+    parent_->evaluated_[reduce_window] = std::move(result);
+    return Status::OK();
+  }
+
+  Status HandleSlice(HloInstruction* slice) override {
+    auto operand = slice->operand(0);
+    const Shape& shape = slice->shape();
+    TF_ASSIGN_OR_RETURN(auto inferred_return_shape,
+                        ShapeInference::InferSliceShape(
+                            operand->shape(), slice->slice_starts(),
+                            slice->slice_limits(), slice->slice_strides()));
+    TF_RET_CHECK(ShapeUtil::Compatible(shape, inferred_return_shape))
+        << "return shape set to: " << ShapeUtil::HumanString(shape)
+        << " but is inferred to be: "
+        << ShapeUtil::HumanString(inferred_return_shape);
+
+    const int64 rank = ShapeUtil::Rank(operand->shape());
+    const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand);
+    auto func = [&](tensorflow::gtl::ArraySlice<int64> out_index) {
+      DimensionVector operand_index(rank);
+      for (int64 i = 0; i < rank; ++i) {
+        operand_index[i] =
+            slice->slice_starts(i) + out_index[i] * slice->slice_strides(i);
+      }
+      return operand_literal.Get<ReturnT>(operand_index);
+    };
+
+    auto result = Literal::CreateFromDimensions(
+        shape.element_type(), AsInt64Slice(shape.dimensions()));
+    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(func));
+    parent_->evaluated_[slice] = std::move(result);
+    return Status::OK();
+  }
+
+  // Enable CLZ only for int32 and uint32.
+  template <
+      typename NativeT,
+      typename std::enable_if<
+          (std::is_floating_point<NativeT>::value ||
+           std::is_integral<NativeT>::value || is_complex_t<NativeT>::value) &&
+          !(std::is_same<NativeT, uint32>::value ||
+            std::is_same<NativeT, int32>::value)>::type* = nullptr>
+  Status HandleClz(HloInstruction* clz) {
+    return InvalidArgument("Unsupported type for Clz");
+  }
+
+  template <typename NativeT,
+            typename std::enable_if<
+                std::is_same<NativeT, uint32>::value ||
+                std::is_same<NativeT, int32>::value>::type* = nullptr>
+  Status HandleClz(HloInstruction* clz) {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[clz],
+                        ElementWiseUnaryOp(clz, [](ElementwiseT elem_operand) {
+                          return 31 - tensorflow::Log2Floor(elem_operand);
+                        }));
+    return Status::OK();
+  }
+
+  Status HandleClz(HloInstruction* clz) override {
+    return HandleClz<ElementwiseT>(clz);
+  }
+
+  template <typename NativeT, typename std::enable_if<std::is_floating_point<
+                                  NativeT>::value>::type* = nullptr>
+  Status HandleSin(HloInstruction* sin) {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[sin],
+                        ElementWiseUnaryOp(sin, [](ElementwiseT elem_operand) {
+                          return std::sin(elem_operand);
+                        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<std::is_integral<NativeT>::value ||
+                              is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleSin(HloInstruction* sin) {
+    return InvalidArgument("Unsupported type for Sin");
+  }
+
+  Status HandleSin(HloInstruction* sin) override {
+    return HandleSin<ElementwiseT>(sin);
+  }
+
+  template <typename NativeT, typename std::enable_if<std::is_floating_point<
+                                  NativeT>::value>::type* = nullptr>
+  Status HandleCos(HloInstruction* cos) {
+    TF_ASSIGN_OR_RETURN(parent_->evaluated_[cos],
+                        ElementWiseUnaryOp(cos, [](ElementwiseT elem_operand) {
+                          return std::cos(elem_operand);
+                        }));
+    return Status::OK();
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<std::is_integral<NativeT>::value ||
+                              is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleCos(HloInstruction* cos) {
+    return InvalidArgument("Unsupported type for Cos");
+  }
+
+  Status HandleCos(HloInstruction* cos) override {
+    return HandleCos<ElementwiseT>(cos);
+  }
+
+  template <typename NativeT, typename std::enable_if<std::is_same<
+                                  float, NativeT>::value>::type* = nullptr>
+  Status HandleReducePrecision(HloInstruction* reduce_precision) {
+    TF_ASSIGN_OR_RETURN(
+        parent_->evaluated_[reduce_precision],
+        ElementWiseUnaryOp(reduce_precision, [reduce_precision](
+                                                 ElementwiseT elem) {
+          uint32_t value_as_int = tensorflow::bit_cast<uint32_t>(elem);
+          const uint32_t mantissa_bits = reduce_precision->mantissa_bits();
+          const uint32_t exponent_bits = reduce_precision->exponent_bits();
+
+          // Code is based on the CPU/GPU implementation in LLVM-emitting code.
+          //
+          // Bits in float type:
+          //   mantissa : bits [0:22]
+          //   exponent : bits [23:30]
+          //   sign     : bits [31]
+          if (mantissa_bits < 23) {
+            const uint32_t last_mantissa_bit_mask = 1u << (23 - mantissa_bits);
+
+            // Compute rounding bias for round-to-nearest with ties to even.
+            // This is equal to a base value of 0111... plus one bit if the last
+            // remaining mantissa bit is 1.
+            const uint32_t base_rounding_bias =
+                (last_mantissa_bit_mask >> 1) - 1;
+            const uint32_t x_last_mantissa_bit =
+                (value_as_int & last_mantissa_bit_mask) >> (23 - mantissa_bits);
+            const uint32_t x_rounding_bias =
+                x_last_mantissa_bit + base_rounding_bias;
+
+            // Add rounding bias, and mask out truncated bits.  Note that the
+            // case where adding the rounding bias overflows into the exponent
+            // bits is correct; the non-masked mantissa bits will all be zero,
+            // and the exponent will be incremented by one.
+            const uint32_t truncation_mask = ~(last_mantissa_bit_mask - 1);
+            value_as_int = value_as_int + x_rounding_bias;
+            value_as_int = value_as_int & truncation_mask;
+          }
+          if (exponent_bits < 8) {
+            // Masks for f32 values.
+            const uint32_t f32_sign_bit_mask = 1u << 31;
+            const uint32_t f32_exp_bits_mask = 0xffu << 23;
+
+            // An exponent of 2^(n-1)-1 -- that is, 0111... with the zero in the
+            // most- significant bit -- is equal to 1.0f for all exponent sizes.
+            // Adding 2^(n-1)-1 to this gives us the highest non-infinite
+            // exponent for a bit- size of n, and subtracting 2^(n-1)-1 from
+            // this gives us the lowest' exponent (corresponding to 0.0f).
+            //
+            // Thus, the f32 exponent corresponding to the highest non-infinite
+            // exponent for a bit size of n is (2^7-1) + 2^(n-1)-1, and the f32
+            // exponent corresponding to the lowest exponent for a bit size of n
+            // is (2^7-1) - 2^(n-1)-1.
+            //
+            // Note that we have already checked that exponents_bits >= 1.
+            const uint32_t f32_exponent_bias = (1 << 7) - 1;
+            const uint32_t reduced_exponent_bias =
+                (1 << (exponent_bits - 1)) - 1;
+            const uint32_t reduced_max_exponent =
+                f32_exponent_bias + reduced_exponent_bias;
+            const uint32_t reduced_min_exponent =
+                f32_exponent_bias - reduced_exponent_bias;
+
+            // Do we overflow or underflow?
+            const uint32_t x_exponent = value_as_int & f32_exp_bits_mask;
+            const bool x_overflows = x_exponent > (reduced_max_exponent << 23);
+            const bool x_underflows =
+                x_exponent <= (reduced_min_exponent << 23);
+
+            // Compute appropriately-signed values of zero and infinity.
+            const uint32_t x_signed_zero = value_as_int & f32_sign_bit_mask;
+            const uint32_t x_signed_inf = x_signed_zero | f32_exp_bits_mask;
+
+            // Force to zero or infinity if overflow or underflow.  (Note that
+            // this truncates all denormal values to zero, rather than rounding
+            // them.)
+            value_as_int = x_overflows ? x_signed_inf : value_as_int;
+            value_as_int = x_underflows ? x_signed_zero : value_as_int;
+          }
+
+          float reduced_result = tensorflow::bit_cast<float>(value_as_int);
+          if (std::isnan(elem)) {
+            reduced_result = mantissa_bits > 0
+                                 ? elem
+                                 : std::numeric_limits<float>::infinity();
+          }
+          return reduced_result;
+        }));
+    return Status::OK();
+  }
+
+  template <typename NativeT, typename std::enable_if<std::is_same<
+                                  double, NativeT>::value>::type* = nullptr>
+  Status HandleReducePrecision(HloInstruction* reduce_precision) {
+    return InvalidArgument("Double not supported for reduce precision");
+  }
+
+  template <
+      typename NativeT,
+      typename std::enable_if<std::is_integral<NativeT>::value ||
+                              is_complex_t<NativeT>::value>::type* = nullptr>
+  Status HandleReducePrecision(HloInstruction* reduce_precision) {
+    return InvalidArgument("Unsupported type for reduce precision");
+  }
+
+  Status HandleReducePrecision(HloInstruction* reduce_precision) override {
+    return HandleReducePrecision<ElementwiseT>(reduce_precision);
+  }
+
+ private:
+  // Creates a vector of multipliers which can be used to create a linear index
+  // into shape.
+  //
+  // Given the multidimensional index {i1, ..., iN} and
+  // M = MakeDimMultipliers(shape), the corresponding linear index LI is simply
+  //
+  //   LI = i1 * M[1] + i2 * M[2] + ... + iN * M[N].
+  //
+  // This lets you calculate LI given the multidimensional indices in any order.
+  static DimensionVector MakeDimMultipliers(const Shape& shape) {
+    DimensionVector v(ShapeUtil::Rank(shape));
+    int64 scale = 1;
+    for (auto dim : LayoutUtil::MinorToMajor(shape)) {
+      v[dim] = scale;
+      scale *= shape.dimensions(dim);
+    }
+    return v;
+  }
+
+  // For one particular placement of a window in a base shape (the placement is
+  // represented as `window_count_index`), iterates inside the window.
+  // Translates the window index into base index. If the base index is within
+  // bound, call `f` with the base index.
+  static void IterateThroughWindow(
+      const Shape& window_shape, const Window& window, const Shape& base_shape,
+      const tensorflow::gtl::ArraySlice<int64>& window_count_index,
+      const std::function<void(const std::vector<int64>&)>& f) {
+    const int64 rank = ShapeUtil::Rank(base_shape);
+    DimensionVector window_index(rank);
+    std::fill(window_index.begin(), window_index.end(), 0);
+    do {
+      std::vector<int64> base_index(rank);
+      bool out_of_bound = false;
+      for (int64 i = 0; i < rank; ++i) {
+        base_index[i] = window_count_index[i] * window.dimensions(i).stride() +
+                        window_index[i] - window.dimensions(i).padding_low();
+        if (base_index[i] < 0 || base_index[i] >= base_shape.dimensions(i)) {
+          out_of_bound = true;
+          break;
+        }
+      }
+      if (!out_of_bound) {
+        f(base_index);
+      }
+    } while (IndexUtil::BumpIndices(window_shape, &window_index));
+  }
+
+  template <typename IndexT>
+  StatusOr<std::unique_ptr<Literal>> DynamicSlice(
+      const Literal& operand_literal, const Literal& start_indices_literal,
+      const Shape& result_shape) {
+    auto start_indices_typed = start_indices_literal.data<IndexT>();
+    std::vector<int64> start(start_indices_typed.begin(),
+                             start_indices_typed.end());
+
+    std::vector<int64> operand_indices(start.size());
+
+    auto result = Literal::CreateFromShape(result_shape);
+    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(
+        [&](tensorflow::gtl::ArraySlice<int64> multi_index) {
+          for (int64 i = 0; i < operand_indices.size(); ++i) {
+            CHECK_GE(multi_index[i] + start[i], 0);
+            // Mod is only used here to be consistent with the existing
+            // backends' behavior.
+            operand_indices[i] = (multi_index[i] + start[i]) %
+                                 operand_literal.shape().dimensions(i);
+          }
+
+          auto result = operand_literal.Get<ReturnT>(operand_indices);
+          return result;
+        }));
+
+    return std::move(result);
+  }
+
+  template <typename IndexT>
+  StatusOr<std::unique_ptr<Literal>> DynamicUpdateSlice(
+      const Literal& operand_literal, const Literal& update_literal,
+      const Literal& start_indices_literal) {
+    auto result = operand_literal.CloneToUnique();
+    auto start_indices_typed = start_indices_literal.data<IndexT>();
+    const auto rank = ShapeUtil::Rank(result->shape());
+    std::vector<int64> start(rank, 0);
+    for (int64 i = 0; i < rank; ++i) {
+      // All other implementations currently wrap-around the index, so this
+      // should do so as well.
+      start[i] = (start_indices_typed[i] % result->shape().dimensions(i));
+      start[i] += (start[i] < 0) * result->shape().dimensions(i);
+    }
+    std::vector<int64> result_index(rank, 0);
+
+    auto func = [&](tensorflow::gtl::ArraySlice<int64> update_index) {
+      std::transform(update_index.begin(), update_index.end(), start.begin(),
+                     result_index.begin(), std::plus<int64>());
+      // Same as above, wrap-around only to match other implementations'
+      // semantics.
+      std::transform(result_index.begin(), result_index.end(),
+                     result->shape().dimensions().begin(), result_index.begin(),
+                     std::modulus<int64>());
+      result->Set<ReturnT>(result_index,
+                           update_literal.Get<ReturnT>(update_index));
+      return true;
+    };
+
+    std::vector<int64> base(update_literal.shape().dimensions_size(), 0);
+    std::vector<int64> step(update_literal.shape().dimensions_size(), 1);
+    ShapeUtil::ForEachIndex(update_literal.shape(), base,
+                            AsInt64Slice(update_literal.shape().dimensions()),
+                            step, func);
+
+    return std::move(result);
+  }
+
+  StatusOr<std::unique_ptr<Literal>> ElementWiseUnaryOp(
+      HloInstruction* instruction,
+      const std::function<ElementwiseT(ElementwiseT)>& unary_op) {
+    const Literal& operand_literal =
+        parent_->GetEvaluatedLiteralFor(instruction->operand(0));
+    TF_ASSIGN_OR_RETURN(
+        auto result_literal,
+        (HloEvaluator::ElementWiseUnaryOpImpl<ReturnT, ReturnT>(
+            instruction, ConvertUnaryFunction(unary_op), operand_literal)));
+
+    return std::move(result_literal);
+  }
+
+  StatusOr<std::unique_ptr<Literal>> ElementWiseBinaryOp(
+      HloInstruction* instruction,
+      const std::function<ElementwiseT(ElementwiseT, ElementwiseT)>&
+          binary_op) {
+    const auto shape = instruction->shape();
+    const auto* lhs = instruction->operand(0);
+    const auto* rhs = instruction->operand(1);
+
+    // TODO(b/35950897, b/27796129): add DCHECK back once implicit broadcast
+    // is removed.
+    if (!(ShapeUtil::SameDimensions(shape, rhs->shape()) &&
+          ShapeUtil::SameDimensions(lhs->shape(), rhs->shape()))) {
+      return Unimplemented(
+          "Implicit broadcasting is currently unsupported in HLO evaluator "
+          "Shape Mismatch: %s vs %s vs %s: ",
+          ShapeUtil::HumanString(shape).c_str(),
+          ShapeUtil::HumanString(lhs->shape()).c_str(),
+          ShapeUtil::HumanString(rhs->shape()).c_str());
+    }
+
+    const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs);
+    const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs);
+
+    auto result = Literal::CreateFromShape(shape);
+
+    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(
+        [&](tensorflow::gtl::ArraySlice<int64> multi_index) {
+          return ConvertBinaryFunction(binary_op)(
+              lhs_literal.Get<ReturnT>(multi_index),
+              rhs_literal.Get<ReturnT>(multi_index));
+        }));
+    return std::move(result);
+  }
+
+  template <typename LhsType, typename RhsType, typename EhsType>
+  StatusOr<std::unique_ptr<Literal>> ElementwiseTernaryOp(
+      HloInstruction* instruction,
+      const std::function<ReturnT(LhsType, RhsType, EhsType)>& ternary_op) {
+    const auto shape = instruction->shape();
+    const auto* lhs = instruction->operand(0);
+    const auto* rhs = instruction->operand(1);
+    const auto* ehs = instruction->operand(2);
+
+    // TODO(b/35950897, b/27796129): add DCHECK back once implicit
+    // broadcast is removed.
+    if (!(ShapeUtil::SameDimensions(shape, lhs->shape()) &&
+          ShapeUtil::SameDimensions(lhs->shape(), rhs->shape()) &&
+          ShapeUtil::SameDimensions(rhs->shape(), ehs->shape()))) {
+      return Unimplemented(
+          "Implicit broadcasting is currently unsupported in HLO evaluator "
+          "Shape Mismatch: %s vs %s vs %s vs %s: ",
+          ShapeUtil::HumanString(shape).c_str(),
+          ShapeUtil::HumanString(lhs->shape()).c_str(),
+          ShapeUtil::HumanString(rhs->shape()).c_str(),
+          ShapeUtil::HumanString(ehs->shape()).c_str());
+    }
+
+    const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs);
+    const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs);
+    const Literal& ehs_literal = parent_->GetEvaluatedLiteralFor(ehs);
+
+    auto result = Literal::CreateFromShape(shape);
+
+    TF_RETURN_IF_ERROR(result->Populate<ReturnT>(
+        [&](tensorflow::gtl::ArraySlice<int64> multi_index) {
+          return ternary_op(lhs_literal.Get<LhsType>(multi_index),
+                            rhs_literal.Get<RhsType>(multi_index),
+                            ehs_literal.Get<EhsType>(multi_index));
+        }));
+
+    return std::move(result);
+  }
+
+  template <typename NativeT>
+  static bool IsShiftOutOfBounds(NativeT rhs) {
+    typedef typename std::make_unsigned<NativeT>::type UnsignedT;
+    UnsignedT lhs_size_unsigned = sizeof(NativeT) * CHAR_BIT;
+    UnsignedT rhs_unsigned = static_cast<UnsignedT>(rhs);
+    return rhs_unsigned >= lhs_size_unsigned;
+  }
+
+  HloEvaluator* parent_;
+};
+
+// These extern templates prevent users of this class from implicitly
+// instantiating it.  We explicitly instantiate this class in the various
+// hlo_evaluator_typed_visitor*.cc files.
+extern template class HloEvaluatorTypedVisitor<bool>;
+extern template class HloEvaluatorTypedVisitor<uint8>;
+extern template class HloEvaluatorTypedVisitor<uint32>;
+extern template class HloEvaluatorTypedVisitor<uint64>;
+extern template class HloEvaluatorTypedVisitor<int8>;
+extern template class HloEvaluatorTypedVisitor<int32>;
+extern template class HloEvaluatorTypedVisitor<int64>;
+extern template class HloEvaluatorTypedVisitor<Eigen::half, float>;
+extern template class HloEvaluatorTypedVisitor<float>;
+extern template class HloEvaluatorTypedVisitor<double>;
+extern template class HloEvaluatorTypedVisitor<complex64>;
+extern template class HloEvaluatorTypedVisitor<bfloat16, float>;
+
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_TYPED_VISITOR_H_
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_bfloat16.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_bfloat16.cc
new file mode 100644
index 0000000000..39c352dfb9
--- /dev/null
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_bfloat16.cc
@@ -0,0 +1,22 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
+
+namespace xla {
+template class HloEvaluatorTypedVisitor<bfloat16, float>;
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_bool.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_bool.cc
new file mode 100644
index 0000000000..289b40fa06
--- /dev/null
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_bool.cc
@@ -0,0 +1,22 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
+
+namespace xla {
+template class HloEvaluatorTypedVisitor<bool>;
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_complex64.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_complex64.cc
new file mode 100644
index 0000000000..9cb4eb921f
--- /dev/null
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_complex64.cc
@@ -0,0 +1,22 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
+
+namespace xla {
+template class HloEvaluatorTypedVisitor<complex64>;
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_double.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_double.cc
new file mode 100644
index 0000000000..5e6252fbf8
--- /dev/null
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_double.cc
@@ -0,0 +1,22 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
+
+namespace xla {
+template class HloEvaluatorTypedVisitor<double>;
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_float.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_float.cc
new file mode 100644
index 0000000000..ee793ae77b
--- /dev/null
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_float.cc
@@ -0,0 +1,22 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
+
+namespace xla {
+template class HloEvaluatorTypedVisitor<float>;
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_half.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_half.cc
new file mode 100644
index 0000000000..038d9d39e4
--- /dev/null
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_half.cc
@@ -0,0 +1,22 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
+
+namespace xla {
+template class HloEvaluatorTypedVisitor<Eigen::half, float>;
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int32.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int32.cc
new file mode 100644
index 0000000000..b1952ca619
--- /dev/null
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int32.cc
@@ -0,0 +1,22 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
+
+namespace xla {
+template class HloEvaluatorTypedVisitor<int32>;
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int64.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int64.cc
new file mode 100644
index 0000000000..0cbaffb40b
--- /dev/null
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int64.cc
@@ -0,0 +1,22 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
+
+namespace xla {
+template class HloEvaluatorTypedVisitor<int64>;
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int8.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int8.cc
new file mode 100644
index 0000000000..6f4bf2a392
--- /dev/null
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_int8.cc
@@ -0,0 +1,22 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
+
+namespace xla {
+template class HloEvaluatorTypedVisitor<int8>;
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint32.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint32.cc
new file mode 100644
index 0000000000..10235447e0
--- /dev/null
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint32.cc
@@ -0,0 +1,22 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
+
+namespace xla {
+template class HloEvaluatorTypedVisitor<uint32>;
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint64.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint64.cc
new file mode 100644
index 0000000000..8abeaa6ffc
--- /dev/null
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint64.cc
@@ -0,0 +1,22 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
+
+namespace xla {
+template class HloEvaluatorTypedVisitor<uint64>;
+}  // namespace xla
diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint8.cc b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint8.cc
new file mode 100644
index 0000000000..6dabd1c176
--- /dev/null
+++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor_uint8.cc
@@ -0,0 +1,22 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h"
+
+#include "tensorflow/compiler/xla/service/hlo_evaluator.h"
+
+namespace xla {
+template class HloEvaluatorTypedVisitor<uint8>;
+}  // namespace xla
-- 
GitLab


From dfae6ff29e95345c7c6c0ef50fd5f45bd458cfdc Mon Sep 17 00:00:00 2001
From: Alexandre Passos <apassos@google.com>
Date: Mon, 7 May 2018 14:29:03 -0700
Subject: [PATCH 324/395] Fix resource variable in cond gradient.

PiperOrigin-RevId: 195722449
---
 tensorflow/python/BUILD                   |  1 +
 tensorflow/python/ops/control_flow_ops.py |  7 +++++++
 tensorflow/python/ops/gradients_impl.py   |  3 ++-
 tensorflow/python/ops/gradients_test.py   | 15 +++++++++++++++
 4 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 087b89b125..4057e37681 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -1762,6 +1762,7 @@ py_library(
         ":logging_ops_gen",
         ":math_ops",
         ":platform",
+        ":resource_variable_ops_gen",
         ":sparse_tensor",
         ":tensor_array_ops",
         ":tf_should_use",
diff --git a/tensorflow/python/ops/control_flow_ops.py b/tensorflow/python/ops/control_flow_ops.py
index 07d4ff7b02..5f60dab6ac 100644
--- a/tensorflow/python/ops/control_flow_ops.py
+++ b/tensorflow/python/ops/control_flow_ops.py
@@ -43,6 +43,7 @@ from tensorflow.python.ops import gen_array_ops
 from tensorflow.python.ops import gen_control_flow_ops
 from tensorflow.python.ops import gen_data_flow_ops
 from tensorflow.python.ops import gen_logging_ops
+from tensorflow.python.ops import gen_resource_variable_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import tensor_array_ops
 # go/tf-wildcard-import
@@ -1433,6 +1434,8 @@ def ZerosLikeOutsideLoop(op, index):
   """Create zeros_like for the specified output of an op."""
   val = op.outputs[index]
   if not util.IsSwitch(op):
+    if val.dtype == dtypes.resource:
+      return array_ops.zeros(gen_resource_variable_ops.variable_shape(val))
     return array_ops.zeros_like(val, optimize=False)
   else:
     op_ctxt = op._get_control_flow_context()
@@ -1441,6 +1444,10 @@ def ZerosLikeOutsideLoop(op, index):
       pred = op_ctxt.pred
       branch = op_ctxt.branch
       switch_val = switch(op.inputs[0], pred)[1 - branch]
+      if val.dtype == dtypes.resource:
+        with ops.control_dependencies([switch_val]):
+          return array_ops.zeros(
+              gen_resource_variable_ops.variable_shape(switch_val))
       zeros_shape = array_ops.shape_internal(switch_val, optimize=False)
       # Ensure ops created within array_ops.zeros are dominated by switch in
       # cond context.
diff --git a/tensorflow/python/ops/gradients_impl.py b/tensorflow/python/ops/gradients_impl.py
index 1448151fef..a6b1e6df54 100644
--- a/tensorflow/python/ops/gradients_impl.py
+++ b/tensorflow/python/ops/gradients_impl.py
@@ -297,7 +297,8 @@ def _DefaultGradYs(grad_ys,
 def _IsTrainable(tensor):
   dtype = dtypes.as_dtype(tensor.dtype)
   return dtype.base_dtype in (dtypes.float16, dtypes.float32, dtypes.float64,
-                              dtypes.complex64, dtypes.complex128)
+                              dtypes.complex64, dtypes.complex128,
+                              dtypes.resource)
 
 
 def _IsBackpropagatable(tensor):
diff --git a/tensorflow/python/ops/gradients_test.py b/tensorflow/python/ops/gradients_test.py
index 5e8b8822ef..e729950201 100644
--- a/tensorflow/python/ops/gradients_test.py
+++ b/tensorflow/python/ops/gradients_test.py
@@ -944,6 +944,21 @@ class CustomGradientTest(test_util.TensorFlowTestCase):
       # Smoke test to ensure numpy inputs are accepted
       F(x)
 
+  def testRVGradientsDynamicCond(self):
+    with self.test_session():
+      alpha = resource_variable_ops.ResourceVariable(
+          np.random.random((1,)),
+          dtype="float32")
+
+      conditional = array_ops.placeholder_with_default(True, shape=())
+      output = control_flow_ops.cond(
+          conditional, lambda: alpha * 2, lambda: alpha * 3)
+
+      g, = gradients_impl.gradients(output, alpha)
+      variables.global_variables_initializer().run()
+      self.assertAllEqual(g.eval(), [2.0])
+      self.assertAllEqual(g.eval(feed_dict={conditional: False}), [3.0])
+
 
 if __name__ == "__main__":
   googletest.main()
-- 
GitLab


From 84986720ee8c07112356bbcc0911629e9c4dafb6 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 14:34:11 -0700
Subject: [PATCH 325/395] Allow output has a different shape from input in the
 image.transform (#17011).

PiperOrigin-RevId: 195723288
---
 tensorflow/contrib/image/kernels/image_ops.cc | 33 ++++++++---
 tensorflow/contrib/image/kernels/image_ops.h  |  2 +-
 tensorflow/contrib/image/ops/image_ops.cc     | 55 +++++++++++++++++--
 .../python/kernel_tests/image_ops_test.py     | 30 ++++++++++
 .../contrib/image/python/ops/image_ops.py     | 49 +++++++++++------
 5 files changed, 139 insertions(+), 30 deletions(-)

diff --git a/tensorflow/contrib/image/kernels/image_ops.cc b/tensorflow/contrib/image/kernels/image_ops.cc
index c2e32da133..575c2004fb 100644
--- a/tensorflow/contrib/image/kernels/image_ops.cc
+++ b/tensorflow/contrib/image/kernels/image_ops.cc
@@ -70,6 +70,7 @@ class ImageProjectiveTransform : public OpKernel {
   void Compute(OpKernelContext* ctx) override {
     const Tensor& images_t = ctx->input(0);
     const Tensor& transform_t = ctx->input(1);
+    const Tensor& shape_t = ctx->input(2);
     OP_REQUIRES(ctx, images_t.shape().dims() == 4,
                 errors::InvalidArgument("Input images must have rank 4"));
     OP_REQUIRES(ctx,
@@ -80,11 +81,28 @@ class ImageProjectiveTransform : public OpKernel {
                      ProjectiveGenerator<Device, T>::kNumParameters),
                 errors::InvalidArgument(
                     "Input transform should be num_images x 8 or 1 x 8"));
-    auto images = images_t.tensor<T, 4>();
-    auto transform = transform_t.matrix<float>();
+    OP_REQUIRES(ctx, shape_t.dims() == 1,
+                errors::InvalidArgument("output shape must be 1-dimensional",
+                                        shape_t.shape().DebugString()));
+    OP_REQUIRES(ctx, shape_t.NumElements() == 2,
+                errors::InvalidArgument("output shape must have two elements",
+                                        shape_t.shape().DebugString()));
+    auto Svec = shape_t.vec<int32>();
+    int32 out_height = Svec(0);
+    int32 out_width = Svec(1);
+    OP_REQUIRES(ctx, out_height > 0 && out_width > 0,
+                errors::InvalidArgument("output dimensions must be positive"));
+
     Tensor* output_t;
-    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, images_t.shape(), &output_t));
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(
+                            0,
+                            TensorShape({images_t.dim_size(0), out_height,
+                                         out_width, images_t.dim_size(3)}),
+                            &output_t));
     auto output = output_t->tensor<T, 4>();
+    auto images = images_t.tensor<T, 4>();
+    auto transform = transform_t.matrix<float>();
+
     (FillProjectiveTransform<Device, T>(interpolation_))(
         ctx->eigen_device<Device>(), &output, images, transform);
   }
@@ -127,10 +145,11 @@ TF_CALL_double(DECLARE_FUNCTOR);
 
 }  // end namespace functor
 
-#define REGISTER(TYPE)                                        \
-  REGISTER_KERNEL_BUILDER(Name("ImageProjectiveTransform")    \
-                              .Device(DEVICE_GPU)             \
-                              .TypeConstraint<TYPE>("dtype"), \
+#define REGISTER(TYPE)                                       \
+  REGISTER_KERNEL_BUILDER(Name("ImageProjectiveTransform")   \
+                              .Device(DEVICE_GPU)            \
+                              .TypeConstraint<TYPE>("dtype") \
+                              .HostMemory("output_shape"),   \
                           ImageProjectiveTransform<GPUDevice, TYPE>)
 
 TF_CALL_uint8(REGISTER);
diff --git a/tensorflow/contrib/image/kernels/image_ops.h b/tensorflow/contrib/image/kernels/image_ops.h
index ad50133061..2320329b92 100644
--- a/tensorflow/contrib/image/kernels/image_ops.h
+++ b/tensorflow/contrib/image/kernels/image_ops.h
@@ -161,7 +161,7 @@ struct FillProjectiveTransform {
   void operator()(const Device& device, OutputType* output,
                   const InputType& images,
                   const TransformsType& transform) const {
-    output->device(device) = images.generate(
+    output->device(device) = output->generate(
         ProjectiveGenerator<Device, T>(images, transform, interpolation_));
   }
 };
diff --git a/tensorflow/contrib/image/ops/image_ops.cc b/tensorflow/contrib/image/ops/image_ops.cc
index ebdcaea7ab..fb62507174 100644
--- a/tensorflow/contrib/image/ops/image_ops.cc
+++ b/tensorflow/contrib/image/ops/image_ops.cc
@@ -19,9 +19,56 @@ limitations under the License.
 
 namespace tensorflow {
 
+using shape_inference::DimensionHandle;
 using shape_inference::InferenceContext;
 using shape_inference::ShapeHandle;
 
+namespace {
+
+// Sets output[0] to shape [batch_dim,height,width,channel_dim], where
+// height and width come from the size_tensor.
+Status SetOutputToSizedImage(InferenceContext* c, DimensionHandle batch_dim,
+                             int size_input_idx, DimensionHandle channel_dim) {
+  // Verify shape of size input.
+  ShapeHandle size;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(size_input_idx), 1, &size));
+  DimensionHandle unused;
+  TF_RETURN_IF_ERROR(c->WithValue(c->Dim(size, 0), 2, &unused));
+
+  // Get size values from the size tensor.
+  const Tensor* size_tensor = c->input_tensor(size_input_idx);
+  DimensionHandle width;
+  DimensionHandle height;
+  if (size_tensor == nullptr) {
+    width = c->UnknownDim();
+    height = c->UnknownDim();
+  } else {
+    // TODO(petewarden) - Remove once we have constant evaluation in C++ only.
+    if (size_tensor->dtype() != DT_INT32) {
+      return errors::InvalidArgument(
+          "Bad size input type for SetOutputToSizedImage: Expected DT_INT32 "
+          "but got ",
+          DataTypeString(size_tensor->dtype()), " for input #", size_input_idx,
+          " in ", c->DebugString());
+    }
+    auto vec = size_tensor->vec<int32>();
+    height = c->MakeDim(vec(0));
+    width = c->MakeDim(vec(1));
+  }
+  c->set_output(0, c->MakeShape({batch_dim, height, width, channel_dim}));
+  return Status::OK();
+}
+
+// TODO(qyu): Move this to core/framework/common_shape_fns.h
+Status ResizeShapeFn(InferenceContext* c) {
+  ShapeHandle input;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 4, &input));
+  return SetOutputToSizedImage(c, c->Dim(input, 0), 2 /* size_input_idx */,
+                               c->Dim(input, 3));
+}
+
+}  // namespace
+
 // TODO(ringwalt): Add a "fill_mode" argument with "constant", "mirror", etc.
 // TODO(ringwalt): Add a "fill_constant" argument for constant mode (default 0).
 // TODO(ringwalt): Add an "output_shape" argument. This is sufficient to
@@ -29,13 +76,11 @@ using shape_inference::ShapeHandle;
 REGISTER_OP("ImageProjectiveTransform")
     .Input("images: dtype")
     .Input("transforms: float32")
+    .Input("output_shape: int32")
     .Attr("dtype: {uint8, int32, int64, float32, float64}")
     .Attr("interpolation: string")
     .Output("transformed_images: dtype")
-    .SetShapeFn([](InferenceContext* c) {
-      c->set_output(0, c->input(0));
-      return Status::OK();
-    })
+    .SetShapeFn(ResizeShapeFn)
     .Doc(R"doc(
 Applies the given transform to each of the images.
 
@@ -49,7 +94,7 @@ If one row of `transforms` is `[a0, a1, a2, b0, b1, b2, c0, c1]`, then it maps
 the *output* point `(x, y)` to a transformed *input* point
 `(x', y') = ((a0 x + a1 y + a2) / k, (b0 x + b1 y + b2) / k)`, where
 `k = c0 x + c1 y + 1`. If the transformed point lays outside of the input
-image, the output pixel is set to 0. The output is the same size as the input,
+image, the output pixel is set to 0.
 
 images: 4D `Tensor`, input image(s) in NHWC format.
 transforms: 2D `Tensor`, projective transform(s) to apply to the image(s).
diff --git a/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py b/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py
index b50177ae56..c0151d320f 100644
--- a/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py
+++ b/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py
@@ -195,10 +195,40 @@ class ImageOpsTest(test_util.TensorFlowTestCase):
           x_init_value=test_image)
       self.assertLess(left_err, 1e-10)
 
+  def _test_grad_different_shape(self, input_shape, output_shape):
+    with self.test_session():
+      test_image_shape = input_shape
+      test_image = np.random.randn(*test_image_shape)
+      test_image_tensor = constant_op.constant(
+          test_image, shape=test_image_shape)
+      test_transform = image_ops.angles_to_projective_transforms(
+          np.pi / 2, 4, 4)
+
+      if len(output_shape) == 2:
+        resize_shape = output_shape
+      elif len(output_shape) == 3:
+        resize_shape = output_shape[0:2]
+      elif len(output_shape) == 4:
+        resize_shape = output_shape[1:3]
+      output = image_ops.transform(
+          images=test_image_tensor,
+          transforms=test_transform,
+          output_shape=resize_shape)
+      left_err = gradient_checker.compute_gradient_error(
+          test_image_tensor,
+          test_image_shape,
+          output,
+          output_shape,
+          x_init_value=test_image)
+      self.assertLess(left_err, 1e-10)
+
   def test_grad(self):
     self._test_grad([16, 16])
     self._test_grad([4, 12, 12])
     self._test_grad([3, 4, 12, 12])
+    self._test_grad_different_shape([16, 16], [8, 8])
+    self._test_grad_different_shape([4, 12, 3], [8, 24, 3])
+    self._test_grad_different_shape([3, 4, 12, 3], [3, 8, 24, 3])
 
 
 class BipartiteMatchTest(test_util.TensorFlowTestCase):
diff --git a/tensorflow/contrib/image/python/ops/image_ops.py b/tensorflow/contrib/image/python/ops/image_ops.py
index cd984c8054..192571ced8 100644
--- a/tensorflow/contrib/image/python/ops/image_ops.py
+++ b/tensorflow/contrib/image/python/ops/image_ops.py
@@ -23,6 +23,7 @@ from tensorflow.python.framework import common_shapes
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import linalg_ops
@@ -212,7 +213,11 @@ def translations_to_projective_transforms(translations, name=None):
         axis=1)
 
 
-def transform(images, transforms, interpolation="NEAREST", name=None):
+def transform(images,
+              transforms,
+              interpolation="NEAREST",
+              output_shape=None,
+              name=None):
   """Applies the given transform(s) to the image(s).
 
   Args:
@@ -229,6 +234,10 @@ def transform(images, transforms, interpolation="NEAREST", name=None):
        the transform mapping input points to output points. Note that gradients
        are not backpropagated into transformation parameters.
     interpolation: Interpolation mode. Supported values: "NEAREST", "BILINEAR".
+    output_shape: Output dimesion after the transform, [height, width].
+       If None, output is the same size as input image.
+
+    name: The name of the op.
 
   Returns:
     Image(s) with the same type and shape as `images`, with the given
@@ -237,6 +246,7 @@ def transform(images, transforms, interpolation="NEAREST", name=None):
 
   Raises:
     TypeError: If `image` is an invalid type.
+    ValueError: If output shape is not 1-D int32 Tensor.
   """
   with ops.name_scope(name, "transform"):
     image_or_images = ops.convert_to_tensor(images, name="images")
@@ -255,6 +265,17 @@ def transform(images, transforms, interpolation="NEAREST", name=None):
     else:
       raise TypeError("Images should have rank between 2 and 4.")
 
+    if output_shape is None:
+      output_shape = tensor_util.constant_value(
+          array_ops.shape(images)[1:3]) or array_ops.shape(images)[1:3]
+
+    output_shape = ops.convert_to_tensor(
+        output_shape, dtypes.int32, name="output_shape")
+
+    if not output_shape.get_shape().is_compatible_with([2]):
+      raise ValueError("output_shape must be a 1-D Tensor of 2 elements: "
+                       "new_height, new_width")
+
     if len(transform_or_transforms.get_shape()) == 1:
       transforms = transform_or_transforms[None]
     elif transform_or_transforms.get_shape().ndims is None:
@@ -264,8 +285,12 @@ def transform(images, transforms, interpolation="NEAREST", name=None):
       transforms = transform_or_transforms
     else:
       raise TypeError("Transforms should have rank 1 or 2.")
+
     output = gen_image_ops.image_projective_transform(
-        images, transforms, interpolation=interpolation.upper())
+        images,
+        output_shape=output_shape,
+        transforms=transforms,
+        interpolation=interpolation.upper())
     if len(image_or_images.get_shape()) == 2:
       return output[0, :, :, 0]
     elif len(image_or_images.get_shape()) == 3:
@@ -375,14 +400,6 @@ def _image_projective_transform_grad(op, grad):
 
   if image_or_images.dtype.base_dtype not in _IMAGE_DTYPES:
     raise TypeError("Invalid dtype %s." % image_or_images.dtype)
-  if len(image_or_images.get_shape()) == 2:
-    images = image_or_images[None, :, :, None]
-  elif len(image_or_images.get_shape()) == 3:
-    images = image_or_images[None, :, :, :]
-  elif len(image_or_images.get_shape()) == 4:
-    images = image_or_images
-  else:
-    raise TypeError("Images should have rank between 2 and 4")
   if len(transform_or_transforms.get_shape()) == 1:
     transforms = transform_or_transforms[None]
   elif len(transform_or_transforms.get_shape()) == 2:
@@ -395,13 +412,11 @@ def _image_projective_transform_grad(op, grad):
   inverse = linalg_ops.matrix_inverse(transforms)
   transforms = matrices_to_flat_transforms(inverse)
   output = gen_image_ops.image_projective_transform(
-      grad, transforms, interpolation=interpolation)
-  if len(image_or_images.get_shape()) == 2:
-    return [output[0, :, :, 0], None]
-  elif len(image_or_images.get_shape()) == 3:
-    return [output[0, :, :, :], None]
-  else:
-    return [output, None]
+      images=grad,
+      transforms=transforms,
+      output_shape=array_ops.shape(image_or_images)[1:3],
+      interpolation=interpolation)
+  return [output, None, None]
 
 
 def bipartite_match(distance_mat,
-- 
GitLab


From 80c7ebc32324d689ae3feb17a0cb4cf32d736f19 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 14:55:26 -0700
Subject: [PATCH 326/395] Disable automated testing of
 tensorflow/compiler/tests:extract_image_patches_op_test_cpu_ondemand

A recent change has made this test flaky.

PiperOrigin-RevId: 195726647
---
 tensorflow/compiler/tests/BUILD | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index a94b298f87..aaea83ae9c 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -300,6 +300,10 @@ tf_xla_py_test(
     name = "extract_image_patches_op_test",
     size = "small",
     srcs = ["extract_image_patches_op_test.py"],
+    tags = [
+        "manual",
+        "notap",
+    ],
     deps = [
         ":xla_test",
         "//tensorflow/python:array_ops",
-- 
GitLab


From 860452f3346e3a450782046081dc4d3263544344 Mon Sep 17 00:00:00 2001
From: Igor Ganichev <iga@google.com>
Date: Mon, 7 May 2018 15:15:01 -0700
Subject: [PATCH 327/395] Replace references to TensorInfo with XlaTensor

PiperOrigin-RevId: 195730139
---
 tensorflow/compiler/jit/xla_tensor.h | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/jit/xla_tensor.h b/tensorflow/compiler/jit/xla_tensor.h
index 922a918973..6b29c82ec1 100644
--- a/tensorflow/compiler/jit/xla_tensor.h
+++ b/tensorflow/compiler/jit/xla_tensor.h
@@ -54,7 +54,7 @@ class XlaTensor {
   // Some Tensors can have complex on-device shapes, including tuple shapes. To
   // manage the memory for these tensors a ShapedBuffer may be required.
 
-  // Return true if this TensorInfo contains a ShapedBuffer.
+  // Return true if this XlaTensor contains a ShapedBuffer.
   bool has_shaped_buffer() const { return shaped_buffer_ != nullptr; }
   // Return the contained ShapedBuffer.
   // REQUIRES: has_shaped_buffer()
@@ -62,7 +62,7 @@ class XlaTensor {
     CHECK(has_shaped_buffer());
     return *shaped_buffer_;
   }
-  // Mutates the TensorInfo to set the ShapedBuffer.
+  // Mutates the XlaTensor to set the ShapedBuffer.
   void set_shaped_buffer(xla::ScopedShapedBuffer shaped_buffer) {
     shaped_buffer_ =
         xla::MakeUnique<xla::ScopedShapedBuffer>(std::move(shaped_buffer));
@@ -72,7 +72,7 @@ class XlaTensor {
   // in on-demand mode to avoid re-copying values from the device if we know the
   // host value already.
 
-  // Return true if this TensorInfo contains a host tensor.
+  // Return true if this XlaTensor contains a host tensor.
   bool has_host_tensor() const { return host_tensor_ != nullptr; }
   // Return the contained host tensor.
   // REQUIRES: has_host_tensor()
-- 
GitLab


From de177c4c73a2eb3e72709a76940c1cc50341e18c Mon Sep 17 00:00:00 2001
From: Pavithra Vijay <psv@google.com>
Date: Mon, 7 May 2018 15:16:59 -0700
Subject: [PATCH 328/395] Add support for tf.data.Dataset iterators in model
 training/eval methods in eager-mode

PiperOrigin-RevId: 195730534
---
 .../keras/_impl/keras/engine/base_layer.py    |   2 +-
 .../_impl/keras/engine/sequential_test.py     |  39 +-
 .../keras/_impl/keras/engine/training.py      | 234 +++--
 .../_impl/keras/engine/training_arrays.py     |   4 +-
 .../_impl/keras/engine/training_eager.py      | 932 +++++++++++++-----
 .../keras/_impl/keras/engine/training_test.py |  96 +-
 .../_impl/keras/engine/training_utils.py      |  91 +-
 .../_impl/keras/model_subclassing_test.py     |  21 +
 8 files changed, 1049 insertions(+), 370 deletions(-)

diff --git a/tensorflow/python/keras/_impl/keras/engine/base_layer.py b/tensorflow/python/keras/_impl/keras/engine/base_layer.py
index 3af4eaabe9..16ee2952b2 100644
--- a/tensorflow/python/keras/_impl/keras/engine/base_layer.py
+++ b/tensorflow/python/keras/_impl/keras/engine/base_layer.py
@@ -1658,7 +1658,7 @@ class DeferredTensor(object):
   """Tensor-like object used to build graphs of layers in Eager mode.
 
   When calling a layer on a DeferredTensor, the layer will not perform any
-  computation and will simply perfom shape inference to return new
+  computation and will simply perform shape inference to return new
   DeferredTensors with appropriate shape information. Thus DeferredTensor
   behaves like a graph-mode Tensor when manipulated by layers.
   """
diff --git a/tensorflow/python/keras/_impl/keras/engine/sequential_test.py b/tensorflow/python/keras/_impl/keras/engine/sequential_test.py
index 8aba16aef3..a90ad131a5 100644
--- a/tensorflow/python/keras/_impl/keras/engine/sequential_test.py
+++ b/tensorflow/python/keras/_impl/keras/engine/sequential_test.py
@@ -20,8 +20,11 @@ from __future__ import print_function
 
 import numpy as np
 
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import context
 from tensorflow.python.framework import test_util as tf_test_util
 from tensorflow.python.keras._impl import keras
+from tensorflow.python.ops import array_ops
 from tensorflow.python.platform import test
 from tensorflow.python.training import rmsprop
 
@@ -75,7 +78,7 @@ class TestSequential(test.TestCase):
       model.pop()
 
   @tf_test_util.run_in_graph_and_eager_modes()
-  def test_sequential_deferred_build(self):
+  def test_sequential_deferred_build_with_np_arrays(self):
     num_hidden = 5
     input_dim = 3
     batch_size = 5
@@ -99,6 +102,40 @@ class TestSequential(test.TestCase):
                      [None, num_classes])
     self.assertEqual(len(model.weights), 2 * 2)
 
+  @tf_test_util.run_in_graph_and_eager_modes()
+  def test_sequential_deferred_build_with_dataset_iterators(self):
+    if not context.executing_eagerly():
+      # TODO(psv/fchollet): Add support for this use case in graph mode.
+      return
+    num_hidden = 5
+    input_dim = 3
+    num_classes = 2
+    num_samples = 50
+    steps_per_epoch = 10
+
+    model = keras.models.Sequential()
+    # We don't specify the input shape.
+    model.add(keras.layers.Dense(num_hidden))
+    model.add(keras.layers.Dense(num_classes))
+    model.compile(loss='mse', optimizer=rmsprop.RMSPropOptimizer(1e-3))
+    self.assertEqual(len(model.layers), 2)
+    self.assertEqual(len(model.weights), 0)
+    self.assertFalse(model.built)
+
+    x = array_ops.ones((num_samples, input_dim))
+    y = array_ops.zeros((num_samples, num_classes))
+    dataset = dataset_ops.Dataset.from_tensor_slices((x, y))
+    dataset = dataset.repeat(100)
+    dataset = dataset.batch(10)
+    iterator = dataset.make_one_shot_iterator()
+
+    model.fit(iterator, epochs=1, steps_per_epoch=steps_per_epoch)
+    self.assertTrue(model.built)
+    self.assertEqual(model.inputs[0].get_shape().as_list(), [None, input_dim])
+    self.assertEqual(model.outputs[0].get_shape().as_list(),
+                     [None, num_classes])
+    self.assertEqual(len(model.weights), 2 * 2)
+
   @tf_test_util.run_in_graph_and_eager_modes()
   def test_invalid_use_cases(self):
     # Added objects must be layer instances
diff --git a/tensorflow/python/keras/_impl/keras/engine/training.py b/tensorflow/python/keras/_impl/keras/engine/training.py
index 5f9b3e8c7d..c7623d2b52 100644
--- a/tensorflow/python/keras/_impl/keras/engine/training.py
+++ b/tensorflow/python/keras/_impl/keras/engine/training.py
@@ -18,11 +18,13 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import weakref
 import numpy as np
 
 from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.data.ops import iterator_ops
 from tensorflow.python.eager import context
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_util
 from tensorflow.python.keras._impl.keras import backend as K
@@ -106,6 +108,11 @@ class Model(Network):
   ```
   """
 
+  def __init__(self, *args, **kwargs):
+    super(Model, self).__init__(*args, **kwargs)
+    # Create a cache for iterator get_next op.
+    self._iterator_get_next = weakref.WeakKeyDictionary()
+
   def compile(self,
               optimizer,
               loss=None,
@@ -623,12 +630,23 @@ class Model(Network):
           **kwargs)
     self._post_build_cleanup()
 
+  def _get_iterator_get_next_tensors(self, iterator):
+    get_next_op = self._iterator_get_next.get(iterator, None)
+    if get_next_op is None:
+      get_next_op = iterator.get_next()
+      self._iterator_get_next[iterator] = get_next_op
+    return get_next_op
+
   def _standardize_user_data(self,
                              x,
                              y=None,
                              sample_weight=None,
                              class_weight=None,
-                             batch_size=None):
+                             batch_size=None,
+                             check_steps=False,
+                             steps_name='steps',
+                             steps=None,
+                             validation_split=0):
     """Runs validation checks on input and target data passed by the user.
 
     Also standardizes the data to lists of arrays, in order.
@@ -660,6 +678,16 @@ class Model(Network):
         to, as conveyed by `y`.
       batch_size: Integer batch size. If provided, it is used to run additional
         validation checks on stateful models.
+      check_steps: boolean, True if we want to check for validity of `steps` and
+        False, otherwise. For example, when we are standardizing one batch of
+        data for train_on_batch/predict_on_batch/test_on_batch APIs, `steps`
+        value is not required and we should not check for its validity in these
+        cases.
+      steps_name: The public API's parameter name for `steps`.
+      steps: Integer or `None`. Total number of steps (batches of samples) to
+        execute.
+      validation_split: Float between 0 and 1.
+        Fraction of the training data to be used as validation data.
 
     Returns:
       A tuple of 3 lists: input arrays, target arrays, sample-weight arrays.
@@ -671,33 +699,54 @@ class Model(Network):
       ValueError: In case of invalid user-provided data.
       RuntimeError: If the model was never compiled.
     """
-    # First, we build/compile the model on the fly if necessary.
     if isinstance(x, dataset_ops.Dataset):
       raise ValueError('You passed a `Dataset` instance to your model (%s), '
                        'which is not supported. Instead, pass an `Iterator`, '
                        'which you can obtain e.g. via '
                        '`dataset.make_one_shot_iterator()` (the exact method '
                        'to use will depend on your specific dataset).' % x)
-    if isinstance(x, iterator_ops.Iterator):
-      if y is not None:
-        raise ValueError('You passed a dataset iterator (%s) as input `x` to '
-                         'your model. In that case, you should not specify '
-                         'a target (`y`) argument, since the dataset iterator '
-                         'generates both input data and target data. '
-                         'Received: %s' % (x, y))
-      if not context.executing_eagerly():
-        x, y = x.get_next()
-        # TODO(fchollet): handle case of `get_next` not returning 2 tensors?
-      else:
-        # TODO(psv): implement this. The way to support it will be to typecheck
-        # for `iterator` before `_standardize_user_data` is called and redirect
-        # to new training/eval functions in `training_eager.py`. The model
-        # may need to get built using the specs of the data from the first batch
-        # drawn from the iterator.
-        raise ValueError('Dataset iterators are not supported '
-                         'with eager execution yet.')
 
+    # Validates `steps` argument based on x's type.
+    if check_steps:
+      training_utils.check_steps_argument(x, steps, steps_name)
+
+    is_x_eager_iterator = isinstance(x, iterator_ops.EagerIterator)
+    is_x_iterator = isinstance(x, iterator_ops.Iterator)
+
+    # Validate user inputs when data is given as a dataset iterator.
+    if is_x_iterator or is_x_eager_iterator:
+      training_utils.validate_iterator_input(x, y, sample_weight,
+                                             validation_split)
+
+    # For eager iterators, when we have to process multiple batches of samples,
+    # we will standardize the data when we actually loop over iterator and get
+    # the batches. For now, we just return the iterator as is.
+    if is_x_eager_iterator and steps is not None:
+      return x, y, sample_weight
+
+    # If input data is a dataset iterator in graph mode or if it is an eager
+    # iterator and only one batch of samples is required, we fetch the data
+    # tensors from the iterator and then standardize them.
+    if is_x_iterator or is_x_eager_iterator:
+      try:
+        if is_x_iterator:
+          next_element = self._get_iterator_get_next_tensors(x)
+        else:
+          next_element = x.get_next()
+      except errors.OutOfRangeError:
+        raise RuntimeError('Your dataset iterator ran out of data; '
+                           'Make sure that your dataset can generate '
+                           'required number of samples.')
+
+      if not isinstance(next_element, (list, tuple)) or len(next_element) != 2:
+        raise ValueError('Please provide data as a list or tuple of 2 elements '
+                         ' - input and target pair. Received %s' % next_element)
+      x, y = next_element
+
+    # First, we build/compile the model on the fly if necessary.
     all_inputs = []
+    is_build_called = False
+    is_compile_called = False
     if not self.built:
       # We need to use `x` to set the model inputs.
       # We type-check that `x` and `y` are either single arrays
@@ -720,6 +769,7 @@ class Model(Network):
       # If values, then in symbolic-mode placeholders will be created
       # to match the value shapes.
       if not self.inputs:
+        is_build_called = True
         self._set_inputs(x)
 
     if y is not None:
@@ -736,6 +786,7 @@ class Model(Network):
             raise ValueError('Please provide as model targets either a single '
                              'array or a list of arrays. '
                              'You passed: y=' + str(y))
+          all_inputs += list(y)
         elif isinstance(y, dict):
           raise ValueError('Please do not pass a dictionary as model targets.')
         else:
@@ -743,14 +794,10 @@ class Model(Network):
             raise ValueError('Please provide as model targets either a single '
                              'array or a list of arrays. '
                              'You passed: y=' + str(y))
+          all_inputs.append(y)
 
         # Typecheck that all inputs are *either* value *or* symbolic.
         # TODO(fchollet): this check could be removed in Eager mode?
-        if y is not None:
-          if isinstance(y, (list, tuple)):
-            all_inputs += list(y)
-          else:
-            all_inputs.append(y)
         if any(tensor_util.is_tensor(v) for v in all_inputs):
           if not all(tensor_util.is_tensor(v) for v in all_inputs):
             raise ValueError('Do not pass inputs that mix Numpy arrays and '
@@ -764,17 +811,22 @@ class Model(Network):
           if not isinstance(y, (list, tuple)):
             y = [y]
           target_tensors = [v for v in y if tensor_util.is_tensor(v)]
+        is_compile_called = True
         self.compile(optimizer=self.optimizer,
                      loss=self.loss,
                      metrics=self.metrics,
                      loss_weights=self.loss_weights,
                      target_tensors=target_tensors)
 
-    # If `x` and `y` were all symbolic, then no model should not be fed any
-    # inputs and targets.
+    # In graph mode, if we had just set inputs and targets as symbolic tensors
+    # by invoking build and compile on the model respectively, we do not have to
+    # feed anything to the model. Model already has input and target data as
+    # part of the graph.
     # Note: in this case, `any` and `all` are equivalent since we disallow
     # mixed symbolic/value inputs.
-    if any(tensor_util.is_tensor(v) for v in all_inputs):
+    if (not context.executing_eagerly() and is_build_called and
+        is_compile_called and
+        any(tensor_util.is_tensor(v) for v in all_inputs)):
       return [], [], []
 
     # What follows is input validation and standardization to list format,
@@ -904,7 +956,12 @@ class Model(Network):
       if isinstance(inputs, list):
         assert len(inputs) == 1
         inputs = inputs[0]
-      self.build(input_shape=(None,) + inputs.shape[1:])
+
+      if tensor_util.is_tensor(inputs):
+        input_shape = (None,) + tuple(inputs.get_shape().as_list()[1:])
+      else:
+        input_shape = (None,) + inputs.shape[1:]
+      self.build(input_shape=input_shape)
     elif context.executing_eagerly():
       self._eager_set_inputs(inputs)
     else:
@@ -931,12 +988,18 @@ class Model(Network):
     # On-the-fly setting of model inputs/outputs as DeferredTensors,
     # to keep track of number of inputs and outputs and their ndim.
     if isinstance(inputs, (list, tuple)):
-      dummy_output_values = self.call(
-          [ops.convert_to_tensor(v, dtype=K.floatx()) for v in inputs])
+      if tensor_util.is_tensor(inputs[0]):
+        dummy_output_values = self.call(inputs)
+      else:
+        dummy_output_values = self.call(
+            [ops.convert_to_tensor(v, dtype=K.floatx()) for v in inputs])
       dummy_input_values = list(inputs)
     else:
-      dummy_output_values = self.call(
-          ops.convert_to_tensor(inputs, dtype=K.floatx()))
+      if tensor_util.is_tensor(inputs):
+        dummy_output_values = self.call(inputs)
+      else:
+        dummy_output_values = self.call(
+            ops.convert_to_tensor(inputs, dtype=K.floatx()))
       dummy_input_values = [inputs]
     if isinstance(dummy_output_values, (list, tuple)):
       dummy_output_values = list(dummy_output_values)
@@ -1071,7 +1134,7 @@ class Model(Network):
         batch_size: Integer or `None`.
             Number of samples per gradient update.
             If unspecified, `batch_size` will default to 32.
-            Do not specify the `batch_size` is your data is in the
+            Do not specify the `batch_size` if your data is in the
             form of symbolic tensors or dataset iterators (since they generate
             batches).
         epochs: Integer. Number of epochs to train the model.
@@ -1094,7 +1157,8 @@ class Model(Network):
             the loss and any model metrics
             on this data at the end of each epoch.
             The validation data is selected from the last samples
-            in the `x` and `y` data provided, before shuffling.
+            in the `x` and `y` data provided, before shuffling. This argument is
+            not supported when `x` is a dataset iterator.
         validation_data: Data on which to evaluate
             the loss and any model metrics at the end of each epoch.
             The model will not be trained on this data.
@@ -1124,7 +1188,8 @@ class Model(Network):
             `(samples, sequence_length)`,
             to apply a different weight to every timestep of every sample.
             In this case you should make sure to specify
-            `sample_weight_mode="temporal"` in `compile()`.
+            `sample_weight_mode="temporal"` in `compile()`. This argument is not
+            supported when `x` is a dataset iterator.
         initial_epoch: Integer.
             Epoch at which to start training
             (useful for resuming a previous training run).
@@ -1165,21 +1230,23 @@ class Model(Network):
       epochs = kwargs.pop('nb_epoch')
     if kwargs:
       raise TypeError('Unrecognized keyword arguments: ' + str(kwargs))
-    if x is None and y is None and steps_per_epoch is None:
-      raise ValueError('If fitting from data tensors, '
-                       'you should specify the `steps_per_epoch` '
-                       'argument.')
 
-    # Validate user data.
+    # Validate and standardize user data.
     x, y, sample_weights = self._standardize_user_data(
         x,
         y,
         sample_weight=sample_weight,
         class_weight=class_weight,
-        batch_size=batch_size)
+        batch_size=batch_size,
+        check_steps=True,
+        steps_name='steps_per_epoch',
+        steps=steps_per_epoch,
+        validation_split=validation_split)
+
     # Prepare validation data.
     if validation_data:
-      if isinstance(validation_data, iterator_ops.Iterator):
+      if (isinstance(validation_data, iterator_ops.Iterator) or
+          isinstance(validation_data, iterator_ops.EagerIterator)):
         val_x = validation_data
         val_y = None
         val_sample_weight = None
@@ -1196,11 +1263,13 @@ class Model(Network):
             'or alternatively it could be a dataset iterator. However we '
             'received `validation_data=%s`' % validation_data)
 
+      # Validate and standardize validation data.
       val_x, val_y, val_sample_weights = self._standardize_user_data(
           val_x,
           val_y,
           sample_weight=val_sample_weight,
-          batch_size=batch_size)
+          batch_size=batch_size,
+          steps=validation_steps)
 
     elif validation_split and 0. < validation_split < 1.:
       if training_utils.has_symbolic_tensors(x):
@@ -1229,6 +1298,7 @@ class Model(Network):
           inputs=x,
           targets=y,
           sample_weights=sample_weights,
+          class_weight=class_weight,
           batch_size=batch_size,
           epochs=epochs,
           verbose=verbose,
@@ -1300,7 +1370,8 @@ class Model(Network):
             `(samples, sequence_length)`,
             to apply a different weight to every timestep of every sample.
             In this case you should make sure to specify
-            `sample_weight_mode="temporal"` in `compile()`.
+            `sample_weight_mode="temporal"` in `compile()`. This argument is not
+            supported when `x` is a dataset iterator.
         steps: Integer or `None`.
             Total number of steps (batches of samples)
             before declaring the evaluation round finished.
@@ -1318,17 +1389,16 @@ class Model(Network):
     # Backwards compatibility.
     if batch_size is None and steps is None:
       batch_size = 32
-    if x is None and y is None and steps is None:
-      raise ValueError('If evaluating from data tensors, '
-                       'you should specify the `steps` '
-                       'argument.')
 
-    # Validate user data.
+    # Validate and standardize user data.
     x, y, sample_weights = self._standardize_user_data(
         x,
         y,
         sample_weight=sample_weight,
-        batch_size=batch_size)
+        batch_size=batch_size,
+        check_steps=True,
+        steps_name='steps',
+        steps=steps)
 
     if context.executing_eagerly():
       return training_eager.test_loop(
@@ -1345,7 +1415,12 @@ class Model(Network):
     Computation is done in batches.
 
     Arguments:
-        x: Input samples, as Numpy array(s) or tensor(s).
+         x: Input samples. It could be:
+          - A Numpy array (or array-like), or a list of arrays
+            (in case the model has multiple inputs).
+          - A TensorFlow tensor, or a list of tensors
+            (in case the model has multiple inputs).
+          - A `tf.data` dataset iterator.
         batch_size: Integer or `None`.
             Number of samples per gradient update.
             If unspecified, `batch_size` will default to 32.
@@ -1369,11 +1444,10 @@ class Model(Network):
     # Backwards compatibility.
     if batch_size is None and steps is None:
       batch_size = 32
-    if x is None and steps is None:
-      raise ValueError('If predicting from data tensors, '
-                       'you should specify the `steps` '
-                       'argument.')
-    x, _, _ = self._standardize_user_data(x)
+
+    # Validate and standardize user data.
+    x, _, _ = self._standardize_user_data(
+        x, check_steps=True, steps_name='steps', steps=steps)
 
     if context.executing_eagerly():
       return training_eager.predict_loop(
@@ -1406,7 +1480,9 @@ class Model(Network):
             with shape (samples, sequence_length),
             to apply a different weight to every timestep of every sample.
             In this case you should make sure to specify
-            sample_weight_mode="temporal" in compile().
+            sample_weight_mode="temporal" in compile(). This argument is not
+            supported when `x` is a dataset iterator.
+
         class_weight: Optional dictionary mapping
             class indices (integers) to
             a weight (float) to apply to the model's loss for the samples
@@ -1424,11 +1500,9 @@ class Model(Network):
     Raises:
       ValueError: In case of invalid user-provided arguments.
     """
+    # Validate and standardize user data.
     x, y, sample_weights = self._standardize_user_data(
-        x,
-        y,
-        sample_weight=sample_weight,
-        class_weight=class_weight)
+        x, y, sample_weight=sample_weight, class_weight=class_weight)
 
     if context.executing_eagerly():
       outputs = training_eager.train_on_batch(
@@ -1470,7 +1544,8 @@ class Model(Network):
             with shape (samples, sequence_length),
             to apply a different weight to every timestep of every sample.
             In this case you should make sure to specify
-            sample_weight_mode="temporal" in compile().
+            sample_weight_mode="temporal" in compile(). This argument is not
+            supported when `x` is a dataset iterator.
 
     Returns:
         Scalar test loss (if the model has a single output and no metrics)
@@ -1481,6 +1556,7 @@ class Model(Network):
     Raises:
         ValueError: In case of invalid user-provided arguments.
     """
+    # Validate and standardize user data.
     x, y, sample_weights = self._standardize_user_data(
         x, y, sample_weight=sample_weight)
 
@@ -1503,23 +1579,34 @@ class Model(Network):
     """Returns predictions for a single batch of samples.
 
     Arguments:
-        x: Input samples, as Numpy array(s) or tensor(s).
+        x: Input data. It could be:
+          - A Numpy array (or array-like), or a list of arrays
+            (in case the model has multiple inputs).
+          - A TensorFlow tensor, or a list of tensors
+            (in case the model has multiple inputs).
+          - A `tf.data` dataset iterator.
 
     Returns:
         Numpy array(s) of predictions.
 
+    Raises:
+        ValueError: In case of mismatch between given number of inputs and
+          expectations of the model.
     """
-    x, _, _ = self._standardize_user_data(x)
-
+    # Validate and standardize user data.
+    inputs, _, _ = self._standardize_user_data(x)
     if context.executing_eagerly():
-      inputs = [ops.convert_to_tensor(val, dtype=K.floatx()) for val in x]
+      if not isinstance(inputs, iterator_ops.EagerIterator):
+        inputs = [
+            ops.convert_to_tensor(val, dtype=K.floatx()) for val in inputs
+        ]
       return self(inputs)  # pylint: disable=not-callable
 
     if not context.executing_eagerly():
       if self.uses_learning_phase and not isinstance(K.learning_phase(), int):
-        ins = x + [0]
+        ins = inputs + [0]
       else:
-        ins = x
+        ins = inputs
 
       self._make_predict_function()
       outputs = self.predict_function(ins)
@@ -1631,8 +1718,7 @@ class Model(Network):
                             steps_per_epoch=10000, epochs=10)
     ```
     Raises:
-        ValueError: In case the generator yields
-            data in an invalid format.
+        ValueError: In case the generator yields data in an invalid format.
     """
     if not self.built and not self._is_graph_network:
       raise NotImplementedError(
@@ -1697,8 +1783,7 @@ class Model(Network):
         ValueError: in case of invalid arguments.
 
     Raises:
-        ValueError: In case the generator yields
-            data in an invalid format.
+        ValueError: In case the generator yields data in an invalid format.
     """
     if not self.built and not self._is_graph_network:
       raise NotImplementedError(
@@ -1751,8 +1836,7 @@ class Model(Network):
         Numpy array(s) of predictions.
 
     Raises:
-        ValueError: In case the generator yields
-            data in an invalid format.
+        ValueError: In case the generator yields data in an invalid format.
     """
     if not self.built and not self._is_graph_network:
       raise NotImplementedError(
diff --git a/tensorflow/python/keras/_impl/keras/engine/training_arrays.py b/tensorflow/python/keras/_impl/keras/engine/training_arrays.py
index 4164cae864..12e74ef51d 100644
--- a/tensorflow/python/keras/_impl/keras/engine/training_arrays.py
+++ b/tensorflow/python/keras/_impl/keras/engine/training_arrays.py
@@ -108,8 +108,8 @@ def fit_loop(model,
   do_validation = False
   if val_inputs:
     do_validation = True
-    if verbose and inputs and hasattr(inputs[0], 'shape') and hasattr(
-        val_inputs[0], 'shape'):
+    if (steps_per_epoch is None and verbose and inputs and
+        hasattr(inputs[0], 'shape') and hasattr(val_inputs[0], 'shape')):
       print('Train on %d samples, validate on %d samples' %
             (inputs[0].shape[0], val_inputs[0].shape[0]))
   if validation_steps:
diff --git a/tensorflow/python/keras/_impl/keras/engine/training_eager.py b/tensorflow/python/keras/_impl/keras/engine/training_eager.py
index 34adeb7599..526ae65321 100644
--- a/tensorflow/python/keras/_impl/keras/engine/training_eager.py
+++ b/tensorflow/python/keras/_impl/keras/engine/training_eager.py
@@ -23,7 +23,9 @@ import copy
 
 import numpy as np
 
+from tensorflow.python.data.ops import iterator_ops
 from tensorflow.python.eager.backprop import GradientTape
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_util
 from tensorflow.python.keras._impl.keras import backend
@@ -177,6 +179,550 @@ def _model_loss(model, inputs, targets, sample_weights=None, training=False):
   return outs, total_loss, loss_metrics
 
 
+def iterator_fit_loop(model,
+                      inputs,
+                      class_weight,
+                      steps_per_epoch,
+                      callback_model,
+                      out_labels,
+                      epoch_logs,
+                      val_inputs=None,
+                      val_targets=None,
+                      val_sample_weights=None,
+                      epochs=1,
+                      verbose=1,
+                      callbacks=None,
+                      callback_metrics=None,
+                      validation_steps=None,
+                      do_validation=False):
+  """Fit function for eager execution when input is given as dataset iterator.
+
+  Updates the given epoch logs.
+
+  Arguments:
+      model: Instance of the `Model`.
+      inputs: Input dataset iterator.
+      class_weight: Optional class-weight array to weight the importance of
+          samples in `inputs` based on the class they belong to, as conveyed by
+          the targets from the `inputs` iterator.
+      steps_per_epoch: Total number of steps (batches of samples)
+          before declaring one epoch finished and starting the
+          next epoch.
+      callback_model: Instance of `Model` to callback.
+      out_labels: Output labels generated from model metric names.
+      epoch_logs: Dictionary of logs from every epoch.
+      val_inputs: Input data for validation.
+      val_targets: Target data for validation.
+      val_sample_weights: Sample weight data for validation.
+      epochs: Number of times to iterate over the data
+      verbose: Verbosity mode, 0, 1 or 2
+      callbacks: List of callbacks to be called during training
+      callback_metrics: List of strings, the display names of the metrics
+          passed to the callbacks. They should be the
+          concatenation of list the display names of the outputs of
+           `f` and the list of display names of the outputs of `f_val`.
+      validation_steps: Number of steps to run validation for (only if doing
+        validation from data tensors). Ignored with default value of `None`.
+      do_validation: Boolean value indicating whether we should do validation.
+
+  Raises:
+      ValueError: In case of mismatch between given number of inputs and
+        expectations of the model.
+  """
+  assert isinstance(inputs, iterator_ops.EagerIterator)
+  for step_index in range(steps_per_epoch):
+    batch_logs = {}
+    batch_logs['batch'] = step_index
+    batch_logs['size'] = 1
+    callbacks.on_batch_begin(step_index, batch_logs)
+
+    # Get data from the iterator.
+    try:
+      next_element = inputs.get_next()
+    except errors.OutOfRangeError:
+      logging.warning(
+          'Your dataset iterator ran out of data; '
+          'interrupting training. Make sure that your dataset'
+          ' can generate at least `steps_per_epoch * epochs` '
+          'batches (in this case, %d batches).' % steps_per_epoch * epochs)
+      break
+
+    if not isinstance(next_element, (list, tuple)) or len(next_element) != 2:
+      raise ValueError('Please provide data as a list or tuple of 2 elements '
+                       ' - input and target pair. Received %s' % next_element)
+    x, y = next_element
+
+    # Validate and standardize data.
+    x, y, sample_weights = model._standardize_user_data(
+        x, y, class_weight=class_weight)
+    if sample_weights:
+      sample_weights = [
+          ops.convert_to_tensor(val, dtype=backend.floatx())
+          if val is not None else None for val in sample_weights
+      ]
+
+    if step_index == 0 and not callback_metrics:
+      out_labels = model.metrics_names
+      if do_validation:
+        callback_metrics = copy.copy(out_labels) + [
+            'val_' + n for n in out_labels
+        ]
+      else:
+        callback_metrics = copy.copy(out_labels)
+      callbacks.set_params({
+          'epochs': epochs,
+          'steps': steps_per_epoch,
+          'verbose': verbose,
+          'do_validation': do_validation,
+          'metrics': callback_metrics or [],
+      })
+
+    # Train model.
+    outs, loss, loss_metrics = _process_single_batch(
+        model, x, y, sample_weights=sample_weights, training=True)
+    if not isinstance(outs, list):
+      outs = [outs]
+
+    # Calculate metrics.
+    for l, o in zip(out_labels, outs):
+      batch_logs[l] = o
+    # Required for eager execution
+    metrics_results = _eager_metrics_fn(model, outs, y)
+    batch_logs['loss'] = tensor_util.constant_value(backend.mean(loss))
+
+    for k, v in zip(model.metrics_names,
+                    [backend.mean(loss)] + loss_metrics + metrics_results):
+      batch_logs[k] = tensor_util.constant_value(v)
+    callbacks.on_batch_end(step_index, batch_logs)
+    if callback_model.stop_training:
+      break
+
+    if step_index == steps_per_epoch - 1:
+      if do_validation:
+        val_outs = test_loop(
+            model,
+            val_inputs,
+            val_targets,
+            sample_weights=val_sample_weights,
+            steps=validation_steps,
+            verbose=0)
+        if not isinstance(val_outs, list):
+          val_outs = [val_outs]
+        # Same labels assumed.
+        for l, o in zip(out_labels, val_outs):
+          epoch_logs['val_' + l] = o
+
+
+def batch_fit_loop(model,
+                   inputs,
+                   targets,
+                   epoch_logs,
+                   index_array,
+                   out_labels,
+                   callback_model,
+                   batch_size,
+                   sample_weights=None,
+                   val_inputs=None,
+                   val_targets=None,
+                   val_sample_weights=None,
+                   callbacks=None,
+                   shuffle=True,
+                   num_train_samples=None,
+                   do_validation=False):
+  """Fit function for eager execution when input is given as arrays or tensors.
+
+  Updates the given epoch logs.
+
+  Arguments:
+      model: Instance of the `Model`.
+      inputs: List of input arrays.
+      targets: List of target arrays.
+      epoch_logs: Dictionary of logs from every epoch.
+      index_array: Index array generated from number of training samples.
+      out_labels: Output labels generated from model metric names.
+      callback_model: Instance of `Model` to callback.
+      batch_size: Integer batch size or None if unknown.
+      sample_weights: Optional list of sample weight arrays.
+      val_inputs: Input data for validation.
+      val_targets: Target data for validation.
+      val_sample_weights: Sample weight data for validation.
+      callbacks: List of callbacks to be called during training.
+      shuffle: Whether to shuffle the data at the beginning of each epoch.
+      num_train_samples: Integer number of training samples.
+      do_validation: Boolean value indicating whether we should do validation.
+  """
+  # TODO(psv): Create a dataset iterator instead of manually creating batches
+  # here and in batch_test_loop, batch_predict_loop.
+  if shuffle == 'batch':
+    index_array = model._batch_shuffle(index_array, batch_size)
+  elif shuffle:
+    np.random.shuffle(index_array)
+
+  batches = generic_utils.make_batches(num_train_samples, batch_size)
+
+  for batch_index, (batch_start, batch_end) in enumerate(batches):
+    batch_ids = index_array[batch_start:batch_end]
+    inputs_batch = slice_arrays(inputs, batch_ids, contiguous=not shuffle)
+    targets_batch = slice_arrays(targets, batch_ids, contiguous=not shuffle)
+    if sample_weights:
+      sample_weights_batch = slice_arrays(
+          sample_weights, batch_ids, contiguous=not shuffle)
+    else:
+      sample_weights_batch = None
+    batch_logs = {}
+    batch_logs['batch'] = batch_index
+    batch_logs['size'] = len(batch_ids)
+
+    callbacks.on_batch_begin(batch_index, batch_logs)
+
+    inputs_batch = [
+        ops.convert_to_tensor(val, dtype=backend.floatx())
+        for val in inputs_batch
+    ]
+    targets_batch = [
+        ops.convert_to_tensor(val, dtype=backend.floatx())
+        for val in targets_batch
+    ]
+    if sample_weights:
+      sample_weights_batch = [
+          ops.convert_to_tensor(val, dtype=backend.floatx())
+          if val is not None else None for val in sample_weights_batch
+      ]
+
+    outs, loss, loss_metrics = _process_single_batch(
+        model,
+        inputs_batch,
+        targets_batch,
+        sample_weights=sample_weights_batch,
+        training=True)
+
+    if not isinstance(outs, list):
+      outs = [outs]
+
+    for l, o in zip(out_labels, outs):
+      batch_logs[l] = o
+    # Required for eager execution
+    metrics_results = _eager_metrics_fn(model, outs, targets_batch)
+    batch_logs['loss'] = tensor_util.constant_value(backend.mean(loss))
+
+    for k, v in zip(model.metrics_names,
+                    [backend.mean(loss)] + loss_metrics + metrics_results):
+      batch_logs[k] = tensor_util.constant_value(v)
+    callbacks.on_batch_end(batch_index, batch_logs)
+    if callback_model.stop_training:
+      break
+
+    if batch_index == len(batches) - 1:  # Last batch.
+      if do_validation:
+        val_outs = test_loop(
+            model,
+            val_inputs,
+            val_targets,
+            sample_weights=val_sample_weights,
+            batch_size=batch_size,
+            verbose=0)
+        if not isinstance(val_outs, list):
+          val_outs = [val_outs]
+        # Same labels assumed.
+        for l, o in zip(out_labels, val_outs):
+          epoch_logs['val_' + l] = o
+
+
+def iterator_test_loop(model, inputs, steps, verbose=0):
+  """Test function for eager execution when input is given as dataset iterator.
+
+  Arguments:
+      model: Model instance that is being evaluated in Eager mode.
+      inputs: Input dataset iterator.
+      steps: Total number of steps (batches of samples) before declaring
+      predictions finished.
+      verbose: Verbosity mode.
+
+  Returns:
+      Scalar loss (if the model has a single output and no metrics)
+      or list of scalars (if the model has multiple outputs
+      and/or metrics). The attribute `model.metrics_names` will give you
+      the display labels for the scalar outputs.
+
+  Raises:
+      ValueError: In case of mismatch between given number of inputs and
+        expectations of the model.
+  """
+  assert isinstance(inputs, iterator_ops.EagerIterator)
+  outs = []
+  num_samples = 0
+  if verbose == 1:
+    progbar = generic_utils.Progbar(target=steps)
+  for step_index in range(steps):
+    # Get data from the iterator.
+    try:
+      next_element = inputs.get_next()
+    except errors.OutOfRangeError:
+      logging.warning(
+          'Your dataset iterator ran out of data interrupting testing. '
+          'Make sure that your dataset can generate at least `steps` batches '
+          '(in this case, %d batches).', steps)
+      break
+
+    if not isinstance(next_element, (list, tuple)) or len(next_element) != 2:
+      raise ValueError('Please provide data as a list or tuple of 2 elements '
+                       ' - input and target pair. Received %s' % next_element)
+    x, y = next_element
+
+    # Validate and standardize data.
+    x, y, sample_weights = model._standardize_user_data(x, y)
+
+    # Calculate model output, loss values.
+    loss_outs, loss, loss_metrics = _model_loss(
+        model, x, y, sample_weights=sample_weights, training=False)
+    metrics_results = _eager_metrics_fn(model, loss_outs, y)
+    batch_outs = []
+    for _, v in zip(model.metrics_names,
+                    [backend.mean(loss)] + loss_metrics + metrics_results):
+      batch_outs.append(tensor_util.constant_value(v))
+
+    # Get current step size.
+    if isinstance(x, list):
+      step_size = x[0].get_shape().as_list()[0]
+    else:
+      step_size = x.get_shape().as_list()[0]
+
+    # Accumulate results in output array.
+    if not isinstance(batch_outs, list):
+      batch_outs = [batch_outs]
+    if step_index == 0:
+      for _ in enumerate(batch_outs):
+        outs.append(0.)
+    for i, batch_out in enumerate(batch_outs):
+      outs[i] += batch_out * step_size
+
+    # Calculate sample size.
+    num_samples += step_size
+    if verbose == 1:
+      progbar.update(step_index + 1)
+
+    for i in range(len(outs)):
+      outs[i] /= num_samples
+    if len(outs) == 1:
+      return outs[0]
+    return outs
+
+
+def batch_test_loop(model,
+                    inputs,
+                    targets,
+                    batch_size,
+                    sample_weights=None,
+                    verbose=0):
+  """Test function for eager execution when input is given as arrays or tensors.
+
+  Arguments:
+      model: Model instance that is being evaluated in Eager mode.
+      inputs: List of input arrays.
+      targets: List of target arrays.
+      batch_size: Integer batch size.
+      sample_weights: Optional list of sample weight arrays.
+      verbose: Verbosity mode.
+
+  Returns:
+      Scalar loss (if the model has a single output and no metrics)
+      or list of scalars (if the model has multiple outputs
+      and/or metrics). The attribute `model.metrics_names` will give you
+      the display labels for the scalar outputs.
+  """
+  outs = []
+  feed_data = inputs + targets
+  if sample_weights:
+    feed_data += sample_weights
+  num_samples = training_utils.check_num_samples(
+      feed_data, batch_size=batch_size)
+  if verbose == 1:
+    progbar = generic_utils.Progbar(target=num_samples)
+  batches = generic_utils.make_batches(num_samples, batch_size)
+  index_array = np.arange(num_samples)
+  for batch_index, (batch_start, batch_end) in enumerate(batches):
+    batch_ids = index_array[batch_start:batch_end]
+    inputs_batch = slice_arrays(inputs, batch_ids)
+    targets_batch = slice_arrays(targets, batch_ids)
+    if sample_weights:
+      sample_weights_batch = slice_arrays(sample_weights, batch_ids)
+    else:
+      sample_weights_batch = None
+
+    inputs_batch = [
+        ops.convert_to_tensor(val, dtype=backend.floatx())
+        for val in inputs_batch
+    ]
+    targets_batch = [
+        ops.convert_to_tensor(val, dtype=backend.floatx())
+        for val in targets_batch
+    ]
+    if sample_weights:
+      sample_weights_batch = [
+          ops.convert_to_tensor(val, dtype=backend.floatx())
+          if val is not None else None for val in sample_weights_batch
+      ]
+
+    loss_outs, loss, loss_metrics = _model_loss(
+        model,
+        inputs_batch,
+        targets_batch,
+        sample_weights=sample_weights_batch,
+        training=False)
+    metrics_results = _eager_metrics_fn(model, loss_outs, targets_batch)
+    batch_outs = []
+    for _, v in zip(model.metrics_names,
+                    [backend.mean(loss)] + loss_metrics + metrics_results):
+      batch_outs.append(tensor_util.constant_value(v))
+
+    if isinstance(batch_outs, list):
+      if batch_index == 0:
+        for _ in enumerate(batch_outs):
+          outs.append(0.)
+      for i, batch_out in enumerate(batch_outs):
+        outs[i] += batch_out * len(batch_ids)
+    else:
+      if batch_index == 0:
+        outs.append(0.)
+      outs[0] += batch_outs * len(batch_ids)
+
+    if verbose == 1:
+      progbar.update(batch_end)
+
+  for i in range(len(outs)):
+    outs[i] /= num_samples
+  if len(outs) == 1:
+    return outs[0]
+  return outs
+
+
+def iterator_predict_loop(model, inputs, steps, verbose=0):
+  """Predict function for eager execution when input is dataset iterator.
+
+  Arguments:
+      model: Instance of `Model`.
+      inputs: Input dataset iterator.
+      steps: Total number of steps (batches of samples) before declaring
+          `_predict_loop` finished.
+      verbose: Verbosity mode.
+
+  Returns:
+      Array of predictions (if the model has a single output)
+      or list of arrays of predictions (if the model has multiple outputs).
+
+  Raises:
+      ValueError: In case of mismatch between given number of inputs and
+        expectations of the model.
+  """
+  assert isinstance(inputs, iterator_ops.EagerIterator)
+  outs = []
+  if verbose == 1:
+    progbar = generic_utils.Progbar(target=steps)
+  for step_index in range(steps):
+    # Get data from the iterator.
+    try:
+      next_element = inputs.get_next()
+    except errors.OutOfRangeError:
+      logging.warning(
+          'Your dataset iterator ran out of data; '
+          'interrupting prediction. Make sure that your '
+          'dataset can generate at least `steps` '
+          'batches (in this case, %d batches).', steps)
+      break
+
+    if not isinstance(next_element, (list, tuple)) or len(next_element) != 2:
+      raise ValueError(
+          'Please provide data as a list or tuple of 2 elements '
+          ' - input and target pair. Received %s. We do not use the '
+          '`target` value here.' % next_element)
+    x, _ = next_element
+
+    # Validate and standardize data.
+    x, _, _ = model._standardize_user_data(x)
+
+    if model._expects_training_arg:
+      batch_outs = model.call(x[0] if len(x) == 1 else x, training=False)
+    else:
+      batch_outs = model.call(x[0] if len(x) == 1 else x)
+    if not isinstance(batch_outs, list):
+      batch_outs = [batch_outs]
+
+    # We collect the results from every step and then concatenate them once
+    # in the end. This is an expensive process. We are doing this because we
+    # do not know the number of samples beforehand.
+    if step_index == 0:
+      for _ in batch_outs:
+        outs.append([])
+    for i, batch_out in enumerate(batch_outs):
+      outs[i].append(backend.get_value(batch_out))
+
+    if verbose == 1:
+      progbar.update(step_index + 1)
+  for i, out in enumerate(outs):
+    outs[i] = np.concatenate(tuple(out), axis=0)
+  if len(outs) == 1:
+    return outs[0]
+  return outs
+
+
+def batch_predict_loop(model, inputs, batch_size, verbose=0):
+  """Predict function for eager execution when input is arrays or tensors.
+
+  Arguments:
+      model: Instance of `Model`.
+      inputs: List of input arrays.
+      batch_size: Integer batch size.
+      verbose: Verbosity mode.
+
+  Returns:
+      Array of predictions (if the model has a single output)
+      or list of arrays of predictions (if the model has multiple outputs).
+  """
+  outs = []
+  num_samples = training_utils.check_num_samples(inputs, batch_size)
+  if verbose == 1:
+    progbar = generic_utils.Progbar(target=num_samples)
+  batches = generic_utils.make_batches(num_samples, batch_size)
+  index_array = np.arange(num_samples)
+  for batch_index, (batch_start, batch_end) in enumerate(batches):
+    batch_ids = index_array[batch_start:batch_end]
+    inputs_batch = slice_arrays(inputs, batch_ids)
+
+    inputs_batch = [
+        ops.convert_to_tensor(val, dtype=backend.floatx())
+        for val in inputs_batch
+    ]
+
+    if len(inputs_batch) == 1:
+      if model._expects_training_arg:
+        batch_outs = model.call(inputs_batch[0], training=False)
+      else:
+        batch_outs = model.call(inputs_batch[0])
+    else:
+      if model._expects_training_arg:
+        batch_outs = model.call(inputs_batch, training=False)
+      else:
+        batch_outs = model.call(inputs_batch)
+
+    if not isinstance(batch_outs, list):
+      batch_outs = [batch_outs]
+    if batch_index == 0:
+      # Pre-allocate the results arrays.
+      for batch_out in batch_outs:
+        dims = batch_out.shape[1:].dims
+        dims_list = [d.value for d in dims]
+        shape = (num_samples,) + tuple(dims_list)
+        outs.append(np.zeros(shape, dtype=batch_out.dtype.as_numpy_dtype))
+    for i, batch_out in enumerate(batch_outs):
+      outs[i][batch_start:batch_end] = batch_out
+    if verbose == 1:
+      progbar.update(batch_end)
+
+  if len(outs) == 1:
+    return outs[0]
+  return outs
+
+
 def slice_arrays(arrays, indices, contiguous=True):
   """Slices batches out of provided arrays (workaround for eager tensors).
 
@@ -268,19 +814,24 @@ def train_on_batch(model, inputs, targets, sample_weights=None):
   Returns:
       total loss and the loss associated with each output.
   """
-  inputs = [
-      ops.convert_to_tensor(val, dtype=backend.floatx()) for val in inputs]
-  targets = [
-      ops.convert_to_tensor(val, dtype=backend.floatx()) for val in targets]
-  sample_weights = [
-      ops.convert_to_tensor(val, dtype=backend.floatx())
-      if val is not None else None for val in sample_weights]
+  if len(inputs) and not tensor_util.is_tensor(inputs[0]):
+    inputs = [
+        ops.convert_to_tensor(val, dtype=backend.floatx()) for val in inputs
+    ]
+    targets = [
+        ops.convert_to_tensor(val, dtype=backend.floatx()) for val in targets
+    ]
+  if sample_weights:
+    sample_weights = [
+        ops.convert_to_tensor(val, dtype=backend.floatx())
+        if val is not None else None for val in sample_weights
+    ]
+
   outs, loss, _ = _process_single_batch(
       model, inputs, targets, sample_weights=sample_weights, training=True)
   if not isinstance(outs, list):
     outs = [outs]
-  metrics_results = _eager_metrics_fn(
-      model, outs, targets)
+  metrics_results = _eager_metrics_fn(model, outs, targets)
   if not isinstance(loss, list):
     loss = [loss]
   return loss + metrics_results
@@ -298,48 +849,55 @@ def test_on_batch(model, inputs, targets, sample_weights=None):
   Returns:
       total loss, loss and metrics associated with each output.
   """
-  inputs = [
-      ops.convert_to_tensor(val, dtype=backend.floatx()) for val in inputs]
-  targets = [
-      ops.convert_to_tensor(val, dtype=backend.floatx()) for val in targets]
-  sample_weights = [
-      ops.convert_to_tensor(val, dtype=backend.floatx())
-      if val is not None else None for val in sample_weights]
-  outs, loss, loss_metrics = _process_single_batch(
+  if len(inputs) and not tensor_util.is_tensor(inputs[0]):
+    inputs = [
+        ops.convert_to_tensor(val, dtype=backend.floatx()) for val in inputs
+    ]
+    targets = [
+        ops.convert_to_tensor(val, dtype=backend.floatx()) for val in targets
+    ]
+  if sample_weights:
+    sample_weights = [
+        ops.convert_to_tensor(val, dtype=backend.floatx())
+        if val is not None else None for val in sample_weights
+    ]
+  outs, loss, loss_metrics = _model_loss(
       model, inputs, targets, sample_weights=sample_weights, training=False)
   if not isinstance(outs, list):
     outs = [outs]
-  metrics_results = _eager_metrics_fn(
-      model, outs, targets)
+  metrics_results = _eager_metrics_fn(model, outs, targets)
   if not isinstance(loss, list):
     loss = [loss]
   return loss + loss_metrics + metrics_results
 
 
-def fit_loop(
-    model,
-    inputs,
-    targets,
-    sample_weights=None,
-    val_inputs=None,
-    val_targets=None,
-    val_sample_weights=None,
-    batch_size=None,
-    epochs=100,
-    verbose=1,
-    callbacks=None,
-    shuffle=True,
-    callback_metrics=None,
-    initial_epoch=0,
-    steps_per_epoch=None,
-    validation_steps=None):
-  """Abstract fit function for eager execution.
+def fit_loop(model,
+             inputs,
+             targets,
+             sample_weights=None,
+             class_weight=None,
+             val_inputs=None,
+             val_targets=None,
+             val_sample_weights=None,
+             batch_size=None,
+             epochs=1,
+             verbose=1,
+             callbacks=None,
+             shuffle=True,
+             callback_metrics=None,
+             initial_epoch=0,
+             steps_per_epoch=None,
+             validation_steps=None):
+  """Fit function for eager execution.
 
   Arguments:
       model: Instance of the model that is being executed in Eager mode.
       inputs: List of input arrays.
       targets: List of target arrays.
       sample_weights: Optional list of sample weight arrays.
+      class_weight: Optional class-weight array to weight the importance of
+          samples in `inputs` based on the class they belong to, as conveyed by
+          `targets`.
       val_inputs: Input data for validation.
       val_targets: Target data for validation.
       val_sample_weights: Sample weight data for validation.
@@ -366,47 +924,40 @@ def fit_loop(
   Raises:
     ValueError: In case of invalid argument values.
   """
-  if not batch_size:
-    raise ValueError('With eager execution, `batch_size` should be specified.')
-  if steps_per_epoch or validation_steps:
-    raise ValueError('With eager execution, `steps_per_epoch` and '
-                     '`validation_steps` are not valid arguments '
-                     '(set `batch_size` instead).')
-  # Required for Eager mode
+  # Required for eager execution
   with backend.learning_phase_scope(1):
     do_validation = False
     if val_inputs:
       do_validation = True
-      if (verbose and inputs and hasattr(inputs[0], 'shape') and
-          hasattr(val_inputs[0], 'shape')):
+      if (steps_per_epoch is None and verbose and inputs and
+          hasattr(inputs[0], 'shape') and hasattr(val_inputs[0], 'shape')):
         print('Train on %d samples, validate on %d samples' %
               (inputs[0].shape[0], val_inputs[0].shape[0]))
-    if validation_steps:
-      if steps_per_epoch is None:
-        raise ValueError('Can only use `validation_steps` when doing step-wise '
-                         'training, i.e. `steps_per_epoch` must be set.')
-      do_validation = True
 
-    out_labels = model.metrics_names
-    if do_validation:
-      callback_metrics = copy.copy(out_labels) + [
-          'val_' + n for n in out_labels
-      ]
-    else:
-      callback_metrics = copy.copy(out_labels)
+    num_train_samples = None
+    out_labels = None
+    if steps_per_epoch is None or model._is_compiled:
+      out_labels = model.metrics_names
+      if do_validation:
+        callback_metrics = copy.copy(out_labels) + [
+            'val_' + n for n in out_labels
+        ]
+      else:
+        callback_metrics = copy.copy(out_labels)
 
-    if sample_weights:
-      feed_data = inputs + targets + sample_weights
-    else:
-      feed_data = inputs + targets
-    num_train_samples = training_utils.check_num_samples(
-        feed_data,
-        batch_size=batch_size,
-        steps=steps_per_epoch,
-        steps_name='steps_per_epoch')
+    if steps_per_epoch is None:
+      if sample_weights:
+        feed_data = inputs + targets + sample_weights
+      else:
+        feed_data = inputs + targets
+      num_train_samples = training_utils.check_num_samples(
+          feed_data,
+          batch_size=batch_size,
+          steps=steps_per_epoch,
+          steps_name='steps_per_epoch')
 
-    if num_train_samples is not None:
-      index_array = np.arange(num_train_samples)
+      if num_train_samples is not None:
+        index_array = np.arange(num_train_samples)
 
     model.history = cbks.History()
     callbacks = [cbks.BaseLogger()] + (callbacks or []) + [model.history]
@@ -441,6 +992,8 @@ def fit_loop(
     for cbk in callbacks:
       if not val_inputs:
         cbk.validation_data = []
+      elif isinstance(val_inputs, iterator_ops.EagerIterator):
+        cbk.validation_data = val_inputs
       elif val_sample_weights:
         cbk.validation_data = val_inputs + val_targets + val_sample_weights
       else:
@@ -449,87 +1002,48 @@ def fit_loop(
     for epoch in range(initial_epoch, epochs):
       callbacks.on_epoch_begin(epoch)
       epoch_logs = {}
-      if shuffle == 'batch':
-        index_array = model._batch_shuffle(index_array, batch_size)
-      elif shuffle:
-        np.random.shuffle(index_array)
-
-      batches = generic_utils.make_batches(num_train_samples, batch_size)
-
-      for batch_index, (batch_start, batch_end) in enumerate(batches):
-        batch_ids = index_array[batch_start:batch_end]
-        try:
-          inputs_batch = slice_arrays(inputs, batch_ids,
-                                      contiguous=not shuffle)
-          targets_batch = slice_arrays(targets, batch_ids,
-                                       contiguous=not shuffle)
-          if sample_weights:
-            sample_weights_batch = slice_arrays(sample_weights, batch_ids,
-                                                contiguous=not shuffle)
-          else:
-            sample_weights_batch = None
-        except TypeError:
-          raise TypeError('TypeError while preparing batch. '
-                          'If using HDF5 input data, '
-                          'pass shuffle="batch".')
-        batch_logs = {}
-        batch_logs['batch'] = batch_index
-        batch_logs['size'] = len(batch_ids)
-
-        callbacks.on_batch_begin(batch_index, batch_logs)
-
-        inputs_batch = [
-            ops.convert_to_tensor(val, dtype=backend.floatx())
-            for val in inputs_batch]
-        targets_batch = [
-            ops.convert_to_tensor(val, dtype=backend.floatx())
-            for val in targets_batch]
-        if sample_weights:
-          sample_weights_batch = [
-              ops.convert_to_tensor(val, dtype=backend.floatx())
-              if val is not None else None
-              for val in sample_weights_batch]
-
-        outs, loss, loss_metrics = _process_single_batch(
+
+      if steps_per_epoch is not None:
+        iterator_fit_loop(
             model,
-            inputs_batch,
-            targets_batch,
-            sample_weights=sample_weights_batch,
-            training=True)
-
-        if not isinstance(outs, list):
-          outs = [outs]
-
-        for l, o in zip(out_labels, outs):
-          batch_logs[l] = o
-        # Required for Eager mode
-        metrics_results = _eager_metrics_fn(model, outs, targets_batch)
-        batch_logs['loss'] = tensor_util.constant_value(backend.mean(loss))
-
-        for k, v in zip(model.metrics_names,
-                        [backend.mean(loss)] + loss_metrics + metrics_results):
-          batch_logs[k] = tensor_util.constant_value(v)
-        callbacks.on_batch_end(batch_index, batch_logs)
-        if callback_model.stop_training:
-          break
-
-        if batch_index == len(batches) - 1:  # Last batch.
-          if do_validation:
-            val_outs = test_loop(
-                model, val_inputs, val_targets,
-                sample_weights=val_sample_weights,
-                batch_size=batch_size,
-                verbose=0)
-            if not isinstance(val_outs, list):
-              val_outs = [val_outs]
-            # Same labels assumed.
-            for l, o in zip(out_labels, val_outs):
-              epoch_logs['val_' + l] = o
+            inputs,
+            class_weight,
+            steps_per_epoch=steps_per_epoch,
+            callback_model=callback_model,
+            out_labels=out_labels,
+            epoch_logs=epoch_logs,
+            val_inputs=val_inputs,
+            val_targets=val_targets,
+            val_sample_weights=val_sample_weights,
+            epochs=epochs,
+            verbose=verbose,
+            callbacks=callbacks,
+            callback_metrics=callback_metrics,
+            validation_steps=validation_steps,
+            do_validation=do_validation)
+      else:
+        batch_fit_loop(
+            model,
+            inputs,
+            targets,
+            epoch_logs=epoch_logs,
+            index_array=index_array,
+            out_labels=out_labels,
+            callback_model=callback_model,
+            batch_size=batch_size,
+            sample_weights=sample_weights,
+            val_inputs=val_inputs,
+            val_targets=val_targets,
+            val_sample_weights=val_sample_weights,
+            callbacks=callbacks,
+            shuffle=shuffle,
+            num_train_samples=num_train_samples,
+            do_validation=do_validation)
       callbacks.on_epoch_end(epoch, epoch_logs)
       if callback_model.stop_training:
         break
-    callbacks.on_train_end()
-    return model.history
+  callbacks.on_train_end()
+  return model.history
 
 
 def test_loop(model, inputs, targets,
@@ -537,7 +1051,7 @@ def test_loop(model, inputs, targets,
               batch_size=None,
               verbose=0,
               steps=None):
-  """Abstract method to loop over some data in batches.
+  """Test function for eager execution.
 
   Arguments:
       model: Model instance that is being evaluated in Eager mode.
@@ -557,77 +1071,26 @@ def test_loop(model, inputs, targets,
       the display labels for the scalar outputs.
   """
   with backend.learning_phase_scope(0):
-    feed_data = inputs + targets
-    if sample_weights:
-      feed_data += sample_weights
-    num_samples = training_utils.check_num_samples(
-        feed_data, batch_size=batch_size, steps=steps, steps_name='steps')
-    outs = []
-    if verbose == 1:
-      progbar = generic_utils.Progbar(target=num_samples)
-    batches = generic_utils.make_batches(num_samples, batch_size)
-    index_array = np.arange(num_samples)
-    for batch_index, (batch_start, batch_end) in enumerate(batches):
-      batch_ids = index_array[batch_start:batch_end]
-      inputs_batch = slice_arrays(inputs, batch_ids)
-      targets_batch = slice_arrays(targets, batch_ids)
-      if sample_weights:
-        sample_weights_batch = slice_arrays(sample_weights, batch_ids)
-      else:
-        sample_weights_batch = None
-
-      inputs_batch = [
-          ops.convert_to_tensor(val, dtype=backend.floatx())
-          for val in inputs_batch]
-      targets_batch = [
-          ops.convert_to_tensor(val, dtype=backend.floatx())
-          for val in targets_batch]
-      if sample_weights:
-        sample_weights_batch = [
-            ops.convert_to_tensor(val, dtype=backend.floatx())
-            if val is not None else None
-            for val in sample_weights_batch]
-
-      loss_outs, loss, loss_metrics = _model_loss(
+    if steps is not None:
+      return iterator_test_loop(model, inputs, steps, verbose=verbose)
+    else:
+      return batch_test_loop(
           model,
-          inputs_batch,
-          targets_batch,
-          sample_weights=sample_weights_batch,
-          training=False)
-      metrics_results = _eager_metrics_fn(model, loss_outs, targets_batch)
-      batch_outs = []
-      for _, v in zip(model.metrics_names,
-                      [backend.mean(loss)] + loss_metrics + metrics_results):
-        batch_outs.append(tensor_util.constant_value(v))
-
-      if isinstance(batch_outs, list):
-        if batch_index == 0:
-          for batch_out in enumerate(batch_outs):
-            outs.append(0.)
-        for i, batch_out in enumerate(batch_outs):
-          outs[i] += batch_out * len(batch_ids)
-      else:
-        if batch_index == 0:
-          outs.append(0.)
-        outs[0] += batch_outs * len(batch_ids)
-
-      if verbose == 1:
-        progbar.update(batch_end)
-    for i in range(len(outs)):
-      outs[i] /= num_samples
-    if len(outs) == 1:
-      return outs[0]
-    return outs
+          inputs,
+          targets,
+          batch_size=batch_size,
+          sample_weights=sample_weights,
+          verbose=verbose)
 
 
 def predict_loop(model, inputs,
                  batch_size=32,
                  verbose=0,
                  steps=None):
-  """Abstract method to loop over some data in batches.
+  """Predict function for eager execution.
 
   Arguments:
-      model:
+      model: Instance of `Model`.
       inputs: List of input arrays.
       batch_size: integer batch size.
       verbose: verbosity mode.
@@ -641,49 +1104,8 @@ def predict_loop(model, inputs,
       (if the model has multiple outputs).
   """
   with backend.learning_phase_scope(0):
-    num_samples = training_utils.check_num_samples(
-        inputs, batch_size, steps, 'steps')
-    if verbose == 1:
-      if steps is not None:
-        progbar = generic_utils.Progbar(target=steps)
-      else:
-        progbar = generic_utils.Progbar(target=num_samples)
-
-    outs = []
-    batches = generic_utils.make_batches(num_samples, batch_size)
-    index_array = np.arange(num_samples)
-    for batch_index, (batch_start, batch_end) in enumerate(batches):
-      batch_ids = index_array[batch_start:batch_end]
-      inputs_batch = slice_arrays(inputs, batch_ids)
-
-      inputs_batch = [
-          ops.convert_to_tensor(val, dtype=backend.floatx())
-          for val in inputs_batch]
-
-      if len(inputs_batch) == 1:
-        if model._expects_training_arg:
-          batch_outs = model.call(inputs_batch[0], training=False)
-        else:
-          batch_outs = model.call(inputs_batch[0])
-      else:
-        if model._expects_training_arg:
-          batch_outs = model.call(inputs_batch, training=False)
-        else:
-          batch_outs = model.call(inputs_batch)
-
-      if not isinstance(batch_outs, list):
-        batch_outs = [batch_outs]
-      if batch_index == 0:
-        # Pre-allocate the results arrays.
-        for batch_out in batch_outs:
-          dims = batch_out.shape[1:].dims
-          dims_list = [d.value for d in dims]
-          shape = (num_samples,) + tuple(dims_list)
-          outs.append(np.zeros(shape, dtype=batch_out.dtype.as_numpy_dtype))
-      for i, batch_out in enumerate(batch_outs):
-        outs[i][batch_start:batch_end] = batch_out
-      if verbose == 1:
-        progbar.update(batch_end)
-    if len(outs) == 1:
-      return outs[0]
-    return outs
+    if steps is not None:
+      return iterator_predict_loop(model, inputs, steps, verbose=verbose)
+    else:
+      return batch_predict_loop(
+          model, inputs, batch_size=batch_size, verbose=verbose)
diff --git a/tensorflow/python/keras/_impl/keras/engine/training_test.py b/tensorflow/python/keras/_impl/keras/engine/training_test.py
index 58011a1412..cc2386a5bd 100644
--- a/tensorflow/python/keras/_impl/keras/engine/training_test.py
+++ b/tensorflow/python/keras/_impl/keras/engine/training_test.py
@@ -24,6 +24,7 @@ import unittest
 import numpy as np
 
 from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_shape
 from tensorflow.python.framework import test_util as tf_test_util
 from tensorflow.python.keras._impl import keras
@@ -1340,16 +1341,12 @@ class TestTrainingWithDataTensors(test.TestCase):
                                  output_a_np)
 
       # test fit
-      out = model.fit(None,
-                      output_a_np, epochs=1, batch_size=10)
-      out = model.fit(None,
-                      output_a_np, epochs=1, batch_size=10)
+      _ = model.fit(None, output_a_np, epochs=1, steps_per_epoch=3)
+      _ = model.fit(None, output_a_np, epochs=1, steps_per_epoch=3)
 
       # test evaluate
-      out = model.evaluate(None,
-                           output_a_np, batch_size=10)
-      out = model.evaluate(None,
-                           output_a_np, batch_size=10)
+      _ = model.evaluate(None, output_a_np, steps=3)
+      _ = model.evaluate(None, output_a_np, steps=3)
 
       # test predict
       out = model.predict(None, steps=3)
@@ -1383,16 +1380,12 @@ class TestTrainingWithDataTensors(test.TestCase):
                                  output_a_np)
 
       # test fit
-      out = model.fit(None,
-                      output_a_np, epochs=1, batch_size=10)
-      out = model.fit(None,
-                      output_a_np, epochs=1, batch_size=10)
+      _ = model.fit(None, output_a_np, epochs=1, steps_per_epoch=10)
+      _ = model.fit(None, output_a_np, epochs=1, steps_per_epoch=10)
 
       # test evaluate
-      out = model.evaluate(None,
-                           output_a_np, batch_size=10)
-      out = model.evaluate(None,
-                           output_a_np, batch_size=10)
+      _ = model.evaluate(None, output_a_np, steps=10)
+      _ = model.evaluate(None, output_a_np, steps=10)
 
       # test predict
       out = model.predict(None, steps=3)
@@ -1715,40 +1708,56 @@ class TestTrainingWithDataTensors(test.TestCase):
 
 class TestTrainingWithDatasetIterators(test.TestCase):
 
+  @tf_test_util.run_in_graph_and_eager_modes()
   def test_training_and_eval_methods_on_iterators_single_io(self):
     with self.test_session():
       x = keras.layers.Input(shape=(3,), name='input')
       y = keras.layers.Dense(4, name='dense')(x)
       model = keras.Model(x, y)
 
-      optimizer = 'rmsprop'
+      optimizer = RMSPropOptimizer(learning_rate=0.001)
       loss = 'mse'
       metrics = ['mae']
       model.compile(optimizer, loss, metrics=metrics)
 
-      inputs = np.zeros((10, 3))
-      targets = np.zeros((10, 4))
+      inputs = np.zeros((10, 3), dtype=np.float32)
+      targets = np.zeros((10, 4), dtype=np.float32)
       dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets))
       dataset = dataset.repeat(100)
       dataset = dataset.batch(10)
       iterator = dataset.make_one_shot_iterator()
 
-      model.fit(iterator, epochs=1, steps_per_epoch=2, verbose=0)
-      model.evaluate(iterator, steps=2, verbose=0)
+      model.fit(iterator, epochs=1, steps_per_epoch=2, verbose=1)
+      model.evaluate(iterator, steps=2, verbose=1)
       model.predict(iterator, steps=2)
       model.train_on_batch(iterator)
       model.test_on_batch(iterator)
+      model.predict_on_batch(iterator)
+
       # Test with validation data
       model.fit(iterator,
                 epochs=1, steps_per_epoch=2, verbose=0,
                 validation_data=iterator, validation_steps=2)
       # Test with validation split
-      with self.assertRaisesRegexp(ValueError,
-                                   'you cannot use `validation_split`'):
+      with self.assertRaisesRegexp(
+          ValueError, '`validation_split` argument is not supported '
+          'when input `x` is a dataset iterator'):
         model.fit(iterator,
                   epochs=1, steps_per_epoch=2, verbose=0,
                   validation_split=0.5, validation_steps=2)
 
+      # Test with sample weight.
+      sample_weight = np.random.random((10,))
+      with self.assertRaisesRegexp(
+          ValueError, '`sample_weight` argument is not supported '
+          'when input `x` is a dataset iterator'):
+        model.fit(
+            iterator,
+            epochs=1,
+            steps_per_epoch=2,
+            verbose=0,
+            sample_weight=sample_weight)
+
       # Test invalid usage
       with self.assertRaisesRegexp(ValueError,
                                    'Instead, pass an `Iterator`'):
@@ -1759,19 +1768,54 @@ class TestTrainingWithDatasetIterators(test.TestCase):
         model.fit(iterator, iterator,
                   epochs=1, steps_per_epoch=2, verbose=0)
 
+      with self.assertRaisesRegexp(
+          ValueError, 'you should specify the `steps_per_epoch` argument'):
+        model.fit(iterator, epochs=1, verbose=0)
+      with self.assertRaisesRegexp(ValueError,
+                                   'you should specify the `steps` argument'):
+        model.evaluate(iterator, verbose=0)
+      with self.assertRaisesRegexp(ValueError,
+                                   'you should specify the `steps` argument'):
+        model.predict(iterator, verbose=0)
+
+  def test_get_next_op_created_once(self):
+    with self.test_session():
+      x = keras.layers.Input(shape=(3,), name='input')
+      y = keras.layers.Dense(4, name='dense')(x)
+      model = keras.Model(x, y)
+
+      optimizer = RMSPropOptimizer(learning_rate=0.001)
+      loss = 'mse'
+      metrics = ['mae']
+      model.compile(optimizer, loss, metrics=metrics)
+
+      inputs = np.zeros((10, 3), dtype=np.float32)
+      targets = np.zeros((10, 4), dtype=np.float32)
+      dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets))
+      dataset = dataset.repeat(100)
+      dataset = dataset.batch(10)
+      iterator = dataset.make_one_shot_iterator()
+
+      model.fit(iterator, epochs=1, steps_per_epoch=2, verbose=1)
+      # Finalize graph to make sure we are not appending another iterator
+      # get_next op in the graph.
+      ops.get_default_graph().finalize()
+      model.fit(iterator, epochs=1, steps_per_epoch=2, verbose=1)
+
+  @tf_test_util.run_in_graph_and_eager_modes()
   def test_iterators_running_out_of_data(self):
     with self.test_session():
       x = keras.layers.Input(shape=(3,), name='input')
       y = keras.layers.Dense(4, name='dense')(x)
       model = keras.Model(x, y)
 
-      optimizer = 'rmsprop'
+      optimizer = RMSPropOptimizer(learning_rate=0.001)
       loss = 'mse'
       metrics = ['mae']
       model.compile(optimizer, loss, metrics=metrics)
 
-      inputs = np.zeros((10, 3))
-      targets = np.zeros((10, 4))
+      inputs = np.zeros((10, 3), dtype=np.float32)
+      targets = np.zeros((10, 4), dtype=np.float32)
       dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets))
       dataset = dataset.repeat(2)
       dataset = dataset.batch(10)
diff --git a/tensorflow/python/keras/_impl/keras/engine/training_utils.py b/tensorflow/python/keras/_impl/keras/engine/training_utils.py
index 662938f421..04d80c891f 100644
--- a/tensorflow/python/keras/_impl/keras/engine/training_utils.py
+++ b/tensorflow/python/keras/_impl/keras/engine/training_utils.py
@@ -22,6 +22,7 @@ import copy
 
 import numpy as np
 
+from tensorflow.python.data.ops import iterator_ops
 from tensorflow.python.eager import context
 from tensorflow.python.framework import tensor_util
 from tensorflow.python.keras._impl.keras import backend as K
@@ -65,14 +66,7 @@ def check_num_samples(ins,
   if steps is not None and batch_size is not None:
     raise ValueError(
         'If ' + steps_name + ' is set, the `batch_size` must be None.')
-
-  if not ins or has_symbolic_tensors(ins):
-    if steps is None:
-      raise ValueError('If your data is in the form of symbolic tensors, '
-                       'you should specify the `' + steps_name + '` argument '
-                       '(instead of the `batch_size` argument, '
-                       'because symbolic tensors are expected to produce '
-                       'batches of input data).')
+  if check_steps_argument(ins, steps, steps_name):
     return None
   if hasattr(ins[0], 'shape'):
     return int(ins[0].shape[0])
@@ -551,8 +545,11 @@ def standardize_weights(y,
 
 
 def has_symbolic_tensors(ls):
-  return (any(tensor_util.is_tensor(v) for v in ls)
-          and not context.executing_eagerly())
+  if context.executing_eagerly():
+    return False
+  if isinstance(ls, (list, tuple)):
+    return any(tensor_util.is_tensor(v) for v in ls)
+  return tensor_util.is_tensor(ls)
 
 
 def populate_metric_names(model):
@@ -614,3 +611,77 @@ def add_metric_name(model, metric_name, index):
     metric_name = '%s_%d' % (base_metric_name, j)
     j += 1
   model.metrics_names.append(metric_name)
+
+
+def validate_iterator_input(x, y, sample_weight, validation_split=None):
+  """Validates user input arguments when a dataset iterator is passed.
+
+  Arguments:
+    x: Input data. A `tf.data` dataset iterator.
+    y: Target data. It could be either Numpy array(s) or TensorFlow tensor(s).
+        Expected to be `None` when `x` is a dataset iterator.
+    sample_weight: An optional sample-weight array passed by the user to
+        weight the importance of each sample in `x`. Expected to be `None` when
+        `x` is a dataset iterator
+    validation_split: Float between 0 and 1. Fraction of the training data to
+        be used as validation data. Expected to be `None` when `x` is a dataset
+        iterator.
+
+  Raises:
+    ValueError: if argument `y` or `sample_weight` or `validation_split` are
+        provided by user.
+  """
+  if y is not None:
+    raise ValueError('You passed a dataset iterator (%s) as input `x` to '
+                     'your model. In that case, you should not specify '
+                     'a target (`y`) argument, since the dataset iterator '
+                     'generates both input data and target data. '
+                     'Received: %s' % (x, y))
+  if sample_weight is not None:
+    raise ValueError('`sample_weight` argument is not supported when input'
+                     ' `x` is a dataset iterator. '
+                     'Received: x=%s, sample_weight=%s' % (x, sample_weight))
+  if validation_split is not None and validation_split != 0.0:
+    raise ValueError(
+        '`validation_split` argument is not supported when '
+        'input `x` is a dataset iterator. '
+        'Received: x=%s, validation_split=%f' % (x, validation_split))
+
+
+def check_steps_argument(input_data, steps, steps_name):
+  """Validates `steps` argument based on input data's type.
+
+  The cases when `steps` value must be provided are when
+    1. input data passed is an iterator.
+    2. model was built on top of symbolic tensors, input data is not
+       required and is `None`.
+    3. input data passed is a symbolic tensor.
+
+  Arguments:
+      input_data: Input data. Can be Numpy array(s) or TensorFlow tensor(s) or
+        tf.data.Dataset iterator or `None`.
+      steps: Integer or `None`. Total number of steps (batches of samples) to
+        execute.
+      steps_name: The public API's parameter name for `steps`.
+
+  Returns:
+    boolean, True if `steps` argument is required, else False.
+
+  Raises:
+      ValueError: if `steps` argument is required for given input data type
+        but not provided.
+  """
+
+  is_x_iterator = (
+      isinstance(input_data, iterator_ops.Iterator) or
+      isinstance(input_data, iterator_ops.EagerIterator))
+
+  if (input_data is None or is_x_iterator or has_symbolic_tensors(input_data) or
+      (isinstance(input_data, list) and not input_data)):
+    if steps is None:
+      input_type_str = 'iterators' if is_x_iterator else 'data tensors'
+      raise ValueError('When using {input_type} as input to a model, you should'
+                       ' specify the `{steps_name}` argument.'.format(
+                           input_type=input_type_str, steps_name=steps_name))
+    return True
+  return False
diff --git a/tensorflow/python/keras/_impl/keras/model_subclassing_test.py b/tensorflow/python/keras/_impl/keras/model_subclassing_test.py
index 3f850e57aa..1e88dc09fb 100644
--- a/tensorflow/python/keras/_impl/keras/model_subclassing_test.py
+++ b/tensorflow/python/keras/_impl/keras/model_subclassing_test.py
@@ -23,6 +23,7 @@ import os
 import numpy as np
 import six
 
+from tensorflow.python.data.ops import dataset_ops
 from tensorflow.python.eager import context
 from tensorflow.python.framework import tensor_shape
 from tensorflow.python.framework import test_util
@@ -250,6 +251,26 @@ class ModelSubclassingTest(test.TestCase):
       model.fit([x1, x2], [y1, y2], epochs=2, steps_per_epoch=10, verbose=0)
       _ = model.evaluate(steps=10, verbose=0)
 
+  @test_util.run_in_graph_and_eager_modes()
+  def test_single_io_workflow_with_dataset_iterators(self):
+    num_classes = 2
+    num_samples = 10
+    input_dim = 50
+
+    with self.test_session():
+      model = SimpleTestModel(num_classes=num_classes, use_dp=True, use_bn=True)
+      model.compile(loss='mse', optimizer=RMSPropOptimizer(learning_rate=0.001))
+
+      x = np.ones((num_samples, input_dim))
+      y = np.zeros((num_samples, num_classes))
+      dataset = dataset_ops.Dataset.from_tensor_slices((x, y))
+      dataset = dataset.repeat(100)
+      dataset = dataset.batch(10)
+      iterator = dataset.make_one_shot_iterator()
+
+      model.fit(iterator, epochs=2, steps_per_epoch=10, verbose=0)
+      _ = model.evaluate(iterator, steps=10, verbose=0)
+
   def test_multi_io_workflow_with_numpy_arrays_and_custom_placeholders(self):
 
     num_classes = (2, 3)
-- 
GitLab


From ba0584a5da7e0ff59486bd77b63eab417fbff352 Mon Sep 17 00:00:00 2001
From: Michael Case <mikecase@google.com>
Date: Mon, 7 May 2018 15:20:49 -0700
Subject: [PATCH 329/395] Fix TypeError in update_version.py

PiperOrigin-RevId: 195731183
---
 tensorflow/tools/ci_build/update_version.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/tools/ci_build/update_version.py b/tensorflow/tools/ci_build/update_version.py
index 9ddb219048..00bfcfd49b 100755
--- a/tensorflow/tools/ci_build/update_version.py
+++ b/tensorflow/tools/ci_build/update_version.py
@@ -250,7 +250,7 @@ def update_md_files(old_version, new_version):
 
   # Update any links to colab notebooks.
   def colab_url(version):
-    version_string = "%d.%d.%d" % (version.major, version.minor, version.patch)
+    version_string = "%s.%s.%s" % (version.major, version.minor, version.patch)
     prefix = "https://colab.research.google.com/github/tensorflow/models/blob/r"
     return prefix + version_string + "/"
 
-- 
GitLab


From 3a2f1cfb73fa6a21eba077485bdc08aa05646ad1 Mon Sep 17 00:00:00 2001
From: Michael Case <mikecase@google.com>
Date: Mon, 7 May 2018 15:24:02 -0700
Subject: [PATCH 330/395] Internal Change.

PiperOrigin-RevId: 195731675
---
 .../estimator/python/estimator/head.py        |  26 ++-
 .../contrib/tpu/python/tpu/tpu_estimator.py   |  26 +--
 tensorflow/python/estimator/canned/dnn.py     |  68 +++++-
 .../estimator/canned/dnn_testing_utils.py     |  20 +-
 tensorflow/python/estimator/canned/head.py    | 216 ++++++++++++++----
 .../python/estimator/canned/head_test.py      |  92 ++++++++
 tensorflow/python/estimator/model_fn.py       |  51 +++++
 7 files changed, 406 insertions(+), 93 deletions(-)

diff --git a/tensorflow/contrib/estimator/python/estimator/head.py b/tensorflow/contrib/estimator/python/estimator/head.py
index 5d19bf4714..109fdd3883 100644
--- a/tensorflow/contrib/estimator/python/estimator/head.py
+++ b/tensorflow/contrib/estimator/python/estimator/head.py
@@ -560,10 +560,10 @@ class _MultiLabelHead(head_lib._Head):  # pylint:disable=protected-access
         weights=weights,
         processed_labels=processed_labels)
 
-  def create_estimator_spec(
+  def _create_tpu_estimator_spec(
       self, features, mode, logits, labels=None, optimizer=None,
       train_op_fn=None, regularization_losses=None):
-    """Returns an `EstimatorSpec`.
+    """Returns an `model_fn._TPUEstimatorSpec`.
 
     Args:
       features: Input `dict` of `Tensor` or `SparseTensor` objects.
@@ -586,7 +586,7 @@ class _MultiLabelHead(head_lib._Head):  # pylint:disable=protected-access
         `loss_reduction=SUM_OVER_NONZERO_WEIGHTS` when creating the head to
         avoid scaling errors.
     Returns:
-      `EstimatorSpec`.
+      `model_fn._TPUEstimatorSpec`.
     Raises:
       ValueError: If both `train_op_fn` and `optimizer` are `None` in TRAIN
         mode, or if both are set.
@@ -606,7 +606,7 @@ class _MultiLabelHead(head_lib._Head):  # pylint:disable=protected-access
         classifier_output = head_lib._classification_output(  # pylint:disable=protected-access
             scores=probabilities, n_classes=self._n_classes,
             label_vocabulary=self._label_vocabulary)
-        return model_fn.EstimatorSpec(
+        return model_fn._TPUEstimatorSpec(  # pylint:disable=protected-access
             mode=model_fn.ModeKeys.PREDICT,
             predictions=predictions,
             export_outputs={
@@ -629,16 +629,18 @@ class _MultiLabelHead(head_lib._Head):  # pylint:disable=protected-access
 
       # Eval.
       if mode == model_fn.ModeKeys.EVAL:
-        return model_fn.EstimatorSpec(
+        return model_fn._TPUEstimatorSpec(  # pylint:disable=protected-access
             mode=model_fn.ModeKeys.EVAL,
             predictions=predictions,
             loss=regularized_training_loss,
-            eval_metric_ops=self._eval_metric_ops(
-                labels=processed_labels,
-                probabilities=probabilities,
-                weights=weights,
-                unreduced_loss=unreduced_loss,
-                regularization_loss=regularization_loss))
+            eval_metrics=head_lib._create_eval_metrics_tuple(  # pylint:disable=protected-access
+                self._eval_metric_ops, {
+                    'labels': processed_labels,
+                    'probabilities': probabilities,
+                    'weights': weights,
+                    'unreduced_loss': unreduced_loss,
+                    'regularization_loss': regularization_loss,
+                }))
 
       # Train.
       if optimizer is not None:
@@ -672,7 +674,7 @@ class _MultiLabelHead(head_lib._Head):  # pylint:disable=protected-access
         summary.scalar(
             head_lib._summary_key(self._name, keys.LOSS_REGULARIZATION),  # pylint:disable=protected-access
             regularization_loss)
-    return model_fn.EstimatorSpec(
+    return model_fn._TPUEstimatorSpec(  # pylint:disable=protected-access
         mode=model_fn.ModeKeys.TRAIN,
         predictions=predictions,
         loss=regularized_training_loss,
diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
index a69bfa9a20..a624eceed9 100644
--- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
+++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py
@@ -175,17 +175,7 @@ class _SIGNAL(object):
   STOP = -2
 
 
-class TPUEstimatorSpec(
-    collections.namedtuple('TPUEstimatorSpec', [
-        'mode',
-        'predictions',
-        'loss',
-        'train_op',
-        'eval_metrics',
-        'export_outputs',
-        'scaffold_fn',
-        'host_call'
-    ])):
+class TPUEstimatorSpec(model_fn_lib._TPUEstimatorSpec):  # pylint: disable=protected-access
   """Ops and objects returned from a `model_fn` and passed to `TPUEstimator`.
 
   See `EstimatorSpec` for `mode`, 'predictions, 'loss', 'train_op', and
@@ -1156,7 +1146,7 @@ class _ModelFnWrapper(object):
           self._call_model_fn(features, labels))
       loss, train_op = estimator_spec.loss, estimator_spec.train_op
 
-      if isinstance(estimator_spec, TPUEstimatorSpec):
+      if isinstance(estimator_spec, model_fn_lib._TPUEstimatorSpec):  # pylint: disable=protected-access
         captured_scaffold_fn.capture(estimator_spec.scaffold_fn)
       else:
         captured_scaffold_fn.capture(None)
@@ -1165,8 +1155,8 @@ class _ModelFnWrapper(object):
       # outfeed.
       with ops.control_dependencies([train_op]):
         host_call_outfeed_ops = []
-        if (isinstance(estimator_spec, TPUEstimatorSpec) and
-            estimator_spec.host_call is not None):
+        if (isinstance(estimator_spec, model_fn_lib._TPUEstimatorSpec)  # pylint: disable=protected-access
+            and estimator_spec.host_call is not None):
           host_call.record({'host_call': estimator_spec.host_call})
           host_call_outfeed_ops = host_call.create_enqueue_op()
         with ops.control_dependencies(host_call_outfeed_ops):
@@ -1209,7 +1199,7 @@ class _ModelFnWrapper(object):
       features, labels = inputs.features_and_labels()
 
       tpu_estimator_spec = self._call_model_fn(features, labels)
-      if not isinstance(tpu_estimator_spec, TPUEstimatorSpec):
+      if not isinstance(tpu_estimator_spec, model_fn_lib._TPUEstimatorSpec):  # pylint: disable=protected-access
         raise RuntimeError(
             'estimator_spec used by TPU evaluation must have type'
             '`TPUEstimatorSpec`. Got {}'.format(type(tpu_estimator_spec)))
@@ -1254,7 +1244,7 @@ class _ModelFnWrapper(object):
 
       tpu_estimator_spec = self._call_model_fn(
           features, labels, is_export_mode=False)
-      if not isinstance(tpu_estimator_spec, TPUEstimatorSpec):
+      if not isinstance(tpu_estimator_spec, model_fn_lib._TPUEstimatorSpec):  # pylint: disable=protected-access
         raise RuntimeError(
             'estimator_spec used by TPU prediction must have type'
             '`TPUEstimatorSpec`. Got {}'.format(type(tpu_estimator_spec)))
@@ -1316,7 +1306,7 @@ class _ModelFnWrapper(object):
 
     estimator_spec = self._model_fn(features=features, **kwargs)
     if (self._ctx.is_running_on_cpu(is_export_mode) and
-        isinstance(estimator_spec, TPUEstimatorSpec)):
+        isinstance(estimator_spec, model_fn_lib._TPUEstimatorSpec)):  # pylint: disable=protected-access
       # The estimator_spec will be passed to `Estimator` directly, which expects
       # type `EstimatorSpec`.
       return estimator_spec.as_estimator_spec()
@@ -1325,7 +1315,7 @@ class _ModelFnWrapper(object):
 
   def _verify_estimator_spec(self, estimator_spec):
     """Validates the estimator_spec."""
-    if isinstance(estimator_spec, TPUEstimatorSpec):
+    if isinstance(estimator_spec, model_fn_lib._TPUEstimatorSpec):  # pylint: disable=protected-access
       return estimator_spec
 
     err_msg = '{} returned by EstimatorSpec is not supported in TPUEstimator.'
diff --git a/tensorflow/python/estimator/canned/dnn.py b/tensorflow/python/estimator/canned/dnn.py
index 973a6ec747..e7fbf8eb72 100644
--- a/tensorflow/python/estimator/canned/dnn.py
+++ b/tensorflow/python/estimator/canned/dnn.py
@@ -151,6 +151,59 @@ def _dnn_model_fn(features,
   Returns:
     An `EstimatorSpec` instance.
 
+  Raises:
+    ValueError: If features has the wrong type.
+  """
+  tpu_estimator_spec = _tpu_dnn_model_fn(
+      features=features,
+      labels=labels,
+      mode=mode,
+      head=head,
+      hidden_units=hidden_units,
+      feature_columns=feature_columns,
+      optimizer=optimizer,
+      activation_fn=activation_fn,
+      dropout=dropout,
+      input_layer_partitioner=input_layer_partitioner,
+      config=config)
+  return tpu_estimator_spec.as_estimator_spec()
+
+
+def _tpu_dnn_model_fn(features,
+                      labels,
+                      mode,
+                      head,
+                      hidden_units,
+                      feature_columns,
+                      optimizer='Adagrad',
+                      activation_fn=nn.relu,
+                      dropout=None,
+                      input_layer_partitioner=None,
+                      config=None):
+  """Deep Neural Net model_fn for TPUEstimator.
+
+  Args:
+    features: dict of `Tensor`.
+    labels: `Tensor` of shape [batch_size, 1] or [batch_size] labels of
+      dtype `int32` or `int64` in the range `[0, n_classes)`.
+    mode: Defines whether this is training, evaluation or prediction.
+      See `ModeKeys`.
+    head: A `head_lib._Head` instance.
+    hidden_units: Iterable of integer number of hidden units per layer.
+    feature_columns: Iterable of `feature_column._FeatureColumn` model inputs.
+    optimizer: String, `tf.Optimizer` object, or callable that creates the
+      optimizer to use for training. If not specified, will use the Adagrad
+      optimizer with a default learning rate of 0.05.
+    activation_fn: Activation function applied to each layer.
+    dropout: When not `None`, the probability we will drop out a given
+      coordinate.
+    input_layer_partitioner: Partitioner for input layer. Defaults
+      to `min_max_variable_partitioner` with `min_slice_size` 64 << 20.
+    config: `RunConfig` object to configure the runtime settings.
+
+  Returns:
+    A `model_fn.TPUEstimatorSpec` instance.
+
   Raises:
     ValueError: If features has the wrong type.
   """
@@ -182,7 +235,7 @@ def _dnn_model_fn(features,
         input_layer_partitioner=input_layer_partitioner)
     logits = logit_fn(features=features, mode=mode)
 
-    return head.create_estimator_spec(
+    return head._create_tpu_estimator_spec(  # pylint: disable=protected-access
         features=features,
         mode=mode,
         labels=labels,
@@ -320,17 +373,8 @@ class DNNClassifier(estimator.Estimator):
       loss_reduction: One of `tf.losses.Reduction` except `NONE`. Describes how
         to reduce training loss over batch. Defaults to `SUM`.
     """
-    if n_classes == 2:
-      head = head_lib._binary_logistic_head_with_sigmoid_cross_entropy_loss(  # pylint: disable=protected-access
-          weight_column=weight_column,
-          label_vocabulary=label_vocabulary,
-          loss_reduction=loss_reduction)
-    else:
-      head = head_lib._multi_class_head_with_softmax_cross_entropy_loss(  # pylint: disable=protected-access
-          n_classes, weight_column=weight_column,
-          label_vocabulary=label_vocabulary,
-          loss_reduction=loss_reduction)
-
+    head = head_lib._binary_logistic_or_multi_class_head(  # pylint: disable=protected-access
+        n_classes, weight_column, label_vocabulary, loss_reduction)
     def _model_fn(features, labels, mode, config):
       """Call the defined shared _dnn_model_fn."""
       return _dnn_model_fn(
diff --git a/tensorflow/python/estimator/canned/dnn_testing_utils.py b/tensorflow/python/estimator/canned/dnn_testing_utils.py
index 62b13c3200..06a648777f 100644
--- a/tensorflow/python/estimator/canned/dnn_testing_utils.py
+++ b/tensorflow/python/estimator/canned/dnn_testing_utils.py
@@ -134,7 +134,7 @@ def mock_head(testcase, hidden_units, logits_dimension, expected_logits):
       hidden_weights_names + hidden_biases_names +
       [LOGITS_WEIGHTS_NAME + '/part_0:0', LOGITS_BIASES_NAME + '/part_0:0'])
 
-  def _create_estimator_spec(
+  def _create_tpu_estimator_spec(
       features, mode, logits, labels, train_op_fn=None, optimizer=None):
     del features, labels  # Not used.
     trainable_vars = ops.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES)
@@ -149,19 +149,29 @@ def mock_head(testcase, hidden_units, logits_dimension, expected_logits):
           train_op = train_op_fn(loss)
         elif optimizer is not None:
           train_op = optimizer.minimize(loss, global_step=None)
-        return model_fn.EstimatorSpec(
+        return model_fn._TPUEstimatorSpec(
             mode=mode, loss=loss, train_op=train_op)
       elif mode == model_fn.ModeKeys.EVAL:
-        return model_fn.EstimatorSpec(mode=mode, loss=array_ops.identity(loss))
+        return model_fn._TPUEstimatorSpec(
+            mode=mode, loss=array_ops.identity(loss))
       elif mode == model_fn.ModeKeys.PREDICT:
-        return model_fn.EstimatorSpec(
+        return model_fn._TPUEstimatorSpec(
             mode=mode, predictions={'logits': array_ops.identity(logits)})
       else:
         testcase.fail('Invalid mode: {}'.format(mode))
 
+  def _create_estimator_spec(
+      features, mode, logits, labels, train_op_fn=None, optimizer=None):
+    tpu_spec = _create_tpu_estimator_spec(
+        features, mode, logits, labels, train_op_fn, optimizer)
+    return tpu_spec.as_estimator_spec()
+
   head = test.mock.NonCallableMagicMock(spec=head_lib._Head)
   head.logits_dimension = logits_dimension
-  head.create_estimator_spec = test.mock.MagicMock(wraps=_create_estimator_spec)
+  head._create_tpu_estimator_spec = test.mock.MagicMock(
+      wraps=_create_tpu_estimator_spec)
+  head.create_estimator_spec = test.mock.MagicMock(
+      wraps=_create_estimator_spec)
 
   return head
 
diff --git a/tensorflow/python/estimator/canned/head.py b/tensorflow/python/estimator/canned/head.py
index 48f448d7f5..232637314d 100644
--- a/tensorflow/python/estimator/canned/head.py
+++ b/tensorflow/python/estimator/canned/head.py
@@ -32,6 +32,7 @@ from tensorflow.python.feature_column import feature_column as feature_column_li
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import check_ops
 from tensorflow.python.ops import control_flow_ops
@@ -69,6 +70,35 @@ def _summary_key(head_name, val):
   return '%s/%s' % (val, head_name) if head_name else val
 
 
+def _create_eval_metrics_tuple(fn, kwargs):
+  """Creates TPU eval metrics tuple.
+
+  Helper function to make eval_metric tuple (eval_metric_fn, fn_kwargs) used
+  by `TPUEstimator`. TPUEstimator requires that `eval_metric_fn` take
+  exclusively Tensor arguments. This helper can help create such a function from
+  a more generic function that can take both Tensor and non-Tensor arguments.
+
+  Args:
+    fn: A eval_metric_fn that takes both Tensor and non-Tensor arguments.
+        This function must return a dict of form
+        {'metric name': (metric_tensor, eval_op)}
+    kwargs: Dict of arguments for `fn`.
+
+  Returns:
+    `eval_metric` tuple that can be passed to a `model_fn._TPUEstimatorSpec`.
+  """
+  tensor_kwargs = {}
+  nontensor_kwargs = {}
+  for k, v in six.iteritems(kwargs):
+    if tensor_util.is_tensor(v):
+      tensor_kwargs[k] = v
+    else:
+      nontensor_kwargs[k] = v
+  def _fn(**tensors):
+    return fn(**dict(nontensor_kwargs, **tensors))
+  return (_fn, tensor_kwargs)
+
+
 class _Head(object):
   """Interface for the head/top of a model.
 
@@ -174,7 +204,6 @@ class _Head(object):
 
   # TODO(b/65403806): By default, collect regularization_losses from
   # GraphKeys.REGULARIZATION_LOSSES collection.
-  @abc.abstractmethod
   def create_estimator_spec(
       self, features, mode, logits, labels=None, optimizer=None,
       train_op_fn=None, regularization_losses=None):
@@ -203,7 +232,47 @@ class _Head(object):
     Returns:
       `EstimatorSpec`.
     """
-    raise NotImplementedError('Calling an abstract method.')
+    try:
+      tpu_estimator_spec = (
+          self._create_tpu_estimator_spec(
+              features, mode, logits, labels, optimizer, train_op_fn,
+              regularization_losses))
+      return tpu_estimator_spec.as_estimator_spec()
+    except NotImplementedError:
+      # Not all subclasses of _Head will have implemented
+      # _create_tpu_estimator_spec. If it is implemented, we can use it to
+      # create our `EstimatorSpec` here.
+      raise NotImplementedError(
+          'Subclasses of _Head must implement `create_estimator_spec()` or '
+          '_create_tpu_estimator_spec().')
+
+  def _create_tpu_estimator_spec(
+      self, features, mode, logits, labels=None, optimizer=None,
+      train_op_fn=None, regularization_losses=None):
+    """Returns `model_fn._TPUEstimatorSpec` that a model_fn can return.
+
+    Args:
+      features: Input `dict` of `Tensor` or `SparseTensor` objects.
+      mode: Estimator's `ModeKeys`.
+      logits: logits `Tensor` to be used by the head.
+      labels: Labels `Tensor`, or `dict` of same.
+      optimizer: `Optimizer` instance to optimize the loss in TRAIN mode.
+        Namely, sets `train_op = optimizer.minimize(loss, global_step)`, which
+        updates variables and increments `global_step`.
+      train_op_fn: Function that takes a scalar loss `Tensor` and returns an op
+        to optimize the model with the loss in TRAIN mode. Used if `optimizer`
+        is `None`. Exactly one of `train_op_fn` and `optimizer` must be set in
+        TRAIN mode. None is allowed in other modes. If you want to optimize loss
+        yourself you can pass `lambda _: tf.no_op()` and then use
+        EstimatorSpec.loss to compute and apply gradients.
+      regularization_losses: A list of additional scalar losses to be added to
+        the training loss, such as regularization losses.
+
+    Returns:
+      A `model_fn._TPUEstimatorSpec' instance.
+    """
+    raise NotImplementedError(
+        'TPUEstimatorSpec not available for this model head.')
 
 
 def _check_dense_labels_match_logits_and_reshape(
@@ -702,10 +771,10 @@ class _MultiClassHeadWithSoftmaxCrossEntropyLoss(_Head):
         weights=weights,
         processed_labels=label_ids)
 
-  def create_estimator_spec(
+  def _create_tpu_estimator_spec(
       self, features, mode, logits, labels=None, optimizer=None,
       train_op_fn=None, regularization_losses=None):
-    """Returns an `EstimatorSpec`.
+    """Returns a `model_fn._TPUEstimatorSpec`.
 
     Args:
       features: Input `dict` of `Tensor` or `SparseTensor` objects.
@@ -727,7 +796,7 @@ class _MultiClassHeadWithSoftmaxCrossEntropyLoss(_Head):
         `loss_reduction=SUM_OVER_NONZERO_WEIGHTS` when creating the head to
         avoid scaling errors.
     Returns:
-      `EstimatorSpec`.
+      A `model_fn._TPUEstimatorSpec` instance.
     Raises:
       ValueError: If both `train_op_fn` and `optimizer` are `None` in TRAIN
         mode, or if both are set.
@@ -761,7 +830,7 @@ class _MultiClassHeadWithSoftmaxCrossEntropyLoss(_Head):
         classifier_output = _classification_output(
             scores=probabilities, n_classes=self._n_classes,
             label_vocabulary=self._label_vocabulary)
-        return model_fn.EstimatorSpec(
+        return model_fn._TPUEstimatorSpec(  # pylint: disable=protected-access
             mode=model_fn.ModeKeys.PREDICT,
             predictions=predictions,
             export_outputs={
@@ -781,16 +850,17 @@ class _MultiClassHeadWithSoftmaxCrossEntropyLoss(_Head):
         regularized_training_loss = training_loss
       # Eval.
       if mode == model_fn.ModeKeys.EVAL:
-        return model_fn.EstimatorSpec(
+        return model_fn._TPUEstimatorSpec(  # pylint: disable=protected-access
             mode=model_fn.ModeKeys.EVAL,
             predictions=predictions,
             loss=regularized_training_loss,
-            eval_metric_ops=self._eval_metric_ops(
-                labels=label_ids,
-                class_ids=class_ids,
-                weights=weights,
-                unreduced_loss=unreduced_loss,
-                regularization_loss=regularization_loss))
+            eval_metrics=_create_eval_metrics_tuple(self._eval_metric_ops, {
+                'labels': label_ids,
+                'class_ids': class_ids,
+                'weights': weights,
+                'unreduced_loss': unreduced_loss,
+                'regularization_loss': regularization_loss
+            }))
 
       # Train.
       if optimizer is not None:
@@ -824,7 +894,7 @@ class _MultiClassHeadWithSoftmaxCrossEntropyLoss(_Head):
         summary.scalar(
             _summary_key(self._name, keys.LOSS_REGULARIZATION),
             regularization_loss)
-    return model_fn.EstimatorSpec(
+    return model_fn._TPUEstimatorSpec(  # pylint: disable=protected-access
         mode=model_fn.ModeKeys.TRAIN,
         predictions=predictions,
         loss=regularized_training_loss,
@@ -1060,7 +1130,7 @@ class _BinaryLogisticHeadWithSigmoidCrossEntropyLoss(_Head):
         weights=weights,
         processed_labels=labels)
 
-  def create_estimator_spec(
+  def _create_tpu_estimator_spec(
       self, features, mode, logits, labels=None, optimizer=None,
       train_op_fn=None, regularization_losses=None):
     """Returns an `EstimatorSpec`.
@@ -1122,7 +1192,7 @@ class _BinaryLogisticHeadWithSigmoidCrossEntropyLoss(_Head):
         classifier_output = _classification_output(
             scores=probabilities, n_classes=2,
             label_vocabulary=self._label_vocabulary)
-        return model_fn.EstimatorSpec(
+        return model_fn._TPUEstimatorSpec(  # pylint: disable=protected-access
             mode=model_fn.ModeKeys.PREDICT,
             predictions=predictions,
             export_outputs={
@@ -1146,18 +1216,22 @@ class _BinaryLogisticHeadWithSigmoidCrossEntropyLoss(_Head):
 
       # Eval.
       if mode == model_fn.ModeKeys.EVAL:
-        return model_fn.EstimatorSpec(
+        return model_fn._TPUEstimatorSpec(  # pylint: disable=protected-access
             mode=model_fn.ModeKeys.EVAL,
             predictions=predictions,
             loss=regularized_training_loss,
-            eval_metric_ops=self._eval_metric_ops(
-                labels=processed_labels,
-                logits=logits,
-                logistic=logistic,
-                class_ids=class_ids,
-                weights=weights,
-                unreduced_loss=unreduced_loss,
-                regularization_loss=regularization_loss))
+            eval_metrics=_create_eval_metrics_tuple(
+                self._eval_metric_ops,
+                {
+                    'labels': processed_labels,
+                    'logits': logits,
+                    'logistic': logistic,
+                    'class_ids': class_ids,
+                    'weights': weights,
+                    'unreduced_loss': unreduced_loss,
+                    'regularization_loss': regularization_loss
+                }
+            ))
 
       # Train.
       if optimizer is not None:
@@ -1190,7 +1264,7 @@ class _BinaryLogisticHeadWithSigmoidCrossEntropyLoss(_Head):
         summary.scalar(
             _summary_key(self._name, keys.LOSS_REGULARIZATION),
             regularization_loss)
-    return model_fn.EstimatorSpec(
+    return model_fn._TPUEstimatorSpec(  # pylint: disable=protected-access
         mode=model_fn.ModeKeys.TRAIN,
         predictions=predictions,
         loss=regularized_training_loss,
@@ -1322,7 +1396,25 @@ class _RegressionHeadWithMeanSquaredErrorLoss(_Head):
         weights=weights,
         processed_labels=labels)
 
-  def create_estimator_spec(
+  def _eval_metric_ops(self, weights, unreduced_loss, regularization_loss):
+    """Returns the Eval metric ops."""
+    keys = metric_keys.MetricKeys
+    # Estimator already adds a metric for loss.
+    eval_metric_ops = {
+        _summary_key(self._name, keys.LOSS_MEAN):
+            metrics_lib.mean(
+                values=unreduced_loss,
+                weights=weights)
+    }
+    if regularization_loss is not None:
+      regularization_loss_key = _summary_key(
+          self._name, keys.LOSS_REGULARIZATION)
+      eval_metric_ops[regularization_loss_key] = metrics_lib.mean(
+          values=regularization_loss,
+          name=keys.LOSS_REGULARIZATION)
+    return eval_metric_ops
+
+  def _create_tpu_estimator_spec(
       self, features, mode, logits, labels=None, optimizer=None,
       train_op_fn=None, regularization_losses=None):
     """Returns an `EstimatorSpec`.
@@ -1348,7 +1440,7 @@ class _RegressionHeadWithMeanSquaredErrorLoss(_Head):
         `loss_reduction=SUM_OVER_NONZERO_WEIGHTS` when creating the head to
         avoid scaling errors.
     Returns:
-      `EstimatorSpec`.
+      A `model_fn._TPUEstimatorSpec` instance.
     Raises:
       ValueError: If both `train_op_fn` and `optimizer` are `None` in TRAIN
         mode, or if both are set.
@@ -1369,7 +1461,7 @@ class _RegressionHeadWithMeanSquaredErrorLoss(_Head):
       if mode == model_fn.ModeKeys.PREDICT:
         regression_output = export_output.RegressionOutput(
             value=predicted_value)
-        return model_fn.EstimatorSpec(
+        return model_fn._TPUEstimatorSpec(  # pylint: disable=protected-access
             mode=model_fn.ModeKeys.PREDICT,
             predictions=predictions,
             export_outputs={
@@ -1390,25 +1482,18 @@ class _RegressionHeadWithMeanSquaredErrorLoss(_Head):
 
       # Eval.
       if mode == model_fn.ModeKeys.EVAL:
-        keys = metric_keys.MetricKeys
-        # Estimator already adds a metric for loss.
-        eval_metric_ops = {
-            _summary_key(self._name, keys.LOSS_MEAN):
-                metrics_lib.mean(
-                    values=unreduced_loss,
-                    weights=weights)
-        }
-        if regularization_loss is not None:
-          regularization_loss_key = _summary_key(
-              self._name, keys.LOSS_REGULARIZATION)
-          eval_metric_ops[regularization_loss_key] = metrics_lib.mean(
-              values=regularization_loss,
-              name=keys.LOSS_REGULARIZATION)
-        return model_fn.EstimatorSpec(
+        return model_fn._TPUEstimatorSpec(  # pylint: disable=protected-access
             mode=model_fn.ModeKeys.EVAL,
             predictions=predictions,
             loss=regularized_training_loss,
-            eval_metric_ops=eval_metric_ops)
+            eval_metrics=_create_eval_metrics_tuple(
+                self._eval_metric_ops,
+                {
+                    'weights': weights,
+                    'unreduced_loss': unreduced_loss,
+                    'regularization_loss': regularization_loss,
+                }
+            ))
 
       # Train.
       if optimizer is not None:
@@ -1441,7 +1526,7 @@ class _RegressionHeadWithMeanSquaredErrorLoss(_Head):
         summary.scalar(
             _summary_key(self._name, keys.LOSS_REGULARIZATION),
             regularization_loss)
-    return model_fn.EstimatorSpec(
+    return model_fn._TPUEstimatorSpec(  # pylint: disable=protected-access
         mode=model_fn.ModeKeys.TRAIN,
         predictions=predictions,
         loss=regularized_training_loss,
@@ -1478,3 +1563,42 @@ def _weights(features, weight_column):
       raise ValueError('Weight column should be castable to float. '
                        'Given dtype: {}'.format(weights.dtype))
     return math_ops.to_float(weights, name='weights')
+
+
+def _binary_logistic_or_multi_class_head(
+    n_classes, weight_column, label_vocabulary, loss_reduction):
+  """Creates either binary or multi-class head.
+
+  Args:
+    n_classes: Number of label classes.
+    weight_column: A string or a `_NumericColumn` created by
+      `tf.feature_column.numeric_column` defining feature column representing
+      weights. It is used to down weight or boost examples during training. It
+      will be multiplied by the loss of the example. If it is a string, it is
+      used as a key to fetch weight tensor from the `features`. If it is a
+      `_NumericColumn`, raw tensor is fetched by key `weight_column.key`,
+      then weight_column.normalizer_fn is applied on it to get weight tensor.
+    label_vocabulary: A list of strings represents possible label values. If
+      given, labels must be string type and have any value in
+      `label_vocabulary`. If it is not given, that means labels are
+      already encoded as integer or float within [0, 1] for `n_classes=2` and
+      encoded as integer values in {0, 1,..., n_classes-1} for `n_classes`>2 .
+      Also there will be errors if vocabulary is not provided and labels are
+      string.
+    loss_reduction: One of `tf.losses.Reduction` except `NONE`. Describes how
+      to reduce training loss over batch. Defaults to `SUM`.
+
+  Returns:
+    `head._Head` instance.
+  """
+  if n_classes == 2:
+    head = _binary_logistic_head_with_sigmoid_cross_entropy_loss(
+        weight_column=weight_column,
+        label_vocabulary=label_vocabulary,
+        loss_reduction=loss_reduction)
+  else:
+    head = _multi_class_head_with_softmax_cross_entropy_loss(
+        n_classes, weight_column=weight_column,
+        label_vocabulary=label_vocabulary,
+        loss_reduction=loss_reduction)
+  return head
diff --git a/tensorflow/python/estimator/canned/head_test.py b/tensorflow/python/estimator/canned/head_test.py
index 32a6339936..ecca3e8b0d 100644
--- a/tensorflow/python/estimator/canned/head_test.py
+++ b/tensorflow/python/estimator/canned/head_test.py
@@ -86,6 +86,98 @@ def _sigmoid(logits):
   return 1 / (1 + np.exp(-logits))
 
 
+class CreateEstimatorSpecTest(test.TestCase):
+
+  class _HeadWithTPUSupport(head_lib._Head):
+    """Head that overrides _create_tpu_estimator_spec."""
+
+    def name(self):
+      return 'HeadWithTPUSupport'
+
+    def logits_dimension(self):
+      return None
+
+    def create_loss(self, features, mode, logits, labels):
+      return None
+
+    def _create_tpu_estimator_spec(self, features, mode, logits, labels=None,
+                                   optimizer=None, train_op_fn=None,
+                                   regularization_losses=None):
+      return model_fn._TPUEstimatorSpec(
+          mode=model_fn.ModeKeys.EVAL,
+          loss=constant_op.constant(0.0, dtype=dtypes.float32))
+
+  class _HeadWithOutTPUSupport(head_lib._Head):
+    """Head that overrides create_estimator_spec."""
+
+    def name(self):
+      return 'HeadWithOutTPUSupport'
+
+    def logits_dimension(self):
+      return None
+
+    def create_loss(self, features, mode, logits, labels):
+      return None
+
+    def create_estimator_spec(self, features, mode, logits, labels=None,
+                              optimizer=None, train_op_fn=None,
+                              regularization_losses=None):
+      return model_fn.EstimatorSpec(
+          mode=model_fn.ModeKeys.EVAL,
+          loss=constant_op.constant(0.0, dtype=dtypes.float32))
+
+  class _InvalidHead(head_lib._Head):
+    """Head that overrides neither estimator_spec functions."""
+
+    def name(self):
+      return 'InvalidHead'
+
+    def logits_dimension(self):
+      return None
+
+    def create_loss(self, features, mode, logits, labels):
+      return None
+
+  def test_head_override_tpu_estimator_spec(self):
+    """Test for `_Head` that overrides _create_tpu_estimator_spec."""
+    head = self._HeadWithTPUSupport()
+
+    tpu_spec = head._create_tpu_estimator_spec(
+        features=None, mode=None, logits=None)
+    self.assertTrue(isinstance(tpu_spec, model_fn._TPUEstimatorSpec))
+    est_spec = head.create_estimator_spec(
+        features=None, mode=None, logits=None)
+    self.assertTrue(isinstance(est_spec, model_fn.EstimatorSpec))
+
+  def test_head_override_estimator_spec(self):
+    """Test for `_Head` that overrides create_estimator_spec."""
+    head = self._HeadWithOutTPUSupport()
+
+    with self.assertRaisesRegexp(
+        NotImplementedError,
+        'TPUEstimatorSpec not available for this model head.'):
+      _ = head._create_tpu_estimator_spec(
+          features=None, mode=None, logits=None)
+    est_spec = head.create_estimator_spec(
+        features=None, mode=None, logits=None)
+    self.assertTrue(isinstance(est_spec, model_fn.EstimatorSpec))
+
+  def test_invalid_head_class(self):
+    head = self._InvalidHead()
+
+    with self.assertRaisesRegexp(
+        NotImplementedError,
+        'TPUEstimatorSpec not available for this model head.'):
+      _ = head._create_tpu_estimator_spec(
+          features=None, mode=None, logits=None)
+    with self.assertRaisesRegexp(
+        NotImplementedError,
+        r'Subclasses of _Head must implement `create_estimator_spec\(\)` or '
+        r'_create_tpu_estimator_spec\(\).'):
+      _ = head.create_estimator_spec(
+          features=None, mode=None, logits=None)
+
+
 class MultiClassHeadWithSoftmaxCrossEntropyLoss(test.TestCase):
 
   def setUp(self):
diff --git a/tensorflow/python/estimator/model_fn.py b/tensorflow/python/estimator/model_fn.py
index 4ab2578769..3edf9fe940 100644
--- a/tensorflow/python/estimator/model_fn.py
+++ b/tensorflow/python/estimator/model_fn.py
@@ -334,6 +334,57 @@ class EstimatorSpec(
     return EstimatorSpec(*new_fields)
 
 
+class _TPUEstimatorSpec(collections.namedtuple('TPUEstimatorSpec', [
+    'mode',
+    'predictions',
+    'loss',
+    'train_op',
+    'eval_metrics',
+    'export_outputs',
+    'scaffold_fn',
+    'host_call'])):
+  """Ops and objects returned from a `model_fn` and passed to `TPUEstimator`.
+
+  This is a simplified implementation of `tf.contrib.tpu.EstimatorSpec`. See
+  tensorflow/contrib/tpu/python/tpu/tpu_estimator.py for more detailed
+  documentation.
+  """
+
+  def __new__(cls,
+              mode,
+              predictions=None,
+              loss=None,
+              train_op=None,
+              eval_metrics=None,
+              export_outputs=None,
+              scaffold_fn=None,
+              host_call=None):
+    """Creates a `_TPUEstimatorSpec` instance."""
+    return super(_TPUEstimatorSpec, cls).__new__(cls,
+                                                 mode=mode,
+                                                 predictions=predictions,
+                                                 loss=loss,
+                                                 train_op=train_op,
+                                                 eval_metrics=eval_metrics,
+                                                 export_outputs=export_outputs,
+                                                 scaffold_fn=scaffold_fn,
+                                                 host_call=host_call)
+
+  def as_estimator_spec(self):
+    """Creates an equivalent `EstimatorSpec` used by CPU train/eval."""
+    if not self.eval_metrics:
+      eval_metric_ops = None
+    else:
+      metric_fn, tensors = self.eval_metrics
+      eval_metric_ops = metric_fn(**tensors)
+    return EstimatorSpec(mode=self.mode,
+                         predictions=self.predictions,
+                         loss=self.loss,
+                         train_op=self.train_op,
+                         eval_metric_ops=eval_metric_ops,
+                         export_outputs=self.export_outputs)
+
+
 def _check_is_tensor_or_operation(x, name):
   if not (isinstance(x, ops.Operation) or isinstance(x, ops.Tensor)):
     raise TypeError('{} must be Operation or Tensor, given: {}'.format(name, x))
-- 
GitLab


From fc7f0b296dd53d1b72af21d36d36b6bcc5291ea7 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 15:41:22 -0700
Subject: [PATCH 331/395] Add support for select (via tf.where) support to
 tflite.

PiperOrigin-RevId: 195734246
---
 tensorflow/contrib/lite/builtin_ops.h         |   1 +
 .../lite/g3doc/tf_ops_compatibility.md        |  14 ++
 tensorflow/contrib/lite/kernels/BUILD         |  18 +++
 .../internal/optimized/optimized_ops.h        |  53 +++++++
 .../internal/reference/reference_ops.h        |  52 +++++++
 tensorflow/contrib/lite/kernels/register.cc   |   2 +
 tensorflow/contrib/lite/kernels/select.cc     | 125 +++++++++++++++
 .../contrib/lite/kernels/select_test.cc       | 143 ++++++++++++++++++
 tensorflow/contrib/lite/model.cc              |   3 +-
 tensorflow/contrib/lite/nnapi_delegate.cc     |   1 +
 tensorflow/contrib/lite/schema/schema.fbs     |   5 +
 .../contrib/lite/schema/schema_generated.h    | 124 ++++++++++++++-
 tensorflow/contrib/lite/testing/BUILD         |   1 +
 .../contrib/lite/testing/generate_examples.py |  33 +++-
 .../testing/generated_examples_zip_test.cc    |   1 +
 .../contrib/lite/toco/export_tensorflow.cc    |  16 ++
 .../propagate_array_data_types.cc             |  11 ++
 .../propagate_fixed_sizes.cc                  |  19 ++-
 .../contrib/lite/toco/import_tensorflow.cc    |  15 ++
 tensorflow/contrib/lite/toco/model.h          |  13 ++
 .../contrib/lite/toco/tflite/operator.cc      |   2 +
 .../contrib/lite/toco/tflite/operator_test.cc |   1 +
 tensorflow/contrib/lite/toco/tooling_util.cc  |   3 +
 tensorflow/contrib/lite/toco/types.proto      |   3 +
 24 files changed, 650 insertions(+), 9 deletions(-)
 create mode 100644 tensorflow/contrib/lite/kernels/select.cc
 create mode 100644 tensorflow/contrib/lite/kernels/select_test.cc
 mode change 100755 => 100644 tensorflow/contrib/lite/schema/schema_generated.h

diff --git a/tensorflow/contrib/lite/builtin_ops.h b/tensorflow/contrib/lite/builtin_ops.h
index 778933f569..a038acf284 100644
--- a/tensorflow/contrib/lite/builtin_ops.h
+++ b/tensorflow/contrib/lite/builtin_ops.h
@@ -89,6 +89,7 @@ typedef enum {
   kTfLiteBuiltinGreater = 61,
   kTfLiteBuiltinGreaterEqual = 62,
   kTfLiteBuiltinLessEqual = 63,
+  kTfLiteBuiltinSelect = 64,
 } TfLiteBuiltinOperator;
 
 #ifdef __cplusplus
diff --git a/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md b/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md
index fc57b8f28b..f45fcceb2e 100644
--- a/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md
+++ b/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md
@@ -639,6 +639,20 @@ Outputs {
 }
 ```
 
+**SELECT**
+
+```
+Inputs {
+  0: tensor
+  1: tensor
+  2: tensor
+}
+Outputs {
+  0: tensor that contains the elementwise values of 'tensor 1' if the
+  corresponding value of 'tensor 0' is true or the value of 'tensor 2' if false.
+}
+```
+
 And these are TensorFlow Lite operations that are present but not ready for
 custom models yet:
 
diff --git a/tensorflow/contrib/lite/kernels/BUILD b/tensorflow/contrib/lite/kernels/BUILD
index feab18b5c2..79e3c9f266 100644
--- a/tensorflow/contrib/lite/kernels/BUILD
+++ b/tensorflow/contrib/lite/kernels/BUILD
@@ -164,6 +164,7 @@ cc_library(
         "register.cc",
         "reshape.cc",
         "resize_bilinear.cc",
+        "select.cc",
         "skip_gram.cc",
         "space_to_batch_nd.cc",
         "space_to_depth.cc",
@@ -870,6 +871,23 @@ tf_cc_test(
     ],
 )
 
+tf_cc_test(
+    name = "select_test",
+    size = "small",
+    srcs = [
+        "select_test.cc",
+    ],
+    tags = [
+        "tflite_not_portable_ios",
+    ],
+    deps = [
+        ":builtin_ops",
+        "//tensorflow/contrib/lite:framework",
+        "//tensorflow/contrib/lite/kernels:test_util",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
 filegroup(
     name = "all_files",
     srcs = glob(
diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
index c506c5636c..8ab6f19b71 100644
--- a/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h
@@ -6318,6 +6318,59 @@ inline void TransposeConv(const float* input_data, const Dims<4>& input_dims,
   }
 }
 
+// UNOPTIMIZED COPY of Select from reference_ops.h.
+template <typename D, typename T>
+inline void Select(const D* input_condition_data,
+                   const Dims<4>& input_condition_dims, const T* input_x_data,
+                   const Dims<4>& input_x_dims, const T* input_y_data,
+                   const Dims<4>& input_y_dims, T* output_data,
+                   const Dims<4>& output_dims) {
+  const int64_t batches =
+      MatchingArraySize(input_condition_dims, 3, input_x_dims, 3, input_y_dims,
+                        3, output_dims, 3);
+  const int64_t height =
+      MatchingArraySize(input_condition_dims, 2, input_x_dims, 2, input_y_dims,
+                        2, output_dims, 2);
+  const int64_t width = MatchingArraySize(input_condition_dims, 1, input_x_dims,
+                                          1, input_y_dims, 1, output_dims, 1);
+  const int64_t depth = MatchingArraySize(input_condition_dims, 0, input_x_dims,
+                                          0, input_y_dims, 0, output_dims, 0);
+
+  const int64_t num_elements = batches * height * width * depth;
+  for (int64_t i = 0; i < num_elements; ++i) {
+    output_data[i] =
+        input_condition_data[i] ? input_x_data[i] : input_y_data[i];
+  }
+}
+
+// UNOPTIMIZED COPY of RankOneSelect from reference_ops.h.
+template <typename D, typename T>
+inline void RankOneSelect(const D* input_condition_data,
+                          const Dims<4>& input_condition_dims,
+                          const T* input_x_data, const Dims<4>& input_x_dims,
+                          const T* input_y_data, const Dims<4>& input_y_dims,
+                          T* output_data, const Dims<4>& output_dims) {
+  const int64_t rank = ArraySize(input_condition_dims, 0);
+
+  const int64_t batches =
+      MatchingArraySize(input_x_dims, 3, input_y_dims, 3, output_dims, 3);
+  const int64_t height =
+      MatchingArraySize(input_x_dims, 2, input_y_dims, 2, output_dims, 2);
+  const int64_t width =
+      MatchingArraySize(input_x_dims, 1, input_y_dims, 1, output_dims, 1);
+  const int64_t depth =
+      MatchingArraySize(input_x_dims, 0, input_y_dims, 0, output_dims, 0);
+
+  TFLITE_DCHECK_EQ(rank, batches);
+
+  int64_t offset = 0;
+  int64_t size = depth * height * width;
+  for (int64_t i = 0; i < rank; i++) {
+    const T* input_data = input_condition_data[i] ? input_x_data : input_y_data;
+    memcpy(output_data + offset, input_data + offset, size * sizeof(T));
+  }
+}
+
 }  // namespace optimized_ops
 }  // namespace tflite
 
diff --git a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
index 93dba1cc8e..c3aff1093f 100644
--- a/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
+++ b/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
@@ -3759,6 +3759,58 @@ TFLITE_COMPARISON_OP(Less);
 TFLITE_COMPARISON_OP(LessEqual);
 #undef TFLITE_COMPARISON_OP
 
+template <typename D, typename T>
+inline void Select(const D* input_condition_data,
+                   const Dims<4>& input_condition_dims, const T* input_x_data,
+                   const Dims<4>& input_x_dims, const T* input_y_data,
+                   const Dims<4>& input_y_dims, T* output_data,
+                   const Dims<4>& output_dims) {
+  const int64_t batches =
+      MatchingArraySize(input_condition_dims, 3, input_x_dims, 3, input_y_dims,
+                        3, output_dims, 3);
+  const int64_t height =
+      MatchingArraySize(input_condition_dims, 2, input_x_dims, 2, input_y_dims,
+                        2, output_dims, 2);
+  const int64_t width = MatchingArraySize(input_condition_dims, 1, input_x_dims,
+                                          1, input_y_dims, 1, output_dims, 1);
+  const int64_t depth = MatchingArraySize(input_condition_dims, 0, input_x_dims,
+                                          0, input_y_dims, 0, output_dims, 0);
+
+  const int64_t num_elements = batches * height * width * depth;
+  for (int64_t i = 0; i < num_elements; ++i) {
+    output_data[i] =
+        input_condition_data[i] ? input_x_data[i] : input_y_data[i];
+  }
+}
+
+template <typename D, typename T>
+inline void RankOneSelect(const D* input_condition_data,
+                          const Dims<4>& input_condition_dims,
+                          const T* input_x_data, const Dims<4>& input_x_dims,
+                          const T* input_y_data, const Dims<4>& input_y_dims,
+                          T* output_data, const Dims<4>& output_dims) {
+  const int64_t rank = ArraySize(input_condition_dims, 0);
+
+  const int64_t batches =
+      MatchingArraySize(input_x_dims, 3, input_y_dims, 3, output_dims, 3);
+  const int64_t height =
+      MatchingArraySize(input_x_dims, 2, input_y_dims, 2, output_dims, 2);
+  const int64_t width =
+      MatchingArraySize(input_x_dims, 1, input_y_dims, 1, output_dims, 1);
+  const int64_t depth =
+      MatchingArraySize(input_x_dims, 0, input_y_dims, 0, output_dims, 0);
+
+  TFLITE_DCHECK_EQ(rank, batches);
+
+  int64_t offset = 0;
+  int64_t size = depth * height * width;
+  for (int64_t i = 0; i < rank; i++) {
+    const T* input_data = input_condition_data[i] ? input_x_data : input_y_data;
+    memcpy(output_data + offset, input_data + offset, size * sizeof(T));
+    offset += size;
+  }
+}
+
 }  // namespace reference_ops
 }  // namespace tflite
 
diff --git a/tensorflow/contrib/lite/kernels/register.cc b/tensorflow/contrib/lite/kernels/register.cc
index 40855891a6..5df35aac62 100644
--- a/tensorflow/contrib/lite/kernels/register.cc
+++ b/tensorflow/contrib/lite/kernels/register.cc
@@ -86,6 +86,7 @@ TfLiteRegistration* Register_LESS();
 TfLiteRegistration* Register_LESS_EQUAL();
 TfLiteRegistration* Register_FLOOR();
 TfLiteRegistration* Register_NEG();
+TfLiteRegistration* Register_SELECT();
 
 BuiltinOpResolver::BuiltinOpResolver() {
   AddBuiltin(BuiltinOperator_RELU, Register_RELU());
@@ -153,6 +154,7 @@ BuiltinOpResolver::BuiltinOpResolver() {
   AddBuiltin(BuiltinOperator_LESS_EQUAL, Register_LESS_EQUAL());
   AddBuiltin(BuiltinOperator_FLOOR, Register_FLOOR());
   AddBuiltin(BuiltinOperator_NEG, Register_NEG());
+  AddBuiltin(BuiltinOperator_SELECT, Register_SELECT());
 
   // TODO(andrewharp, ahentz): Move these somewhere more appropriate so that
   // custom ops aren't always included by default.
diff --git a/tensorflow/contrib/lite/kernels/select.cc b/tensorflow/contrib/lite/kernels/select.cc
new file mode 100644
index 0000000000..029ad9a709
--- /dev/null
+++ b/tensorflow/contrib/lite/kernels/select.cc
@@ -0,0 +1,125 @@
+/* Copyright 2018 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.
+==============================================================================*/
+#include "tensorflow/contrib/lite/context.h"
+#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h"
+#include "tensorflow/contrib/lite/kernels/internal/tensor.h"
+#include "tensorflow/contrib/lite/kernels/kernel_util.h"
+#include "tensorflow/contrib/lite/kernels/op_macros.h"
+#include "tensorflow/contrib/lite/string_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace select {
+
+constexpr int kInputTensorCondition = 0;
+constexpr int kInputTensorX = 1;
+constexpr int kInputTensorY = 2;
+constexpr int kOutputTensor = 0;
+
+TfLiteStatus SelectPrepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 3);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  TfLiteTensor* input_condition =
+      GetInput(context, node, kInputTensorCondition);
+  TfLiteTensor* input_x = GetInput(context, node, kInputTensorX);
+  TfLiteTensor* input_y = GetInput(context, node, kInputTensorY);
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+
+  // Input must be bool.
+  TF_LITE_ENSURE(context, input_condition->type == kTfLiteBool);
+
+  // Input tensors must have the same type and size
+  TF_LITE_ENSURE_EQ(context, input_x->type, input_y->type);
+  TF_LITE_ENSURE(context, HaveSameShapes(input_x, input_y));
+  output->type = input_x->type;
+
+  // Either the same shape, or input_condition must be Rank 1 and match over the
+  // first dimension.
+  bool same_shape = HaveSameShapes(input_condition, input_x);
+  if (!same_shape && NumDimensions(input_condition) == 1) {
+    same_shape =
+        SizeOfDimension(input_condition, 0) == SizeOfDimension(input_x, 0);
+  }
+
+  TF_LITE_ENSURE(context, same_shape);
+
+  TfLiteIntArray* output_size = TfLiteIntArrayCopy(input_x->dims);
+  return context->ResizeTensor(context, output, output_size);
+}
+
+TfLiteStatus SelectEval(TfLiteContext* context, TfLiteNode* node) {
+  TfLiteTensor* input_condition =
+      GetInput(context, node, kInputTensorCondition);
+  TfLiteTensor* input_x = GetInput(context, node, kInputTensorX);
+  TfLiteTensor* input_y = GetInput(context, node, kInputTensorY);
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+
+  bool is_rank_one = !HaveSameShapes(input_condition, input_x);
+
+#define TF_LITE_SELECT(type, op)                                          \
+  reference_ops::op(GetTensorData<bool>(input_condition),                 \
+                    GetTensorDims(input_condition),                       \
+                    GetTensorData<type>(input_x), GetTensorDims(input_x), \
+                    GetTensorData<type>(input_y), GetTensorDims(input_y), \
+                    GetTensorData<type>(output), GetTensorDims(output));
+
+#define TF_LITE_SWITCH(type, op)                                               \
+  switch (type) {                                                              \
+    break;                                                                     \
+    case kTfLiteBool:                                                          \
+      TF_LITE_SELECT(bool, op);                                                \
+      break;                                                                   \
+    case kTfLiteFloat32:                                                       \
+      TF_LITE_SELECT(float, op);                                               \
+      break;                                                                   \
+    case kTfLiteUInt8:                                                         \
+      TF_LITE_SELECT(uint8_t, op);                                             \
+      break;                                                                   \
+    case kTfLiteInt32:                                                         \
+      TF_LITE_SELECT(int32_t, op);                                             \
+      break;                                                                   \
+    case kTfLiteInt64:                                                         \
+      TF_LITE_SELECT(int64_t, op);                                             \
+      break;                                                                   \
+    default:                                                                   \
+      context->ReportError(context,                                            \
+                           "Does not support type other than bool|float|int"); \
+      return kTfLiteError;                                                     \
+  }
+
+  if (is_rank_one) {
+    TF_LITE_SWITCH(input_x->type, RankOneSelect);
+  } else {
+    TF_LITE_SWITCH(input_x->type, Select);
+  }
+
+#undef TF_LITE_SELECT
+#undef TF_LITE_SWITCH
+  return kTfLiteOk;
+}
+
+}  // namespace select
+
+TfLiteRegistration* Register_SELECT() {
+  static TfLiteRegistration r = {nullptr, nullptr, select::SelectPrepare,
+                                 select::SelectEval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite
diff --git a/tensorflow/contrib/lite/kernels/select_test.cc b/tensorflow/contrib/lite/kernels/select_test.cc
new file mode 100644
index 0000000000..cfe24a5fc9
--- /dev/null
+++ b/tensorflow/contrib/lite/kernels/select_test.cc
@@ -0,0 +1,143 @@
+/* Copyright 2018 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.
+==============================================================================*/
+#include <gtest/gtest.h>
+#include "tensorflow/contrib/lite/interpreter.h"
+#include "tensorflow/contrib/lite/kernels/register.h"
+#include "tensorflow/contrib/lite/kernels/test_util.h"
+#include "tensorflow/contrib/lite/model.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+class SelectOpModel : public SingleOpModel {
+ public:
+  SelectOpModel(std::initializer_list<int> input1_shape,
+                std::initializer_list<int> input2_shape,
+                std::initializer_list<int> input3_shape,
+                TensorType input_type) {
+    input1_ = AddInput(TensorType_BOOL);
+    input2_ = AddInput(input_type);
+    input3_ = AddInput(input_type);
+    output_ = AddOutput(input_type);
+    SetBuiltinOp(BuiltinOperator_SELECT, BuiltinOptions_SelectOptions,
+                 CreateSelectOptions(builder_).Union());
+    BuildInterpreter({input1_shape, input2_shape, input3_shape});
+  }
+
+  int input1() { return input1_; }
+  int input2() { return input2_; }
+  int input3() { return input3_; }
+
+  template <typename T>
+  std::vector<T> GetOutput() {
+    return ExtractVector<T>(output_);
+  }
+
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ private:
+  int input1_;
+  int input2_;
+  int input3_;
+  int output_;
+};
+
+TEST(SelectOpTest, SelectBool) {
+  SelectOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, {1, 1, 1, 4},
+                      TensorType_BOOL);
+
+  model.PopulateTensor<bool>(model.input1(), {true, false, true, false});
+  model.PopulateTensor<bool>(model.input2(), {false, false, false, false});
+  model.PopulateTensor<bool>(model.input3(), {true, true, true, true});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput<bool>(),
+              ElementsAreArray({false, true, false, true}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(SelectOpTest, SelectFloat) {
+  SelectOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, {1, 1, 1, 4},
+                      TensorType_FLOAT32);
+
+  model.PopulateTensor<bool>(model.input1(), {true, false, true, false});
+  model.PopulateTensor<float>(model.input2(), {0.1, 0.2, 0.3, 0.4});
+  model.PopulateTensor<float>(model.input3(), {0.5, 0.6, 0.7, 0.8});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput<float>(), ElementsAreArray({0.1, 0.6, 0.3, 0.8}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(SelectOpTest, SelectUInt8) {
+  SelectOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, {1, 1, 1, 4},
+                      TensorType_UINT8);
+
+  model.PopulateTensor<bool>(model.input1(), {false, true, false, false});
+  model.PopulateTensor<uint8>(model.input2(), {1, 2, 3, 4});
+  model.PopulateTensor<uint8>(model.input3(), {5, 6, 7, 8});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput<uint8>(), ElementsAreArray({5, 2, 7, 8}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(SelectOpTest, SelectInt32) {
+  SelectOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, {1, 1, 1, 4},
+                      TensorType_INT32);
+
+  model.PopulateTensor<bool>(model.input1(), {false, true, false, false});
+  model.PopulateTensor<int32>(model.input2(), {1, 2, 3, 4});
+  model.PopulateTensor<int32>(model.input3(), {5, 6, 7, 8});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput<int32>(), ElementsAreArray({5, 2, 7, 8}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(SelectOpTest, RankOneSelectInt32) {
+  SelectOpModel model({2}, {2, 1, 2, 1}, {2, 1, 2, 1}, TensorType_INT32);
+
+  model.PopulateTensor<bool>(model.input1(), {false, true});
+  model.PopulateTensor<int32>(model.input2(), {1, 2, 3, 4});
+  model.PopulateTensor<int32>(model.input3(), {5, 6, 7, 8});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput<int32>(), ElementsAreArray({5, 6, 3, 4}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({2, 1, 2, 1}));
+}
+
+TEST(SelectOpTest, RankZeroSelectInt32) {
+  SelectOpModel model({1}, {1, 2, 2, 1}, {1, 2, 2, 1}, TensorType_INT32);
+
+  model.PopulateTensor<bool>(model.input1(), {false});
+  model.PopulateTensor<int32>(model.input2(), {1, 2, 3, 4});
+  model.PopulateTensor<int32>(model.input3(), {5, 6, 7, 8});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput<int32>(), ElementsAreArray({5, 6, 7, 8}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 2, 2, 1}));
+}
+
+}  // namespace
+}  // namespace tflite
+
+int main(int argc, char** argv) {
+  ::tflite::LogToStderr();
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}
diff --git a/tensorflow/contrib/lite/model.cc b/tensorflow/contrib/lite/model.cc
index 21c2181377..e89036ce73 100644
--- a/tensorflow/contrib/lite/model.cc
+++ b/tensorflow/contrib/lite/model.cc
@@ -675,7 +675,8 @@ TfLiteStatus ParseOpData(const Operator* op, BuiltinOperator op_type,
     case BuiltinOperator_GREATER:
     case BuiltinOperator_GREATER_EQUAL:
     case BuiltinOperator_LESS:
-    case BuiltinOperator_LESS_EQUAL: {
+    case BuiltinOperator_LESS_EQUAL:
+    case BuiltinOperator_SELECT: {
       break;
     }
     case BuiltinOperator_DELEGATE: {
diff --git a/tensorflow/contrib/lite/nnapi_delegate.cc b/tensorflow/contrib/lite/nnapi_delegate.cc
index e903af87b7..6a231dc6bc 100644
--- a/tensorflow/contrib/lite/nnapi_delegate.cc
+++ b/tensorflow/contrib/lite/nnapi_delegate.cc
@@ -377,6 +377,7 @@ void AddOpsAndParams(tflite::Interpreter* interpreter,
       case tflite::BuiltinOperator_LESS:
       case tflite::BuiltinOperator_LESS_EQUAL:
       case tflite::BuiltinOperator_NEG:
+      case tflite::BuiltinOperator_SELECT:
         FATAL("Op code %d is currently not delegated to NNAPI", builtin);
         nn_op_type = -1;  // set to invalid
         break;
diff --git a/tensorflow/contrib/lite/schema/schema.fbs b/tensorflow/contrib/lite/schema/schema.fbs
index 3ec91e505d..9de6180874 100644
--- a/tensorflow/contrib/lite/schema/schema.fbs
+++ b/tensorflow/contrib/lite/schema/schema.fbs
@@ -141,6 +141,7 @@ enum BuiltinOperator : byte {
   GREATER = 61,
   GREATER_EQUAL = 62,
   LESS_EQUAL = 63,
+  SELECT = 64,
 }
 
 // Options for the builtin operators.
@@ -191,6 +192,7 @@ union BuiltinOptions {
   GreaterOptions,
   GreaterEqualOptions,
   LessEqualOptions,
+  SelectOptions,
 }
 
 enum Padding : byte { SAME, VALID }
@@ -431,6 +433,9 @@ table LessEqualOptions {
 table NegOptions {
 }
 
+table SelectOptions {
+}
+
 // An OperatorCode can be an enum value (BuiltinOperator) if the operator is a
 // builtin, or a string if the operator is custom.
 table OperatorCode {
diff --git a/tensorflow/contrib/lite/schema/schema_generated.h b/tensorflow/contrib/lite/schema/schema_generated.h
old mode 100755
new mode 100644
index c6e4dab454..a2f0c8cdd2
--- a/tensorflow/contrib/lite/schema/schema_generated.h
+++ b/tensorflow/contrib/lite/schema/schema_generated.h
@@ -169,6 +169,9 @@ struct LessEqualOptionsT;
 struct NegOptions;
 struct NegOptionsT;
 
+struct SelectOptions;
+struct SelectOptionsT;
+
 struct OperatorCode;
 struct OperatorCodeT;
 
@@ -292,11 +295,12 @@ enum BuiltinOperator {
   BuiltinOperator_GREATER = 61,
   BuiltinOperator_GREATER_EQUAL = 62,
   BuiltinOperator_LESS_EQUAL = 63,
+  BuiltinOperator_SELECT = 64,
   BuiltinOperator_MIN = BuiltinOperator_ADD,
-  BuiltinOperator_MAX = BuiltinOperator_LESS_EQUAL
+  BuiltinOperator_MAX = BuiltinOperator_SELECT
 };
 
-inline BuiltinOperator (&EnumValuesBuiltinOperator())[63] {
+inline BuiltinOperator (&EnumValuesBuiltinOperator())[64] {
   static BuiltinOperator values[] = {
     BuiltinOperator_ADD,
     BuiltinOperator_AVERAGE_POOL_2D,
@@ -360,7 +364,8 @@ inline BuiltinOperator (&EnumValuesBuiltinOperator())[63] {
     BuiltinOperator_PADV2,
     BuiltinOperator_GREATER,
     BuiltinOperator_GREATER_EQUAL,
-    BuiltinOperator_LESS_EQUAL
+    BuiltinOperator_LESS_EQUAL,
+    BuiltinOperator_SELECT
   };
   return values;
 }
@@ -431,6 +436,7 @@ inline const char **EnumNamesBuiltinOperator() {
     "GREATER",
     "GREATER_EQUAL",
     "LESS_EQUAL",
+    "SELECT",
     nullptr
   };
   return names;
@@ -489,11 +495,12 @@ enum BuiltinOptions {
   BuiltinOptions_GreaterOptions = 44,
   BuiltinOptions_GreaterEqualOptions = 45,
   BuiltinOptions_LessEqualOptions = 46,
+  BuiltinOptions_SelectOptions = 47,
   BuiltinOptions_MIN = BuiltinOptions_NONE,
-  BuiltinOptions_MAX = BuiltinOptions_LessEqualOptions
+  BuiltinOptions_MAX = BuiltinOptions_SelectOptions
 };
 
-inline BuiltinOptions (&EnumValuesBuiltinOptions())[47] {
+inline BuiltinOptions (&EnumValuesBuiltinOptions())[48] {
   static BuiltinOptions values[] = {
     BuiltinOptions_NONE,
     BuiltinOptions_Conv2DOptions,
@@ -541,7 +548,8 @@ inline BuiltinOptions (&EnumValuesBuiltinOptions())[47] {
     BuiltinOptions_PadV2Options,
     BuiltinOptions_GreaterOptions,
     BuiltinOptions_GreaterEqualOptions,
-    BuiltinOptions_LessEqualOptions
+    BuiltinOptions_LessEqualOptions,
+    BuiltinOptions_SelectOptions
   };
   return values;
 }
@@ -595,6 +603,7 @@ inline const char **EnumNamesBuiltinOptions() {
     "GreaterOptions",
     "GreaterEqualOptions",
     "LessEqualOptions",
+    "SelectOptions",
     nullptr
   };
   return names;
@@ -793,6 +802,10 @@ template<> struct BuiltinOptionsTraits<LessEqualOptions> {
   static const BuiltinOptions enum_value = BuiltinOptions_LessEqualOptions;
 };
 
+template<> struct BuiltinOptionsTraits<SelectOptions> {
+  static const BuiltinOptions enum_value = BuiltinOptions_SelectOptions;
+};
+
 struct BuiltinOptionsUnion {
   BuiltinOptions type;
   void *value;
@@ -1192,6 +1205,14 @@ struct BuiltinOptionsUnion {
     return type == BuiltinOptions_LessEqualOptions ?
       reinterpret_cast<const LessEqualOptionsT *>(value) : nullptr;
   }
+  SelectOptionsT *AsSelectOptions() {
+    return type == BuiltinOptions_SelectOptions ?
+      reinterpret_cast<SelectOptionsT *>(value) : nullptr;
+  }
+  const SelectOptionsT *AsSelectOptions() const {
+    return type == BuiltinOptions_SelectOptions ?
+      reinterpret_cast<const SelectOptionsT *>(value) : nullptr;
+  }
 };
 
 bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *obj, BuiltinOptions type);
@@ -4319,6 +4340,46 @@ inline flatbuffers::Offset<NegOptions> CreateNegOptions(
 
 flatbuffers::Offset<NegOptions> CreateNegOptions(flatbuffers::FlatBufferBuilder &_fbb, const NegOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
 
+struct SelectOptionsT : public flatbuffers::NativeTable {
+  typedef SelectOptions TableType;
+  SelectOptionsT() {
+  }
+};
+
+struct SelectOptions FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
+  typedef SelectOptionsT NativeTableType;
+  bool Verify(flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           verifier.EndTable();
+  }
+  SelectOptionsT *UnPack(const flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  void UnPackTo(SelectOptionsT *_o, const flatbuffers::resolver_function_t *_resolver = nullptr) const;
+  static flatbuffers::Offset<SelectOptions> Pack(flatbuffers::FlatBufferBuilder &_fbb, const SelectOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
+};
+
+struct SelectOptionsBuilder {
+  flatbuffers::FlatBufferBuilder &fbb_;
+  flatbuffers::uoffset_t start_;
+  explicit SelectOptionsBuilder(flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  SelectOptionsBuilder &operator=(const SelectOptionsBuilder &);
+  flatbuffers::Offset<SelectOptions> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = flatbuffers::Offset<SelectOptions>(end);
+    return o;
+  }
+};
+
+inline flatbuffers::Offset<SelectOptions> CreateSelectOptions(
+    flatbuffers::FlatBufferBuilder &_fbb) {
+  SelectOptionsBuilder builder_(_fbb);
+  return builder_.Finish();
+}
+
+flatbuffers::Offset<SelectOptions> CreateSelectOptions(flatbuffers::FlatBufferBuilder &_fbb, const SelectOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher = nullptr);
+
 struct OperatorCodeT : public flatbuffers::NativeTable {
   typedef OperatorCode TableType;
   BuiltinOperator builtin_code;
@@ -4574,6 +4635,9 @@ struct Operator FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
   const LessEqualOptions *builtin_options_as_LessEqualOptions() const {
     return builtin_options_type() == BuiltinOptions_LessEqualOptions ? static_cast<const LessEqualOptions *>(builtin_options()) : nullptr;
   }
+  const SelectOptions *builtin_options_as_SelectOptions() const {
+    return builtin_options_type() == BuiltinOptions_SelectOptions ? static_cast<const SelectOptions *>(builtin_options()) : nullptr;
+  }
   const flatbuffers::Vector<uint8_t> *custom_options() const {
     return GetPointer<const flatbuffers::Vector<uint8_t> *>(VT_CUSTOM_OPTIONS);
   }
@@ -4784,6 +4848,10 @@ template<> inline const LessEqualOptions *Operator::builtin_options_as<LessEqual
   return builtin_options_as_LessEqualOptions();
 }
 
+template<> inline const SelectOptions *Operator::builtin_options_as<SelectOptions>() const {
+  return builtin_options_as_SelectOptions();
+}
+
 struct OperatorBuilder {
   flatbuffers::FlatBufferBuilder &fbb_;
   flatbuffers::uoffset_t start_;
@@ -6525,6 +6593,29 @@ inline flatbuffers::Offset<NegOptions> CreateNegOptions(flatbuffers::FlatBufferB
       _fbb);
 }
 
+inline SelectOptionsT *SelectOptions::UnPack(const flatbuffers::resolver_function_t *_resolver) const {
+  auto _o = new SelectOptionsT();
+  UnPackTo(_o, _resolver);
+  return _o;
+}
+
+inline void SelectOptions::UnPackTo(SelectOptionsT *_o, const flatbuffers::resolver_function_t *_resolver) const {
+  (void)_o;
+  (void)_resolver;
+}
+
+inline flatbuffers::Offset<SelectOptions> SelectOptions::Pack(flatbuffers::FlatBufferBuilder &_fbb, const SelectOptionsT* _o, const flatbuffers::rehasher_function_t *_rehasher) {
+  return CreateSelectOptions(_fbb, _o, _rehasher);
+}
+
+inline flatbuffers::Offset<SelectOptions> CreateSelectOptions(flatbuffers::FlatBufferBuilder &_fbb, const SelectOptionsT *_o, const flatbuffers::rehasher_function_t *_rehasher) {
+  (void)_rehasher;
+  (void)_o;
+  struct _VectorArgs { flatbuffers::FlatBufferBuilder *__fbb; const SelectOptionsT* __o; const flatbuffers::rehasher_function_t *__rehasher; } _va = { &_fbb, _o, _rehasher}; (void)_va;
+  return tflite::CreateSelectOptions(
+      _fbb);
+}
+
 inline OperatorCodeT *OperatorCode::UnPack(const flatbuffers::resolver_function_t *_resolver) const {
   auto _o = new OperatorCodeT();
   UnPackTo(_o, _resolver);
@@ -6892,6 +6983,10 @@ inline bool VerifyBuiltinOptions(flatbuffers::Verifier &verifier, const void *ob
       auto ptr = reinterpret_cast<const LessEqualOptions *>(obj);
       return verifier.VerifyTable(ptr);
     }
+    case BuiltinOptions_SelectOptions: {
+      auto ptr = reinterpret_cast<const SelectOptions *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
     default: return false;
   }
 }
@@ -7094,6 +7189,10 @@ inline void *BuiltinOptionsUnion::UnPack(const void *obj, BuiltinOptions type, c
       auto ptr = reinterpret_cast<const LessEqualOptions *>(obj);
       return ptr->UnPack(resolver);
     }
+    case BuiltinOptions_SelectOptions: {
+      auto ptr = reinterpret_cast<const SelectOptions *>(obj);
+      return ptr->UnPack(resolver);
+    }
     default: return nullptr;
   }
 }
@@ -7284,6 +7383,10 @@ inline flatbuffers::Offset<void> BuiltinOptionsUnion::Pack(flatbuffers::FlatBuff
       auto ptr = reinterpret_cast<const LessEqualOptionsT *>(value);
       return CreateLessEqualOptions(_fbb, ptr, _rehasher).Union();
     }
+    case BuiltinOptions_SelectOptions: {
+      auto ptr = reinterpret_cast<const SelectOptionsT *>(value);
+      return CreateSelectOptions(_fbb, ptr, _rehasher).Union();
+    }
     default: return 0;
   }
 }
@@ -7474,6 +7577,10 @@ inline BuiltinOptionsUnion::BuiltinOptionsUnion(const BuiltinOptionsUnion &u) FL
       value = new LessEqualOptionsT(*reinterpret_cast<LessEqualOptionsT *>(u.value));
       break;
     }
+    case BuiltinOptions_SelectOptions: {
+      value = new SelectOptionsT(*reinterpret_cast<SelectOptionsT *>(u.value));
+      break;
+    }
     default:
       break;
   }
@@ -7711,6 +7818,11 @@ inline void BuiltinOptionsUnion::Reset() {
       delete ptr;
       break;
     }
+    case BuiltinOptions_SelectOptions: {
+      auto ptr = reinterpret_cast<SelectOptionsT *>(value);
+      delete ptr;
+      break;
+    }
     default: break;
   }
   value = nullptr;
diff --git a/tensorflow/contrib/lite/testing/BUILD b/tensorflow/contrib/lite/testing/BUILD
index 6749e63552..f89c0d28d3 100644
--- a/tensorflow/contrib/lite/testing/BUILD
+++ b/tensorflow/contrib/lite/testing/BUILD
@@ -64,6 +64,7 @@ gen_zipped_test_files(
         "sub.zip",
         "topk.zip",
         "transpose.zip",
+        "where.zip",
     ],
 )
 
diff --git a/tensorflow/contrib/lite/testing/generate_examples.py b/tensorflow/contrib/lite/testing/generate_examples.py
index 7a658d43d3..05d099a82c 100644
--- a/tensorflow/contrib/lite/testing/generate_examples.py
+++ b/tensorflow/contrib/lite/testing/generate_examples.py
@@ -2242,10 +2242,41 @@ def make_neg_tests(zip_path):
   make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs)
 
 
+def make_where_tests(zip_path):
+  """Make a set of tests to do where."""
+
+  test_parameters = [{
+      "input_dtype": [tf.float32, tf.int32],
+      "input_shape_set": [([1, 2, 3, 4], [1, 2, 3, 4]),],
+  }]
+
+  def build_graph(parameters):
+    """Build the where op testing graph."""
+    input_value1 = tf.placeholder(
+        dtype=parameters["input_dtype"],
+        name="input2",
+        shape=parameters["input_shape_set"][0])
+    input_value2 = tf.placeholder(
+        dtype=parameters["input_dtype"],
+        name="input3",
+        shape=parameters["input_shape_set"][1])
+    less = tf.less(input_value1, input_value2)
+    out = tf.where(less, input_value1, input_value2)
+    return [input_value1, input_value2], [out]
+
+  def build_inputs(parameters, sess, inputs, outputs):
+    input_value1 = create_tensor_data(parameters["input_dtype"],
+                                      parameters["input_shape_set"][0])
+    input_value2 = create_tensor_data(parameters["input_dtype"],
+                                      parameters["input_shape_set"][1])
+    return [input_value1, input_value2], sess.run(
+        outputs, feed_dict=dict(zip(inputs, [input_value1, input_value2])))
+
+  make_zip_of_tests(zip_path, test_parameters, build_graph, build_inputs)
+
 # Toco binary path provided by the generate rule.
 bin_path = None
 
-
 def main(unused_args):
   global bin_path
   def mkdir_if_not_exist(x):
diff --git a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
index 2ce14f3b38..49762bdfe7 100644
--- a/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
+++ b/tensorflow/contrib/lite/testing/generated_examples_zip_test.cc
@@ -289,6 +289,7 @@ INSTANTIATE_TESTS(squeeze)
 INSTANTIATE_TESTS(strided_slice)
 INSTANTIATE_TESTS(sub)
 INSTANTIATE_TESTS(transpose)
+INSTANTIATE_TESTS(where)
 
 }  // namespace testing
 }  // namespace tflite
diff --git a/tensorflow/contrib/lite/toco/export_tensorflow.cc b/tensorflow/contrib/lite/toco/export_tensorflow.cc
index 53df1987b3..f5157149af 100644
--- a/tensorflow/contrib/lite/toco/export_tensorflow.cc
+++ b/tensorflow/contrib/lite/toco/export_tensorflow.cc
@@ -1674,6 +1674,19 @@ void ConvertTensorFlowMaximumOperator(const Model& model,
   (*sub_op->mutable_attr())["T"].set_type(data_type);
 }
 
+void ConvertSelectOperator(const Model& model, const SelectOperator& src_op,
+                           GraphDef* tensorflow_graph) {
+  auto* sub_op = tensorflow_graph->add_node();
+  sub_op->set_op("Select");
+  sub_op->set_name(src_op.outputs[0]);
+  CHECK_EQ(src_op.inputs.size(), 3);
+  *sub_op->add_input() = src_op.inputs[0];
+  *sub_op->add_input() = src_op.inputs[1];
+  *sub_op->add_input() = src_op.inputs[2];
+  const auto data_type = GetTensorFlowDataType(model, src_op.inputs[1]);
+  (*sub_op->mutable_attr())["T"].set_type(data_type);
+}
+
 void ConvertTopKV2Operator(const Model& model, const TopKV2Operator& src_op,
                            GraphDef* tensorflow_graph) {
   auto* topk_op = tensorflow_graph->add_node();
@@ -1914,6 +1927,9 @@ void ConvertOperator(const Model& model, const Operator& src_op,
     ConvertComparisonOperator(model, src_op, "Less", tensorflow_graph);
   } else if (src_op.type == OperatorType::kTensorFlowLessEqual) {
     ConvertComparisonOperator(model, src_op, "LessEqual", tensorflow_graph);
+  } else if (src_op.type == OperatorType::kSelect) {
+    ConvertSelectOperator(model, static_cast<const SelectOperator&>(src_op),
+                          tensorflow_graph);
   } else {
     LOG(FATAL) << "Unhandled operator type " << OperatorTypeName(src_op.type);
   }
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc b/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc
index c1cf79f626..6342cf3e8a 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc
@@ -152,6 +152,17 @@ bool PropagateArrayDataTypes::Run(Model* model, std::size_t op_index) {
       // Yield on ExpandDim until it is converted to Reshape
       return false;
     }
+    case OperatorType::kSelect: {
+      // Select produces outputs with the same type as their 2nd input
+      CHECK_EQ(op->inputs.size(), 3);
+      const ArrayDataType data_type_x =
+          model->GetArray(op->inputs[1]).data_type;
+      const ArrayDataType data_type_y =
+          model->GetArray(op->inputs[2]).data_type;
+      CHECK(data_type_x == data_type_y);
+      SetDataTypeForAllOutputs(model, op, data_type_x);
+      break;
+    }
     default: {
       // These operators produce outputs with the same type as their 1st input
       CHECK_GT(op->inputs.size(), 0);
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
index a081abea55..52b739c5e2 100644
--- a/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_fixed_sizes.cc
@@ -529,6 +529,21 @@ void ProcessSimpleBinaryOperator(Model* model, Operator* op) {
                                   &output_array);
 }
 
+void ProcessSelectOperator(Model* model, SelectOperator* op) {
+  // Yield until all input dims have been resolved.
+  for (const auto& input : op->inputs) {
+    const auto& input_array = model->GetArray(input);
+    if (!input_array.has_shape()) {
+      return;
+    }
+  }
+
+  // Select's output matches the second and third output.
+  const auto& input1_array = model->GetArray(op->inputs[1]);
+  auto& output_array = model->GetArray(op->outputs[0]);
+  output_array.copy_shape(input1_array.shape());
+}
+
 void ProcessAddNOperator(Model* model, Operator* op) {
   // Yield until all input dims have been resolved.
   //
@@ -1570,7 +1585,9 @@ bool PropagateFixedSizes::Run(Model* model, std::size_t op_index) {
     case OperatorType::kMean:
       ProcessTensorFlowReductionOperator(model, op);
       break;
-
+    case OperatorType::kSelect:
+      ProcessSelectOperator(model, static_cast<SelectOperator*>(op));
+      break;
     case OperatorType::kSlice:
       ProcessSliceOperator(model, static_cast<SliceOperator*>(op));
       break;
diff --git a/tensorflow/contrib/lite/toco/import_tensorflow.cc b/tensorflow/contrib/lite/toco/import_tensorflow.cc
index 532fcdd808..52757ca748 100644
--- a/tensorflow/contrib/lite/toco/import_tensorflow.cc
+++ b/tensorflow/contrib/lite/toco/import_tensorflow.cc
@@ -1344,6 +1344,19 @@ void ConvertUnsupportedOperator(const NodeDef& node,
   }
 }
 
+void ConvertSelectOperator(const NodeDef& node,
+                           const TensorFlowImportFlags& tf_import_flags,
+                           Model* model) {
+  CheckInputsCount(node, tf_import_flags, 3);
+
+  auto* op = new SelectOperator;
+  for (const auto& input : node.input()) {
+    op->inputs.push_back(input);
+  }
+  op->outputs.push_back(node.name());
+  model->operators.emplace_back(op);
+}
+
 void ConvertStridedSliceOperator(const NodeDef& node,
                                  const TensorFlowImportFlags& tf_import_flags,
                                  Model* model) {
@@ -2254,6 +2267,8 @@ Status ImportTensorFlowNode(const tensorflow::NodeDef& node,
     ConvertDynamicStitchOperator(node, tf_import_flags, model);
   } else if (node.op() == "RandomUniform") {
     ConvertRandomUniform(node, tf_import_flags, model);
+  } else if (node.op() == "Select") {
+    ConvertSelectOperator(node, tf_import_flags, model);
   } else {
     ConvertUnsupportedOperator(node, tf_import_flags, model);
   }
diff --git a/tensorflow/contrib/lite/toco/model.h b/tensorflow/contrib/lite/toco/model.h
index 7ee7841511..47f8db5978 100644
--- a/tensorflow/contrib/lite/toco/model.h
+++ b/tensorflow/contrib/lite/toco/model.h
@@ -133,6 +133,7 @@ enum class OperatorType {
   // instead of being given as plain constant arrays. So we need to insert
   // special nodes in the graph to shuffle axes.
   kReorderAxes,
+  kSelect,
 };
 
 // Helper to deal with TensorFlow arrays using a different ordering of
@@ -1087,6 +1088,18 @@ struct NegOperator : Operator {
   NegOperator() : Operator(OperatorType::kNeg) {}
 };
 
+// Element-wise select operator choosing elements from inputs[1] or input[2]
+//
+// Inputs:
+//  inputs[0]: required: boolean mask per index
+//  inputs[1]: required: tensor of values if true
+//  inputs[2]: required: tensor of values if false
+//
+//  TensorFlow equivalent: Select
+struct SelectOperator : Operator {
+  SelectOperator() : Operator(OperatorType::kSelect) {}
+};
+
 // Element-wise reciprocal-square-root (x^-0.5) operator.
 //
 // Inputs:
diff --git a/tensorflow/contrib/lite/toco/tflite/operator.cc b/tensorflow/contrib/lite/toco/tflite/operator.cc
index a008e63351..90e24aa104 100644
--- a/tensorflow/contrib/lite/toco/tflite/operator.cc
+++ b/tensorflow/contrib/lite/toco/tflite/operator.cc
@@ -924,6 +924,8 @@ std::vector<std::unique_ptr<BaseOperator>> BuildOperatorList() {
   ops.emplace_back(new SimpleOperator<TensorFlowLessEqualOperator>(
       "LESS_EQUAL", OperatorType::kTensorFlowLessEqual));
   ops.emplace_back(new SimpleOperator<NegOperator>("NEG", OperatorType::kNeg));
+  ops.emplace_back(
+      new SimpleOperator<SelectOperator>("SELECT", OperatorType::kSelect));
 
   return ops;
 }
diff --git a/tensorflow/contrib/lite/toco/tflite/operator_test.cc b/tensorflow/contrib/lite/toco/tflite/operator_test.cc
index 2b6c32b07c..a4fff9974a 100644
--- a/tensorflow/contrib/lite/toco/tflite/operator_test.cc
+++ b/tensorflow/contrib/lite/toco/tflite/operator_test.cc
@@ -116,6 +116,7 @@ TEST_F(OperatorTest, SimpleOperators) {
   CheckSimpleOperator<TensorFlowLessOperator>("LESS",
                                               OperatorType::kTensorFlowLess);
   CheckSimpleOperator<NegOperator>("NEG", OperatorType::kNeg);
+  CheckSimpleOperator<SelectOperator>("SELECT", OperatorType::kSelect);
 }
 
 TEST_F(OperatorTest, BuiltinAdd) {
diff --git a/tensorflow/contrib/lite/toco/tooling_util.cc b/tensorflow/contrib/lite/toco/tooling_util.cc
index f82bb33535..1f56fe5c83 100644
--- a/tensorflow/contrib/lite/toco/tooling_util.cc
+++ b/tensorflow/contrib/lite/toco/tooling_util.cc
@@ -391,6 +391,7 @@ const char* OperatorTypeName(OperatorType type) {
     HANDLE_OPERATORTYPENAME_CASE(Exp)
     HANDLE_OPERATORTYPENAME_CASE(DynamicPartition)
     HANDLE_OPERATORTYPENAME_CASE(DynamicStitch)
+    HANDLE_OPERATORTYPENAME_CASE(Select)
     default:
       LOG(FATAL) << "Unhandled op type";
 #undef HANDLE_OPERATORTYPENAME_CASE
@@ -2097,6 +2098,8 @@ ArrayDataType ConvertIODataTypeToArrayDataType(IODataType type) {
       return ArrayDataType::kInt32;
     case INT64:
       return ArrayDataType::kInt64;
+    case BOOL:
+      return ArrayDataType::kBool;
     default:
       return ArrayDataType::kNone;
   }
diff --git a/tensorflow/contrib/lite/toco/types.proto b/tensorflow/contrib/lite/toco/types.proto
index 03bd6150bc..421667a83c 100644
--- a/tensorflow/contrib/lite/toco/types.proto
+++ b/tensorflow/contrib/lite/toco/types.proto
@@ -37,4 +37,7 @@ enum IODataType {
 
   // Int16, quantized
   QUANTIZED_INT16 = 6;
+
+  // Boolean
+  BOOL = 7;
 }
-- 
GitLab


From 37b8860e302d73845e74e1bfb6c3cb59207f2d77 Mon Sep 17 00:00:00 2001
From: Michael Kuperstein <mkuper@google.com>
Date: Mon, 7 May 2018 15:41:52 -0700
Subject: [PATCH 332/395] [XLA] Fix a "we're we're" in the operation semantics.

PiperOrigin-RevId: 195734316
---
 tensorflow/docs_src/performance/xla/operation_semantics.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/docs_src/performance/xla/operation_semantics.md b/tensorflow/docs_src/performance/xla/operation_semantics.md
index f530fe1206..21e4c71a60 100644
--- a/tensorflow/docs_src/performance/xla/operation_semantics.md
+++ b/tensorflow/docs_src/performance/xla/operation_semantics.md
@@ -1049,8 +1049,8 @@ For a more intuitive description, see the "Informal Description" section below.
 :                  :                         : from.                           :
 |`gather_indices`  | `ComputationDataHandle` | Tensor containing the starting  |
 :                  :                         : indices of the slices we're     :
-:                  :                         : we're stitching together into   :
-:                  :                         : the output tensor.              :
+:                  :                         : stitching together into the     :
+:                  :                         : output tensor.                  :
 |`index_vector_dim`  | `int64`               | The dimension in                |
 :                  :                         : `gather_indices` that contains  :
 :                  :                         : the starting indices.           :
-- 
GitLab


From 4a9beef315c3e456e7f087b5b3205df99f4a0876 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 15:47:57 -0700
Subject: [PATCH 333/395] Add EvaluateNodes to tests:
 RemoveIdentityTransposesMultipleOutputs,
 RemoveTransposesWithControlDependency, CombineBitcasts,
 CombineAndRemoveBitcasts, RemoveRedundantCast

PiperOrigin-RevId: 195735234
---
 .../optimizers/arithmetic_optimizer_test.cc   | 45 +++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
index 741cc135a1..067adb359c 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer_test.cc
@@ -1166,6 +1166,11 @@ TEST_F(ArithmeticOptimizerTest, RemoveIdentityTransposesMultipleOutputs) {
   item.fetch = {"outputs"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto x_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({8, 12, 28, 28}));
+  item.feed = {{"inputs", x_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, item.feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyRemoveIdentityTranspose(&optimizer);
@@ -1178,6 +1183,10 @@ TEST_F(ArithmeticOptimizerTest, RemoveIdentityTransposesMultipleOutputs) {
       EXPECT_EQ(node.input(2), "Split:2");
     }
   }
+
+  auto tensors = EvaluateNodes(output, item.fetch, item.feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorNear<float>(tensors_expected[0], tensors[0], 1e-6);
 }
 
 TEST_F(ArithmeticOptimizerTest, RemoveTransposesWithControlDependency) {
@@ -1194,6 +1203,11 @@ TEST_F(ArithmeticOptimizerTest, RemoveTransposesWithControlDependency) {
   item.fetch = {"outputs"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto x_t = GenerateRandomTensor<DT_FLOAT>(TensorShape({2, 3}));
+  item.feed = {{"Placeholder", x_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, item.feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyRemoveIdentityTranspose(&optimizer);
@@ -1204,6 +1218,10 @@ TEST_F(ArithmeticOptimizerTest, RemoveTransposesWithControlDependency) {
   EXPECT_EQ(2, outputs_node->input_size());
   EXPECT_EQ(outputs_node->input(0), "outputs_const");
   EXPECT_EQ(outputs_node->input(1), "^Placeholder");
+
+  auto tensors = EvaluateNodes(output, item.fetch, item.feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorNear<float>(tensors_expected[0], tensors[0], 1e-6);
 }
 
 TEST_F(ArithmeticOptimizerTest, NotRemoveTransposes) {
@@ -1450,6 +1468,11 @@ TEST_F(ArithmeticOptimizerTest, CombineBitcasts) {
   item.fetch = {"outputs"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto x_t = GenerateRandomTensor<DT_UINT8>(TensorShape({2, 3}));
+  item.feed = {{"inputs", x_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, item.feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyRemoveRedundantBitcast(&optimizer);
@@ -1461,6 +1484,10 @@ TEST_F(ArithmeticOptimizerTest, CombineBitcasts) {
   EXPECT_EQ(3, output.node_size());
   EXPECT_EQ(1, CountOpNodes(output, "Bitcast"));
   EXPECT_TRUE(IsNodesDirectlyConnected(node_map, "inputs", "bc2"));
+
+  auto tensors = EvaluateNodes(output, item.fetch, item.feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorEqual<int8>(tensors_expected[0], tensors[0]);
 }
 
 TEST_F(ArithmeticOptimizerTest, CombineAndRemoveBitcasts) {
@@ -1475,6 +1502,11 @@ TEST_F(ArithmeticOptimizerTest, CombineAndRemoveBitcasts) {
   item.fetch = {"outputs"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto x_t = GenerateRandomTensor<DT_INT8>(TensorShape({2, 3}));
+  item.feed = {{"inputs", x_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, item.feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyRemoveRedundantBitcast(&optimizer);
@@ -1486,6 +1518,10 @@ TEST_F(ArithmeticOptimizerTest, CombineAndRemoveBitcasts) {
   EXPECT_EQ(2, output.node_size());
   EXPECT_EQ(0, CountOpNodes(output, "Bitcast"));
   EXPECT_TRUE(IsNodesDirectlyConnected(node_map, "inputs", "outputs"));
+
+  auto tensors = EvaluateNodes(output, item.fetch, item.feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorEqual<int8>(tensors_expected[0], tensors[0]);
 }
 
 TEST_F(ArithmeticOptimizerTest, RemoveRedundantCast) {
@@ -1499,6 +1535,11 @@ TEST_F(ArithmeticOptimizerTest, RemoveRedundantCast) {
   item.fetch = {"outputs"};
   TF_CHECK_OK(s.ToGraphDef(&item.graph));
 
+  auto x_t = GenerateRandomTensor<DT_INT8>(TensorShape({2, 3}));
+  item.feed = {{"inputs", x_t}};
+  auto tensors_expected = EvaluateNodes(item.graph, item.fetch, item.feed);
+  EXPECT_EQ(1, tensors_expected.size());
+
   GraphDef output;
   ArithmeticOptimizer optimizer;
   EnableOnlyRemoveRedundantCast(&optimizer);
@@ -1510,6 +1551,10 @@ TEST_F(ArithmeticOptimizerTest, RemoveRedundantCast) {
   EXPECT_EQ(2, output.node_size());
   EXPECT_EQ(0, CountOpNodes(output, "Cast"));
   EXPECT_TRUE(IsNodesDirectlyConnected(node_map, "inputs", "outputs"));
+
+  auto tensors = EvaluateNodes(output, item.fetch, item.feed);
+  EXPECT_EQ(1, tensors.size());
+  test::ExpectTensorEqual<int8>(tensors_expected[0], tensors[0]);
 }
 
 TEST_F(ArithmeticOptimizerTest, AddOpsRewrite_AddOpsOfIdenticalShape) {
-- 
GitLab


From 27e6ab7c8b33d7f5e5795d31226b596ec70642fd Mon Sep 17 00:00:00 2001
From: Derek Murray <mrry@google.com>
Date: Mon, 7 May 2018 15:51:05 -0700
Subject: [PATCH 334/395] [Remote functions] Only set the default runner
 *after* resolving the remote FLR.

Previously, if the `runner` was not specified for a function
execution, we would immediately set it to the default runner of the
*local* FLR, even if the function was to be executed remotely. This
change postpones the resolution of the default runner until after the
function invocation has been routed to the FLR that will actually
execute it.

As a result, we avoid the pathological case where a GPU device using a
private threadpool (TF_GPU_THREAD_MODE=gpu_private) ends up running
all of the ops for the CPU-side input pipeline on the private
threadpool.

PiperOrigin-RevId: 195735734
---
 tensorflow/core/common_runtime/function.cc | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/common_runtime/function.cc b/tensorflow/core/common_runtime/function.cc
index a6f637b488..bf05f6f1d9 100644
--- a/tensorflow/core/common_runtime/function.cc
+++ b/tensorflow/core/common_runtime/function.cc
@@ -795,16 +795,16 @@ void FunctionLibraryRuntimeImpl::Run(const Options& opts, Handle handle,
     };
   }
 
-  if (run_opts.runner == nullptr) {
-    run_opts.runner = &default_runner_;
-  }
-  DCHECK(run_opts.runner != nullptr);
-
   if (!parent_->IsInstantiatedOnDevice(device_name_, handle)) {
     parent_->Run(run_opts, handle, args, rets, done);
     return;
   }
 
+  if (run_opts.runner == nullptr) {
+    run_opts.runner = &default_runner_;
+  }
+  DCHECK(run_opts.runner != nullptr);
+
   Executor::Args* exec_args = new Executor::Args;
   // Inherit the step_id from the caller.
   exec_args->step_id = run_opts.step_id;
-- 
GitLab


From 94b0b2fbce60100c4fe81bf92f5c927626ed66b6 Mon Sep 17 00:00:00 2001
From: Blake Hechtman <blakehechtman@google.com>
Date: Mon, 7 May 2018 15:58:29 -0700
Subject: [PATCH 335/395] [XLA] Make post order a possible schedule as it
 sometimes uses less memory than the DFS or list scheduler and it is very
 simple.

PiperOrigin-RevId: 195736916
---
 .../compiler/xla/service/cpu/cpu_compiler.cc  |  3 ++-
 .../compiler/xla/service/hlo_scheduling.cc    | 26 ++++++++++++++++++-
 .../compiler/xla/service/hlo_scheduling.h     |  6 +++++
 tensorflow/compiler/xla/tests/BUILD           |  1 +
 4 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
index 91ed6e427a..3d2e24ca14 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
+++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc
@@ -535,7 +535,8 @@ StatusOr<std::unique_ptr<Executable>> CpuCompiler::RunBackend(
   // and reduced memory usage (as compared to using DependencyHloOrdering).
   TF_ASSIGN_OR_RETURN(
       SequentialHloOrdering::HloModuleSequence module_sequence,
-      CreateMemoryMinimizingSequence(*module, BufferSizeBytesFunction()));
+      CreateMemoryMinimizingSequence(*module, BufferSizeBytesFunction(),
+                                     DFSMemoryScheduler));
 
   // Run buffer analysis on the HLO graph. This analysis figures out which
   // temporary buffers are required to run the computation.
diff --git a/tensorflow/compiler/xla/service/hlo_scheduling.cc b/tensorflow/compiler/xla/service/hlo_scheduling.cc
index 1a767628f6..23ace5afea 100644
--- a/tensorflow/compiler/xla/service/hlo_scheduling.cc
+++ b/tensorflow/compiler/xla/service/hlo_scheduling.cc
@@ -430,6 +430,15 @@ StatusOr<std::vector<const HloInstruction*>> ListMemoryScheduler(
   return ListScheduler::Run(computation, points_to_analysis, size_function);
 }
 
+StatusOr<std::vector<const HloInstruction*>> PostOrderMemoryScheduler(
+    const HloComputation& computation,
+    const TuplePointsToAnalysis& points_to_analysis,
+    const LogicalBuffer::SizeFunction& size_function) {
+  const auto& post_order = computation.MakeInstructionPostOrder();
+  return std::vector<const HloInstruction*>{post_order.begin(),
+                                            post_order.end()};
+}
+
 StatusOr<std::vector<const HloInstruction*>> DefaultMemoryScheduler(
     const HloComputation& computation,
     const TuplePointsToAnalysis& points_to_analysis,
@@ -459,7 +468,22 @@ StatusOr<std::vector<const HloInstruction*>> DefaultMemoryScheduler(
                                   size_function));
   VLOG(2) << "Min-memory dfs sequence: " << HumanReadableNumBytes(dfs_memory);
 
-  if (list_memory <= dfs_memory) {
+  TF_ASSIGN_OR_RETURN(
+      std::vector<const HloInstruction*> post_order_sequence,
+      PostOrderMemoryScheduler(computation, points_to_analysis, size_function));
+  TF_ASSIGN_OR_RETURN(
+      const int64 post_order_memory,
+      MinimumMemoryForComputation(computation, post_order_sequence,
+                                  points_to_analysis, size_function));
+  VLOG(2) << "Min-memory post order sequence: "
+          << HumanReadableNumBytes(post_order_memory);
+
+  if (post_order_memory < std::min(list_memory, dfs_memory)) {
+    VLOG(2) << "Chose min-memory post_order sequence: "
+            << HumanReadableNumBytes(post_order_memory);
+    return post_order_sequence;
+
+  } else if (list_memory <= dfs_memory) {
     VLOG(2) << "Chose min-memory list sequence: "
             << HumanReadableNumBytes(list_memory);
     return list_sequence;
diff --git a/tensorflow/compiler/xla/service/hlo_scheduling.h b/tensorflow/compiler/xla/service/hlo_scheduling.h
index 068e68383d..fcb006f818 100644
--- a/tensorflow/compiler/xla/service/hlo_scheduling.h
+++ b/tensorflow/compiler/xla/service/hlo_scheduling.h
@@ -55,6 +55,12 @@ StatusOr<std::vector<const HloInstruction*>> DFSMemoryScheduler(
     const TuplePointsToAnalysis& points_to_analysis,
     const LogicalBuffer::SizeFunction& size_function);
 
+// Naive Post Order scheduler
+StatusOr<std::vector<const HloInstruction*>> PostOrderMemoryScheduler(
+    const HloComputation& computation,
+    const TuplePointsToAnalysis& points_to_analysis,
+    const LogicalBuffer::SizeFunction& size_function);
+
 // The default scheduling algorithm. Runs both the list scheduler
 // and the DFS scheduler, and chooses whichever returns a lower min-memory,
 // not accounting for fragmentation.
diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD
index 0571ff5055..1c29abcb80 100644
--- a/tensorflow/compiler/xla/tests/BUILD
+++ b/tensorflow/compiler/xla/tests/BUILD
@@ -1868,6 +1868,7 @@ xla_test(
 xla_test(
     name = "local_client_execute_test",
     srcs = ["local_client_execute_test.cc"],
+    shard_count = 30,
     tags = ["optonly"],
     deps = [
         "//tensorflow/compiler/xla:literal_util",
-- 
GitLab


From 5802096c267c805f6a69798aac10aefef759bb9f Mon Sep 17 00:00:00 2001
From: Akshay Agrawal <akshayka@google.com>
Date: Mon, 7 May 2018 16:16:24 -0700
Subject: [PATCH 336/395] Refactor TensorArray to avoid copies and memory
 allocations when executing eagerly.

With this change, writes to TensorArrays when eager execution is enabled take O(1) time instead of O(n). Additionally, whereas writing to a TensorArray when constructing a graph results in allocating a new Python TensorArray object, writing to a TensorArray with eager enabled no longer performs that allocation (graph construction uses these allocations to ensure correctness of control flow and gradients, but this isn't necessary when executing eagerly). Finally, this change also removes the artificial write-once semantics of TensorArrays when executing eagerly.

PiperOrigin-RevId: 195739572
---
 .../kernel_tests/tensor_array_ops_test.py     |   1 -
 tensorflow/python/ops/tensor_array_ops.py     | 196 ++++++++----------
 2 files changed, 81 insertions(+), 116 deletions(-)

diff --git a/tensorflow/python/kernel_tests/tensor_array_ops_test.py b/tensorflow/python/kernel_tests/tensor_array_ops_test.py
index 918bbd38ed..c0b36f143d 100644
--- a/tensorflow/python/kernel_tests/tensor_array_ops_test.py
+++ b/tensorflow/python/kernel_tests/tensor_array_ops_test.py
@@ -438,7 +438,6 @@ class TensorArrayTest(test.TestCase):
           "Tried to read from index 3 but array size is: 3"):
         self.evaluate(ta.read(3))
 
-  @test_util.run_in_graph_and_eager_modes()
   def testTensorArrayWriteMultipleFails(self):
     with self.test_session(use_gpu=True):
       ta = tensor_array_ops.TensorArray(
diff --git a/tensorflow/python/ops/tensor_array_ops.py b/tensorflow/python/ops/tensor_array_ops.py
index d2f45ce37b..cc92da4fd7 100644
--- a/tensorflow/python/ops/tensor_array_ops.py
+++ b/tensorflow/python/ops/tensor_array_ops.py
@@ -20,6 +20,7 @@ from __future__ import division
 from __future__ import print_function
 
 import contextlib
+import weakref
 
 from tensorflow.python.eager import context
 from tensorflow.python.framework import constant_op
@@ -395,69 +396,8 @@ class _GraphTensorArray(object):
 # pylint: enable=protected-access
 
 
-# pylint: disable=protected-access
-def _eager_write_no_copy(ta, index, value):
-  """Writes value into an _EagerTensorArray without creating a new TensorArray.
-
-  Args:
-    ta: _EagerTensorArray into which to write value.
-    index: 0-D.  int32 scalar with the index to write to.
-    value: N-D.  Tensor of type `dtype`.  The Tensor to write to this index.
-
-  Raises:
-    errors_impl.AlreadyExistsError: attempting to overwrite an entry.
-    errors_impl.InvalidArgumentError: value dtype does not match `ta`'s dtype.
-    errors_impl.OutOfRangeError: `index` is out of bounds.
-    ValueError: shape of `value` is not consistent with inferred shape.
-  """
-
-  if isinstance(index, ops.EagerTensor):
-    index = index.numpy()
-
-  if index < 0:
-    raise errors_impl.OutOfRangeError(
-        None, None,
-        "Writing to negative indices (index %d) is not allowed." % index)
-
-  tensor_array = ta._tensor_array
-  size = len(tensor_array)
-  if index >= size:
-    if not ta._dynamic_size:
-      raise errors_impl.OutOfRangeError(
-          None, None,
-          "Tried to write to index %d but array is not resizeable and size "
-          "is: %d" % (index, size))
-    tensor_array.extend([None for _ in range(index - size + 1)])
-
-  if not isinstance(value, ops.EagerTensor):
-    value = constant_op.constant(value)
-
-  if ta._infer_shape:
-    if ta._element_shape is None:
-      ta._element_shape = value.shape
-    elif ta._element_shape != value.shape:
-      raise ValueError("Incompatible shape for value (%s), expected (%s)" %
-                       (value.shape.as_list(), ta._element_shape.as_list()))
-
-  if ta._dtype != value.dtype:
-    raise errors_impl.InvalidArgumentError(
-        None, None,
-        "TensorArray dtype is %s but Op is trying to write dtype %s" %
-        (ta._dtype.name, value.dtype.name))
-
-  if ta._tensor_array[index] is not None:
-    raise errors_impl.AlreadyExistsError(
-        None, None,
-        "Could not write to TensorArray index %d because it has already been "
-        "written to." % index)
-
-  tensor_array[index] = value
-
-# pylint: enable=protected-access
-
-
 class _EagerTensorArray(object):
-  """Eager-mode implementation of TensorArray.
+  """Eager-compatible implementation of TensorArray.
   """
 
   def __init__(self,
@@ -472,7 +412,7 @@ class _EagerTensorArray(object):
                element_shape=None,
                colocate_with_first_write_call=True,
                name=None):
-    """Constructs an Eager mode TensorArray.
+    """Constructs a TensorArray compatible with eager execution.
 
     Args:
       dtype: (required) data type of the TensorArray.
@@ -495,16 +435,19 @@ class _EagerTensorArray(object):
       ValueError: handle or flow are supplied, or if size is not supplied.
     """
 
-    del (flow, tensor_array_name, name)  # not meaningful in Eager
+    del (flow, tensor_array_name, name)  # Unused.
 
     if handle is not None:
-      raise ValueError("TensorArray handles are not supported in Eager mode.")
+      raise ValueError("TensorArray handles are not supported when eager "
+                       "execution is enabled.")
     if size is None:
-      raise ValueError("Size must be declared for TensorArrays in Eager mode.")
+      raise ValueError("Size must be declared for TensorArrays when eager "
+                       "execution is enabled.")
 
-    # These attributes are not meaningful in Eager, but some library functions
-    # (e.g., those in control_flow_ops.py) access them to create new tensor
-    # arrays; as such, we define them for the sake of compatibility.
+    # These attributes are not meaningful when eager is enabled, but some
+    # library functions (e.g., those in control_flow_ops.py) access them to
+    # create new tensor arrays; as such, we define them for the sake of
+    # compatibility.
     self._handle = None
     # we assign a dummy value to _flow in case other code assumes it to be
     # a Tensor
@@ -525,7 +468,7 @@ class _EagerTensorArray(object):
 
   @property
   def flow(self):
-    """Flows are not meaningful in Eager; this exists for compatibility."""
+    """For compatibility; flows are not meaningful when eager is enabled."""
     return self._flow
 
   @property
@@ -534,42 +477,22 @@ class _EagerTensorArray(object):
 
   @property
   def handle(self):
-    """Handles are not meaningful in Eager; this exists for compatibility."""
+    """For compatibility; handles are not meaningful when eager is enabled."""
     return self._handle
 
-  def _identity_without_array(self):
-    """Returns a new TensorArray with the same properties as this Eager one.
-
-    NB: Does not set the underlying _tensor_array attribute.
-    """
-    ta = TensorArray(
-        dtype=self._dtype,
-        size=len(self._tensor_array),
-        dynamic_size=self._dynamic_size,
-        clear_after_read=self._clear_after_read,
-        handle=self._handle,
-        flow=self._flow,
-        infer_shape=self._infer_shape,
-        element_shape=self._element_shape,
-        colocate_with_first_write_call=self._colocate_with_first_write_call)
-    ta._implementation._previously_read_indices = self._previously_read_indices  # pylint: disable=protected-access
-    return ta
-
   def identity(self):
     """See TensorArray."""
-    ta = self._identity_without_array()
-    ta._implementation._tensor_array = [t for t in self._tensor_array]  # pylint: disable=protected-access
-    return ta
+    return self.parent()
 
   def grad(self, source, flow=None, name=None):
     raise NotImplementedError(
-        "TensorArray.grad is not supported in Eager mode; Eager's gradient "
-        "implementation does not use/need this function to compute gradients "
-        "of operations that use TensorArrays.")
+        "TensorArray.grad is not supported when executing eagerly; eager's "
+        "gradient implementation does not use/need this function to compute "
+        "gradients of operations that use TensorArrays.")
 
   def read(self, index, name=None):
     """See TensorArray."""
-    del name  # not meaningful in Eager mode
+    del name  # not meaningful when executing eagerly.
 
     if isinstance(index, ops.EagerTensor):
       index = index.numpy()
@@ -600,12 +523,58 @@ class _EagerTensorArray(object):
       self._previously_read_indices.append(index)
     return tensor
 
+  def _write(self, index, value):
+    """Writes `value` into index named by `index`.
+
+    Args:
+      index: 0-D.  int32 scalar with the index to write to.
+      value: N-D.  Tensor of type `dtype`.  The `Tensor` to write to `index`.
+
+    Raises:
+      errors_impl.InvalidArgumentError: `value` dtype does not match dtype.
+      errors_impl.OutOfRangeError: `index` is out of bounds.
+      ValueError: shape of `value` is not consistent with inferred shape.
+    """
+
+    if isinstance(index, ops.EagerTensor):
+      index = index.numpy()
+
+    if index < 0:
+      raise errors_impl.OutOfRangeError(
+          None, None,
+          "Writing to negative indices (index %d) is not allowed." % index)
+
+    size = len(self._tensor_array)
+    if index >= size:
+      if not self._dynamic_size:
+        raise errors_impl.OutOfRangeError(
+            None, None,
+            "Tried to write to index %d but array is not resizeable and size "
+            "is: %d" % (index, size))
+      self._tensor_array.extend([None for _ in range(index - size + 1)])
+
+    if not isinstance(value, ops.EagerTensor):
+      value = constant_op.constant(value)
+
+    if self._infer_shape:
+      if self._element_shape is None:
+        self._element_shape = value.shape
+      elif self._element_shape != value.shape:
+        raise ValueError("Incompatible shape for value (%s), expected (%s)" %
+                         (value.shape.as_list(), self._element_shape.as_list()))
+
+    if self._dtype != value.dtype:
+      raise errors_impl.InvalidArgumentError(
+          None, None,
+          "TensorArray dtype is %s but Op is trying to write dtype %s" %
+          (self._dtype.name, value.dtype.name))
+    self._tensor_array[index] = value
+
   def write(self, index, value, name=None):
     """See TensorArray."""
-    del name  # not meaningful in Eager mode
-    ta = self.identity()
-    _eager_write_no_copy(ta._implementation, index, value)  # pylint: disable=protected-access
-    return ta
+    del name  # not meaningful when executing eagerly.
+    self._write(index, value)
+    return self.parent()
 
   def _maybe_zero(self, ix):
     val = self._tensor_array[ix]
@@ -623,7 +592,7 @@ class _EagerTensorArray(object):
 
   def gather(self, indices, name=None):
     """See TensorArray."""
-    del name  # not meaningful in Eager mode
+    del name  # not meaningful when executing eagerly.
     return array_ops.stack([self._maybe_zero(i) for i in indices.numpy()])
 
   def concat(self, name=None):
@@ -651,17 +620,15 @@ class _EagerTensorArray(object):
       raise ValueError(
           "Cannot unstack %d tensors into a TensorArray of static size %d" %
           (len(tensors), len(self._tensor_array)))
-    ta = self._identity_without_array()
-    ta._implementation._tensor_array = tensors  # pylint: disable=protected-access
-    return ta
+    self._tensor_array = tensors
+    return self.parent()
 
   def scatter(self, indices, value, name=None):
     """See TensorArray."""
-    del name  # unused in Eager
-    ta = self.identity()
+    del name  # not meaningful when executing eagerly.
     for index, val in zip(indices.numpy(), array_ops.unstack(value)):
-      _eager_write_no_copy(ta._implementation, index, val)  # pylint: disable=protected-access
-    return ta
+      self._write(index, val)  # pylint: disable=protected-access
+    return self.parent()
 
   def split(self, value, lengths, name=None):
     """See TensorArray."""
@@ -690,20 +657,17 @@ class _EagerTensorArray(object):
           "dynamically resizeable" % (len(self._tensor_array),
                                       lengths.shape[0]))
     else:
-      ta = self._identity_without_array()
-      tensor_array = array_ops.split(value, lengths, name=name)
-      ta._implementation._tensor_array = tensor_array  # pylint: disable=protected-access
-      return ta
+      self._tensor_array = array_ops.split(value, lengths, name=name)
+      return self.parent()
 
   def size(self, name=None):
     """See TensorArray."""
-    del name  # not meaningful in Eager mode
+    del name  # not meaningful when executing eagerly.
     return constant_op.constant(len(self._tensor_array))
 
   def close(self, name=None):
-    del name  # not meaningful in Eager mode
+    del name  # not meaningful when executing eagerly.
     del self._tensor_array[:]
-    return
 
 
 # TensorArray is designed to hide an underlying implementation object
@@ -789,6 +753,8 @@ class TensorArray(object):
         colocate_with_first_write_call=colocate_with_first_write_call,
         name=name)
 
+    self._implementation.parent = weakref.ref(self)
+
   @property
   def flow(self):
     """The flow `Tensor` forcing ops leading to this TensorArray state."""
-- 
GitLab


From 6e1784b6b4e0542de0ac3ebd790633c6db9cfe46 Mon Sep 17 00:00:00 2001
From: Skye Wanderman-Milne <skyewm@google.com>
Date: Mon, 7 May 2018 16:16:32 -0700
Subject: [PATCH 337/395] ShapeRefiner fix: some variant-type tensors have
 handle data.

ShapeRefiner::AddNode() would only propagate handle data for
DT_RESOURCE tensors, but not DT_VARIANT. The Python shape inference
logic in common_shapes.py handled this correct, which is why we didn't
notice this earlier. In particular, list ops use DT_VARIANT with
handle data.
PiperOrigin-RevId: 195739586
---
 tensorflow/core/common_runtime/shape_refiner.cc | 13 ++++++-------
 tensorflow/python/kernel_tests/list_ops_test.py |  1 +
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensorflow/core/common_runtime/shape_refiner.cc b/tensorflow/core/common_runtime/shape_refiner.cc
index 06dbe04986..a0772713d4 100644
--- a/tensorflow/core/common_runtime/shape_refiner.cc
+++ b/tensorflow/core/common_runtime/shape_refiner.cc
@@ -232,13 +232,12 @@ Status ShapeRefiner::AddNode(const Node* node) {
     input_nodes[e->dst_input()] = input;
     input_shapes[e->dst_input()] = c->output(e->src_output());
 
-    // Only propagate handle data of edges which are carrying resource handles.
-    if (e->src()->output_type(e->src_output()) == DT_RESOURCE) {
-      const auto* in_v = c->output_handle_shapes_and_types(e->src_output());
-      if (in_v != nullptr) {
-        input_handle_shapes_and_types[e->dst_input()].reset(
-            new std::vector<ShapeAndType>(*in_v));
-      }
+    const auto* in_v = c->output_handle_shapes_and_types(e->src_output());
+    if (in_v != nullptr) {
+      DataType input_type = e->src()->output_type(e->src_output());
+      DCHECK(input_type == DT_RESOURCE || input_type == DT_VARIANT);
+      input_handle_shapes_and_types[e->dst_input()].reset(
+          new std::vector<ShapeAndType>(*in_v));
     }
   }
 
diff --git a/tensorflow/python/kernel_tests/list_ops_test.py b/tensorflow/python/kernel_tests/list_ops_test.py
index 098f9724a2..49855200c2 100644
--- a/tensorflow/python/kernel_tests/list_ops_test.py
+++ b/tensorflow/python/kernel_tests/list_ops_test.py
@@ -43,6 +43,7 @@ def scalar_shape():
   return ops.convert_to_tensor([], dtype=dtypes.int32)
 
 
+@test_util.with_c_shapes
 class ListOpsTest(test_util.TensorFlowTestCase):
 
   @test_util.run_in_graph_and_eager_modes()
-- 
GitLab


From 97fea64e69fbac87867343d92d2b47a2a582a79f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 16:31:07 -0700
Subject: [PATCH 338/395] Reorder executor NodeItem variable length data
 section so that all multi-byte aligned types precede all byte-aligned types
 so that alignment is satisfied without padding.

PiperOrigin-RevId: 195741712
---
 tensorflow/core/common_runtime/executor.cc | 26 ++++++++++------------
 1 file changed, 12 insertions(+), 14 deletions(-)

diff --git a/tensorflow/core/common_runtime/executor.cc b/tensorflow/core/common_runtime/executor.cc
index e389eb9b2a..7d63626b95 100644
--- a/tensorflow/core/common_runtime/executor.cc
+++ b/tensorflow/core/common_runtime/executor.cc
@@ -272,9 +272,9 @@ struct NodeItem {
   // (uint8 is enough for DataType).
   //   EdgeInfo            out_edges[num_out_edges];
   //   AllocatorAttributes output_attr[num_outputs];
+  //   int                 forward_from[num_outputs];
   //   uint8               input_type[num_inputs];
   //   uint8               output_type[num_outputs];
-  //   int                 forward_from[num_outputs];
 
   // Return pointer to variable length section.
   char* var() const {
@@ -289,22 +289,20 @@ struct NodeItem {
     return reinterpret_cast<AllocatorAttributes*>(var() + sizeof(EdgeInfo) *
                                                               num_output_edges);
   }
+  int* forward_from_base() const {
+    return reinterpret_cast<int*>(var() + sizeof(EdgeInfo) * num_output_edges +
+                                  sizeof(AllocatorAttributes) * num_outputs);
+  }
   uint8* input_type_base() const {
-    return reinterpret_cast<uint8*>(var() +
-                                    sizeof(EdgeInfo) * num_output_edges +
-                                    sizeof(AllocatorAttributes) * num_outputs);
+    return reinterpret_cast<uint8*>(
+        var() + sizeof(EdgeInfo) * num_output_edges +
+        sizeof(AllocatorAttributes) * num_outputs + sizeof(int) * num_outputs);
   }
   uint8* output_type_base() const {
     return reinterpret_cast<uint8*>(
         var() + sizeof(EdgeInfo) * num_output_edges +
-        sizeof(AllocatorAttributes) * num_outputs + sizeof(uint8) * num_inputs);
-  }
-
-  int* forward_from_base() const {
-    return reinterpret_cast<int*>(var() + sizeof(EdgeInfo) * num_output_edges +
-                                  sizeof(AllocatorAttributes) * num_outputs +
-                                  sizeof(uint8) * num_inputs +
-                                  sizeof(uint8) * num_outputs);
+        sizeof(AllocatorAttributes) * num_outputs + sizeof(int) * num_outputs +
+        sizeof(uint8) * num_inputs);
   }
 
   TF_DISALLOW_COPY_AND_ASSIGN(NodeItem);
@@ -481,9 +479,9 @@ size_t GraphView::NodeItemBytes(const Node* n) {
       sizeof(NodeItem)                             // Fixed
       + num_output_edges * sizeof(EdgeInfo)        // output_edges[...]
       + num_outputs * sizeof(AllocatorAttributes)  // output_attr[...]
+      + num_outputs * sizeof(int)                  // forward_from[num_outputs]
       + num_inputs * sizeof(uint8)                 // input_type[num_inputs]
-      + num_outputs * sizeof(uint8)                // output_type[num_outputs]
-      + num_outputs * sizeof(int);                 // forward_from[num_outputs]
+      + num_outputs * sizeof(uint8);               // output_type[num_outputs]
   static constexpr size_t kItemAlignment = sizeof(NodeItem*);
   static_assert(kItemAlignment % alignof(NodeItem) == 0,
                 "NodeItem must be aligned with kItemAlignment");
-- 
GitLab


From f64b16aa146aada0b2d20cafc0036a71f7460228 Mon Sep 17 00:00:00 2001
From: Billy Lamberta <blamb@google.com>
Date: Mon, 7 May 2018 16:38:02 -0700
Subject: [PATCH 339/395] Add TFX section. Add Ecosystem page and dropdown
 menu.

PiperOrigin-RevId: 195742728
---
 tensorflow/docs_src/deploy/index.md | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensorflow/docs_src/deploy/index.md b/tensorflow/docs_src/deploy/index.md
index 61edba04b4..3322004189 100644
--- a/tensorflow/docs_src/deploy/index.md
+++ b/tensorflow/docs_src/deploy/index.md
@@ -15,3 +15,7 @@ the following documents:
     out-of-the-box integration with TensorFlow models.
     [Source code for TensorFlow Serving](https://github.com/tensorflow/serving)
     is available on GitHub.
+
+[TensorFlow Extended (TFX)](/tfx) is an end-to-end machine learning platform for
+TensorFlow. Implemented at Google, we've open sourced some TFX libraries with the
+rest of the system to come.
-- 
GitLab


From a72ee2f74061cdd72f1197eed4c90a8216d39d74 Mon Sep 17 00:00:00 2001
From: Alexandre Passos <apassos@google.com>
Date: Mon, 7 May 2018 16:49:44 -0700
Subject: [PATCH 340/395] Fast-path to VarHandleOp

PiperOrigin-RevId: 195744374
---
 tensorflow/core/framework/resource_mgr.h | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/framework/resource_mgr.h b/tensorflow/core/framework/resource_mgr.h
index c84ea3b034..3cc17e1ca6 100644
--- a/tensorflow/core/framework/resource_mgr.h
+++ b/tensorflow/core/framework/resource_mgr.h
@@ -338,6 +338,9 @@ class ResourceHandleOp : public OpKernel {
  private:
   string container_;
   string name_;
+  mutex mutex_;
+  Tensor resource_ GUARDED_BY(mutex_);
+  std::atomic<bool> initialized_{false};
 };
 
 // Registers a kernel for an op which produces a handle to a resource of the
@@ -511,10 +514,17 @@ ResourceHandleOp<T>::ResourceHandleOp(OpKernelConstruction* context)
 
 template <typename T>
 void ResourceHandleOp<T>::Compute(OpKernelContext* ctx) {
-  Tensor* output = nullptr;
-  OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({}), &output));
-  output->scalar<ResourceHandle>()() =
-      MakeResourceHandle<T>(ctx, container_, name_);
+  if (!initialized_.load()) {
+    mutex_lock ml(mutex_);
+    AllocatorAttributes attr;
+    attr.set_on_host(true);
+    OP_REQUIRES_OK(ctx, ctx->allocate_temp(DT_RESOURCE, TensorShape({}),
+                                           &resource_, attr));
+    resource_.scalar<ResourceHandle>()() =
+        MakeResourceHandle<T>(ctx, container_, name_);
+    initialized_.store(true);
+  }
+  ctx->set_output(0, resource_);
 }
 
 }  //  end namespace tensorflow
-- 
GitLab


From 3964bdeef88cb9f7824bbfc8ca4f44c7a4bd4dbd Mon Sep 17 00:00:00 2001
From: Sanjoy Das <sanjoy@google.com>
Date: Mon, 7 May 2018 16:55:10 -0700
Subject: [PATCH 341/395] Delete kTransposeDot (it is no longer in use)

PiperOrigin-RevId: 195745124
---
 .../xla/service/cpu/cpu_layout_assignment.cc  | 10 ++---
 .../xla/service/cpu/dot_op_emitter.cc         | 13 ------
 .../compiler/xla/service/cpu/ir_emitter.cc    |  2 -
 .../service/cpu/parallel_task_assignment.cc   |  2 +-
 .../xla/service/gpu/ir_emission_utils.cc      |  4 --
 .../compiler/xla/service/hlo_instruction.cc   | 23 +---------
 .../compiler/xla/service/hlo_instruction.h    |  1 -
 .../xla/service/hlo_instruction_test.cc       | 11 +----
 .../compiler/xla/service/liveness_util.cc     | 22 ++++------
 .../xla/service/liveness_util_test.cc         | 42 -------------------
 10 files changed, 15 insertions(+), 115 deletions(-)

diff --git a/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc b/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc
index e8117377e6..6c642080c3 100644
--- a/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc
+++ b/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment.cc
@@ -139,13 +139,9 @@ Status CpuLayoutAssignment::AddBackendConstraints(
       Shape lhs_shape(RowMajorShape(lhs_instruction->shape()));
       TF_RETURN_IF_ERROR(constraints->SetOperandLayout(lhs_shape, dot, 0));
 
-      // dot is a kDot or a kTransposeDot fusion node.  In the latter case, if
-      // it represents X @ X, it may have just one operand.
-      if (dot->operand_count() > 1) {
-        const HloInstruction* rhs_instruction = dot->operand(1);
-        Shape rhs_shape(RowMajorShape(rhs_instruction->shape()));
-        TF_RETURN_IF_ERROR(constraints->SetOperandLayout(rhs_shape, dot, 1));
-      }
+      const HloInstruction* rhs_instruction = dot->operand(1);
+      Shape rhs_shape(RowMajorShape(rhs_instruction->shape()));
+      TF_RETURN_IF_ERROR(constraints->SetOperandLayout(rhs_shape, dot, 1));
 
       // Set layouts of the instructions' shapes.
       TF_RETURN_IF_ERROR(constraints->SetInstructionLayout(output_shape, dot));
diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
index e5ac2a33e1..8db4a0650d 100644
--- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
+++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc
@@ -1098,19 +1098,6 @@ bool PotentiallyImplementedAsEigenDot(const HloInstruction& hlo) {
     }
   }
 
-  if (hlo.opcode() == HloOpcode::kFusion &&
-      hlo.fusion_kind() == HloInstruction::FusionKind::kTransposeDot &&
-      hlo.fused_expression_root()->opcode() == HloOpcode::kDot) {
-    auto* dot = hlo.fused_expression_root();
-    const Shape& lhs_shape = dot->operand(0)->shape();
-    const Shape& rhs_shape = dot->operand(1)->shape();
-    if (ShapeUtil::HasZeroElements(lhs_shape) ||
-        ShapeUtil::HasZeroElements(rhs_shape)) {
-      return false;
-    }
-    return true;
-  }
-
   return false;
 }
 
diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
index 12f50e00b5..55e5aa5063 100644
--- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
+++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc
@@ -2077,8 +2077,6 @@ static const HloInstruction* StripTranspose(const HloInstruction& hlo) {
 }
 
 Status IrEmitter::HandleFusion(HloInstruction* fusion) {
-  CHECK_NE(fusion->fusion_kind(), HloInstruction::FusionKind::kTransposeDot);
-
   auto* root = fusion->fused_expression_root();
   if (llvm_ir::CanEmitFusedDynamicUpdateSliceInPlace(fusion, assignment_)) {
     VLOG(3) << "HandleFusion FusedDynamicUpdateSliceInPlace";
diff --git a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc
index fb28280fad..47e8405ff2 100644
--- a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc
+++ b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc
@@ -127,7 +127,7 @@ int64 ParallelTaskAssignment::GetTargetParallelTaskCount(
   // Currently, we do not assign parallel tasks to instructions with at least
   // one of the following properties:
   // *) Internal threading (library calls to kConv, kDot, kFft, kCustomCall).
-  // *) Emit custom loops (kSelectAndScatter, FusionKind::kTransposeDot).
+  // *) Emit custom loops (kSelectAndScatter).
   // *) Operations that are not thread safe (like infeed and rng).
   // *) Tuple-shaped.
   // TODO(b/27458679) Parallelize instructions which are skipped here.
diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
index 777345722c..96199035b9 100644
--- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
+++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
@@ -85,10 +85,6 @@ bool ImplementedAsGemm(const HloInstruction& hlo) {
     }
   }
 
-  if (hlo.opcode() == HloOpcode::kFusion) {
-    CHECK_NE(hlo.fusion_kind(), HloInstruction::FusionKind::kTransposeDot);
-  }
-
   if (hlo.opcode() == HloOpcode::kFusion &&
       hlo.fusion_kind() == HloInstruction::FusionKind::kOutput &&
       hlo.fused_expression_root()->opcode() == HloOpcode::kMultiply) {
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.cc b/tensorflow/compiler/xla/service/hlo_instruction.cc
index f9189077a1..857cd39adb 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.cc
+++ b/tensorflow/compiler/xla/service/hlo_instruction.cc
@@ -793,23 +793,11 @@ HloInstruction::CreateBroadcastSequence(
   return instruction;
 }
 
-// We put the fusion kind into the instruction's name for transpose-dot fusions,
-// since those fusions are really just describing a type of dot rather than
-// generating a novel computation.
-static string FusionNodeName(HloInstruction::FusionKind fusion_kind) {
-  switch (fusion_kind) {
-    case HloInstruction::FusionKind::kTransposeDot:
-      return "dot_fusion";
-    default:
-      return "fusion";
-  }
-}
-
 /* static */ std::unique_ptr<HloInstruction> HloInstruction::CreateFusion(
     const Shape& shape, FusionKind fusion_kind, HloInstruction* fused_root) {
   auto instruction = WrapUnique(new HloInstruction(HloOpcode::kFusion, shape));
   instruction->fusion_kind_ = fusion_kind;
-  instruction->name_ = FusionNodeName(fusion_kind);
+  instruction->name_ = "fusion";
   instruction->set_parent(fused_root->parent());
   instruction->set_metadata(fused_root->metadata());
   instruction->CloneAndFuseInternal(fused_root);
@@ -825,7 +813,7 @@ static string FusionNodeName(HloInstruction::FusionKind fusion_kind) {
     instruction->AppendOperand(operand);
   }
   instruction->fusion_kind_ = fusion_kind;
-  instruction->name_ = FusionNodeName(fusion_kind);
+  instruction->name_ = "fusion";
   instruction->called_computations_.push_back(fusion_computation);
   fusion_computation->SetFusionInstruction(instruction.get());
   return instruction;
@@ -2442,8 +2430,6 @@ string HloInstruction::ToCategory() const {
         return "input fusion";
       case FusionKind::kOutput:
         return "output fusion";
-      case FusionKind::kTransposeDot:
-        return "dot";
       case FusionKind::kCustom:
         return "custom fusion";
     }
@@ -3226,8 +3212,6 @@ string ToString(HloInstruction::FusionKind kind) {
       return "kInput";
     case HloInstruction::FusionKind::kOutput:
       return "kOutput";
-    case HloInstruction::FusionKind::kTransposeDot:
-      return "kTransposeDot";
     case HloInstruction::FusionKind::kCustom:
       return "kCustom";
   }
@@ -3244,9 +3228,6 @@ StatusOr<HloInstruction::FusionKind> StringToFusionKind(
   if (kind_name == "kOutput") {
     return HloInstruction::FusionKind::kOutput;
   }
-  if (kind_name == "kTransposeDot") {
-    return HloInstruction::FusionKind::kTransposeDot;
-  }
   if (kind_name == "kCustom") {
     return HloInstruction::FusionKind::kCustom;
   }
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h
index 0bf2c589e4..14be58d069 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction.h
+++ b/tensorflow/compiler/xla/service/hlo_instruction.h
@@ -177,7 +177,6 @@ class HloInstruction {
     kOutput,        // Op's output is fused into the op itself.
                     // REQUIRES: At least one operand buffer must be able
                     // to alias the output buffer.
-    kTransposeDot,  // Fused into a dot with transposed operands.
     kCustom,        // Custom category for backend-specific fusions that
                     // do not match any of the more specific ones.
   };
diff --git a/tensorflow/compiler/xla/service/hlo_instruction_test.cc b/tensorflow/compiler/xla/service/hlo_instruction_test.cc
index 5b65b1152c..909cdc0b62 100644
--- a/tensorflow/compiler/xla/service/hlo_instruction_test.cc
+++ b/tensorflow/compiler/xla/service/hlo_instruction_test.cc
@@ -1102,7 +1102,7 @@ TEST_F(HloInstructionTest, CloneOfFusionPreservesShape) {
   auto module = CreateNewModule();
   auto* computation = module->AddEntryComputation(builder.Build());
   HloInstruction* fusion = computation->CreateFusionInstruction(
-      {dot, reshape}, HloInstruction::FusionKind::kTransposeDot);
+      {dot, reshape}, HloInstruction::FusionKind::kLoop);
 
   auto fusion2 = fusion->Clone();
   const HloInstruction* root = fusion->fused_expression_root();
@@ -1169,7 +1169,7 @@ TEST_F(HloInstructionTest, NestedFusionEquality) {
   auto computation = module->AddEntryComputation(builder.Build());
 
   auto nested_fusion = computation->CreateFusionInstruction(
-      {dot, b_t}, HloInstruction::FusionKind::kTransposeDot);
+      {dot, b_t}, HloInstruction::FusionKind::kLoop);
 
   auto fusion = computation->CreateFusionInstruction(
       {add, nested_fusion}, HloInstruction::FusionKind::kOutput);
@@ -1246,13 +1246,6 @@ TEST_F(HloInstructionTest, Stringification) {
 
   auto module = CreateNewModule();
   auto* computation = module->AddEntryComputation(builder.Build());
-  HloInstruction* fusion = computation->CreateFusionInstruction(
-      {dot, reshape}, HloInstruction::FusionKind::kTransposeDot);
-
-  EXPECT_EQ(
-      fusion->ToString(options),
-      "%dot_fusion = f32[5,20]{1,0} fusion(f32[5,10]{1,0} %x, "
-      "f32[20,10]{1,0} %y), kind=kTransposeDot, calls=%fused_computation");
 
   HloInstruction* loop = builder.AddInstruction(
       HloInstruction::CreateWhile(sout, computation, computation, x));
diff --git a/tensorflow/compiler/xla/service/liveness_util.cc b/tensorflow/compiler/xla/service/liveness_util.cc
index 68c99256a2..79dfd1e409 100644
--- a/tensorflow/compiler/xla/service/liveness_util.cc
+++ b/tensorflow/compiler/xla/service/liveness_util.cc
@@ -173,9 +173,9 @@ bool HasUniqueFusedUseOfOperandAt(
 // (2) Is a loop fusion instruction where the only use of 'operand' at 'index'
 //     in the set 'user.fused_instructions' is a DynamicUpdateSlice fused root
 //     at operand 0. Or...
-// (3) Is a kDot -> kAdd (or fused kTransposeDot -> kAdd) output fusion
-//     instruction where the only use of 'operand' at 'index' in the set
-//     'user.fused_instructions' is a kAdd fused root at operand 0 or 1. Or...
+// (3) Is a kDot -> kAdd output fusion instruction where the only use of
+//     'operand' at 'index' in the set 'user.fused_instructions' is a kAdd fused
+//     root at operand 0 or 1. Or...
 // (4) The 'user' of 'operand' is DynamicUpdateSlice or While at operand index
 //     0.
 //
@@ -209,17 +209,13 @@ bool CanShareOperandBufferWithUser(
                user->fused_expression_root()->opcode() == HloOpcode::kAdd) {
       // Output fusion with kAdd fused root.
 
-      // Check if one operand of kAdd fused root is either kDot, or nested
-      // kFusion of kind kTransposeDot.
+      // Check if one operand of kAdd fused root is kDot or kConvolution.
       auto* add = user->fused_expression_root();
       auto add_operand_it =
           std::find_if(add->operands().begin(), add->operands().end(),
                        [&](HloInstruction* operand) {
                          return operand->opcode() == HloOpcode::kConvolution ||
-                                operand->opcode() == HloOpcode::kDot ||
-                                (operand->opcode() == HloOpcode::kFusion &&
-                                 operand->fusion_kind() ==
-                                     HloInstruction::FusionKind::kTransposeDot);
+                                operand->opcode() == HloOpcode::kDot;
                        });
       if (add_operand_it == add->operands().end()) {
         return false;
@@ -314,17 +310,13 @@ bool CanShareOperandBufferWithUser(HloInstruction* operand,
                user->fused_expression_root()->opcode() == HloOpcode::kAdd) {
       // Output fusion with kAdd fused root.
 
-      // Check if one operand of kAdd fused root is either kDot, or nested
-      // kFusion of kind kTransposeDot.
+      // Check if one operand of kAdd fused root is kDot, or kConvolution.
       auto* add = user->fused_expression_root();
       auto add_operand_it =
           std::find_if(add->operands().begin(), add->operands().end(),
                        [&](HloInstruction* operand) {
                          return operand->opcode() == HloOpcode::kConvolution ||
-                                operand->opcode() == HloOpcode::kDot ||
-                                (operand->opcode() == HloOpcode::kFusion &&
-                                 operand->fusion_kind() ==
-                                     HloInstruction::FusionKind::kTransposeDot);
+                                operand->opcode() == HloOpcode::kDot;
                        });
       if (add_operand_it == add->operands().end()) {
         return false;
diff --git a/tensorflow/compiler/xla/service/liveness_util_test.cc b/tensorflow/compiler/xla/service/liveness_util_test.cc
index f8b309488e..c01b52df62 100644
--- a/tensorflow/compiler/xla/service/liveness_util_test.cc
+++ b/tensorflow/compiler/xla/service/liveness_util_test.cc
@@ -303,48 +303,6 @@ TEST_F(CanShareOperandBufferWithUserTest, FusedDotAdd) {
                                             *dataflow_analysis_));
 }
 
-TEST_F(CanShareOperandBufferWithUserTest, FusedTransposeDotAdd) {
-  auto builder = HloComputation::Builder(TestName());
-  Shape data_shape = ShapeUtil::MakeShape(F32, {2, 2});
-
-  auto a = builder.AddInstruction(HloInstruction::CreateConstant(
-      Literal::CreateR2<float>({{1.0, 0.0}, {0.0, 1.0}})));
-  auto b = builder.AddInstruction(HloInstruction::CreateConstant(
-      Literal::CreateR2<float>({{2.0, 2.0}, {2.0, 2.0}})));
-  auto b_t = builder.AddInstruction(
-      HloInstruction::CreateTranspose(data_shape, b, {1, 0}));
-
-  DotDimensionNumbers dot_dnums;
-  dot_dnums.add_lhs_contracting_dimensions(1);
-  dot_dnums.add_rhs_contracting_dimensions(0);
-  auto dot = builder.AddInstruction(
-      HloInstruction::CreateDot(data_shape, a, b_t, dot_dnums));
-
-  auto one = builder.AddInstruction(
-      HloInstruction::CreateConstant(Literal::CreateR0<float>(1.0)));
-  auto add_operand = builder.AddInstruction(
-      HloInstruction::CreateBroadcast(data_shape, one, {1}));
-
-  auto add = builder.AddInstruction(HloInstruction::CreateBinary(
-      data_shape, HloOpcode::kAdd, dot, add_operand));
-
-  BuildModule(builder.Build());
-
-  auto nested_fusion = computation_->CreateFusionInstruction(
-      {dot, b_t}, HloInstruction::FusionKind::kTransposeDot);
-
-  auto fusion = computation_->CreateFusionInstruction(
-      {add, nested_fusion}, HloInstruction::FusionKind::kOutput);
-  RunAnalysis();
-
-  // Output fused transpose-dot-add should be share buffer with 'add_operand'.
-  EXPECT_TRUE(CanShareOperandBufferWithUser(add_operand, {}, fusion, {},
-                                            *points_to_analysis_));
-
-  EXPECT_TRUE(CanShareOperandBufferWithUser(add_operand, {}, fusion, {},
-                                            *dataflow_analysis_));
-}
-
 TEST_F(CanShareOperandBufferWithUserTest, OutputFusionCantAliasOperandBuffer) {
   auto builder = HloComputation::Builder(TestName());
   Shape data_shape = ShapeUtil::MakeShape(F32, {2, 2});
-- 
GitLab


From db63348bf14d911f2eebeb418a0b570b65b64f92 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar@google.com>
Date: Mon, 7 May 2018 16:59:41 -0700
Subject: [PATCH 342/395] Add test with tf.cond.

PiperOrigin-RevId: 195745718
---
 tensorflow/compiler/aot/tests/BUILD           | 14 +++++++++
 .../compiler/aot/tests/make_test_graphs.py    | 29 ++++++++++++-------
 .../aot/tests/test_graph_tfcond.config.pbtxt  | 20 +++++++++++++
 .../compiler/aot/tests/tfcompile_test.cc      | 26 +++++++++++++++++
 4 files changed, 79 insertions(+), 10 deletions(-)
 create mode 100644 tensorflow/compiler/aot/tests/test_graph_tfcond.config.pbtxt

diff --git a/tensorflow/compiler/aot/tests/BUILD b/tensorflow/compiler/aot/tests/BUILD
index 222e26810a..fd2cf2b67d 100644
--- a/tensorflow/compiler/aot/tests/BUILD
+++ b/tensorflow/compiler/aot/tests/BUILD
@@ -15,6 +15,7 @@ test_suite(
         ":test_graph_tfadd_with_ckpt_saver_test",
         ":test_graph_tfadd_with_ckpt_test",
         ":test_graph_tfassert_eq_test",
+        ":test_graph_tfcond_test",
         ":test_graph_tffunction_test",
         ":test_graph_tfgather_test",
         ":test_graph_tfmatmul_test",
@@ -55,6 +56,7 @@ genrule(
         "test_graph_tfadd_with_ckpt_saver.pb",
         "test_graph_tfadd_with_ckpt_saver.saver",
         "test_graph_tfassert_eq.pb",
+        "test_graph_tfcond.pb",
         "test_graph_tffunction.pb",
         "test_graph_tfgather.pb",
         "test_graph_tfmatmul.pb",
@@ -118,6 +120,17 @@ tf_library(
     ],
 )
 
+tf_library(
+    name = "test_graph_tfcond",
+    testonly = 1,
+    config = "test_graph_tfcond.config.pbtxt",
+    cpp_class = "CondComp",
+    graph = "test_graph_tfcond.pb",
+    tags = [
+        "manual",
+    ],
+)
+
 tf_library(
     name = "test_graph_tffunction",
     testonly = 1,
@@ -194,6 +207,7 @@ tf_cc_test(
         ":test_graph_tfadd_with_ckpt",
         ":test_graph_tfadd_with_ckpt_saver",
         ":test_graph_tfassert_eq",
+        ":test_graph_tfcond",
         ":test_graph_tffunction",
         ":test_graph_tfgather",
         ":test_graph_tfmatmul",
diff --git a/tensorflow/compiler/aot/tests/make_test_graphs.py b/tensorflow/compiler/aot/tests/make_test_graphs.py
index 67767f55da..9ec7df163b 100644
--- a/tensorflow/compiler/aot/tests/make_test_graphs.py
+++ b/tensorflow/compiler/aot/tests/make_test_graphs.py
@@ -78,6 +78,22 @@ def tfadd_with_ckpt_saver(out_dir):
       f.write(saver.as_saver_def().SerializeToString())
 
 
+def tfassert_eq(_):
+  x = array_ops.placeholder(dtypes.int32, name='x_hold')
+  y = array_ops.placeholder(dtypes.int32, name='y_hold')
+  control_flow_ops.Assert(
+      math_ops.equal(x, y), ['Expected x == y.'], name='assert_eq')
+  math_ops.add(x, math_ops.negative(y), name='x_y_diff')
+
+
+def tfcond(_):
+  p = array_ops.placeholder(dtypes.bool, name='p_hold')
+  x = array_ops.placeholder(dtypes.int32, name='x_hold')
+  y = array_ops.placeholder(dtypes.int32, name='y_hold')
+  z = control_flow_ops.cond(p, lambda: x, lambda: y)
+  array_ops.identity(z, name='result')
+
+
 def tfgather(_):
   params = array_ops.placeholder(dtypes.float32, name='params')
   indices = array_ops.placeholder(dtypes.int32, name='indices')
@@ -126,14 +142,6 @@ def tfsplits(_):
   array_ops.identity(y, name='result')
 
 
-def tfassert_eq(_):
-  x = array_ops.placeholder(dtypes.int32, name='x_hold')
-  y = array_ops.placeholder(dtypes.int32, name='y_hold')
-  control_flow_ops.Assert(
-      math_ops.equal(x, y), ['Expected x == y.'], name='assert_eq')
-  math_ops.add(x, math_ops.negative(y), name='x_y_diff')
-
-
 def write_graph(build_graph, out_dir):
   """Build a graph using build_graph and write it out."""
   g = ops.Graph()
@@ -148,12 +156,13 @@ def main(_):
   write_graph(tfadd, FLAGS.out_dir)
   write_graph(tfadd_with_ckpt, FLAGS.out_dir)
   write_graph(tfadd_with_ckpt_saver, FLAGS.out_dir)
+  write_graph(tfassert_eq, FLAGS.out_dir)
+  write_graph(tfcond, FLAGS.out_dir)
+  write_graph(tffunction, FLAGS.out_dir)
   write_graph(tfgather, FLAGS.out_dir)
   write_graph(tfmatmul, FLAGS.out_dir)
   write_graph(tfmatmulandadd, FLAGS.out_dir)
-  write_graph(tffunction, FLAGS.out_dir)
   write_graph(tfsplits, FLAGS.out_dir)
-  write_graph(tfassert_eq, FLAGS.out_dir)
 
 
 if __name__ == '__main__':
diff --git a/tensorflow/compiler/aot/tests/test_graph_tfcond.config.pbtxt b/tensorflow/compiler/aot/tests/test_graph_tfcond.config.pbtxt
new file mode 100644
index 0000000000..94a01ad4ab
--- /dev/null
+++ b/tensorflow/compiler/aot/tests/test_graph_tfcond.config.pbtxt
@@ -0,0 +1,20 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed {
+  id { node_name: "p_hold" }
+  shape {}
+}
+feed {
+  id { node_name: "x_hold" }
+  shape {
+    dim { size: 1 }
+  }
+}
+feed {
+  id { node_name: "y_hold" }
+  shape {
+    dim { size: 1 }
+  }
+}
+fetch {
+  id { node_name: "result" }
+}
diff --git a/tensorflow/compiler/aot/tests/tfcompile_test.cc b/tensorflow/compiler/aot/tests/tfcompile_test.cc
index 27ba42b31f..309a991fc1 100644
--- a/tensorflow/compiler/aot/tests/tfcompile_test.cc
+++ b/tensorflow/compiler/aot/tests/tfcompile_test.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include "tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt.h"
 #include "tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt_saver.h"
 #include "tensorflow/compiler/aot/tests/test_graph_tfassert_eq.h"
+#include "tensorflow/compiler/aot/tests/test_graph_tfcond.h"
 #include "tensorflow/compiler/aot/tests/test_graph_tffunction.h"
 #include "tensorflow/compiler/aot/tests/test_graph_tfgather.h"
 #include "tensorflow/compiler/aot/tests/test_graph_tfmatmul.h"
@@ -150,6 +151,31 @@ TEST(TFCompileTest, AddWithCkptSaver) {
   EXPECT_EQ(add_const.result0_data(), add_const.results()[0]);
 }
 
+TEST(TFCompileTest, Cond) {
+  CondComp cond;
+  EXPECT_EQ(cond.arg0_data(), cond.args()[0]);
+  EXPECT_EQ(cond.arg1_data(), cond.args()[1]);
+  EXPECT_EQ(cond.arg2_data(), cond.args()[2]);
+  cond.arg1() = 10;
+  cond.arg2() = 20;
+  {
+    cond.arg0() = true;
+    const int32 expected_result = cond.arg1();
+    EXPECT_TRUE(cond.Run());
+    EXPECT_EQ(cond.result0(), expected_result);
+    EXPECT_EQ(cond.result0_data()[0], expected_result);
+    EXPECT_EQ(cond.result0_data(), cond.results()[0]);
+  }
+  {
+    cond.arg0() = false;
+    const int32 expected_result = cond.arg2();
+    EXPECT_TRUE(cond.Run());
+    EXPECT_EQ(cond.result0(), expected_result);
+    EXPECT_EQ(cond.result0_data()[0], expected_result);
+    EXPECT_EQ(cond.result0_data(), cond.results()[0]);
+  }
+}
+
 TEST(TFCompileTest, Gather) {
   GatherComp gather;
   EXPECT_EQ(gather.arg0_data(), gather.args()[0]);
-- 
GitLab


From b67d8b278d48a046491b42eccbd5c5c23975d054 Mon Sep 17 00:00:00 2001
From: Blake Hechtman <blakehechtman@google.com>
Date: Mon, 7 May 2018 17:00:27 -0700
Subject: [PATCH 343/395] Internal change

PiperOrigin-RevId: 195745819
---
 tensorflow/compiler/xla/tests/BUILD | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD
index 1c29abcb80..b982cf0dbc 100644
--- a/tensorflow/compiler/xla/tests/BUILD
+++ b/tensorflow/compiler/xla/tests/BUILD
@@ -1867,6 +1867,8 @@ xla_test(
 
 xla_test(
     name = "local_client_execute_test",
+    # TODO(b/79375911): Test times out in LLVM at normal size.
+    size = "large",
     srcs = ["local_client_execute_test.cc"],
     shard_count = 30,
     tags = ["optonly"],
-- 
GitLab


From 482ed8eb666d8bc1e5c3f47e5c1e61cc19e0fdb1 Mon Sep 17 00:00:00 2001
From: Skye Wanderman-Milne <skyewm@google.com>
Date: Mon, 7 May 2018 17:21:39 -0700
Subject: [PATCH 344/395] Raise an error if we try to take the gradient wrt to
 the initial value of a loop variable.

Fixes #14101

PiperOrigin-RevId: 195748688
---
 .../kernel_tests/control_flow_ops_py_test.py  | 17 ++++++++
 tensorflow/python/ops/gradients_impl.py       | 39 +++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/tensorflow/python/kernel_tests/control_flow_ops_py_test.py b/tensorflow/python/kernel_tests/control_flow_ops_py_test.py
index 77e6f5f1a0..843759fed0 100644
--- a/tensorflow/python/kernel_tests/control_flow_ops_py_test.py
+++ b/tensorflow/python/kernel_tests/control_flow_ops_py_test.py
@@ -1847,6 +1847,23 @@ class ControlFlowTest(test.TestCase):
       r = control_flow_ops.cond(math_ops.less(1, 2), fn1, lambda: x)
       self.assertAllClose(9.0, r.eval(feed_dict={x: 1.0}))
 
+  def testGradInWhileWrtInitialLoopVal(self):
+    with self.test_session():
+      x = array_ops.placeholder(dtypes.float32, shape=(), name="x")
+      y = x + 1
+
+      def body(i, v):
+        z = v * 2
+        return i + 1, gradients_impl.gradients(z, x)[0]
+
+      with self.assertRaisesRegexp(
+          ValueError,
+          "Cannot compute gradient inside while loop with respect to op 'x'. "
+          "We do not support taking the gradient wrt or through the initial "
+          "value of a loop variable. Gradients can be computed through "
+          "loop invariants or wrt the input parameters to the loop body."):
+        control_flow_ops.while_loop(lambda i, x: i < 3, body, [0, y])
+
   def testWhileGradInWhile(self):
     with self.test_session():
       n = ops.convert_to_tensor(1.0, name="n")
diff --git a/tensorflow/python/ops/gradients_impl.py b/tensorflow/python/ops/gradients_impl.py
index a6b1e6df54..069b5a4308 100644
--- a/tensorflow/python/ops/gradients_impl.py
+++ b/tensorflow/python/ops/gradients_impl.py
@@ -418,6 +418,30 @@ def _MaybeCompile(scope, op, func, grad_fn):
     return grad_fn()
 
 
+def _RaiseNoGradWrtInitialLoopValError(op, from_ops):
+  """Raises an error if we backprop through a loop var."""
+  # Find the nearest 'to_op' reachable from 'op' to provide a more helpful error
+  # message.
+  target_op = None
+  queue = collections.deque([op])
+  visited = set()
+  while queue:
+    curr_op = queue.popleft()
+    if curr_op in visited: continue
+    visited.add(curr_op)
+    if curr_op in from_ops:
+      target_op = curr_op
+      break
+    queue.extend(t.op for t in curr_op.inputs)
+  assert target_op
+  raise ValueError(
+      "Cannot compute gradient inside while loop with respect to op '%s'. "
+      "We do not support taking the gradient wrt or through the initial value "
+      "of a loop variable. Gradients can be computed through loop invariants "
+      "or wrt the input parameters to the loop body."
+      % target_op.name)
+
+
 @tf_export("gradients")
 def gradients(ys,
               xs,
@@ -630,6 +654,21 @@ def _GradientsHelper(ys, xs, grad_ys, name, colocate_gradients_with_ops,
                   (op.name, op.type))
         if loop_state:
           loop_state.EnterGradWhileContext(op, before=False)
+
+        # NOTE(skyewm): We don't support computing gradients wrt a loop variable
+        # unless it's within the context of a single iteration (i.e. the
+        # gradient is wrt to the loop parameter in the body function, not wrt or
+        # through the initial value). This means if we're in a while loop
+        # context, we should never see a switch node from this context.
+        # pylint: disable=protected-access
+        if (control_flow_util.IsSwitch(op) and
+            op._control_flow_context is not None and
+            op._control_flow_context.IsWhileContext() and
+            op._control_flow_context ==
+            ops.get_default_graph()._get_control_flow_context()):
+          _RaiseNoGradWrtInitialLoopValError(op, from_ops)
+        # pylint: enable=protected-access
+
         if (grad_fn or is_func_call) and has_out_grads:
           # NOTE: If _AggregatedGrads didn't compute a value for the i'th
           # output, it means that the cost does not depend on output[i],
-- 
GitLab


From b6906d19bacffa25fec074216a5c281e5689ef03 Mon Sep 17 00:00:00 2001
From: Igor Ganichev <iga@google.com>
Date: Mon, 7 May 2018 17:21:53 -0700
Subject: [PATCH 345/395] Make eager functions runable on TPU

PiperOrigin-RevId: 195748721
---
 tensorflow/compiler/jit/BUILD                 |  22 ++
 .../compiler/jit/create_xla_launch_op.cc      | 206 ++++++++++++++----
 .../compiler/jit/create_xla_launch_op.h       |  35 +++
 .../compiler/jit/create_xla_launch_op_test.cc | 144 ++++++++++++
 .../compiler/jit/kernels/xla_launch_op.cc     |  90 ++++++--
 .../compiler/jit/kernels/xla_launch_op.h      |  51 +++--
 .../compiler/jit/xla_compile_on_demand_op.cc  |   3 +-
 tensorflow/compiler/jit/xla_launch_util.cc    |  18 +-
 tensorflow/compiler/jit/xla_launch_util.h     |  15 +-
 tensorflow/compiler/tests/BUILD               |   4 +
 tensorflow/compiler/tests/eager_test.py       | 112 +++++++++-
 .../python/examples/resnet50/resnet50_test.py |  55 +++--
 tensorflow/python/eager/function.py           | 127 +++++++----
 13 files changed, 718 insertions(+), 164 deletions(-)
 create mode 100644 tensorflow/compiler/jit/create_xla_launch_op.h
 create mode 100644 tensorflow/compiler/jit/create_xla_launch_op_test.cc

diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD
index 07136d6a74..e942b46086 100644
--- a/tensorflow/compiler/jit/BUILD
+++ b/tensorflow/compiler/jit/BUILD
@@ -261,6 +261,7 @@ cc_library(
     name = "create_xla_launch_op",
     srcs = [
         "create_xla_launch_op.cc",
+        "create_xla_launch_op.h",
     ],
     deps = [
         ":common",
@@ -270,6 +271,27 @@ cc_library(
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+    ],
+    alwayslink = 1,
+)
+
+tf_cc_test(
+    name = "create_xla_launch_op_test",
+    srcs = [
+        "create_xla_launch_op.h",
+        "create_xla_launch_op_test.cc",
+    ],
+    deps = [
+        ":create_xla_launch_op",
+        "//tensorflow/core:core_cpu_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:session_options",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
     ],
 )
 
diff --git a/tensorflow/compiler/jit/create_xla_launch_op.cc b/tensorflow/compiler/jit/create_xla_launch_op.cc
index 18d901323f..6ac84dc19c 100644
--- a/tensorflow/compiler/jit/create_xla_launch_op.cc
+++ b/tensorflow/compiler/jit/create_xla_launch_op.cc
@@ -12,6 +12,7 @@ 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.
 ==============================================================================*/
+#include "tensorflow/compiler/jit/create_xla_launch_op.h"
 
 #include "tensorflow/compiler/jit/defs.h"
 #include "tensorflow/compiler/jit/kernels/xla_launch_op.h"
@@ -25,78 +26,189 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
-// Givens a NodeDef 'ndef' and the function library runtime 'flr', if
-// 'ndef' is a call to a compilable function defined in 'flr', returns OK
-// and fills in 'kernel' with a XlaLaunchOp kernel which computes the
-// node. Otherwise, returns a non-OK.
+// Utility which searches for values in a sorted list by scanning over it once.
+// No matter how many times ScanForValue is called, the list is scanned at most
+// once. However, if a call to ScanForValue skips over a value, that value is
+// not revisited in future calls to ScanForValue, so callers must take
+// care to order their calls.
 //
-// This routine is here so that FunctionLibraryRuntime can jit a
-// specific function call as requested.
-Status CreateXlaLaunchOp(FunctionLibraryRuntime* flr, const NodeDef& ndef,
-                         std::unique_ptr<OpKernel>* kernel) {
-  bool xla_compile = false;
-  if (!flr->GetFunctionLibraryDefinition()
-           ->GetAttr(ndef, kXlaCompileAttr, &xla_compile)
-           .ok() ||
-      !xla_compile) {
-    // Not marked as _XlaCompile=true.
-    return errors::InvalidArgument("No ", kXlaCompileAttr, " for ", ndef.op());
+// Useful for merging multiple sorted lists in O(n) time.
+class SinglePassSearch {
+ public:
+  // Creates a SinglePassSearch object that can be used to search in `values`.
+  // Does not take ownership of `values`. `values` must outlive this.
+  // `values` must be sorted.
+  explicit SinglePassSearch(const std::vector<int>* values)
+      : current_index_(0), values_(values) {}
+
+  // Scans forward in the vector looking for "value", updating the internal
+  // position in to the vector.
+  // Returns true iff the vector contains the given value at or after current
+  // position.
+  // Not thread-safe.
+  bool ScanForValue(int value) {
+    while (current_index_ < values_->size() &&
+           (*values_)[current_index_] <= value) {
+      if ((*values_)[current_index_] == value) {
+        current_index_++;
+        return true;
+      }
+      current_index_++;
+    }
+    return false;
   }
-  // Make sure that kernels have been registered on the JIT device.
-  XlaOpRegistry::RegisterCompilationKernels();
-  if (!IsCompilable(flr, ndef)) {
-    // ndef is calling a function that XLA can't compile.
-    return errors::InvalidArgument("Not compilable: ", ndef.ShortDebugString());
+
+ private:
+  int current_index_;
+  const std::vector<int>* values_;
+};
+
+Status CompilationRequested(const FunctionLibraryRuntime& flr,
+                            const NodeDef& node_def) {
+  bool xla_compile = false;
+  // Check if op is marked _XlaCompile=true.
+  Status status = flr.GetFunctionLibraryDefinition()->GetAttr(
+      node_def, kXlaCompileAttr, &xla_compile);
+  if (!status.ok() || !xla_compile) {
+    if (VLOG_IS_ON(3)) {
+      if (!status.ok()) {
+        VLOG(3) << "No " << kXlaCompileAttr << " attr defined for "
+                << node_def.op() << ". status=" << status.ToString();
+      } else {
+        VLOG(3) << node_def.op() << " is explicitly marked not to be compiled";
+      }
+    }
+    return Status(error::INVALID_ARGUMENT, "");
   }
+  return Status::OK();
+}
+
+// Given a FunctionLibraryRuntime and a NodeDef calling a function in the
+// runtime, returns this function's body in `fbody` as well as the indices
+// of its constant and resource arguments.
+// `fbody` is owned by `flr`.
+// `constant_arg_indices` and `resource_arg_indices` should be empty vector.
+// They are sorted in ascending order on this function's return.
+Status GetBodyAndConstantsAndResources(FunctionLibraryRuntime* flr,
+                                       const NodeDef& node_def,
+                                       const FunctionBody** fbody,
+                                       std::vector<int>* constant_arg_indices,
+                                       std::vector<int>* resource_arg_indices) {
   FunctionLibraryRuntime::Handle handle;
-  // If ndef is not instantiable, e.g., the function does not exist,
+  // If node_def is not instantiable, e.g., the function does not exist,
   // simply bail out.
   TF_RETURN_IF_ERROR(
-      flr->Instantiate(ndef.op(), AttrSlice(&ndef.attr()), &handle));
-  const FunctionBody* fbody = flr->GetFunctionBody(handle);
-  CHECK(fbody);  // Can't be nullptr since we just instantiated it.
-  std::vector<bool> const_args(fbody->arg_types.size());
+      flr->Instantiate(node_def.op(), AttrSlice(&node_def.attr()), &handle));
+  *fbody = flr->GetFunctionBody(handle);
+  CHECK(*fbody);  // Can't be nullptr since we just instantiated it.
+  const DataTypeVector& arg_types = (*fbody)->arg_types;
+  std::vector<bool> const_args(arg_types.size());
   // If we can't analyze the const args. Bail out.
-  TF_RETURN_IF_ERROR(BackwardsConstAnalysis(*(fbody->graph), &const_args));
+  TF_RETURN_IF_ERROR(BackwardsConstAnalysis(*((*fbody)->graph), &const_args));
 
   for (int i = 0; i < const_args.size(); ++i) {
     if (const_args[i]) {
-      // There is a const arg. Bail out.
-      return errors::InvalidArgument("Const arg: ", i, " in ",
-                                     DebugString(fbody->fdef));
+      constant_arg_indices->push_back(i);
+    }
+  }
+
+  // There can be hundreds of resource variables. Reserve the space for them.
+  // We don't reserve for constants above as they are usually few.
+  resource_arg_indices->reserve(arg_types.size());
+  for (int i = 0; i < arg_types.size(); ++i) {
+    if (arg_types[i] == DT_RESOURCE) {
+      resource_arg_indices->push_back(i);
     }
   }
 
-  NodeDef launch_def;
-  launch_def.set_name(ndef.name());
-  launch_def.set_op("_XlaLaunch");
-  launch_def.set_device(flr->device()->name());
-  AddNodeAttr("Tconstants", DataTypeVector{}, &launch_def);
-  AddNodeAttr("Nresources", 0, &launch_def);
-  AddNodeAttr("Targs", fbody->arg_types, &launch_def);
-  AddNodeAttr("Tresults", fbody->ret_types, &launch_def);
-  NameAttrList func;
-  func.set_name(ndef.op());
-  *(func.mutable_attr()) = ndef.attr();
-  AddNodeAttr("function", func, &launch_def);
-
-  // TODO(b/32387911): Handles the host memory types across function
-  // calls properly. For now, we assume all inputs and outputs are on
-  // the device memory.
+  return Status::OK();
+}
+
+}  // namespace
+
+Status CreateXlaLaunchOp(FunctionLibraryRuntime* flr, const NodeDef& node_def,
+                         std::unique_ptr<OpKernel>* kernel) {
+  TF_RETURN_IF_ERROR(CompilationRequested(*flr, node_def));
+
+  VLOG(3) << "Creating XlaLaunchOp for " << node_def.DebugString();
+
+  // Make sure that kernels have been registered on the JIT device.
+  XlaOpRegistry::RegisterCompilationKernels();
+  if (!IsCompilable(flr, node_def)) {
+    // node_def is calling a function that XLA can't compile.
+    return errors::InvalidArgument("Not compilable: ",
+                                   node_def.ShortDebugString());
+  }
+
+  // Get function body, constant args, and resource args.
+  const FunctionBody* fbody = nullptr;
+  std::vector<int> constant_arg_indices;
+  std::vector<int> resource_arg_indices;
+  TF_RETURN_IF_ERROR(GetBodyAndConstantsAndResources(
+      flr, node_def, &fbody, &constant_arg_indices, &resource_arg_indices));
+
+  // Set input and output memory types.
   MemoryTypeVector input_memory_types(fbody->arg_types.size(), DEVICE_MEMORY);
+  // These indices are used only for optimization purposes. They allow us
+  // to loop over constant_arg_indices and resource_arg_indices only once
+  // while iterating over all the function arguments checking if it is a
+  // resource or a constant.
+  // The reason we optimized this code is because functions can have a lot of
+  // captured arguments. For example, the backward pass of ResNet50 takes in all
+  // 214 variables and a similar number of activations.
+  SinglePassSearch constants_search(&constant_arg_indices);
+  SinglePassSearch resources_search(&resource_arg_indices);
+  for (int i = 0; i < fbody->arg_types.size(); ++i) {
+    if (resources_search.ScanForValue(i) || constants_search.ScanForValue(i)) {
+      // Compile-time constants and resource handles are expected to be in
+      // host memory.
+      input_memory_types[i] = HOST_MEMORY;
+    }
+  }
+  // One might wonder, about the case where a compile-time constant argument
+  // (which must be in host memory) is also used as an input into an op,
+  // e.g. Add, that expects its inputs in device memory. Here is how it
+  // works now.
+  // First, what do we mean by "op expects an input in XYZ memory"?
+  // There are two types of "ops" here: the tf2xla kernel and the HLO
+  // computation it builds. The tf2xla kernel needs to retrieve the actual
+  // numeric value of the compile-time constant tensors, so it really expects
+  // them to be on in host memory. However, for other inputs, it refers to them
+  // using xla::ComputationDataHandle, which is just a symbolic handle that
+  // xla::ComputationBuilder assigns. How does this handle gets assigned for
+  // constant arguments? Even constant arguments get an _Arg node in the graph
+  // instatiated for Function compilation. The tf2xla kernel for constant _Arg
+  // nodes takes the constant value, converts it to XlaLiteral, and feeds it
+  // to xla::ComputationBuilder.ConstantLiteral, which returns the handle. This
+  // constant XlaLiteral is included in the HLO graph, and subsequently, in
+  // the actual executable, which is copied to the device before being
+  // executed. Thus, when this executable runs, the constant is available in
+  // device memory.
+
+  // XlaLaunch kernel keeps all outputs (including constants, which it copies),
+  // in device memory
   MemoryTypeVector output_memory_types(fbody->ret_types.size(), DEVICE_MEMORY);
 
+  // Create the kernel.
+  NameAttrList function;
+  function.set_name(node_def.op());
+  *(function.mutable_attr()) = node_def.attr();
+
   Device* dev = flr->device();
   Status s;
   OpKernelConstruction construction(
       DeviceType(dev->device_type()), dev,
-      dev->GetAllocator(AllocatorAttributes()), &launch_def,
+      dev->GetAllocator(AllocatorAttributes()), &node_def,
       &fbody->fdef.signature(), flr, fbody->arg_types, input_memory_types,
       fbody->ret_types, output_memory_types, flr->graph_def_version(), &s);
-  kernel->reset(new XlaLocalLaunchOp(&construction));
+
+  *kernel = absl::make_unique<XlaLocalLaunchBase>(
+      &construction, constant_arg_indices, resource_arg_indices, function);
   return s;
 }
 
+namespace {
+
 bool RegisterLaunchOpCreator() {
   RegisterDefaultCustomKernelCreator(CreateXlaLaunchOp);
   return true;
diff --git a/tensorflow/compiler/jit/create_xla_launch_op.h b/tensorflow/compiler/jit/create_xla_launch_op.h
new file mode 100644
index 0000000000..98a22e3515
--- /dev/null
+++ b/tensorflow/compiler/jit/create_xla_launch_op.h
@@ -0,0 +1,35 @@
+/* Copyright 2018 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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_JIT_CREATE_XLA_LAUNCH_OP_H_
+#define TENSORFLOW_COMPILER_JIT_CREATE_XLA_LAUNCH_OP_H_
+
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+
+class FunctionLibraryRuntime;
+class OpKernel;
+
+// Given a NodeDef 'node_def' and the function library runtime 'flr', if
+// 'node_def' is a call to a compilable function defined in 'flr', returns OK
+// and fills in 'kernel' with a XlaLaunchOp kernel which computes the
+// node. Otherwise, returns a non-OK.
+Status CreateXlaLaunchOp(FunctionLibraryRuntime* flr, const NodeDef& node_def,
+                         std::unique_ptr<OpKernel>* kernel);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_JIT_CREATE_XLA_LAUNCH_OP_H_
diff --git a/tensorflow/compiler/jit/create_xla_launch_op_test.cc b/tensorflow/compiler/jit/create_xla_launch_op_test.cc
new file mode 100644
index 0000000000..c222824eda
--- /dev/null
+++ b/tensorflow/compiler/jit/create_xla_launch_op_test.cc
@@ -0,0 +1,144 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/jit/create_xla_launch_op.h"
+
+#include "tensorflow/core/common_runtime/device_factory.h"
+#include "tensorflow/core/common_runtime/function.h"
+#include "tensorflow/core/framework/function_testlib.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/public/session_options.h"
+#include "tensorflow/core/public/version.h"
+
+namespace tensorflow {
+
+NodeDef ToNodeDef(const string& text) {
+  NodeDef node_def;
+  EXPECT_TRUE(protobuf::TextFormat::MergeFromString(text, &node_def));
+  return node_def;
+}
+
+// Create a FunctionDef that takes one resource and one regular param
+FunctionDef XTimesY() {
+  return FunctionDefHelper::Define(
+      // Name
+      "XTimesY",
+      // Args
+      {"x: float", "y: resource"},
+      // Return values
+      {"z: float"},
+      // Attr def
+      {},
+      // Nodes
+      {
+          {{"y0"}, "ReadVariableOp", {"y"}, {{"dtype", DT_FLOAT}}},
+          {{"z"}, "Mul", {"x", "y0"}, {{"T", DT_FLOAT}}},
+      });
+}
+
+class CreateXlaLaunchOpTest : public ::testing::Test {
+ protected:
+  void Init(const std::vector<FunctionDef>& flib) {
+    SessionOptions options;
+    auto* device_count = options.config.mutable_device_count();
+    device_count->insert({"CPU", 1});
+    TF_CHECK_OK(DeviceFactory::AddDevices(
+        options, "/job:localhost/replica:0/task:0", &devices_));
+
+    FunctionDefLibrary proto;
+    for (const auto& fdef : flib) {
+      *(proto.add_function()) = fdef;
+    }
+    lib_def_ = absl::make_unique<FunctionLibraryDefinition>(
+        OpRegistry::Global(), proto);
+    OptimizerOptions opts;
+    device_mgr_ = absl::make_unique<DeviceMgr>(devices_);
+    pflr_ = absl::make_unique<ProcessFunctionLibraryRuntime>(
+        device_mgr_.get(), Env::Default(), TF_GRAPH_DEF_VERSION, lib_def_.get(),
+        opts, /*default_thread_pool=*/nullptr, /*cluster_flr=*/nullptr);
+    flr_ = pflr_->GetFLR("/job:localhost/replica:0/task:0/cpu:0");
+  }
+
+  FunctionLibraryRuntime* flr_;
+  std::vector<Device*> devices_;
+  std::unique_ptr<DeviceMgr> device_mgr_;
+  std::unique_ptr<FunctionLibraryDefinition> lib_def_;
+  std::unique_ptr<ProcessFunctionLibraryRuntime> pflr_;
+
+  std::unique_ptr<OpKernel> kernel_;
+};
+
+AttrValue BoolAttr(bool b) {
+  AttrValue v;
+  v.set_b(b);
+  return v;
+}
+
+TEST_F(CreateXlaLaunchOpTest, OneFloatOneResourceArgument) {
+  FunctionDef fdef = XTimesY();
+  (*fdef.mutable_attr())["_XlaCompile"] = BoolAttr(true);
+  Init({fdef});
+
+  Status status = CreateXlaLaunchOp(
+      flr_, ToNodeDef(R"pb(
+        name: 'XTimesY' op: 'XTimesY' input: 'a' input: 'b'
+      )pb"), &kernel_);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+
+  EXPECT_EQ("XTimesY", kernel_->name());
+  EXPECT_EQ("XTimesY", kernel_->type_string());
+
+  EXPECT_EQ(2, kernel_->num_inputs());
+  EXPECT_EQ(DT_FLOAT, kernel_->input_type(0));
+  EXPECT_EQ(DT_RESOURCE, kernel_->input_type(1));
+  EXPECT_EQ(DEVICE_MEMORY, kernel_->input_memory_types()[0]);
+  EXPECT_EQ(HOST_MEMORY, kernel_->input_memory_types()[1]);
+
+  EXPECT_EQ(1, kernel_->num_outputs());
+  EXPECT_EQ(DT_FLOAT, kernel_->output_type(0));
+  EXPECT_EQ(DEVICE_MEMORY, kernel_->output_memory_types()[0]);
+}
+
+TEST_F(CreateXlaLaunchOpTest, FailsIfXlaCompileAttrNotSet) {
+  FunctionDef fdef = XTimesY();
+  Init({fdef});
+
+  Status status = CreateXlaLaunchOp(flr_, ToNodeDef(R"proto(
+                                      name: 'XTimesY'
+                                      op: 'XTimesY'
+                                      input: 'a'
+                                      input: 'b'
+                                    )proto"), &kernel_);
+  EXPECT_TRUE(errors::IsInvalidArgument(status)) << status.ToString();
+}
+
+TEST_F(CreateXlaLaunchOpTest, FailsIfXlaCompileAttrIsSetToFalse) {
+  FunctionDef fdef = XTimesY();
+  (*fdef.mutable_attr())["_XlaCompile"] = BoolAttr(false);
+  Init({fdef});
+
+  Status status = CreateXlaLaunchOp(flr_, ToNodeDef(R"proto(
+                                      name: 'XTimesY'
+                                      op: 'XTimesY'
+                                      input: 'a'
+                                      input: 'b'
+                                    )proto"), &kernel_);
+  EXPECT_TRUE(errors::IsInvalidArgument(status)) << status.ToString();
+}
+
+}  // namespace tensorflow
diff --git a/tensorflow/compiler/jit/kernels/xla_launch_op.cc b/tensorflow/compiler/jit/kernels/xla_launch_op.cc
index 049d170fa4..86a9fd3b8e 100644
--- a/tensorflow/compiler/jit/kernels/xla_launch_op.cc
+++ b/tensorflow/compiler/jit/kernels/xla_launch_op.cc
@@ -39,15 +39,15 @@ limitations under the License.
 
 namespace tensorflow {
 
-XlaLocalLaunchOp::XlaLocalLaunchOp(OpKernelConstruction* ctx)
-    : OpKernel(ctx), device_type_(ctx->device_type()) {
-  const NameAttrList* func;
-  OP_REQUIRES_OK(ctx, ctx->GetAttr("function", &func));
-  function_ = *func;
-  DataTypeVector constant_types;
-  OP_REQUIRES_OK(ctx, ctx->GetAttr("Tconstants", &constant_types));
-  num_constant_args_ = constant_types.size();
-  OP_REQUIRES_OK(ctx, ctx->GetAttr("Nresources", &num_resource_args_));
+XlaLocalLaunchBase::XlaLocalLaunchBase(OpKernelConstruction* ctx,
+                                       const std::vector<int>& constants,
+                                       const std::vector<int>& resources,
+                                       const NameAttrList& function)
+    : OpKernel(ctx),
+      constants_(constants),
+      resources_(resources),
+      device_type_(ctx->device_type()),
+      function_(function) {
   if (device_type_ == DeviceType(DEVICE_CPU)) {
     platform_id_ = se::host::kHostPlatformId;
   } else if (device_type_ == DeviceType(DEVICE_GPU)) {
@@ -57,8 +57,8 @@ XlaLocalLaunchOp::XlaLocalLaunchOp(OpKernelConstruction* ctx)
   }
 }
 
-Status XlaLocalLaunchOp::BuildCompilationCache(OpKernelContext* ctx,
-                                               XlaCompilationCache** cache) {
+Status XlaLocalLaunchBase::BuildCompilationCache(OpKernelContext* ctx,
+                                                 XlaCompilationCache** cache) {
   const XlaDevice::Metadata* metadata;
   Status s = XlaDevice::GetMetadata(ctx, &metadata);
   if (s.ok()) {
@@ -90,8 +90,8 @@ Status XlaLocalLaunchOp::BuildCompilationCache(OpKernelContext* ctx,
   return Status::OK();
 }
 
-void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) {
-  VLOG(1) << "XlaLocalLaunchOp::Compute "
+void XlaLocalLaunchBase::Compute(OpKernelContext* ctx) {
+  VLOG(1) << "XlaLocalLaunchOpBase::Compute "
           << Canonicalize(function_.name(), AttrSlice(&function_.attr()));
   // We store information about the JIT-compiled XLA computation
   // in the ResourceMgr.
@@ -124,7 +124,7 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) {
   }
 
   std::map<int, OptionalTensor> variables =
-      SnapshotResourceVariables(ctx, num_resource_args_);
+      SnapshotResourceVariables(ctx, resources_);
 
   xla::LocalClient* client = static_cast<xla::LocalClient*>(cache->client());
 
@@ -161,7 +161,7 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) {
   xla::LocalExecutable* executable;
 
   std::map<int, Tensor> constant_args;
-  for (int i = 0; i < num_constant_args_; ++i) {
+  for (int i : constants_) {
     constant_args.insert({i, ctx->input(i)});
   }
   OP_REQUIRES_OK(ctx, cache->Compile(options, function_, constant_args,
@@ -170,8 +170,8 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) {
 
   VLOG(1) << "Executing XLA Computation...";
 
-  XlaComputationLaunchContext launch_context(
-      num_resource_args_, client, xla_allocator, allocate_xla_tensors);
+  XlaComputationLaunchContext launch_context(client, xla_allocator,
+                                             allocate_xla_tensors);
   launch_context.PopulateInputs(ctx, kernel, variables);
 
   // Execute the computation.
@@ -194,6 +194,62 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) {
   VLOG(1) << "Done";
 }
 
+namespace {
+
+// OP_REQUIRES_OK_RETURN is the same as OP_REQUIRES_OK except that
+// in error case, it returns RET instead of void.
+#define OP_REQUIRES_OK_RETURN(CTX, RET, ...)                \
+  do {                                                      \
+    ::tensorflow::Status _s(__VA_ARGS__);                   \
+    if (!TF_PREDICT_TRUE(_s.ok())) {                        \
+      (CTX)->CtxFailureWithWarning(__FILE__, __LINE__, _s); \
+      return RET;                                           \
+    }                                                       \
+  } while (0)
+
+// Helper static functions to construct parameters for
+// XlaLocalLaunchBase constructor from OpKernelConstruction.
+std::vector<int> ConstantsVector(OpKernelConstruction* ctx) {
+  DataTypeVector constant_types;
+  OP_REQUIRES_OK_RETURN(ctx, std::vector<int>(),
+                        ctx->GetAttr("Tconstants", &constant_types));
+  std::vector<int> constants(constant_types.size());
+  std::iota(constants.begin(), constants.end(), 0);
+  return constants;
+}
+
+std::vector<int> ResourcesVector(OpKernelConstruction* ctx) {
+  DataTypeVector constant_types;
+  OP_REQUIRES_OK_RETURN(ctx, std::vector<int>(),
+                        ctx->GetAttr("Tconstants", &constant_types));
+
+  DataTypeVector arg_types;
+  OP_REQUIRES_OK_RETURN(ctx, std::vector<int>(),
+                        ctx->GetAttr("Targs", &arg_types));
+
+  int num_resources;
+  OP_REQUIRES_OK_RETURN(ctx, std::vector<int>(),
+                        ctx->GetAttr("Nresources", &num_resources));
+
+  std::vector<int> resources(num_resources);
+  std::iota(resources.begin(), resources.end(),
+            constant_types.size() + arg_types.size());
+  return resources;
+}
+
+NameAttrList FunctionAttr(OpKernelConstruction* ctx) {
+  const NameAttrList* func;
+  OP_REQUIRES_OK_RETURN(ctx, NameAttrList(), ctx->GetAttr("function", &func));
+  return *func;
+}
+
+#undef OP_REQUIRES_OK_RETURN
+}  // namespace
+
+XlaLocalLaunchOp::XlaLocalLaunchOp(OpKernelConstruction* ctx)
+    : XlaLocalLaunchBase(ctx, ConstantsVector(ctx), ResourcesVector(ctx),
+                         FunctionAttr(ctx)) {}
+
 XlaLocalLaunchOp::~XlaLocalLaunchOp() {
   VLOG(1) << "XlaLocalLaunchOp destroyed";
 }
diff --git a/tensorflow/compiler/jit/kernels/xla_launch_op.h b/tensorflow/compiler/jit/kernels/xla_launch_op.h
index 8f8e646f0f..8dfc4b382d 100644
--- a/tensorflow/compiler/jit/kernels/xla_launch_op.h
+++ b/tensorflow/compiler/jit/kernels/xla_launch_op.h
@@ -26,6 +26,41 @@ limitations under the License.
 
 namespace tensorflow {
 
+// XlaLocalLaunchBase is almost the same as XlaLocalLaunchOp.
+// The only difference is that it does not require arguments to follow
+// the "constants, then regular args, then resources" order.
+// It takes vectors of constant and resource arguments explicitly.
+// It does not have corresponding OpDef because it is never present
+// in the GraphDef.
+// Currently, it is used by eager runtime. FunctionLibraryRuntime creates
+// this kernel when asked to create a kernel for an XLA-compiled function.
+class XlaLocalLaunchBase : public OpKernel {
+ public:
+  XlaLocalLaunchBase(OpKernelConstruction* ctx,
+                     const std::vector<int>& constants,
+                     const std::vector<int>& resources,
+                     const NameAttrList& function);
+  XlaLocalLaunchBase(const XlaLocalLaunchBase&) = delete;
+  XlaLocalLaunchBase& operator=(const XlaLocalLaunchBase&) = delete;
+  ~XlaLocalLaunchBase() override = default;
+
+  void Compute(OpKernelContext* ctx) override;
+
+ protected:
+  // Builds a XlaCompilationCache class suitable for the current device.
+  Status BuildCompilationCache(OpKernelContext* ctx,
+                               XlaCompilationCache** cache);
+
+  // Indexes of compile-time constant inputs
+  std::vector<int> constants_;
+  // Indexes of resource inputs
+  std::vector<int> resources_;
+
+  DeviceType device_type_;
+  NameAttrList function_;
+  se::Platform::Id platform_id_;
+};
+
 // XlaLocalLaunchOp is used to replace a region of the TensorFlow graph
 // which will be compiled and executed using XLA.  The XlaLocalLaunchOp is
 // responsible for handling interactions with the TensorFlow executor.
@@ -35,26 +70,12 @@ namespace tensorflow {
 // XlaLocalLaunchOp uses xla::LocalClient::Compile() and
 // xla::LocalExecutable::Run(), and passes arguments into/out of XLA in device
 // memory.
-class XlaLocalLaunchOp : public OpKernel {
+class XlaLocalLaunchOp : public XlaLocalLaunchBase {
  public:
   explicit XlaLocalLaunchOp(OpKernelConstruction* ctx);
   ~XlaLocalLaunchOp() override;
 
-  void Compute(OpKernelContext* ctx) override;
-
  private:
-  // Builds a XlaCompilationCache class suitable for the current device.
-  Status BuildCompilationCache(OpKernelContext* ctx,
-                               XlaCompilationCache** compiler);
-
-  DeviceType device_type_;
-  NameAttrList function_;
-  int num_constant_args_;
-  // Number of resource variable arguments.
-  int num_resource_args_;
-
-  se::Platform::Id platform_id_;
-
   TF_DISALLOW_COPY_AND_ASSIGN(XlaLocalLaunchOp);
 };
 
diff --git a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc
index 60458f6f33..6b83cf67ff 100644
--- a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc
+++ b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc
@@ -48,13 +48,12 @@ Status XlaCompileOnDemandOp::Run(OpKernelContext* ctx,
                                  const XlaCompiler::CompilationResult* result,
                                  xla::LocalExecutable* executable) {
   std::map<int, OptionalTensor> variables = GetVariables(ctx);
-  int64 num_resource_args = variables.size();
 
   xla::LocalClient* client = metadata.client();
 
   // Builds an XLA allocator for the device.
   XlaComputationLaunchContext launch_context(
-      num_resource_args, client, client->backend().memory_allocator(), true);
+      client, client->backend().memory_allocator(), true);
 
   launch_context.PopulateInputs(ctx, result, variables);
 
diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc
index 33e53612b9..0223f97a03 100644
--- a/tensorflow/compiler/jit/xla_launch_util.cc
+++ b/tensorflow/compiler/jit/xla_launch_util.cc
@@ -38,14 +38,13 @@ using xla::ScopedShapedBuffer;
 using xla::ShapedBuffer;
 }  // anonymous namespace
 
-std::map<int, OptionalTensor> SnapshotResourceVariables(OpKernelContext* ctx,
-                                                        int num_variables) {
+std::map<int, OptionalTensor> SnapshotResourceVariables(
+    OpKernelContext* ctx, const std::vector<int>& variables) {
   std::map<int, OptionalTensor> snapshot;
-  int first_variable = ctx->num_inputs() - num_variables;
-  for (int i = 0; i < num_variables; ++i) {
+  for (int i : variables) {
     Var* variable = nullptr;
-    ResourceHandle handle = HandleFromInput(ctx, first_variable + i);
-    OptionalTensor& tensor = snapshot[first_variable + i];
+    ResourceHandle handle = HandleFromInput(ctx, i);
+    OptionalTensor& tensor = snapshot[i];
     if (LookupResource(ctx, handle, &variable).ok()) {
       tf_shared_lock lock(*variable->mu());
       tensor.name = handle.name();
@@ -112,10 +111,9 @@ ScopedShapedBuffer ExtractSubShapedBuffer(
 using internal::ExtractSubShapedBuffer;
 
 XlaComputationLaunchContext::XlaComputationLaunchContext(
-    int64 num_resource_args, xla::LocalClient* client,
-    xla::DeviceMemoryAllocator* xla_allocator, bool allocate_xla_tensors)
-    : num_resource_args_(num_resource_args),
-      client_(client),
+    xla::LocalClient* client, xla::DeviceMemoryAllocator* xla_allocator,
+    bool allocate_xla_tensors)
+    : client_(client),
       xla_allocator_(xla_allocator),
       allocate_xla_tensors_(allocate_xla_tensors) {}
 
diff --git a/tensorflow/compiler/jit/xla_launch_util.h b/tensorflow/compiler/jit/xla_launch_util.h
index 38291b0bd4..a2431253f8 100644
--- a/tensorflow/compiler/jit/xla_launch_util.h
+++ b/tensorflow/compiler/jit/xla_launch_util.h
@@ -31,15 +31,17 @@ limitations under the License.
 namespace tensorflow {
 class XlaAllocator;
 
-// Takes a snapshot of the values of resource variable arguments, which are
-// the last `num_variables` arguments. We snapshot tensors that back
+// Takes a snapshot of the values of resource variable arguments, whose
+// indices are specified in `variables` argument. We snapshot tensors that back
 // resource variables since concurrent updates may modify the shape, and it is
 // important that the shapes used for compilation match the true shapes of the
 // buffers.
 //
-// Returns a map of TensorFlow argument index to resource variable.
-std::map<int, OptionalTensor> SnapshotResourceVariables(OpKernelContext* ctx,
-                                                        int num_variables);
+// Returns a map of TensorFlow argument index to resource variable. If a
+// resource variable is not initialized, the corresponding OptionalTensor
+// will have its `present` field set to false.
+std::map<int, OptionalTensor> SnapshotResourceVariables(
+    OpKernelContext* ctx, const std::vector<int>& variables);
 
 // Adapter class that wraps a Tensorflow allocator as an XLA allocator.
 // Assumes that the Tensorflow allocator permits asynchronous deallocation:
@@ -72,7 +74,7 @@ class XlaComputationLaunchContext {
   // Create a new launch context. 'allocate_xla_tensors' is true if allocated
   // output tensors and variables are always XlaTensors. If false they are
   // assumed to be "normal" device pointers.
-  XlaComputationLaunchContext(int64 num_resource_args, xla::LocalClient* client,
+  XlaComputationLaunchContext(xla::LocalClient* client,
                               xla::DeviceMemoryAllocator* xla_allocator,
                               bool allocate_xla_tensors);
 
@@ -92,7 +94,6 @@ class XlaComputationLaunchContext {
   const std::vector<xla::ShapedBuffer*>& arguments() const { return arg_ptrs_; }
 
  private:
-  int64 num_resource_args_;
   xla::LocalClient* client_;
   xla::DeviceMemoryAllocator* xla_allocator_;
   bool allocate_xla_tensors_;
diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index aaea83ae9c..9791792f29 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -327,7 +327,11 @@ tf_xla_py_test(
         ":xla_test",
         "//tensorflow/python:array_ops",
         "//tensorflow/python:framework_for_generated_wrappers",
+        "//tensorflow/python:layers",
+        "//tensorflow/python:math_ops",
+        "//tensorflow/python:nn",
         "//tensorflow/python:platform_test",
+        "//tensorflow/python/eager:function",
     ],
 )
 
diff --git a/tensorflow/compiler/tests/eager_test.py b/tensorflow/compiler/tests/eager_test.py
index bdd0185dfe..5ab1585f8c 100644
--- a/tensorflow/compiler/tests/eager_test.py
+++ b/tensorflow/compiler/tests/eager_test.py
@@ -24,10 +24,16 @@ from tensorflow.compiler.tests.xla_test import XLATestCase
 from tensorflow.core.protobuf import config_pb2
 from tensorflow.python.eager import backprop
 from tensorflow.python.eager import context
+from tensorflow.python.eager import function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
+from tensorflow.python.layers import convolutional
+from tensorflow.python.layers import pooling
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
 from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.platform import googletest
 
@@ -43,7 +49,7 @@ class EagerTest(XLATestCase):
 
   def testExecuteListOutputLen0(self):
     with self.test_scope():
-      empty = constant_op.constant([], dtype=dtypes.int32)
+      empty = constant_op.constant([], dtype=dtypes.float32)
       result = array_ops.unstack(empty, 0)
       self.assertTrue(isinstance(result, list))
       self.assertEqual(0, len(result))
@@ -51,7 +57,7 @@ class EagerTest(XLATestCase):
   def testExecuteListOutputLen1(self):
     with self.test_scope():
       split_dim = constant_op.constant(1)
-      value = constant_op.constant([[0, 1, 2], [3, 4, 5]])
+      value = constant_op.constant([[0., 1., 2.], [3., 4., 5.]])
       result = array_ops.split(value, 1, axis=split_dim)
       self.assertTrue(isinstance(result, list))
       self.assertEqual(1, len(result))
@@ -60,7 +66,7 @@ class EagerTest(XLATestCase):
   def testExecuteListOutputLen3(self):
     with self.test_scope():
       split_dim = constant_op.constant(1)
-      value = constant_op.constant([[0, 1, 2], [3, 4, 5]])
+      value = constant_op.constant([[0., 1., 2.], [3., 4., 5.]])
       result = array_ops.split(value, 3, axis=split_dim)
       self.assertTrue(isinstance(result, list))
       self.assertEqual(3, len(result))
@@ -131,7 +137,105 @@ class EagerTest(XLATestCase):
     self.assertEqual(2., grads[0][0].numpy())
 
 
-if __name__ == "__main__":
+class EagerFunctionTest(XLATestCase):
+
+  def testBasic(self):
+    with self.test_scope():
+      matmul = function.defun(math_ops.matmul, compiled=True)
+      t = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
+      sq = matmul(t, t, transpose_a=True)
+      self.assertAllEqual(sq.numpy().reshape(-1), [10, 14, 14, 20])
+
+  def testConv(self):
+    if 'GPU' in self.device:
+      # TODO(b/32333178)
+      self.skipTest('Current implementation of RandomStandardNormal kernel '
+                    'is very slow on GPU, and has been blacklisted.')
+    with self.test_scope():
+      data_format = 'channels_last'
+      conv = convolutional.Conv2D(
+          filters=1, kernel_size=2, padding='VALID',
+          data_format=data_format, activation=nn_ops.relu,
+          kernel_initializer=init_ops.ones_initializer(),
+          bias_initializer=init_ops.zeros_initializer())
+      pool = pooling.MaxPooling2D(2, 2, data_format=data_format)
+
+      def model(x):
+        x = conv(x)
+        return pool(x)
+      model = function.defun(model, compiled=True)
+
+      x = array_ops.ones([1, 4, 4, 1])
+      y = model(x)
+      self.assertAllEqual(y.numpy(), [[[[4.]]]])
+
+  def testReadVariable(self):
+    with self.test_scope():
+      v = resource_variable_ops.ResourceVariable(1.0)
+
+      @function.defun(compiled=True)
+      def f():
+        return v.read_value()
+
+      var = f()
+      self.assertEqual(1.0, var.numpy())
+
+  def testUpdateVariable(self):
+    with self.test_scope():
+      v = resource_variable_ops.ResourceVariable(1.0)
+
+      def f(v):
+        v.assign_add(1.0)
+        return v
+
+      f = function.defun(f, compiled=True)
+
+      var = f(v)
+      self.assertEqual(2.0, var.numpy())
+
+  def testAllArgumentKinds(self):
+    """Test a complex function that takes different argument kinds.
+
+    tf2xla machinery that translates, compiles, and runs defuns
+    classifies arguments into: compile-time constants, regular tensors,
+    and resources. This test creates a function with a mix of all these
+    kinds. Moreover, the order of function arguments is intentionally mixed up.
+
+    This also tests the case when the same argument is a compile-time constant
+    as well as used in an operation that normally expects its inputs to be
+    in device memory - addition in this case.
+    """
+    with self.test_scope():
+      def foo(c1, r1, v1, c2, v2, r2):
+        # c1 and c2 are compile-time constants
+        # r1 and r2 are regular tensors
+        # v1 and v2 are resource variables
+        a = c1 + r1
+        b = math_ops.cast(c2, dtypes.float32) + v2
+        c = array_ops.slice(v1, c1, c2)
+        d = r2 * v2
+        return a, b, c, d
+
+      foo = function.defun(foo, compiled=True)
+
+      c1 = [0, 0]
+      c2 = array_ops.ones([2], dtype=dtypes.int32)
+
+      r1 = array_ops.ones([2])
+      r2 = [[2., 2.], [3., 3.]]
+
+      v1 = resource_variable_ops.ResourceVariable([[1., 2.], [3., 4.]])
+      v2 = resource_variable_ops.ResourceVariable([[10., 20.], [30., 40.]])
+
+      a, b, c, d = foo(c1, r1, v1, c2, v2, r2)
+
+      self.assertAllEqual([1, 1], a.numpy())
+      self.assertAllEqual([[11., 21.], [31., 41.]], b.numpy())
+      self.assertAllEqual([[1.]], c.numpy())
+      self.assertAllEqual([[20., 40.], [90., 120.]], d.numpy())
+
+
+if __name__ == '__main__':
   ops.enable_eager_execution(
       config=config_pb2.ConfigProto(log_device_placement=True))
   googletest.main()
diff --git a/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py b/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py
index 8517a3bf7b..b8f352d5f5 100644
--- a/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py
+++ b/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py
@@ -36,9 +36,7 @@ def device_and_data_format():
                                                               'channels_last')
 
 
-def random_batch(batch_size, device_and_format=None):
-  _, data_format = device_and_format or device_and_data_format()
-
+def random_batch(batch_size, data_format):
   shape = (3, 224, 224) if data_format == 'channels_first' else (224, 224, 3)
   shape = (batch_size,) + shape
 
@@ -70,7 +68,7 @@ class ResNet50Test(tf.test.TestCase):
     if defun:
       model.call = tfe.defun(model.call)
     with tf.device(device), tfe.execution_mode(execution_mode):
-      images, _ = random_batch(2)
+      images, _ = random_batch(2, data_format)
       output = model(images, training=False)
       tfe.async_wait()
     self.assertEqual((2, 1000), output.shape)
@@ -91,7 +89,7 @@ class ResNet50Test(tf.test.TestCase):
     device, data_format = device_and_data_format()
     model = resnet50.ResNet50(data_format, include_top=False)
     with tf.device(device):
-      images, _ = random_batch(2)
+      images, _ = random_batch(2, data_format)
       output = model(images, training=False)
     output_shape = ((2, 2048, 1, 1)
                     if data_format == 'channels_first' else (2, 1, 1, 2048))
@@ -101,7 +99,7 @@ class ResNet50Test(tf.test.TestCase):
     device, data_format = device_and_data_format()
     model = resnet50.ResNet50(data_format, include_top=False, pooling='avg')
     with tf.device(device):
-      images, _ = random_batch(2)
+      images, _ = random_batch(2, data_format)
       output = model(images, training=False)
     self.assertEqual((2, 2048), output.shape)
 
@@ -115,7 +113,7 @@ class ResNet50Test(tf.test.TestCase):
         name='t0').as_default(), tf.contrib.summary.always_record_summaries():
       with tf.device(device), tfe.execution_mode(execution_mode):
         optimizer = tf.train.GradientDescentOptimizer(0.1)
-        images, labels = random_batch(2)
+        images, labels = random_batch(2, data_format)
         train_one_step(model, images, labels, optimizer)
         self.assertEqual(320, len(model.variables))
         tfe.async_wait()
@@ -134,7 +132,7 @@ class ResNet50Test(tf.test.TestCase):
     model = resnet50.ResNet50(data_format)
     optimizer = tf.train.GradientDescentOptimizer(0.1)
     with tf.device(device):
-      images, labels = random_batch(2)
+      images, labels = random_batch(2, data_format)
       gc.disable()
       # Warm up. Note that this first run does create significant amounts of
       # garbage to be collected. The hope is that this is a build-only effect,
@@ -202,18 +200,18 @@ class ResNet50Benchmarks(tf.test.Benchmark):
     # which forces a sync. This is a roundabout way, yes.
     tf.constant(1.).cpu()
 
-  def _benchmark_eager_apply(self, label, defun=False, execution_mode=None,
-                             device_and_format=None):
+  def _benchmark_eager_apply(self, label, device_and_format, defun=False,
+                             execution_mode=None, compiled=False):
     with tfe.execution_mode(execution_mode):
-      device, data_format = device_and_format or device_and_data_format()
+      device, data_format = device_and_format
       model = resnet50.ResNet50(data_format)
       if defun:
-        model.call = tfe.defun(model.call)
+        model.call = tfe.defun(model.call, compiled=compiled)
       batch_size = 64
       num_burn = 5
       num_iters = 30
       with tf.device(device):
-        images, _ = random_batch(batch_size, device_and_format)
+        images, _ = random_batch(batch_size, data_format)
         for _ in xrange(num_burn):
           model(images, training=False).cpu()
         if execution_mode:
@@ -227,30 +225,34 @@ class ResNet50Benchmarks(tf.test.Benchmark):
         self._report(label, start, num_iters, device, batch_size, data_format)
 
   def benchmark_eager_apply_sync(self):
-    self._benchmark_eager_apply('eager_apply', defun=False)
+    self._benchmark_eager_apply('eager_apply', device_and_data_format(),
+                                defun=False)
 
   def benchmark_eager_apply_async(self):
     self._benchmark_eager_apply(
-        'eager_apply_async', defun=False, execution_mode=tfe.ASYNC)
+        'eager_apply_async', device_and_data_format(), defun=False,
+        execution_mode=tfe.ASYNC)
 
   def benchmark_eager_apply_with_defun(self):
-    self._benchmark_eager_apply('eager_apply_with_defun', defun=True)
+    self._benchmark_eager_apply('eager_apply_with_defun',
+                                device_and_data_format(), defun=True)
 
   def _benchmark_eager_train(self,
                              label,
                              make_iterator,
+                             device_and_format,
                              defun=False,
                              execution_mode=None,
-                             device_and_format=None):
+                             compiled=False):
     with tfe.execution_mode(execution_mode):
-      device, data_format = device_and_format or device_and_data_format()
+      device, data_format = device_and_format
       for batch_size in self._train_batch_sizes():
-        (images, labels) = random_batch(batch_size, device_and_format)
+        (images, labels) = random_batch(batch_size, data_format)
         num_burn = 3
         num_iters = 10
         model = resnet50.ResNet50(data_format)
         if defun:
-          model.call = tfe.defun(model.call)
+          model.call = tfe.defun(model.call, compiled=compiled)
         optimizer = tf.train.GradientDescentOptimizer(0.1)
 
         with tf.device(device):
@@ -273,18 +275,21 @@ class ResNet50Benchmarks(tf.test.Benchmark):
           self._report(label, start, num_iters, device, batch_size, data_format)
 
   def benchmark_eager_train_sync(self):
-    self._benchmark_eager_train('eager_train', MockIterator, defun=False)
+    self._benchmark_eager_train('eager_train', MockIterator,
+                                device_and_data_format(), defun=False)
 
   def benchmark_eager_train_async(self):
     self._benchmark_eager_train(
         'eager_train_async',
         MockIterator,
+        device_and_data_format(),
         defun=False,
         execution_mode=tfe.ASYNC)
 
   def benchmark_eager_train_with_defun(self):
     self._benchmark_eager_train(
-        'eager_train_with_defun', MockIterator, defun=True)
+        'eager_train_with_defun', MockIterator,
+        device_and_data_format(), defun=True)
 
   def benchmark_eager_train_datasets(self):
 
@@ -294,7 +299,8 @@ class ResNet50Benchmarks(tf.test.Benchmark):
       return tfe.Iterator(ds)
 
     self._benchmark_eager_train(
-        'eager_train_dataset', make_iterator, defun=False)
+        'eager_train_dataset', make_iterator,
+        device_and_data_format(), defun=False)
 
   def benchmark_eager_train_datasets_with_defun(self):
 
@@ -304,7 +310,8 @@ class ResNet50Benchmarks(tf.test.Benchmark):
       return tfe.Iterator(ds)
 
     self._benchmark_eager_train(
-        'eager_train_dataset_with_defun', make_iterator, defun=True)
+        'eager_train_dataset_with_defun', make_iterator,
+        device_and_data_format(), defun=True)
 
 
 if __name__ == '__main__':
diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py
index 741bd2ac9c..60cfacc141 100644
--- a/tensorflow/python/eager/function.py
+++ b/tensorflow/python/eager/function.py
@@ -23,6 +23,7 @@ import collections
 
 import numpy as np
 
+from tensorflow.core.framework import attr_value_pb2
 from tensorflow.core.framework import function_pb2
 from tensorflow.python import pywrap_tensorflow
 from tensorflow.python.eager import context
@@ -225,7 +226,7 @@ def _inference_name(n):
 class _EagerDefinedFunction(object):
   """Function object with the interface of tf _DefinedFunction."""
 
-  def __init__(self, name, graph, operations, inputs, outputs):
+  def __init__(self, name, graph, operations, inputs, outputs, attrs):
     """Initializes an eager defined function.
 
     Args:
@@ -235,6 +236,7 @@ class _EagerDefinedFunction(object):
         which will be in the function
       inputs: the tensors in the graph to be used as inputs to the function
       outputs: the tensors in the graph which will be outputs to the function
+      attrs: dict mapping names of attributes to their AttrValue values
     """
     fn = pywrap_tensorflow.TF_GraphToFunction_wrapper(
         graph._c_graph,  # pylint: disable=protected-access
@@ -246,6 +248,14 @@ class _EagerDefinedFunction(object):
         [],
         None,
         compat.as_str(""))
+
+    for name, attr_value in attrs.items():
+      serialized = attr_value.SerializeToString()
+      # TODO(iga): this creates and deletes a new TF_Status for every attr.
+      # It might be worth creating a convenient way to re-use status.
+      pywrap_tensorflow.TF_FunctionSetAttrValueProto(
+          fn, compat.as_str(name), serialized)
+
     # TODO(apassos) avoid creating a FunctionDef (specially to grab the
     # signature, but also in general it's nice not to depend on it.
     with c_api_util.tf_buffer() as buffer_:
@@ -287,25 +297,6 @@ def _flatten(sequence):
 
 class GraphModeFunction(object):
   """Callable object representing a graph-mode function.
-
-  Args:
-    name: str the name of the created function
-    input_placeholders: list of placeholder values (tensors) to feed when
-      calling the wrapped function.
-    extra_inputs: Tensor inputs this function definition closed over which
-      are passed as arguments. Need to track so gradients are supported
-      correctly.
-    graph: the Graph from which the operations will be pulled. Used as
-      a context when computing gradients.
-    operations: the subset of Operations in the graph used in the function
-      definition.
-    outputs: a flat list of the Tensors in the graph used as outputs to the
-      function
-    func_outputs: a possibly nested python object which will be returned by
-      this function. The Tensors in this structure will be replaced by their
-      corresponding values in outputs.
-    output_shapes: List of shapes of all tensors in outputs
-    variables: (optional) List of variables to watch during function execution.
   """
 
   def __init__(self,
@@ -317,9 +308,36 @@ class GraphModeFunction(object):
                outputs,
                func_outputs,
                output_shapes,
-               variables=None):
+               variables=None,
+               attrs=None):
+    """Initialize a GraphModeFunction.
+
+    Args:
+      name: str the name of the created function
+      input_placeholders: list of placeholder values (tensors) to feed when
+        calling the wrapped function.
+      extra_inputs: Tensor inputs this function definition closed over which
+        are passed as arguments. Need to track so gradients are supported
+        correctly.
+      graph: the Graph from which the operations will be pulled. Used as
+        a context when computing gradients.
+      operations: the subset of Operations in the graph used in the function
+        definition.
+      outputs: a flat list of the Tensors in the graph used as outputs to the
+        function
+      func_outputs: a possibly nested python object which will be returned by
+        this function. The Tensors in this structure will be replaced by their
+        corresponding values in outputs.
+      output_shapes: List of shapes of all tensors in outputs
+      variables: (optional) List of variables to watch during function
+        execution.
+      attrs: (optional) dict mapping names of attributes to their AttrValue
+        values. Attributes in `attrs` will be included in this function's
+        definition.
+    """
+    self._attrs = attrs or {}
     defined_function = _EagerDefinedFunction(
-        name, graph, operations, input_placeholders, outputs)
+        name, graph, operations, input_placeholders, outputs, self._attrs)
     if len(input_placeholders) != len(defined_function.signature.input_arg):
       raise ValueError("Internal error: invalid lengths. %s %s" % (
           len(input_placeholders), len(defined_function.signature.input_arg)))
@@ -372,7 +390,7 @@ class GraphModeFunction(object):
     forward_name = _forward_name(self._func_name)
     self._forward_fdef = _EagerDefinedFunction(
         forward_name, self._graph, self._ops, self._input_placeholders,
-        filtered_outputs + captures)
+        filtered_outputs + captures, self._attrs)
     all_inputs = self._out_grad_placeholders + captures
     # Excluding input ops from the body as we do not intend to execute these
     # operations when the function is executed.
@@ -386,7 +404,7 @@ class GraphModeFunction(object):
     bname = _backward_name(self._func_name)
     self._backward_function = GraphModeFunction(
         bname, all_inputs, [], self._graph, function_def_ops,
-        backward_outputs, in_gradients, output_shapes)
+        backward_outputs, in_gradients, output_shapes, attrs=self._attrs)
 
   def _backprop_call(self, args):
     """Calls the wrapped function and records the result on a tape."""
@@ -560,7 +578,7 @@ def _get_defun_inputs(args):
   return nest.pack_sequence_as(args, ret)
 
 
-def _defun_internal(name, func, args, kwds):
+def _defun_internal(name, func, compiled, args, kwds):
   """Defines and returns graph-mode version of func."""
   graph_key = ops.get_default_graph()._graph_key  # pylint: disable=protected-access
   with context.graph_mode():
@@ -625,9 +643,14 @@ def _defun_internal(name, func, args, kwds):
     for f in tmp_graph._functions.values():  # pylint: disable=protected-access
       # TODO(ashankar): What about the gradient registry?
       _register(f._c_func.func)  # pylint: disable=protected-access
+
+  attrs = {}
+  if compiled:
+    attrs["_XlaCompile"] = attr_value_pb2.AttrValue(b=True)
+
   return GraphModeFunction(
       fname, all_inputs, extra_inputs, tmp_graph, operations, func_def_outputs,
-      func_outputs, output_shapes, variables)
+      func_outputs, output_shapes, variables, attrs)
 
 
 # Defun uses this instead of Tensor as a cache key. Using dtype because
@@ -669,7 +692,7 @@ def _register(fn):
 
 
 # TODO(apassos): better error messages for non-hashable arguments.
-def named_defun(func, name):
+def named_defun(func, name, compiled=False):
   """Defines a function with a given name.
 
   See the documentation for `defun` for more information on the semantics of the
@@ -678,6 +701,7 @@ def named_defun(func, name):
   Args:
     func: the function to be wrapped.
     name: the name given to it.
+    compiled: if true, the framework will attempt to compile func with XLA.
 
   Returns:
     the wrapped function.
@@ -694,13 +718,13 @@ def named_defun(func, name):
 
     if cache_key not in arguments_to_functions:
       arguments_to_functions[cache_key] = _defun_internal(
-          name, func, args, kwds)
+          name, func, compiled, args, kwds)
     return arguments_to_functions[cache_key](*args)
 
   return decorated
 
 
-def defun(func):
+def defun(func=None, compiled=False):
   """Decorator to compile func into graph_mode.
 
   `defun` converts a function that constructs a TensorFlow graph into a function
@@ -743,18 +767,45 @@ def defun(func):
   ```
 
   Args:
-    func: function to be compiled.
+    func: function to be compiled. If `func` is None, returns a
+      decorator that can be invoked with a single argument - `func`. The
+      end result is equivalent to providing all the arguments up front.
+      In other words, defun(compiled=True)(func) is equivalent to
+      defun(func, compiled=True). The former allows the following use case:
+        @tfe.defun(compiled=True)
+        def foo(...):
+          ...
+    compiled: If True, an attempt to compile `func` with XLA will be made.
+      If it fails, function will be run normally. Experimental.
+      Currently, supported only for execution on TPUs.
 
   Returns:
-     A callable that will execute the compiled function (and return zero
-     or more `tf.Tensor` objects).
+     If `func` is not None, returns callable that will execute the compiled
+     function (and return zero or more `tf.Tensor` objects).
+     If `func` is None, returns a decorator that, when invoked with a single
+     `func` argument, returns a callable equivalent to the case above.
   """
   # TODO(apassos): deal with captured global state. Deal with control flow.
-  try:
-    name = func.__name__
-  except AttributeError:
-    name = "function"
-  return tf_decorator.make_decorator(func, named_defun(func, name))
+  def decorated(function):
+    try:
+      name = function.__name__
+    except AttributeError:
+      name = "function"
+    return tf_decorator.make_decorator(
+        function, named_defun(function, name, compiled=compiled))
+
+  # This code path is for the `foo = tfe.defun(foo, ...)` use case
+  if func is not None:
+    return decorated(func)
+
+  # This code path is for the
+  #
+  # @tfe.defun(...)
+  # def foo(...):
+  #    ...
+  #
+  # use case, which is equivalent to `foo = tfe.defun(...)(foo)`
+  return decorated
 
 
 def make_defun_op(func, *args, **kwds):
@@ -806,7 +857,7 @@ def make_defun_op(func, *args, **kwds):
   name = func.__name__
   if any(isinstance(x, ops.EagerTensor) for x in kwds.values()):
     raise ValueError("Tensor keyword arguments are not supported.")
-  return _defun_internal(name, func, args, kwds)
+  return _defun_internal(name, func, False, args, kwds)
 
 
 class AutomaticControlDependencies(object):
-- 
GitLab


From 2585a8181904b39c71fc314940587c02b30a68a6 Mon Sep 17 00:00:00 2001
From: Skye Wanderman-Milne <skyewm@google.com>
Date: Mon, 7 May 2018 17:24:28 -0700
Subject: [PATCH 346/395] Make conv2d_tranpose_test.py work with C API shapes
 enabled.

The C API provides more accurate shape information in many cases.

PiperOrigin-RevId: 195749030
---
 tensorflow/python/kernel_tests/conv2d_transpose_test.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/kernel_tests/conv2d_transpose_test.py b/tensorflow/python/kernel_tests/conv2d_transpose_test.py
index b692d3da60..27804be65c 100644
--- a/tensorflow/python/kernel_tests/conv2d_transpose_test.py
+++ b/tensorflow/python/kernel_tests/conv2d_transpose_test.py
@@ -23,6 +23,7 @@ from six.moves import xrange  # pylint: disable=redefined-builtin
 
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import gradient_checker
 from tensorflow.python.ops import nn_ops
@@ -292,6 +293,7 @@ class Conv2DTransposeTest(test.TestCase):
 
         self.assertAllClose(cache_values, value)
 
+  @test_util.enable_c_shapes
   def testConv2DTransposeShapeInference(self):
     # Test case for 8972
     initializer = random_ops.truncated_normal(
@@ -301,7 +303,8 @@ class Conv2DTransposeTest(test.TestCase):
     f_shape = array_ops.stack([array_ops.shape(x)[0], 10, 5, 5])
     output = nn_ops.conv2d_transpose(
         x, f, f_shape, strides=[1, 1, 1, 1], padding="SAME")
-    self.assertEqual(output.get_shape().as_list(), [None, 10, 5, 5])
+    self.assertEqual(output.get_shape().as_list(), [3, 10, 5, 5])
+
 
 if __name__ == "__main__":
   test.main()
-- 
GitLab


From 1af09b57ef663d4ab0c02a00e2af1f1e2819d32f Mon Sep 17 00:00:00 2001
From: Skye Wanderman-Milne <skyewm@google.com>
Date: Mon, 7 May 2018 17:28:41 -0700
Subject: [PATCH 347/395] Add logic for StridedSlice ops in
 ShapeRefiner::ConstantPartialShape().

This mimics the logic in tensor_util.constant_value_as_shape, allowing
the C++ shape inference code to infer more shapes than it could before.

This change also adds an optional stride argument to InferenceContext::Subshape().

PiperOrigin-RevId: 195749522
---
 .../core/common_runtime/shape_refiner.cc      | 113 ++++++++++++++++--
 .../core/common_runtime/shape_refiner.h       |  14 +++
 .../core/common_runtime/shape_refiner_test.cc | 100 ++++++++++++++++
 tensorflow/core/framework/shape_inference.cc  |  29 ++++-
 tensorflow/core/framework/shape_inference.h   |   7 ++
 5 files changed, 245 insertions(+), 18 deletions(-)

diff --git a/tensorflow/core/common_runtime/shape_refiner.cc b/tensorflow/core/common_runtime/shape_refiner.cc
index a0772713d4..fa4d1eda62 100644
--- a/tensorflow/core/common_runtime/shape_refiner.cc
+++ b/tensorflow/core/common_runtime/shape_refiner.cc
@@ -421,6 +421,28 @@ Status ShapeRefiner::EvaluateConstantTensorForEdge(const Node* node,
                                 kMaxTensorSize, disable_constant_propagation_);
 }
 
+Status ShapeRefiner::EvaluateConstantIntScalarEdge(const Node* node,
+                                                   int dst_idx, bool* evaluated,
+                                                   int64* result) {
+  Tensor scalar;
+  TF_RETURN_IF_ERROR(
+      EvaluateConstantTensorForEdge(node, dst_idx, evaluated, &scalar));
+  if (*evaluated) {
+    DCHECK_EQ(scalar.NumElements(), 1)
+        << "EvaluateConstantIntScalarEdge called on non-scalar edge: "
+        << scalar.NumElements();
+    if (scalar.dtype() == DT_INT32) {
+      *result = scalar.scalar<int32>()();
+    } else {
+      DCHECK_EQ(scalar.dtype(), DT_INT64)
+          << "EvaluateConstantIntScalarEdge called on non-integer edge: "
+          << scalar.dtype();
+      *result = scalar.scalar<int64>()();
+    }
+  }
+  return Status::OK();
+}
+
 Status ShapeRefiner::ConstantPartialShape(InferenceContext* target_context,
                                           const Node* node, int dst_idx,
                                           ShapeHandle* result) {
@@ -471,19 +493,11 @@ Status ShapeRefiner::ConstantPartialShape(InferenceContext* target_context,
     std::vector<DimensionHandle> dims;
     // Pack is concatenating its input scalars to form the shape tensor vector.
     for (int i = 0; i < src_context->num_inputs(); ++i) {
-      Tensor scalar;
-      bool evaluated = false;
-      TF_RETURN_IF_ERROR(EvaluateConstantTensorForEdge(input_edge->src(), i,
-                                                       &evaluated, &scalar));
+      int64 size;
+      bool evaluated;
+      TF_RETURN_IF_ERROR(EvaluateConstantIntScalarEdge(input_edge->src(), i,
+                                                       &evaluated, &size));
       if (evaluated) {
-        int64 size;
-        if (scalar.dtype() == DT_INT32) {
-          size = scalar.scalar<int32>()();
-        } else if (scalar.dtype() == DT_INT64) {
-          size = scalar.scalar<int64>()();
-        } else {
-          return errors::InvalidArgument("Pack input must be int32 or int64");
-        }
         dims.push_back(size < 0 ? target_context->UnknownDim()
                                 : target_context->MakeDim(size));
       } else {
@@ -513,6 +527,9 @@ Status ShapeRefiner::ConstantPartialShape(InferenceContext* target_context,
       TF_RETURN_IF_ERROR(
           target_context->Concatenate(*result, sub_result, result));
     }
+  } else if (src_op == "StridedSlice") {
+    TF_RETURN_IF_ERROR(
+        PartialStridedSliceShape(input_edge->src(), src_context, result));
   } else {
     Tensor t;
     bool evaluated = false;
@@ -524,6 +541,78 @@ Status ShapeRefiner::ConstantPartialShape(InferenceContext* target_context,
   return Status::OK();
 }
 
+Status ShapeRefiner::PartialStridedSliceShape(Node* slice_node,
+                                              InferenceContext* ctx,
+                                              ShapeHandle* result) {
+  // Only attempt to evaluate if begin/end/strides all are scalars.
+  for (int i = 1; i <= 3; ++i) {
+    ShapeHandle input_shape = ctx->input(i);
+    if (ctx->Value(ctx->Dim(input_shape, 0)) != 1) {
+      *result = ctx->UnknownShape();
+      return Status::OK();
+    }
+  }
+
+  int begin_mask, end_mask, ellipsis_mask, new_axis_mask, shrink_axis_mask;
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(slice_node->attrs(), "begin_mask", &begin_mask));
+  TF_RETURN_IF_ERROR(GetNodeAttr(slice_node->attrs(), "end_mask", &end_mask));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(slice_node->attrs(), "ellipsis_mask", &ellipsis_mask));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(slice_node->attrs(), "new_axis_mask", &new_axis_mask));
+  TF_RETURN_IF_ERROR(
+      GetNodeAttr(slice_node->attrs(), "shrink_axis_mask", &shrink_axis_mask));
+
+  // Only attempt to evaluate if there are no special masks set (note that we
+  // can handle begin/end_mask == 1).
+  if (!(begin_mask == 0 || begin_mask == 1) ||
+      !(end_mask == 0 || end_mask == 1) || ellipsis_mask != 0 ||
+      new_axis_mask != 0 || shrink_axis_mask != 0) {
+    *result = ctx->UnknownShape();
+    return Status::OK();
+  }
+
+  bool evaluated;
+  int64 begin;
+  if (begin_mask == 1) {
+    begin = 0;
+  } else {
+    TF_RETURN_IF_ERROR(
+        EvaluateConstantIntScalarEdge(slice_node, 1, &evaluated, &begin));
+    if (!evaluated) {
+      *result = ctx->UnknownShape();
+      return Status::OK();
+    }
+  }
+
+  int64 end;
+  if (end_mask == 1) {
+    end = std::numeric_limits<int64>::max();
+  } else {
+    TF_RETURN_IF_ERROR(
+        EvaluateConstantIntScalarEdge(slice_node, 2, &evaluated, &end));
+    if (!evaluated) {
+      *result = ctx->UnknownShape();
+      return Status::OK();
+    }
+  }
+
+  int64 stride;
+  TF_RETURN_IF_ERROR(
+      EvaluateConstantIntScalarEdge(slice_node, 3, &evaluated, &stride));
+  if (!evaluated) {
+    *result = ctx->UnknownShape();
+    return Status::OK();
+  }
+
+  // Apply stride to input interpreted as a partial shape.
+  ShapeHandle input;
+  TF_RETURN_IF_ERROR(ConstantPartialShape(ctx, slice_node, 0, &input));
+  TF_RETURN_IF_ERROR(ctx->Subshape(input, begin, end, stride, result));
+  return Status::OK();
+}
+
 Status ShapeRefiner::RunShapeFn(const Node* node,
                                 const OpRegistrationData* op_reg_data,
                                 ExtendedInferenceContext* ec) {
diff --git a/tensorflow/core/common_runtime/shape_refiner.h b/tensorflow/core/common_runtime/shape_refiner.h
index d49c4373f0..9c96dcbc20 100644
--- a/tensorflow/core/common_runtime/shape_refiner.h
+++ b/tensorflow/core/common_runtime/shape_refiner.h
@@ -215,9 +215,18 @@ class ShapeRefiner {
                                 bool keep_nested_shapes,
                                 ExtendedInferenceContext* outer_context);
 
+  // Attempts to evaluate the 'dst_idx'-th input to 'node'. If the input edge
+  // value can be evaluated, 'evaluated' is set to true and the value returned
+  // in 'result'. Otherwise 'evaluated' is set to false.
   Status EvaluateConstantTensorForEdge(const Node* node, int dst_idx,
                                        bool* evaluated, Tensor* result);
 
+  // Wrapper around EvaluateConstantTensorForEdge for scalar int32/int64 input
+  // tensors. The caller is responsible for checking that the specified edge is
+  // scalar and int32 or int64.
+  Status EvaluateConstantIntScalarEdge(const Node* node, int dst_idx,
+                                       bool* evaluated, int64* result);
+
   // This function tries to materialize as much information about the 'node''s
   // dst_idx input as a statically computable shape, and the result may be
   // partially known, depending on what is statically inferable.
@@ -243,6 +252,11 @@ class ShapeRefiner {
                               const Node* node, int dst_idx,
                               shape_inference::ShapeHandle* result);
 
+  // Implementation of ConstantPartialShape for StridedSlice nodes.
+  Status PartialStridedSliceShape(Node* slice_node,
+                                  shape_inference::InferenceContext* ctx,
+                                  shape_inference::ShapeHandle* result);
+
   Status RunShapeFn(const Node* node, const OpRegistrationData* op_reg_data,
                     ExtendedInferenceContext* ec);
 
diff --git a/tensorflow/core/common_runtime/shape_refiner_test.cc b/tensorflow/core/common_runtime/shape_refiner_test.cc
index f48638afc0..8b9657eec8 100644
--- a/tensorflow/core/common_runtime/shape_refiner_test.cc
+++ b/tensorflow/core/common_runtime/shape_refiner_test.cc
@@ -60,6 +60,39 @@ class ShapeRefinerTest : public ::testing::Test {
   }
 
   static constexpr int64 kMaxTensorSize = ShapeRefiner::kMaxTensorSize;
+
+  void TestStridedSlice(const PartialTensorShape& input_shape, int begin,
+                        int end, int stride, const char* expected,
+                        int begin_mask = 0, int end_mask = 0,
+                        int ellipsis_mask = 0) {
+    Scope root = Scope::DisabledShapeInferenceScope();
+    auto placeholder =
+        ops::Placeholder(root, DT_INT32, ops::Placeholder::Shape(input_shape));
+    auto input = ops::Shape(root, placeholder);
+    auto begin_op = ops::Const(root, {begin});
+    auto end_op = ops::Const(root, {end});
+    auto stride_op = ops::Const(root, {stride});
+    auto slice = ops::StridedSlice(root, input, begin_op, end_op, stride_op,
+                                   ops::StridedSlice::BeginMask(begin_mask)
+                                       .EndMask(end_mask)
+                                       .EllipsisMask(ellipsis_mask));
+    Node* result;
+    TF_ASSERT_OK(NodeBuilder("test", "TensorAsShapeInt32")
+                     .Input(slice.node())
+                     .Finalize(root.graph(), &result));
+
+    ShapeRefiner m(TF_GRAPH_DEF_VERSION, OpRegistry::Global());
+    TF_ASSERT_OK(m.AddNode(placeholder.node()));
+    TF_ASSERT_OK(m.AddNode(input.node()));
+    TF_ASSERT_OK(m.AddNode(begin_op.node()));
+    TF_ASSERT_OK(m.AddNode(end_op.node()));
+    TF_ASSERT_OK(m.AddNode(stride_op.node()));
+    TF_ASSERT_OK(m.AddNode(slice.node()));
+    TF_ASSERT_OK(m.AddNode(result));
+
+    shape_inference::InferenceContext* ctx = m.GetContext(result);
+    EXPECT_EQ(ctx->DebugString(ctx->output(0)), expected);
+  }
 };
 
 namespace {
@@ -1156,6 +1189,73 @@ TEST_F(ShapeRefinerTest, ConstantValueAsShape_ConcatInvalidDimValue) {
             m.AddNode(result).error_message());
 }
 
+TEST_F(ShapeRefinerTest, ConstantValueAsShape_StridedSlice) {
+  TestStridedSlice(
+      /*input_shape=*/{1, -1, 3, -1, 5},
+      /*begin=*/2,
+      /*end=*/5,
+      /*stride=*/1,
+      /*expected=*/"[3,?,5]");
+}
+
+TEST_F(ShapeRefinerTest, ConstantValueAsShape_StridedSliceNegativeStride) {
+  // clang-format off
+  TestStridedSlice(
+      /*input_shape=*/{1, -1, 3, -1, 5},
+      /*begin=*/10,
+      /*end=*/0,
+      /*stride=*/-1,
+      /*expected=*/"[5,?,3,?]");
+  // clang-format on
+}
+
+TEST_F(ShapeRefinerTest, ConstantValueAsShape_StridedSliceMasks) {
+  TestStridedSlice(
+      /*input_shape=*/{1, -1, 3, -1, 5},
+      /*begin=*/3,
+      /*end=*/4,
+      /*stride=*/1,
+      /*expected=*/"[1,?,3,?,5]",
+      /*begin_mask=*/1,
+      /*end_mask=*/1);
+}
+
+TEST_F(ShapeRefinerTest, ConstantValueAsShape_StridedSliceInvalidMask) {
+  TestStridedSlice(
+      /*input_shape=*/{1, -1, 3},
+      /*begin=*/2,
+      /*end=*/3,
+      /*stride=*/1,
+      /*expected=*/"[?,?,?]",
+      /*begin_mask=*/0,
+      /*end_mask=*/0,
+      /*ellipsis_mask=*/1);
+}
+
+TEST_F(ShapeRefinerTest, ConstantValueAsShape_StridedSliceMulti) {
+  Scope root = Scope::DisabledShapeInferenceScope();
+  auto input = ops::Placeholder(root, DT_INT32);
+  auto begin = ops::Const(root, {0, 0});
+  auto end = ops::Const(root, {2, 2});
+  auto stride = ops::Const(root, {1, 1});
+  auto slice = ops::StridedSlice(root, input, begin, end, stride);
+  Node* result;
+  TF_ASSERT_OK(NodeBuilder("test", "TensorAsShapeInt32")
+                   .Input(slice.node())
+                   .Finalize(root.graph(), &result));
+
+  ShapeRefiner m(TF_GRAPH_DEF_VERSION, OpRegistry::Global());
+  TF_ASSERT_OK(m.AddNode(input.node()));
+  TF_ASSERT_OK(m.AddNode(begin.node()));
+  TF_ASSERT_OK(m.AddNode(end.node()));
+  TF_ASSERT_OK(m.AddNode(stride.node()));
+  TF_ASSERT_OK(m.AddNode(slice.node()));
+  TF_ASSERT_OK(m.AddNode(result));
+
+  shape_inference::InferenceContext* ctx = m.GetContext(result);
+  EXPECT_EQ(ctx->DebugString(ctx->output(0)), "?");
+}
+
 namespace {
 
 // Dummy op to test ShapeRefiner util functions
diff --git a/tensorflow/core/framework/shape_inference.cc b/tensorflow/core/framework/shape_inference.cc
index 2b995e8b5e..3185875e3b 100644
--- a/tensorflow/core/framework/shape_inference.cc
+++ b/tensorflow/core/framework/shape_inference.cc
@@ -605,10 +605,16 @@ Status InferenceContext::Subshape(ShapeHandle s, int64 start,
   return Subshape(s, start, std::numeric_limits<int64>::max() /* end */, out);
 }
 
-Status InferenceContext::Subshape(ShapeHandle s, int64 start_in, int64 end_in,
+Status InferenceContext::Subshape(ShapeHandle s, int64 start, int64 end,
                                   ShapeHandle* out) {
-  int64 start = start_in;
-  int64 end = end_in;
+  return Subshape(s, start, end, 1 /* stride */, out);
+}
+
+Status InferenceContext::Subshape(ShapeHandle s, int64 start, int64 end,
+                                  int64 stride, ShapeHandle* out) {
+  int64 start_in = start;
+  int64 end_in = end;
+
   const int32 rank = Rank(s);
   if (start == 0 && ((RankKnown(s) && end >= rank) ||
                      end == std::numeric_limits<int64>::max())) {
@@ -621,6 +627,9 @@ Status InferenceContext::Subshape(ShapeHandle s, int64 start_in, int64 end_in,
 
   if (start > rank) start = rank;
   if (end > rank) end = rank;
+
+  if (stride < 0 && start == rank) --start;
+
   if (start < 0) {
     start = rank + start;
     if (start < 0) {
@@ -638,16 +647,24 @@ Status InferenceContext::Subshape(ShapeHandle s, int64 start_in, int64 end_in,
                                      ", for shape with rank ", rank);
     }
   }
-  if (start > end) {
+  if (stride > 0 && start > end) {
     *out = nullptr;
     return errors::InvalidArgument(
         "Subshape must have computed start <= end, but is ", start, " and ",
         end, " (computed from start ", start_in, " and end ", end_in,
         " over shape with rank ", rank, ")");
+  } else if (stride < 0 && start < end) {
+    *out = nullptr;
+    return errors::InvalidArgument(
+        "Subshape must have computed start >= end since stride is negative, "
+        "but is ",
+        start, " and ", end, " (computed from start ", start_in, " and end ",
+        end_in, " over shape with rank ", rank, " and stride", stride, ")");
   }
+
   std::vector<DimensionHandle> dims;
-  dims.reserve(end - start);
-  for (int i = start; i < end; ++i) {
+  dims.reserve((end - start) / stride);
+  for (int i = start; stride > 0 ? i < end : i > end; i += stride) {
     dims.push_back(Dim(s, i));
   }
   return ReturnCreatedShape(dims, out);
diff --git a/tensorflow/core/framework/shape_inference.h b/tensorflow/core/framework/shape_inference.h
index 9431a62abe..3f3729dcf9 100644
--- a/tensorflow/core/framework/shape_inference.h
+++ b/tensorflow/core/framework/shape_inference.h
@@ -434,6 +434,13 @@ class InferenceContext {
   Status Subshape(ShapeHandle s, int64 start, int64 end,
                   ShapeHandle* out) TF_MUST_USE_RESULT;
 
+  // Returns in <*out> a sub-shape of <s>, with dimensions [start:end:stride].
+  // <start> and <end> can be negative, to index from the end of the shape.
+  // <start> and <end> are set to the rank of <s> if > rank of <s>.
+  // <stride> can be negative, to reverse the <s>.
+  Status Subshape(ShapeHandle s, int64 start, int64 end, int64 stride,
+                  ShapeHandle* out) TF_MUST_USE_RESULT;
+
   // Returns in <*out> the result of appending the dimensions of <s2> to those
   // of <s1>.
   Status Concatenate(ShapeHandle s1, ShapeHandle s2,
-- 
GitLab


From 9ecbb5574fb86d9f5280315141a11acd47e50dee Mon Sep 17 00:00:00 2001
From: wangsiyu <siyu.wsy@gmail.com>
Date: Tue, 8 May 2018 10:54:04 +0800
Subject: [PATCH 348/395] refine unit test case coding style and move
 _should_add_regularizer function into add_weight

---
 tensorflow/python/layers/base.py      | 23 ++++++++++++-----------
 tensorflow/python/layers/base_test.py |  4 ++--
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/tensorflow/python/layers/base.py b/tensorflow/python/layers/base.py
index f7b2e471b2..78db47681a 100644
--- a/tensorflow/python/layers/base.py
+++ b/tensorflow/python/layers/base.py
@@ -191,6 +191,18 @@ class Layer(base_layer.Layer):
       RuntimeError: If called with partioned variable regularization and
         eager execution is enabled.
     """
+    
+    def _should_add_regularizer(variable, existing_variable_set):
+      result = True
+      if isinstance(variable, tf_variables.PartitionedVariable):
+        for var in variable:
+          if var in existing_variable_set:
+            result = False
+            break
+      else:
+        result = variable not in existing_variable_set
+      return result
+
     init_graph = None
     if not context.executing_eagerly():
       default_graph = ops.get_default_graph()
@@ -354,14 +366,3 @@ def _add_elements_to_collection(elements, collection_list):
     for element in elements:
       if element not in collection_set:
         collection.append(element)
-
-def _should_add_regularizer(variable, existing_variable_set):
-  result = True
-  if isinstance(variable, tf_variables.PartitionedVariable):
-    for var in variable:
-      if var in existing_variable_set:
-        result = False
-        break
-  else:
-    result = variable not in existing_variable_set
-  return result
diff --git a/tensorflow/python/layers/base_test.py b/tensorflow/python/layers/base_test.py
index 361e3de7aa..7158fd42e1 100644
--- a/tensorflow/python/layers/base_test.py
+++ b/tensorflow/python/layers/base_test.py
@@ -99,10 +99,10 @@ class BaseLayerTest(test.TestCase):
   def testReusePartitionedVaraiblesAndRegularizers(self):
     regularizer = lambda x: math_ops.reduce_sum(x) * 1e-3
     partitioner = partitioned_variables.fixed_size_partitioner(3)
-    for i in xrange(2):
+    for reuse in [False, True]:
       with variable_scope.variable_scope(variable_scope.get_variable_scope(),
                                          partitioner=partitioner,
-                                         reuse=False if i == 0 else True):
+                                         reuse=reuse):
         layer = base_layers.Layer(name='my_layer')
         variable = layer.add_variable(
             'reg_part_var', [4, 4],
-- 
GitLab


From 263c094c1d4f9509c4428e97fdd83957d8225c25 Mon Sep 17 00:00:00 2001
From: wangsiyu <siyu.wsy@gmail.com>
Date: Tue, 8 May 2018 12:58:27 +0800
Subject: [PATCH 349/395] eliminate result variable in _should_add_regularizer
 to make code clean

---
 tensorflow/python/layers/base.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/tensorflow/python/layers/base.py b/tensorflow/python/layers/base.py
index 78db47681a..aa416d1ff6 100644
--- a/tensorflow/python/layers/base.py
+++ b/tensorflow/python/layers/base.py
@@ -193,15 +193,12 @@ class Layer(base_layer.Layer):
     """
     
     def _should_add_regularizer(variable, existing_variable_set):
-      result = True
       if isinstance(variable, tf_variables.PartitionedVariable):
         for var in variable:
           if var in existing_variable_set:
-            result = False
-            break
+            return False
       else:
-        result = variable not in existing_variable_set
-      return result
+        return variable not in existing_variable_set
 
     init_graph = None
     if not context.executing_eagerly():
-- 
GitLab


From f32699406f31e0b6a38a15c9a3d580d1ef9d6204 Mon Sep 17 00:00:00 2001
From: wangsiyu <siyu.wsy@gmail.com>
Date: Tue, 8 May 2018 15:46:38 +0800
Subject: [PATCH 350/395] fix bug of return value in _should_add_regularizer
 function and refine code in base_test.py to make it no more than 80 columns

---
 tensorflow/python/layers/base.py      | 2 +-
 tensorflow/python/layers/base_test.py | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/layers/base.py b/tensorflow/python/layers/base.py
index aa416d1ff6..e122d6533c 100644
--- a/tensorflow/python/layers/base.py
+++ b/tensorflow/python/layers/base.py
@@ -197,6 +197,7 @@ class Layer(base_layer.Layer):
         for var in variable:
           if var in existing_variable_set:
             return False
+        return True
       else:
         return variable not in existing_variable_set
 
@@ -240,7 +241,6 @@ class Layer(base_layer.Layer):
             partitioner=partitioner,
             use_resource=use_resource,
             getter=vs.get_variable)
-
         if regularizer:
           if context.executing_eagerly() or _should_add_regularizer(
               variable, existing_variables):
diff --git a/tensorflow/python/layers/base_test.py b/tensorflow/python/layers/base_test.py
index 7158fd42e1..ab49e37b90 100644
--- a/tensorflow/python/layers/base_test.py
+++ b/tensorflow/python/layers/base_test.py
@@ -108,7 +108,8 @@ class BaseLayerTest(test.TestCase):
             'reg_part_var', [4, 4],
             initializer=init_ops.zeros_initializer(),
             regularizer=regularizer)
-    self.assertEqual(len(ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)), 3)
+    self.assertEqual(
+        len(ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)), 3)
 
   def testNoEagerActivityRegularizer(self):
     with context.eager_mode():
-- 
GitLab


From 334d8bfb594caafe5ab7ecaf007b1bf9ca062590 Mon Sep 17 00:00:00 2001
From: wangsiyu <siyu.wsy@gmail.com>
Date: Tue, 8 May 2018 15:50:36 +0800
Subject: [PATCH 351/395] remove type

---
 tensorflow/python/layers/base.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/layers/base.py b/tensorflow/python/layers/base.py
index e122d6533c..aa43a153c2 100644
--- a/tensorflow/python/layers/base.py
+++ b/tensorflow/python/layers/base.py
@@ -241,6 +241,7 @@ class Layer(base_layer.Layer):
             partitioner=partitioner,
             use_resource=use_resource,
             getter=vs.get_variable)
+
         if regularizer:
           if context.executing_eagerly() or _should_add_regularizer(
               variable, existing_variables):
-- 
GitLab


From 7a9e695d82ef75b3619177da245842fdddc3b8a8 Mon Sep 17 00:00:00 2001
From: Brennan Saeta <saeta@google.com>
Date: Mon, 7 May 2018 18:31:47 -0700
Subject: [PATCH 352/395] [tf.data] Move
 tensorflow::dataset::MakeIteratorContext to core/framework

PiperOrigin-RevId: 195756342
---
 tensorflow/core/framework/dataset.cc          | 19 +++++++++++++++++++
 tensorflow/core/framework/dataset.h           |  6 ++++++
 tensorflow/core/kernels/data/dataset_utils.cc | 12 ------------
 tensorflow/core/kernels/data/dataset_utils.h  |  2 --
 4 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/tensorflow/core/framework/dataset.cc b/tensorflow/core/framework/dataset.cc
index 4145ef7bc9..62a9d5751d 100644
--- a/tensorflow/core/framework/dataset.cc
+++ b/tensorflow/core/framework/dataset.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/core/framework/dataset.h"
 
+#include "tensorflow/core/framework/device_base.h"
 #include "tensorflow/core/graph/graph_def_builder.h"
 #include "tensorflow/core/graph/node_builder.h"
 
@@ -269,4 +270,22 @@ const char GraphDatasetBase::kDatasetGraphKey[] = "_DATASET_GRAPH";
 const char GraphDatasetBase::kDatasetGraphOutputNodeKey[] =
     "_DATASET_GRAPH_OUTPUT_NODE";
 
+namespace dataset {
+
+IteratorContext MakeIteratorContext(OpKernelContext* ctx) {
+  IteratorContext::Params params;
+  params.env = ctx->env();
+  params.runner = *(ctx->runner());
+  params.lib = ctx->function_library();
+  // Note: must use reinterpret_cast because function.h forward-declares Device.
+  DeviceBase* device =
+      reinterpret_cast<DeviceBase*>(ctx->function_library()->device());
+  params.allocator_getter = [device](AllocatorAttributes attrs) {
+    return device->GetAllocator(attrs);
+  };
+  return IteratorContext(params);
+}
+
+}  // namespace dataset
+
 }  // namespace tensorflow
diff --git a/tensorflow/core/framework/dataset.h b/tensorflow/core/framework/dataset.h
index 775d9f6eb6..8624af9bf5 100644
--- a/tensorflow/core/framework/dataset.h
+++ b/tensorflow/core/framework/dataset.h
@@ -619,6 +619,12 @@ Status GetDatasetFromVariantTensor(const Tensor& tensor,
 // The ownership of `dataset` is transferred to `tensor`.
 Status StoreDatasetInVariantTensor(DatasetBase* dataset, Tensor* tensor);
 
+namespace dataset {
+
+IteratorContext MakeIteratorContext(OpKernelContext* ctx);
+
+}  // namespace dataset
+
 }  // namespace tensorflow
 
 #endif  // TENSORFLOW_CORE_FRAMEWORK_DATASET_H_
diff --git a/tensorflow/core/kernels/data/dataset_utils.cc b/tensorflow/core/kernels/data/dataset_utils.cc
index 67ddb52d57..c608f9e1c6 100644
--- a/tensorflow/core/kernels/data/dataset_utils.cc
+++ b/tensorflow/core/kernels/data/dataset_utils.cc
@@ -46,18 +46,6 @@ Status MakeIteratorFromInputElement(
   return Status::OK();
 }
 
-IteratorContext MakeIteratorContext(OpKernelContext* ctx) {
-  IteratorContext::Params params;
-  params.env = ctx->env();
-  params.runner = *(ctx->runner());
-  params.lib = ctx->function_library();
-  DeviceBase* device = ctx->function_library()->device();
-  params.allocator_getter = [device](AllocatorAttributes attrs) {
-    return device->GetAllocator(attrs);
-  };
-  return IteratorContext(params);
-}
-
 }  // namespace dataset
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/kernels/data/dataset_utils.h b/tensorflow/core/kernels/data/dataset_utils.h
index e5ca71dd99..6c4191c2be 100644
--- a/tensorflow/core/kernels/data/dataset_utils.h
+++ b/tensorflow/core/kernels/data/dataset_utils.h
@@ -28,8 +28,6 @@ Status MakeIteratorFromInputElement(
     int64 thread_index, CapturedFunction* captured_func, StringPiece prefix,
     std::unique_ptr<IteratorBase>* out_iterator);
 
-IteratorContext MakeIteratorContext(OpKernelContext* ctx);
-
 }  // namespace dataset
 
 }  // namespace tensorflow
-- 
GitLab


From 069f3124eedab44b4e884c3c64ba8d5eccadfe93 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 7 May 2018 19:56:26 -0700
Subject: [PATCH 353/395] Temporarily disable concat rewrite.

PiperOrigin-RevId: 195762860
---
 tensorflow/core/grappler/optimizers/arithmetic_optimizer.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
index 3f9feac55f..1f6f563687 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
@@ -65,7 +65,7 @@ class ArithmeticOptimizer : public GraphOptimizer {
     bool remove_redundant_bitcast = true;
     bool remove_redundant_cast = true;
     bool remove_negation = true;
-    bool hoist_cwise_unary_chains = true;
+    bool hoist_cwise_unary_chains = false;
     bool convert_sqrt_div_to_rsqrt_mul = false;
     bool remove_idempotent = true;
 
-- 
GitLab


From a799cdbe78ca2c2e9c41f2b1bf8a3f57162fbcea Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 02:11:52 -0700
Subject: [PATCH 354/395] Automated g4 rollback of changelist 195748721

PiperOrigin-RevId: 195790581
---
 tensorflow/compiler/jit/BUILD                 |  22 --
 .../compiler/jit/create_xla_launch_op.cc      | 206 ++++--------------
 .../compiler/jit/create_xla_launch_op.h       |  35 ---
 .../compiler/jit/create_xla_launch_op_test.cc | 144 ------------
 .../compiler/jit/kernels/xla_launch_op.cc     |  90 ++------
 .../compiler/jit/kernels/xla_launch_op.h      |  51 ++---
 .../compiler/jit/xla_compile_on_demand_op.cc  |   3 +-
 tensorflow/compiler/jit/xla_launch_util.cc    |  18 +-
 tensorflow/compiler/jit/xla_launch_util.h     |  15 +-
 tensorflow/compiler/tests/BUILD               |   4 -
 tensorflow/compiler/tests/eager_test.py       | 112 +---------
 .../python/examples/resnet50/resnet50_test.py |  55 ++---
 tensorflow/python/eager/function.py           | 127 ++++-------
 13 files changed, 164 insertions(+), 718 deletions(-)
 delete mode 100644 tensorflow/compiler/jit/create_xla_launch_op.h
 delete mode 100644 tensorflow/compiler/jit/create_xla_launch_op_test.cc

diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD
index e942b46086..07136d6a74 100644
--- a/tensorflow/compiler/jit/BUILD
+++ b/tensorflow/compiler/jit/BUILD
@@ -261,7 +261,6 @@ cc_library(
     name = "create_xla_launch_op",
     srcs = [
         "create_xla_launch_op.cc",
-        "create_xla_launch_op.h",
     ],
     deps = [
         ":common",
@@ -271,27 +270,6 @@ cc_library(
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
-        "//tensorflow/core:protos_all_cc",
-    ],
-    alwayslink = 1,
-)
-
-tf_cc_test(
-    name = "create_xla_launch_op_test",
-    srcs = [
-        "create_xla_launch_op.h",
-        "create_xla_launch_op_test.cc",
-    ],
-    deps = [
-        ":create_xla_launch_op",
-        "//tensorflow/core:core_cpu_internal",
-        "//tensorflow/core:framework",
-        "//tensorflow/core:lib",
-        "//tensorflow/core:protos_all_cc",
-        "//tensorflow/core:session_options",
-        "//tensorflow/core:test",
-        "//tensorflow/core:test_main",
-        "//tensorflow/core:testlib",
     ],
 )
 
diff --git a/tensorflow/compiler/jit/create_xla_launch_op.cc b/tensorflow/compiler/jit/create_xla_launch_op.cc
index 6ac84dc19c..18d901323f 100644
--- a/tensorflow/compiler/jit/create_xla_launch_op.cc
+++ b/tensorflow/compiler/jit/create_xla_launch_op.cc
@@ -12,7 +12,6 @@ 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.
 ==============================================================================*/
-#include "tensorflow/compiler/jit/create_xla_launch_op.h"
 
 #include "tensorflow/compiler/jit/defs.h"
 #include "tensorflow/compiler/jit/kernels/xla_launch_op.h"
@@ -26,189 +25,78 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
-// Utility which searches for values in a sorted list by scanning over it once.
-// No matter how many times ScanForValue is called, the list is scanned at most
-// once. However, if a call to ScanForValue skips over a value, that value is
-// not revisited in future calls to ScanForValue, so callers must take
-// care to order their calls.
+// Givens a NodeDef 'ndef' and the function library runtime 'flr', if
+// 'ndef' is a call to a compilable function defined in 'flr', returns OK
+// and fills in 'kernel' with a XlaLaunchOp kernel which computes the
+// node. Otherwise, returns a non-OK.
 //
-// Useful for merging multiple sorted lists in O(n) time.
-class SinglePassSearch {
- public:
-  // Creates a SinglePassSearch object that can be used to search in `values`.
-  // Does not take ownership of `values`. `values` must outlive this.
-  // `values` must be sorted.
-  explicit SinglePassSearch(const std::vector<int>* values)
-      : current_index_(0), values_(values) {}
-
-  // Scans forward in the vector looking for "value", updating the internal
-  // position in to the vector.
-  // Returns true iff the vector contains the given value at or after current
-  // position.
-  // Not thread-safe.
-  bool ScanForValue(int value) {
-    while (current_index_ < values_->size() &&
-           (*values_)[current_index_] <= value) {
-      if ((*values_)[current_index_] == value) {
-        current_index_++;
-        return true;
-      }
-      current_index_++;
-    }
-    return false;
-  }
-
- private:
-  int current_index_;
-  const std::vector<int>* values_;
-};
-
-Status CompilationRequested(const FunctionLibraryRuntime& flr,
-                            const NodeDef& node_def) {
+// This routine is here so that FunctionLibraryRuntime can jit a
+// specific function call as requested.
+Status CreateXlaLaunchOp(FunctionLibraryRuntime* flr, const NodeDef& ndef,
+                         std::unique_ptr<OpKernel>* kernel) {
   bool xla_compile = false;
-  // Check if op is marked _XlaCompile=true.
-  Status status = flr.GetFunctionLibraryDefinition()->GetAttr(
-      node_def, kXlaCompileAttr, &xla_compile);
-  if (!status.ok() || !xla_compile) {
-    if (VLOG_IS_ON(3)) {
-      if (!status.ok()) {
-        VLOG(3) << "No " << kXlaCompileAttr << " attr defined for "
-                << node_def.op() << ". status=" << status.ToString();
-      } else {
-        VLOG(3) << node_def.op() << " is explicitly marked not to be compiled";
-      }
-    }
-    return Status(error::INVALID_ARGUMENT, "");
+  if (!flr->GetFunctionLibraryDefinition()
+           ->GetAttr(ndef, kXlaCompileAttr, &xla_compile)
+           .ok() ||
+      !xla_compile) {
+    // Not marked as _XlaCompile=true.
+    return errors::InvalidArgument("No ", kXlaCompileAttr, " for ", ndef.op());
+  }
+  // Make sure that kernels have been registered on the JIT device.
+  XlaOpRegistry::RegisterCompilationKernels();
+  if (!IsCompilable(flr, ndef)) {
+    // ndef is calling a function that XLA can't compile.
+    return errors::InvalidArgument("Not compilable: ", ndef.ShortDebugString());
   }
-  return Status::OK();
-}
-
-// Given a FunctionLibraryRuntime and a NodeDef calling a function in the
-// runtime, returns this function's body in `fbody` as well as the indices
-// of its constant and resource arguments.
-// `fbody` is owned by `flr`.
-// `constant_arg_indices` and `resource_arg_indices` should be empty vector.
-// They are sorted in ascending order on this function's return.
-Status GetBodyAndConstantsAndResources(FunctionLibraryRuntime* flr,
-                                       const NodeDef& node_def,
-                                       const FunctionBody** fbody,
-                                       std::vector<int>* constant_arg_indices,
-                                       std::vector<int>* resource_arg_indices) {
   FunctionLibraryRuntime::Handle handle;
-  // If node_def is not instantiable, e.g., the function does not exist,
+  // If ndef is not instantiable, e.g., the function does not exist,
   // simply bail out.
   TF_RETURN_IF_ERROR(
-      flr->Instantiate(node_def.op(), AttrSlice(&node_def.attr()), &handle));
-  *fbody = flr->GetFunctionBody(handle);
-  CHECK(*fbody);  // Can't be nullptr since we just instantiated it.
-  const DataTypeVector& arg_types = (*fbody)->arg_types;
-  std::vector<bool> const_args(arg_types.size());
+      flr->Instantiate(ndef.op(), AttrSlice(&ndef.attr()), &handle));
+  const FunctionBody* fbody = flr->GetFunctionBody(handle);
+  CHECK(fbody);  // Can't be nullptr since we just instantiated it.
+  std::vector<bool> const_args(fbody->arg_types.size());
   // If we can't analyze the const args. Bail out.
-  TF_RETURN_IF_ERROR(BackwardsConstAnalysis(*((*fbody)->graph), &const_args));
+  TF_RETURN_IF_ERROR(BackwardsConstAnalysis(*(fbody->graph), &const_args));
 
   for (int i = 0; i < const_args.size(); ++i) {
     if (const_args[i]) {
-      constant_arg_indices->push_back(i);
-    }
-  }
-
-  // There can be hundreds of resource variables. Reserve the space for them.
-  // We don't reserve for constants above as they are usually few.
-  resource_arg_indices->reserve(arg_types.size());
-  for (int i = 0; i < arg_types.size(); ++i) {
-    if (arg_types[i] == DT_RESOURCE) {
-      resource_arg_indices->push_back(i);
+      // There is a const arg. Bail out.
+      return errors::InvalidArgument("Const arg: ", i, " in ",
+                                     DebugString(fbody->fdef));
     }
   }
 
-  return Status::OK();
-}
-
-}  // namespace
-
-Status CreateXlaLaunchOp(FunctionLibraryRuntime* flr, const NodeDef& node_def,
-                         std::unique_ptr<OpKernel>* kernel) {
-  TF_RETURN_IF_ERROR(CompilationRequested(*flr, node_def));
-
-  VLOG(3) << "Creating XlaLaunchOp for " << node_def.DebugString();
-
-  // Make sure that kernels have been registered on the JIT device.
-  XlaOpRegistry::RegisterCompilationKernels();
-  if (!IsCompilable(flr, node_def)) {
-    // node_def is calling a function that XLA can't compile.
-    return errors::InvalidArgument("Not compilable: ",
-                                   node_def.ShortDebugString());
-  }
-
-  // Get function body, constant args, and resource args.
-  const FunctionBody* fbody = nullptr;
-  std::vector<int> constant_arg_indices;
-  std::vector<int> resource_arg_indices;
-  TF_RETURN_IF_ERROR(GetBodyAndConstantsAndResources(
-      flr, node_def, &fbody, &constant_arg_indices, &resource_arg_indices));
-
-  // Set input and output memory types.
+  NodeDef launch_def;
+  launch_def.set_name(ndef.name());
+  launch_def.set_op("_XlaLaunch");
+  launch_def.set_device(flr->device()->name());
+  AddNodeAttr("Tconstants", DataTypeVector{}, &launch_def);
+  AddNodeAttr("Nresources", 0, &launch_def);
+  AddNodeAttr("Targs", fbody->arg_types, &launch_def);
+  AddNodeAttr("Tresults", fbody->ret_types, &launch_def);
+  NameAttrList func;
+  func.set_name(ndef.op());
+  *(func.mutable_attr()) = ndef.attr();
+  AddNodeAttr("function", func, &launch_def);
+
+  // TODO(b/32387911): Handles the host memory types across function
+  // calls properly. For now, we assume all inputs and outputs are on
+  // the device memory.
   MemoryTypeVector input_memory_types(fbody->arg_types.size(), DEVICE_MEMORY);
-  // These indices are used only for optimization purposes. They allow us
-  // to loop over constant_arg_indices and resource_arg_indices only once
-  // while iterating over all the function arguments checking if it is a
-  // resource or a constant.
-  // The reason we optimized this code is because functions can have a lot of
-  // captured arguments. For example, the backward pass of ResNet50 takes in all
-  // 214 variables and a similar number of activations.
-  SinglePassSearch constants_search(&constant_arg_indices);
-  SinglePassSearch resources_search(&resource_arg_indices);
-  for (int i = 0; i < fbody->arg_types.size(); ++i) {
-    if (resources_search.ScanForValue(i) || constants_search.ScanForValue(i)) {
-      // Compile-time constants and resource handles are expected to be in
-      // host memory.
-      input_memory_types[i] = HOST_MEMORY;
-    }
-  }
-  // One might wonder, about the case where a compile-time constant argument
-  // (which must be in host memory) is also used as an input into an op,
-  // e.g. Add, that expects its inputs in device memory. Here is how it
-  // works now.
-  // First, what do we mean by "op expects an input in XYZ memory"?
-  // There are two types of "ops" here: the tf2xla kernel and the HLO
-  // computation it builds. The tf2xla kernel needs to retrieve the actual
-  // numeric value of the compile-time constant tensors, so it really expects
-  // them to be on in host memory. However, for other inputs, it refers to them
-  // using xla::ComputationDataHandle, which is just a symbolic handle that
-  // xla::ComputationBuilder assigns. How does this handle gets assigned for
-  // constant arguments? Even constant arguments get an _Arg node in the graph
-  // instatiated for Function compilation. The tf2xla kernel for constant _Arg
-  // nodes takes the constant value, converts it to XlaLiteral, and feeds it
-  // to xla::ComputationBuilder.ConstantLiteral, which returns the handle. This
-  // constant XlaLiteral is included in the HLO graph, and subsequently, in
-  // the actual executable, which is copied to the device before being
-  // executed. Thus, when this executable runs, the constant is available in
-  // device memory.
-
-  // XlaLaunch kernel keeps all outputs (including constants, which it copies),
-  // in device memory
   MemoryTypeVector output_memory_types(fbody->ret_types.size(), DEVICE_MEMORY);
 
-  // Create the kernel.
-  NameAttrList function;
-  function.set_name(node_def.op());
-  *(function.mutable_attr()) = node_def.attr();
-
   Device* dev = flr->device();
   Status s;
   OpKernelConstruction construction(
       DeviceType(dev->device_type()), dev,
-      dev->GetAllocator(AllocatorAttributes()), &node_def,
+      dev->GetAllocator(AllocatorAttributes()), &launch_def,
       &fbody->fdef.signature(), flr, fbody->arg_types, input_memory_types,
       fbody->ret_types, output_memory_types, flr->graph_def_version(), &s);
-
-  *kernel = absl::make_unique<XlaLocalLaunchBase>(
-      &construction, constant_arg_indices, resource_arg_indices, function);
+  kernel->reset(new XlaLocalLaunchOp(&construction));
   return s;
 }
 
-namespace {
-
 bool RegisterLaunchOpCreator() {
   RegisterDefaultCustomKernelCreator(CreateXlaLaunchOp);
   return true;
diff --git a/tensorflow/compiler/jit/create_xla_launch_op.h b/tensorflow/compiler/jit/create_xla_launch_op.h
deleted file mode 100644
index 98a22e3515..0000000000
--- a/tensorflow/compiler/jit/create_xla_launch_op.h
+++ /dev/null
@@ -1,35 +0,0 @@
-/* Copyright 2018 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.
-==============================================================================*/
-#ifndef TENSORFLOW_COMPILER_JIT_CREATE_XLA_LAUNCH_OP_H_
-#define TENSORFLOW_COMPILER_JIT_CREATE_XLA_LAUNCH_OP_H_
-
-#include "tensorflow/core/framework/node_def.pb.h"
-#include "tensorflow/core/lib/core/status.h"
-
-namespace tensorflow {
-
-class FunctionLibraryRuntime;
-class OpKernel;
-
-// Given a NodeDef 'node_def' and the function library runtime 'flr', if
-// 'node_def' is a call to a compilable function defined in 'flr', returns OK
-// and fills in 'kernel' with a XlaLaunchOp kernel which computes the
-// node. Otherwise, returns a non-OK.
-Status CreateXlaLaunchOp(FunctionLibraryRuntime* flr, const NodeDef& node_def,
-                         std::unique_ptr<OpKernel>* kernel);
-
-}  // namespace tensorflow
-
-#endif  // TENSORFLOW_COMPILER_JIT_CREATE_XLA_LAUNCH_OP_H_
diff --git a/tensorflow/compiler/jit/create_xla_launch_op_test.cc b/tensorflow/compiler/jit/create_xla_launch_op_test.cc
deleted file mode 100644
index c222824eda..0000000000
--- a/tensorflow/compiler/jit/create_xla_launch_op_test.cc
+++ /dev/null
@@ -1,144 +0,0 @@
-/* Copyright 2018 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.
-==============================================================================*/
-
-#include "tensorflow/compiler/jit/create_xla_launch_op.h"
-
-#include "tensorflow/core/common_runtime/device_factory.h"
-#include "tensorflow/core/common_runtime/function.h"
-#include "tensorflow/core/framework/function_testlib.h"
-#include "tensorflow/core/framework/node_def_builder.h"
-#include "tensorflow/core/framework/tensor_testutil.h"
-#include "tensorflow/core/lib/core/status.h"
-#include "tensorflow/core/platform/test.h"
-#include "tensorflow/core/public/session_options.h"
-#include "tensorflow/core/public/version.h"
-
-namespace tensorflow {
-
-NodeDef ToNodeDef(const string& text) {
-  NodeDef node_def;
-  EXPECT_TRUE(protobuf::TextFormat::MergeFromString(text, &node_def));
-  return node_def;
-}
-
-// Create a FunctionDef that takes one resource and one regular param
-FunctionDef XTimesY() {
-  return FunctionDefHelper::Define(
-      // Name
-      "XTimesY",
-      // Args
-      {"x: float", "y: resource"},
-      // Return values
-      {"z: float"},
-      // Attr def
-      {},
-      // Nodes
-      {
-          {{"y0"}, "ReadVariableOp", {"y"}, {{"dtype", DT_FLOAT}}},
-          {{"z"}, "Mul", {"x", "y0"}, {{"T", DT_FLOAT}}},
-      });
-}
-
-class CreateXlaLaunchOpTest : public ::testing::Test {
- protected:
-  void Init(const std::vector<FunctionDef>& flib) {
-    SessionOptions options;
-    auto* device_count = options.config.mutable_device_count();
-    device_count->insert({"CPU", 1});
-    TF_CHECK_OK(DeviceFactory::AddDevices(
-        options, "/job:localhost/replica:0/task:0", &devices_));
-
-    FunctionDefLibrary proto;
-    for (const auto& fdef : flib) {
-      *(proto.add_function()) = fdef;
-    }
-    lib_def_ = absl::make_unique<FunctionLibraryDefinition>(
-        OpRegistry::Global(), proto);
-    OptimizerOptions opts;
-    device_mgr_ = absl::make_unique<DeviceMgr>(devices_);
-    pflr_ = absl::make_unique<ProcessFunctionLibraryRuntime>(
-        device_mgr_.get(), Env::Default(), TF_GRAPH_DEF_VERSION, lib_def_.get(),
-        opts, /*default_thread_pool=*/nullptr, /*cluster_flr=*/nullptr);
-    flr_ = pflr_->GetFLR("/job:localhost/replica:0/task:0/cpu:0");
-  }
-
-  FunctionLibraryRuntime* flr_;
-  std::vector<Device*> devices_;
-  std::unique_ptr<DeviceMgr> device_mgr_;
-  std::unique_ptr<FunctionLibraryDefinition> lib_def_;
-  std::unique_ptr<ProcessFunctionLibraryRuntime> pflr_;
-
-  std::unique_ptr<OpKernel> kernel_;
-};
-
-AttrValue BoolAttr(bool b) {
-  AttrValue v;
-  v.set_b(b);
-  return v;
-}
-
-TEST_F(CreateXlaLaunchOpTest, OneFloatOneResourceArgument) {
-  FunctionDef fdef = XTimesY();
-  (*fdef.mutable_attr())["_XlaCompile"] = BoolAttr(true);
-  Init({fdef});
-
-  Status status = CreateXlaLaunchOp(
-      flr_, ToNodeDef(R"pb(
-        name: 'XTimesY' op: 'XTimesY' input: 'a' input: 'b'
-      )pb"), &kernel_);
-  ASSERT_TRUE(status.ok()) << status.ToString();
-
-  EXPECT_EQ("XTimesY", kernel_->name());
-  EXPECT_EQ("XTimesY", kernel_->type_string());
-
-  EXPECT_EQ(2, kernel_->num_inputs());
-  EXPECT_EQ(DT_FLOAT, kernel_->input_type(0));
-  EXPECT_EQ(DT_RESOURCE, kernel_->input_type(1));
-  EXPECT_EQ(DEVICE_MEMORY, kernel_->input_memory_types()[0]);
-  EXPECT_EQ(HOST_MEMORY, kernel_->input_memory_types()[1]);
-
-  EXPECT_EQ(1, kernel_->num_outputs());
-  EXPECT_EQ(DT_FLOAT, kernel_->output_type(0));
-  EXPECT_EQ(DEVICE_MEMORY, kernel_->output_memory_types()[0]);
-}
-
-TEST_F(CreateXlaLaunchOpTest, FailsIfXlaCompileAttrNotSet) {
-  FunctionDef fdef = XTimesY();
-  Init({fdef});
-
-  Status status = CreateXlaLaunchOp(flr_, ToNodeDef(R"proto(
-                                      name: 'XTimesY'
-                                      op: 'XTimesY'
-                                      input: 'a'
-                                      input: 'b'
-                                    )proto"), &kernel_);
-  EXPECT_TRUE(errors::IsInvalidArgument(status)) << status.ToString();
-}
-
-TEST_F(CreateXlaLaunchOpTest, FailsIfXlaCompileAttrIsSetToFalse) {
-  FunctionDef fdef = XTimesY();
-  (*fdef.mutable_attr())["_XlaCompile"] = BoolAttr(false);
-  Init({fdef});
-
-  Status status = CreateXlaLaunchOp(flr_, ToNodeDef(R"proto(
-                                      name: 'XTimesY'
-                                      op: 'XTimesY'
-                                      input: 'a'
-                                      input: 'b'
-                                    )proto"), &kernel_);
-  EXPECT_TRUE(errors::IsInvalidArgument(status)) << status.ToString();
-}
-
-}  // namespace tensorflow
diff --git a/tensorflow/compiler/jit/kernels/xla_launch_op.cc b/tensorflow/compiler/jit/kernels/xla_launch_op.cc
index 86a9fd3b8e..049d170fa4 100644
--- a/tensorflow/compiler/jit/kernels/xla_launch_op.cc
+++ b/tensorflow/compiler/jit/kernels/xla_launch_op.cc
@@ -39,15 +39,15 @@ limitations under the License.
 
 namespace tensorflow {
 
-XlaLocalLaunchBase::XlaLocalLaunchBase(OpKernelConstruction* ctx,
-                                       const std::vector<int>& constants,
-                                       const std::vector<int>& resources,
-                                       const NameAttrList& function)
-    : OpKernel(ctx),
-      constants_(constants),
-      resources_(resources),
-      device_type_(ctx->device_type()),
-      function_(function) {
+XlaLocalLaunchOp::XlaLocalLaunchOp(OpKernelConstruction* ctx)
+    : OpKernel(ctx), device_type_(ctx->device_type()) {
+  const NameAttrList* func;
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("function", &func));
+  function_ = *func;
+  DataTypeVector constant_types;
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("Tconstants", &constant_types));
+  num_constant_args_ = constant_types.size();
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("Nresources", &num_resource_args_));
   if (device_type_ == DeviceType(DEVICE_CPU)) {
     platform_id_ = se::host::kHostPlatformId;
   } else if (device_type_ == DeviceType(DEVICE_GPU)) {
@@ -57,8 +57,8 @@ XlaLocalLaunchBase::XlaLocalLaunchBase(OpKernelConstruction* ctx,
   }
 }
 
-Status XlaLocalLaunchBase::BuildCompilationCache(OpKernelContext* ctx,
-                                                 XlaCompilationCache** cache) {
+Status XlaLocalLaunchOp::BuildCompilationCache(OpKernelContext* ctx,
+                                               XlaCompilationCache** cache) {
   const XlaDevice::Metadata* metadata;
   Status s = XlaDevice::GetMetadata(ctx, &metadata);
   if (s.ok()) {
@@ -90,8 +90,8 @@ Status XlaLocalLaunchBase::BuildCompilationCache(OpKernelContext* ctx,
   return Status::OK();
 }
 
-void XlaLocalLaunchBase::Compute(OpKernelContext* ctx) {
-  VLOG(1) << "XlaLocalLaunchOpBase::Compute "
+void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) {
+  VLOG(1) << "XlaLocalLaunchOp::Compute "
           << Canonicalize(function_.name(), AttrSlice(&function_.attr()));
   // We store information about the JIT-compiled XLA computation
   // in the ResourceMgr.
@@ -124,7 +124,7 @@ void XlaLocalLaunchBase::Compute(OpKernelContext* ctx) {
   }
 
   std::map<int, OptionalTensor> variables =
-      SnapshotResourceVariables(ctx, resources_);
+      SnapshotResourceVariables(ctx, num_resource_args_);
 
   xla::LocalClient* client = static_cast<xla::LocalClient*>(cache->client());
 
@@ -161,7 +161,7 @@ void XlaLocalLaunchBase::Compute(OpKernelContext* ctx) {
   xla::LocalExecutable* executable;
 
   std::map<int, Tensor> constant_args;
-  for (int i : constants_) {
+  for (int i = 0; i < num_constant_args_; ++i) {
     constant_args.insert({i, ctx->input(i)});
   }
   OP_REQUIRES_OK(ctx, cache->Compile(options, function_, constant_args,
@@ -170,8 +170,8 @@ void XlaLocalLaunchBase::Compute(OpKernelContext* ctx) {
 
   VLOG(1) << "Executing XLA Computation...";
 
-  XlaComputationLaunchContext launch_context(client, xla_allocator,
-                                             allocate_xla_tensors);
+  XlaComputationLaunchContext launch_context(
+      num_resource_args_, client, xla_allocator, allocate_xla_tensors);
   launch_context.PopulateInputs(ctx, kernel, variables);
 
   // Execute the computation.
@@ -194,62 +194,6 @@ void XlaLocalLaunchBase::Compute(OpKernelContext* ctx) {
   VLOG(1) << "Done";
 }
 
-namespace {
-
-// OP_REQUIRES_OK_RETURN is the same as OP_REQUIRES_OK except that
-// in error case, it returns RET instead of void.
-#define OP_REQUIRES_OK_RETURN(CTX, RET, ...)                \
-  do {                                                      \
-    ::tensorflow::Status _s(__VA_ARGS__);                   \
-    if (!TF_PREDICT_TRUE(_s.ok())) {                        \
-      (CTX)->CtxFailureWithWarning(__FILE__, __LINE__, _s); \
-      return RET;                                           \
-    }                                                       \
-  } while (0)
-
-// Helper static functions to construct parameters for
-// XlaLocalLaunchBase constructor from OpKernelConstruction.
-std::vector<int> ConstantsVector(OpKernelConstruction* ctx) {
-  DataTypeVector constant_types;
-  OP_REQUIRES_OK_RETURN(ctx, std::vector<int>(),
-                        ctx->GetAttr("Tconstants", &constant_types));
-  std::vector<int> constants(constant_types.size());
-  std::iota(constants.begin(), constants.end(), 0);
-  return constants;
-}
-
-std::vector<int> ResourcesVector(OpKernelConstruction* ctx) {
-  DataTypeVector constant_types;
-  OP_REQUIRES_OK_RETURN(ctx, std::vector<int>(),
-                        ctx->GetAttr("Tconstants", &constant_types));
-
-  DataTypeVector arg_types;
-  OP_REQUIRES_OK_RETURN(ctx, std::vector<int>(),
-                        ctx->GetAttr("Targs", &arg_types));
-
-  int num_resources;
-  OP_REQUIRES_OK_RETURN(ctx, std::vector<int>(),
-                        ctx->GetAttr("Nresources", &num_resources));
-
-  std::vector<int> resources(num_resources);
-  std::iota(resources.begin(), resources.end(),
-            constant_types.size() + arg_types.size());
-  return resources;
-}
-
-NameAttrList FunctionAttr(OpKernelConstruction* ctx) {
-  const NameAttrList* func;
-  OP_REQUIRES_OK_RETURN(ctx, NameAttrList(), ctx->GetAttr("function", &func));
-  return *func;
-}
-
-#undef OP_REQUIRES_OK_RETURN
-}  // namespace
-
-XlaLocalLaunchOp::XlaLocalLaunchOp(OpKernelConstruction* ctx)
-    : XlaLocalLaunchBase(ctx, ConstantsVector(ctx), ResourcesVector(ctx),
-                         FunctionAttr(ctx)) {}
-
 XlaLocalLaunchOp::~XlaLocalLaunchOp() {
   VLOG(1) << "XlaLocalLaunchOp destroyed";
 }
diff --git a/tensorflow/compiler/jit/kernels/xla_launch_op.h b/tensorflow/compiler/jit/kernels/xla_launch_op.h
index 8dfc4b382d..8f8e646f0f 100644
--- a/tensorflow/compiler/jit/kernels/xla_launch_op.h
+++ b/tensorflow/compiler/jit/kernels/xla_launch_op.h
@@ -26,41 +26,6 @@ limitations under the License.
 
 namespace tensorflow {
 
-// XlaLocalLaunchBase is almost the same as XlaLocalLaunchOp.
-// The only difference is that it does not require arguments to follow
-// the "constants, then regular args, then resources" order.
-// It takes vectors of constant and resource arguments explicitly.
-// It does not have corresponding OpDef because it is never present
-// in the GraphDef.
-// Currently, it is used by eager runtime. FunctionLibraryRuntime creates
-// this kernel when asked to create a kernel for an XLA-compiled function.
-class XlaLocalLaunchBase : public OpKernel {
- public:
-  XlaLocalLaunchBase(OpKernelConstruction* ctx,
-                     const std::vector<int>& constants,
-                     const std::vector<int>& resources,
-                     const NameAttrList& function);
-  XlaLocalLaunchBase(const XlaLocalLaunchBase&) = delete;
-  XlaLocalLaunchBase& operator=(const XlaLocalLaunchBase&) = delete;
-  ~XlaLocalLaunchBase() override = default;
-
-  void Compute(OpKernelContext* ctx) override;
-
- protected:
-  // Builds a XlaCompilationCache class suitable for the current device.
-  Status BuildCompilationCache(OpKernelContext* ctx,
-                               XlaCompilationCache** cache);
-
-  // Indexes of compile-time constant inputs
-  std::vector<int> constants_;
-  // Indexes of resource inputs
-  std::vector<int> resources_;
-
-  DeviceType device_type_;
-  NameAttrList function_;
-  se::Platform::Id platform_id_;
-};
-
 // XlaLocalLaunchOp is used to replace a region of the TensorFlow graph
 // which will be compiled and executed using XLA.  The XlaLocalLaunchOp is
 // responsible for handling interactions with the TensorFlow executor.
@@ -70,12 +35,26 @@ class XlaLocalLaunchBase : public OpKernel {
 // XlaLocalLaunchOp uses xla::LocalClient::Compile() and
 // xla::LocalExecutable::Run(), and passes arguments into/out of XLA in device
 // memory.
-class XlaLocalLaunchOp : public XlaLocalLaunchBase {
+class XlaLocalLaunchOp : public OpKernel {
  public:
   explicit XlaLocalLaunchOp(OpKernelConstruction* ctx);
   ~XlaLocalLaunchOp() override;
 
+  void Compute(OpKernelContext* ctx) override;
+
  private:
+  // Builds a XlaCompilationCache class suitable for the current device.
+  Status BuildCompilationCache(OpKernelContext* ctx,
+                               XlaCompilationCache** compiler);
+
+  DeviceType device_type_;
+  NameAttrList function_;
+  int num_constant_args_;
+  // Number of resource variable arguments.
+  int num_resource_args_;
+
+  se::Platform::Id platform_id_;
+
   TF_DISALLOW_COPY_AND_ASSIGN(XlaLocalLaunchOp);
 };
 
diff --git a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc
index 6b83cf67ff..60458f6f33 100644
--- a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc
+++ b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc
@@ -48,12 +48,13 @@ Status XlaCompileOnDemandOp::Run(OpKernelContext* ctx,
                                  const XlaCompiler::CompilationResult* result,
                                  xla::LocalExecutable* executable) {
   std::map<int, OptionalTensor> variables = GetVariables(ctx);
+  int64 num_resource_args = variables.size();
 
   xla::LocalClient* client = metadata.client();
 
   // Builds an XLA allocator for the device.
   XlaComputationLaunchContext launch_context(
-      client, client->backend().memory_allocator(), true);
+      num_resource_args, client, client->backend().memory_allocator(), true);
 
   launch_context.PopulateInputs(ctx, result, variables);
 
diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc
index 0223f97a03..33e53612b9 100644
--- a/tensorflow/compiler/jit/xla_launch_util.cc
+++ b/tensorflow/compiler/jit/xla_launch_util.cc
@@ -38,13 +38,14 @@ using xla::ScopedShapedBuffer;
 using xla::ShapedBuffer;
 }  // anonymous namespace
 
-std::map<int, OptionalTensor> SnapshotResourceVariables(
-    OpKernelContext* ctx, const std::vector<int>& variables) {
+std::map<int, OptionalTensor> SnapshotResourceVariables(OpKernelContext* ctx,
+                                                        int num_variables) {
   std::map<int, OptionalTensor> snapshot;
-  for (int i : variables) {
+  int first_variable = ctx->num_inputs() - num_variables;
+  for (int i = 0; i < num_variables; ++i) {
     Var* variable = nullptr;
-    ResourceHandle handle = HandleFromInput(ctx, i);
-    OptionalTensor& tensor = snapshot[i];
+    ResourceHandle handle = HandleFromInput(ctx, first_variable + i);
+    OptionalTensor& tensor = snapshot[first_variable + i];
     if (LookupResource(ctx, handle, &variable).ok()) {
       tf_shared_lock lock(*variable->mu());
       tensor.name = handle.name();
@@ -111,9 +112,10 @@ ScopedShapedBuffer ExtractSubShapedBuffer(
 using internal::ExtractSubShapedBuffer;
 
 XlaComputationLaunchContext::XlaComputationLaunchContext(
-    xla::LocalClient* client, xla::DeviceMemoryAllocator* xla_allocator,
-    bool allocate_xla_tensors)
-    : client_(client),
+    int64 num_resource_args, xla::LocalClient* client,
+    xla::DeviceMemoryAllocator* xla_allocator, bool allocate_xla_tensors)
+    : num_resource_args_(num_resource_args),
+      client_(client),
       xla_allocator_(xla_allocator),
       allocate_xla_tensors_(allocate_xla_tensors) {}
 
diff --git a/tensorflow/compiler/jit/xla_launch_util.h b/tensorflow/compiler/jit/xla_launch_util.h
index a2431253f8..38291b0bd4 100644
--- a/tensorflow/compiler/jit/xla_launch_util.h
+++ b/tensorflow/compiler/jit/xla_launch_util.h
@@ -31,17 +31,15 @@ limitations under the License.
 namespace tensorflow {
 class XlaAllocator;
 
-// Takes a snapshot of the values of resource variable arguments, whose
-// indices are specified in `variables` argument. We snapshot tensors that back
+// Takes a snapshot of the values of resource variable arguments, which are
+// the last `num_variables` arguments. We snapshot tensors that back
 // resource variables since concurrent updates may modify the shape, and it is
 // important that the shapes used for compilation match the true shapes of the
 // buffers.
 //
-// Returns a map of TensorFlow argument index to resource variable. If a
-// resource variable is not initialized, the corresponding OptionalTensor
-// will have its `present` field set to false.
-std::map<int, OptionalTensor> SnapshotResourceVariables(
-    OpKernelContext* ctx, const std::vector<int>& variables);
+// Returns a map of TensorFlow argument index to resource variable.
+std::map<int, OptionalTensor> SnapshotResourceVariables(OpKernelContext* ctx,
+                                                        int num_variables);
 
 // Adapter class that wraps a Tensorflow allocator as an XLA allocator.
 // Assumes that the Tensorflow allocator permits asynchronous deallocation:
@@ -74,7 +72,7 @@ class XlaComputationLaunchContext {
   // Create a new launch context. 'allocate_xla_tensors' is true if allocated
   // output tensors and variables are always XlaTensors. If false they are
   // assumed to be "normal" device pointers.
-  XlaComputationLaunchContext(xla::LocalClient* client,
+  XlaComputationLaunchContext(int64 num_resource_args, xla::LocalClient* client,
                               xla::DeviceMemoryAllocator* xla_allocator,
                               bool allocate_xla_tensors);
 
@@ -94,6 +92,7 @@ class XlaComputationLaunchContext {
   const std::vector<xla::ShapedBuffer*>& arguments() const { return arg_ptrs_; }
 
  private:
+  int64 num_resource_args_;
   xla::LocalClient* client_;
   xla::DeviceMemoryAllocator* xla_allocator_;
   bool allocate_xla_tensors_;
diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index 9791792f29..aaea83ae9c 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -327,11 +327,7 @@ tf_xla_py_test(
         ":xla_test",
         "//tensorflow/python:array_ops",
         "//tensorflow/python:framework_for_generated_wrappers",
-        "//tensorflow/python:layers",
-        "//tensorflow/python:math_ops",
-        "//tensorflow/python:nn",
         "//tensorflow/python:platform_test",
-        "//tensorflow/python/eager:function",
     ],
 )
 
diff --git a/tensorflow/compiler/tests/eager_test.py b/tensorflow/compiler/tests/eager_test.py
index 5ab1585f8c..bdd0185dfe 100644
--- a/tensorflow/compiler/tests/eager_test.py
+++ b/tensorflow/compiler/tests/eager_test.py
@@ -24,16 +24,10 @@ from tensorflow.compiler.tests.xla_test import XLATestCase
 from tensorflow.core.protobuf import config_pb2
 from tensorflow.python.eager import backprop
 from tensorflow.python.eager import context
-from tensorflow.python.eager import function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
-from tensorflow.python.layers import convolutional
-from tensorflow.python.layers import pooling
 from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import init_ops
-from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import nn_ops
 from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.platform import googletest
 
@@ -49,7 +43,7 @@ class EagerTest(XLATestCase):
 
   def testExecuteListOutputLen0(self):
     with self.test_scope():
-      empty = constant_op.constant([], dtype=dtypes.float32)
+      empty = constant_op.constant([], dtype=dtypes.int32)
       result = array_ops.unstack(empty, 0)
       self.assertTrue(isinstance(result, list))
       self.assertEqual(0, len(result))
@@ -57,7 +51,7 @@ class EagerTest(XLATestCase):
   def testExecuteListOutputLen1(self):
     with self.test_scope():
       split_dim = constant_op.constant(1)
-      value = constant_op.constant([[0., 1., 2.], [3., 4., 5.]])
+      value = constant_op.constant([[0, 1, 2], [3, 4, 5]])
       result = array_ops.split(value, 1, axis=split_dim)
       self.assertTrue(isinstance(result, list))
       self.assertEqual(1, len(result))
@@ -66,7 +60,7 @@ class EagerTest(XLATestCase):
   def testExecuteListOutputLen3(self):
     with self.test_scope():
       split_dim = constant_op.constant(1)
-      value = constant_op.constant([[0., 1., 2.], [3., 4., 5.]])
+      value = constant_op.constant([[0, 1, 2], [3, 4, 5]])
       result = array_ops.split(value, 3, axis=split_dim)
       self.assertTrue(isinstance(result, list))
       self.assertEqual(3, len(result))
@@ -137,105 +131,7 @@ class EagerTest(XLATestCase):
     self.assertEqual(2., grads[0][0].numpy())
 
 
-class EagerFunctionTest(XLATestCase):
-
-  def testBasic(self):
-    with self.test_scope():
-      matmul = function.defun(math_ops.matmul, compiled=True)
-      t = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
-      sq = matmul(t, t, transpose_a=True)
-      self.assertAllEqual(sq.numpy().reshape(-1), [10, 14, 14, 20])
-
-  def testConv(self):
-    if 'GPU' in self.device:
-      # TODO(b/32333178)
-      self.skipTest('Current implementation of RandomStandardNormal kernel '
-                    'is very slow on GPU, and has been blacklisted.')
-    with self.test_scope():
-      data_format = 'channels_last'
-      conv = convolutional.Conv2D(
-          filters=1, kernel_size=2, padding='VALID',
-          data_format=data_format, activation=nn_ops.relu,
-          kernel_initializer=init_ops.ones_initializer(),
-          bias_initializer=init_ops.zeros_initializer())
-      pool = pooling.MaxPooling2D(2, 2, data_format=data_format)
-
-      def model(x):
-        x = conv(x)
-        return pool(x)
-      model = function.defun(model, compiled=True)
-
-      x = array_ops.ones([1, 4, 4, 1])
-      y = model(x)
-      self.assertAllEqual(y.numpy(), [[[[4.]]]])
-
-  def testReadVariable(self):
-    with self.test_scope():
-      v = resource_variable_ops.ResourceVariable(1.0)
-
-      @function.defun(compiled=True)
-      def f():
-        return v.read_value()
-
-      var = f()
-      self.assertEqual(1.0, var.numpy())
-
-  def testUpdateVariable(self):
-    with self.test_scope():
-      v = resource_variable_ops.ResourceVariable(1.0)
-
-      def f(v):
-        v.assign_add(1.0)
-        return v
-
-      f = function.defun(f, compiled=True)
-
-      var = f(v)
-      self.assertEqual(2.0, var.numpy())
-
-  def testAllArgumentKinds(self):
-    """Test a complex function that takes different argument kinds.
-
-    tf2xla machinery that translates, compiles, and runs defuns
-    classifies arguments into: compile-time constants, regular tensors,
-    and resources. This test creates a function with a mix of all these
-    kinds. Moreover, the order of function arguments is intentionally mixed up.
-
-    This also tests the case when the same argument is a compile-time constant
-    as well as used in an operation that normally expects its inputs to be
-    in device memory - addition in this case.
-    """
-    with self.test_scope():
-      def foo(c1, r1, v1, c2, v2, r2):
-        # c1 and c2 are compile-time constants
-        # r1 and r2 are regular tensors
-        # v1 and v2 are resource variables
-        a = c1 + r1
-        b = math_ops.cast(c2, dtypes.float32) + v2
-        c = array_ops.slice(v1, c1, c2)
-        d = r2 * v2
-        return a, b, c, d
-
-      foo = function.defun(foo, compiled=True)
-
-      c1 = [0, 0]
-      c2 = array_ops.ones([2], dtype=dtypes.int32)
-
-      r1 = array_ops.ones([2])
-      r2 = [[2., 2.], [3., 3.]]
-
-      v1 = resource_variable_ops.ResourceVariable([[1., 2.], [3., 4.]])
-      v2 = resource_variable_ops.ResourceVariable([[10., 20.], [30., 40.]])
-
-      a, b, c, d = foo(c1, r1, v1, c2, v2, r2)
-
-      self.assertAllEqual([1, 1], a.numpy())
-      self.assertAllEqual([[11., 21.], [31., 41.]], b.numpy())
-      self.assertAllEqual([[1.]], c.numpy())
-      self.assertAllEqual([[20., 40.], [90., 120.]], d.numpy())
-
-
-if __name__ == '__main__':
+if __name__ == "__main__":
   ops.enable_eager_execution(
       config=config_pb2.ConfigProto(log_device_placement=True))
   googletest.main()
diff --git a/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py b/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py
index b8f352d5f5..8517a3bf7b 100644
--- a/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py
+++ b/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py
@@ -36,7 +36,9 @@ def device_and_data_format():
                                                               'channels_last')
 
 
-def random_batch(batch_size, data_format):
+def random_batch(batch_size, device_and_format=None):
+  _, data_format = device_and_format or device_and_data_format()
+
   shape = (3, 224, 224) if data_format == 'channels_first' else (224, 224, 3)
   shape = (batch_size,) + shape
 
@@ -68,7 +70,7 @@ class ResNet50Test(tf.test.TestCase):
     if defun:
       model.call = tfe.defun(model.call)
     with tf.device(device), tfe.execution_mode(execution_mode):
-      images, _ = random_batch(2, data_format)
+      images, _ = random_batch(2)
       output = model(images, training=False)
       tfe.async_wait()
     self.assertEqual((2, 1000), output.shape)
@@ -89,7 +91,7 @@ class ResNet50Test(tf.test.TestCase):
     device, data_format = device_and_data_format()
     model = resnet50.ResNet50(data_format, include_top=False)
     with tf.device(device):
-      images, _ = random_batch(2, data_format)
+      images, _ = random_batch(2)
       output = model(images, training=False)
     output_shape = ((2, 2048, 1, 1)
                     if data_format == 'channels_first' else (2, 1, 1, 2048))
@@ -99,7 +101,7 @@ class ResNet50Test(tf.test.TestCase):
     device, data_format = device_and_data_format()
     model = resnet50.ResNet50(data_format, include_top=False, pooling='avg')
     with tf.device(device):
-      images, _ = random_batch(2, data_format)
+      images, _ = random_batch(2)
       output = model(images, training=False)
     self.assertEqual((2, 2048), output.shape)
 
@@ -113,7 +115,7 @@ class ResNet50Test(tf.test.TestCase):
         name='t0').as_default(), tf.contrib.summary.always_record_summaries():
       with tf.device(device), tfe.execution_mode(execution_mode):
         optimizer = tf.train.GradientDescentOptimizer(0.1)
-        images, labels = random_batch(2, data_format)
+        images, labels = random_batch(2)
         train_one_step(model, images, labels, optimizer)
         self.assertEqual(320, len(model.variables))
         tfe.async_wait()
@@ -132,7 +134,7 @@ class ResNet50Test(tf.test.TestCase):
     model = resnet50.ResNet50(data_format)
     optimizer = tf.train.GradientDescentOptimizer(0.1)
     with tf.device(device):
-      images, labels = random_batch(2, data_format)
+      images, labels = random_batch(2)
       gc.disable()
       # Warm up. Note that this first run does create significant amounts of
       # garbage to be collected. The hope is that this is a build-only effect,
@@ -200,18 +202,18 @@ class ResNet50Benchmarks(tf.test.Benchmark):
     # which forces a sync. This is a roundabout way, yes.
     tf.constant(1.).cpu()
 
-  def _benchmark_eager_apply(self, label, device_and_format, defun=False,
-                             execution_mode=None, compiled=False):
+  def _benchmark_eager_apply(self, label, defun=False, execution_mode=None,
+                             device_and_format=None):
     with tfe.execution_mode(execution_mode):
-      device, data_format = device_and_format
+      device, data_format = device_and_format or device_and_data_format()
       model = resnet50.ResNet50(data_format)
       if defun:
-        model.call = tfe.defun(model.call, compiled=compiled)
+        model.call = tfe.defun(model.call)
       batch_size = 64
       num_burn = 5
       num_iters = 30
       with tf.device(device):
-        images, _ = random_batch(batch_size, data_format)
+        images, _ = random_batch(batch_size, device_and_format)
         for _ in xrange(num_burn):
           model(images, training=False).cpu()
         if execution_mode:
@@ -225,34 +227,30 @@ class ResNet50Benchmarks(tf.test.Benchmark):
         self._report(label, start, num_iters, device, batch_size, data_format)
 
   def benchmark_eager_apply_sync(self):
-    self._benchmark_eager_apply('eager_apply', device_and_data_format(),
-                                defun=False)
+    self._benchmark_eager_apply('eager_apply', defun=False)
 
   def benchmark_eager_apply_async(self):
     self._benchmark_eager_apply(
-        'eager_apply_async', device_and_data_format(), defun=False,
-        execution_mode=tfe.ASYNC)
+        'eager_apply_async', defun=False, execution_mode=tfe.ASYNC)
 
   def benchmark_eager_apply_with_defun(self):
-    self._benchmark_eager_apply('eager_apply_with_defun',
-                                device_and_data_format(), defun=True)
+    self._benchmark_eager_apply('eager_apply_with_defun', defun=True)
 
   def _benchmark_eager_train(self,
                              label,
                              make_iterator,
-                             device_and_format,
                              defun=False,
                              execution_mode=None,
-                             compiled=False):
+                             device_and_format=None):
     with tfe.execution_mode(execution_mode):
-      device, data_format = device_and_format
+      device, data_format = device_and_format or device_and_data_format()
       for batch_size in self._train_batch_sizes():
-        (images, labels) = random_batch(batch_size, data_format)
+        (images, labels) = random_batch(batch_size, device_and_format)
         num_burn = 3
         num_iters = 10
         model = resnet50.ResNet50(data_format)
         if defun:
-          model.call = tfe.defun(model.call, compiled=compiled)
+          model.call = tfe.defun(model.call)
         optimizer = tf.train.GradientDescentOptimizer(0.1)
 
         with tf.device(device):
@@ -275,21 +273,18 @@ class ResNet50Benchmarks(tf.test.Benchmark):
           self._report(label, start, num_iters, device, batch_size, data_format)
 
   def benchmark_eager_train_sync(self):
-    self._benchmark_eager_train('eager_train', MockIterator,
-                                device_and_data_format(), defun=False)
+    self._benchmark_eager_train('eager_train', MockIterator, defun=False)
 
   def benchmark_eager_train_async(self):
     self._benchmark_eager_train(
         'eager_train_async',
         MockIterator,
-        device_and_data_format(),
         defun=False,
         execution_mode=tfe.ASYNC)
 
   def benchmark_eager_train_with_defun(self):
     self._benchmark_eager_train(
-        'eager_train_with_defun', MockIterator,
-        device_and_data_format(), defun=True)
+        'eager_train_with_defun', MockIterator, defun=True)
 
   def benchmark_eager_train_datasets(self):
 
@@ -299,8 +294,7 @@ class ResNet50Benchmarks(tf.test.Benchmark):
       return tfe.Iterator(ds)
 
     self._benchmark_eager_train(
-        'eager_train_dataset', make_iterator,
-        device_and_data_format(), defun=False)
+        'eager_train_dataset', make_iterator, defun=False)
 
   def benchmark_eager_train_datasets_with_defun(self):
 
@@ -310,8 +304,7 @@ class ResNet50Benchmarks(tf.test.Benchmark):
       return tfe.Iterator(ds)
 
     self._benchmark_eager_train(
-        'eager_train_dataset_with_defun', make_iterator,
-        device_and_data_format(), defun=True)
+        'eager_train_dataset_with_defun', make_iterator, defun=True)
 
 
 if __name__ == '__main__':
diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py
index 60cfacc141..741bd2ac9c 100644
--- a/tensorflow/python/eager/function.py
+++ b/tensorflow/python/eager/function.py
@@ -23,7 +23,6 @@ import collections
 
 import numpy as np
 
-from tensorflow.core.framework import attr_value_pb2
 from tensorflow.core.framework import function_pb2
 from tensorflow.python import pywrap_tensorflow
 from tensorflow.python.eager import context
@@ -226,7 +225,7 @@ def _inference_name(n):
 class _EagerDefinedFunction(object):
   """Function object with the interface of tf _DefinedFunction."""
 
-  def __init__(self, name, graph, operations, inputs, outputs, attrs):
+  def __init__(self, name, graph, operations, inputs, outputs):
     """Initializes an eager defined function.
 
     Args:
@@ -236,7 +235,6 @@ class _EagerDefinedFunction(object):
         which will be in the function
       inputs: the tensors in the graph to be used as inputs to the function
       outputs: the tensors in the graph which will be outputs to the function
-      attrs: dict mapping names of attributes to their AttrValue values
     """
     fn = pywrap_tensorflow.TF_GraphToFunction_wrapper(
         graph._c_graph,  # pylint: disable=protected-access
@@ -248,14 +246,6 @@ class _EagerDefinedFunction(object):
         [],
         None,
         compat.as_str(""))
-
-    for name, attr_value in attrs.items():
-      serialized = attr_value.SerializeToString()
-      # TODO(iga): this creates and deletes a new TF_Status for every attr.
-      # It might be worth creating a convenient way to re-use status.
-      pywrap_tensorflow.TF_FunctionSetAttrValueProto(
-          fn, compat.as_str(name), serialized)
-
     # TODO(apassos) avoid creating a FunctionDef (specially to grab the
     # signature, but also in general it's nice not to depend on it.
     with c_api_util.tf_buffer() as buffer_:
@@ -297,6 +287,25 @@ def _flatten(sequence):
 
 class GraphModeFunction(object):
   """Callable object representing a graph-mode function.
+
+  Args:
+    name: str the name of the created function
+    input_placeholders: list of placeholder values (tensors) to feed when
+      calling the wrapped function.
+    extra_inputs: Tensor inputs this function definition closed over which
+      are passed as arguments. Need to track so gradients are supported
+      correctly.
+    graph: the Graph from which the operations will be pulled. Used as
+      a context when computing gradients.
+    operations: the subset of Operations in the graph used in the function
+      definition.
+    outputs: a flat list of the Tensors in the graph used as outputs to the
+      function
+    func_outputs: a possibly nested python object which will be returned by
+      this function. The Tensors in this structure will be replaced by their
+      corresponding values in outputs.
+    output_shapes: List of shapes of all tensors in outputs
+    variables: (optional) List of variables to watch during function execution.
   """
 
   def __init__(self,
@@ -308,36 +317,9 @@ class GraphModeFunction(object):
                outputs,
                func_outputs,
                output_shapes,
-               variables=None,
-               attrs=None):
-    """Initialize a GraphModeFunction.
-
-    Args:
-      name: str the name of the created function
-      input_placeholders: list of placeholder values (tensors) to feed when
-        calling the wrapped function.
-      extra_inputs: Tensor inputs this function definition closed over which
-        are passed as arguments. Need to track so gradients are supported
-        correctly.
-      graph: the Graph from which the operations will be pulled. Used as
-        a context when computing gradients.
-      operations: the subset of Operations in the graph used in the function
-        definition.
-      outputs: a flat list of the Tensors in the graph used as outputs to the
-        function
-      func_outputs: a possibly nested python object which will be returned by
-        this function. The Tensors in this structure will be replaced by their
-        corresponding values in outputs.
-      output_shapes: List of shapes of all tensors in outputs
-      variables: (optional) List of variables to watch during function
-        execution.
-      attrs: (optional) dict mapping names of attributes to their AttrValue
-        values. Attributes in `attrs` will be included in this function's
-        definition.
-    """
-    self._attrs = attrs or {}
+               variables=None):
     defined_function = _EagerDefinedFunction(
-        name, graph, operations, input_placeholders, outputs, self._attrs)
+        name, graph, operations, input_placeholders, outputs)
     if len(input_placeholders) != len(defined_function.signature.input_arg):
       raise ValueError("Internal error: invalid lengths. %s %s" % (
           len(input_placeholders), len(defined_function.signature.input_arg)))
@@ -390,7 +372,7 @@ class GraphModeFunction(object):
     forward_name = _forward_name(self._func_name)
     self._forward_fdef = _EagerDefinedFunction(
         forward_name, self._graph, self._ops, self._input_placeholders,
-        filtered_outputs + captures, self._attrs)
+        filtered_outputs + captures)
     all_inputs = self._out_grad_placeholders + captures
     # Excluding input ops from the body as we do not intend to execute these
     # operations when the function is executed.
@@ -404,7 +386,7 @@ class GraphModeFunction(object):
     bname = _backward_name(self._func_name)
     self._backward_function = GraphModeFunction(
         bname, all_inputs, [], self._graph, function_def_ops,
-        backward_outputs, in_gradients, output_shapes, attrs=self._attrs)
+        backward_outputs, in_gradients, output_shapes)
 
   def _backprop_call(self, args):
     """Calls the wrapped function and records the result on a tape."""
@@ -578,7 +560,7 @@ def _get_defun_inputs(args):
   return nest.pack_sequence_as(args, ret)
 
 
-def _defun_internal(name, func, compiled, args, kwds):
+def _defun_internal(name, func, args, kwds):
   """Defines and returns graph-mode version of func."""
   graph_key = ops.get_default_graph()._graph_key  # pylint: disable=protected-access
   with context.graph_mode():
@@ -643,14 +625,9 @@ def _defun_internal(name, func, compiled, args, kwds):
     for f in tmp_graph._functions.values():  # pylint: disable=protected-access
       # TODO(ashankar): What about the gradient registry?
       _register(f._c_func.func)  # pylint: disable=protected-access
-
-  attrs = {}
-  if compiled:
-    attrs["_XlaCompile"] = attr_value_pb2.AttrValue(b=True)
-
   return GraphModeFunction(
       fname, all_inputs, extra_inputs, tmp_graph, operations, func_def_outputs,
-      func_outputs, output_shapes, variables, attrs)
+      func_outputs, output_shapes, variables)
 
 
 # Defun uses this instead of Tensor as a cache key. Using dtype because
@@ -692,7 +669,7 @@ def _register(fn):
 
 
 # TODO(apassos): better error messages for non-hashable arguments.
-def named_defun(func, name, compiled=False):
+def named_defun(func, name):
   """Defines a function with a given name.
 
   See the documentation for `defun` for more information on the semantics of the
@@ -701,7 +678,6 @@ def named_defun(func, name, compiled=False):
   Args:
     func: the function to be wrapped.
     name: the name given to it.
-    compiled: if true, the framework will attempt to compile func with XLA.
 
   Returns:
     the wrapped function.
@@ -718,13 +694,13 @@ def named_defun(func, name, compiled=False):
 
     if cache_key not in arguments_to_functions:
       arguments_to_functions[cache_key] = _defun_internal(
-          name, func, compiled, args, kwds)
+          name, func, args, kwds)
     return arguments_to_functions[cache_key](*args)
 
   return decorated
 
 
-def defun(func=None, compiled=False):
+def defun(func):
   """Decorator to compile func into graph_mode.
 
   `defun` converts a function that constructs a TensorFlow graph into a function
@@ -767,45 +743,18 @@ def defun(func=None, compiled=False):
   ```
 
   Args:
-    func: function to be compiled. If `func` is None, returns a
-      decorator that can be invoked with a single argument - `func`. The
-      end result is equivalent to providing all the arguments up front.
-      In other words, defun(compiled=True)(func) is equivalent to
-      defun(func, compiled=True). The former allows the following use case:
-        @tfe.defun(compiled=True)
-        def foo(...):
-          ...
-    compiled: If True, an attempt to compile `func` with XLA will be made.
-      If it fails, function will be run normally. Experimental.
-      Currently, supported only for execution on TPUs.
+    func: function to be compiled.
 
   Returns:
-     If `func` is not None, returns callable that will execute the compiled
-     function (and return zero or more `tf.Tensor` objects).
-     If `func` is None, returns a decorator that, when invoked with a single
-     `func` argument, returns a callable equivalent to the case above.
+     A callable that will execute the compiled function (and return zero
+     or more `tf.Tensor` objects).
   """
   # TODO(apassos): deal with captured global state. Deal with control flow.
-  def decorated(function):
-    try:
-      name = function.__name__
-    except AttributeError:
-      name = "function"
-    return tf_decorator.make_decorator(
-        function, named_defun(function, name, compiled=compiled))
-
-  # This code path is for the `foo = tfe.defun(foo, ...)` use case
-  if func is not None:
-    return decorated(func)
-
-  # This code path is for the
-  #
-  # @tfe.defun(...)
-  # def foo(...):
-  #    ...
-  #
-  # use case, which is equivalent to `foo = tfe.defun(...)(foo)`
-  return decorated
+  try:
+    name = func.__name__
+  except AttributeError:
+    name = "function"
+  return tf_decorator.make_decorator(func, named_defun(func, name))
 
 
 def make_defun_op(func, *args, **kwds):
@@ -857,7 +806,7 @@ def make_defun_op(func, *args, **kwds):
   name = func.__name__
   if any(isinstance(x, ops.EagerTensor) for x in kwds.values()):
     raise ValueError("Tensor keyword arguments are not supported.")
-  return _defun_internal(name, func, False, args, kwds)
+  return _defun_internal(name, func, args, kwds)
 
 
 class AutomaticControlDependencies(object):
-- 
GitLab


From 392ce20dccefe86b5ef38ef8ac2bf6534ca17cd8 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 03:26:06 -0700
Subject: [PATCH 355/395] Fix a test expectation.

PiperOrigin-RevId: 195796348
---
 tensorflow/compiler/xla/service/instruction_fusion_test.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/compiler/xla/service/instruction_fusion_test.cc b/tensorflow/compiler/xla/service/instruction_fusion_test.cc
index b4b1955fe2..6dd8fa1ab0 100644
--- a/tensorflow/compiler/xla/service/instruction_fusion_test.cc
+++ b/tensorflow/compiler/xla/service/instruction_fusion_test.cc
@@ -126,7 +126,7 @@ TEST_F(InstructionFusionTest, FuseCheapNonDuplicatableOps) {
   EXPECT_EQ(Count(*module, HloOpcode::kFusion), 1) << module->ToString();
 
   // Make sure the add hasn't been duplicated.
-  EXPECT_EQ(Count(*module, HloOpcode::kFusion), 1) << module->ToString();
+  EXPECT_EQ(Count(*module, HloOpcode::kAdd), 1) << module->ToString();
 }
 
 TEST_F(InstructionFusionTest, AvoidDuplicationIfNotAllFusableRecursively) {
-- 
GitLab


From 42115bdf2b9d3bc2d544d19e2c822879cc634379 Mon Sep 17 00:00:00 2001
From: Asim Shankar <ashankar@google.com>
Date: Tue, 8 May 2018 07:28:43 -0700
Subject: [PATCH 356/395] ProfileHandler: Remove unnecessary interface method.
 PiperOrigin-RevId: 195815565

---
 tensorflow/core/common_runtime/profile_handler.h | 16 ----------------
 1 file changed, 16 deletions(-)

diff --git a/tensorflow/core/common_runtime/profile_handler.h b/tensorflow/core/common_runtime/profile_handler.h
index 9d31b1aecb..391dc8c198 100644
--- a/tensorflow/core/common_runtime/profile_handler.h
+++ b/tensorflow/core/common_runtime/profile_handler.h
@@ -29,22 +29,6 @@ class ProfileHandler {
   ProfileHandler() {}
   virtual ~ProfileHandler() {}
 
-  // Records that a miscellaneous activity occurred in the current step.
-  //
-  // Implementations of this method must be thread-safe.
-  //
-  // Args:
-  // - device: The device on which the activity occurred.
-  // - start: The time at which the activity started.
-  // - limit: The time at which the activity finished.
-  // - label: A label for the op, which may be used in visualization.
-  // - op_type: A type string for the op, which may be used in visualization.
-  // - details: A details string, which may be used in visualization.
-  // from time "start" to "limit" with "op_type" and "details".
-  virtual void RecordActivity(const string& device, Microseconds start,
-                              Microseconds limit, StringPiece label,
-                              StringPiece op_type, StringPiece details) = 0;
-
   // Records that a single Op was executed in the current step.
   //
   // Implementations of this method must be thread-safe.
-- 
GitLab


From 0bd1408a2d95d0d30bf9412dc64edc45c71b915f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 07:57:12 -0700
Subject: [PATCH 357/395] Add missing #include for OpResponse. This class
 currently happens to be forward declared by xla.proto.h, but that proto
 doesn't actually need this type anywhere and we are working on removing such
 unneeded forward declarations.

PiperOrigin-RevId: 195818397
---
 tensorflow/compiler/xla/BUILD               | 1 +
 tensorflow/compiler/xla/service_interface.h | 1 +
 2 files changed, 2 insertions(+)

diff --git a/tensorflow/compiler/xla/BUILD b/tensorflow/compiler/xla/BUILD
index 1af9cb6d2a..dbf14f32bc 100644
--- a/tensorflow/compiler/xla/BUILD
+++ b/tensorflow/compiler/xla/BUILD
@@ -99,6 +99,7 @@ cc_library(
     hdrs = ["service_interface.h"],
     visibility = [":friends"],
     deps = [
+        ":xla_data_proto",
         ":xla_proto",
         "//tensorflow/core:lib",
     ],
diff --git a/tensorflow/compiler/xla/service_interface.h b/tensorflow/compiler/xla/service_interface.h
index 5b44c26b7c..4f64fe8f83 100644
--- a/tensorflow/compiler/xla/service_interface.h
+++ b/tensorflow/compiler/xla/service_interface.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define TENSORFLOW_COMPILER_XLA_SERVICE_INTERFACE_H_
 
 #include "tensorflow/compiler/xla/xla.pb.h"
+#include "tensorflow/compiler/xla/xla_data.pb.h"
 #include "tensorflow/core/lib/core/status.h"
 
 namespace xla {
-- 
GitLab


From 07fdb697d33478d7a72d09fc2371fa834e870b83 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 08:04:07 -0700
Subject: [PATCH 358/395] Automated g4 rollback of changelist 195723288

PiperOrigin-RevId: 195819297
---
 tensorflow/contrib/image/kernels/image_ops.cc | 33 +++--------
 tensorflow/contrib/image/kernels/image_ops.h  |  2 +-
 tensorflow/contrib/image/ops/image_ops.cc     | 55 ++-----------------
 .../python/kernel_tests/image_ops_test.py     | 30 ----------
 .../contrib/image/python/ops/image_ops.py     | 49 ++++++-----------
 5 files changed, 30 insertions(+), 139 deletions(-)

diff --git a/tensorflow/contrib/image/kernels/image_ops.cc b/tensorflow/contrib/image/kernels/image_ops.cc
index 575c2004fb..c2e32da133 100644
--- a/tensorflow/contrib/image/kernels/image_ops.cc
+++ b/tensorflow/contrib/image/kernels/image_ops.cc
@@ -70,7 +70,6 @@ class ImageProjectiveTransform : public OpKernel {
   void Compute(OpKernelContext* ctx) override {
     const Tensor& images_t = ctx->input(0);
     const Tensor& transform_t = ctx->input(1);
-    const Tensor& shape_t = ctx->input(2);
     OP_REQUIRES(ctx, images_t.shape().dims() == 4,
                 errors::InvalidArgument("Input images must have rank 4"));
     OP_REQUIRES(ctx,
@@ -81,28 +80,11 @@ class ImageProjectiveTransform : public OpKernel {
                      ProjectiveGenerator<Device, T>::kNumParameters),
                 errors::InvalidArgument(
                     "Input transform should be num_images x 8 or 1 x 8"));
-    OP_REQUIRES(ctx, shape_t.dims() == 1,
-                errors::InvalidArgument("output shape must be 1-dimensional",
-                                        shape_t.shape().DebugString()));
-    OP_REQUIRES(ctx, shape_t.NumElements() == 2,
-                errors::InvalidArgument("output shape must have two elements",
-                                        shape_t.shape().DebugString()));
-    auto Svec = shape_t.vec<int32>();
-    int32 out_height = Svec(0);
-    int32 out_width = Svec(1);
-    OP_REQUIRES(ctx, out_height > 0 && out_width > 0,
-                errors::InvalidArgument("output dimensions must be positive"));
-
-    Tensor* output_t;
-    OP_REQUIRES_OK(ctx, ctx->allocate_output(
-                            0,
-                            TensorShape({images_t.dim_size(0), out_height,
-                                         out_width, images_t.dim_size(3)}),
-                            &output_t));
-    auto output = output_t->tensor<T, 4>();
     auto images = images_t.tensor<T, 4>();
     auto transform = transform_t.matrix<float>();
-
+    Tensor* output_t;
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, images_t.shape(), &output_t));
+    auto output = output_t->tensor<T, 4>();
     (FillProjectiveTransform<Device, T>(interpolation_))(
         ctx->eigen_device<Device>(), &output, images, transform);
   }
@@ -145,11 +127,10 @@ TF_CALL_double(DECLARE_FUNCTOR);
 
 }  // end namespace functor
 
-#define REGISTER(TYPE)                                       \
-  REGISTER_KERNEL_BUILDER(Name("ImageProjectiveTransform")   \
-                              .Device(DEVICE_GPU)            \
-                              .TypeConstraint<TYPE>("dtype") \
-                              .HostMemory("output_shape"),   \
+#define REGISTER(TYPE)                                        \
+  REGISTER_KERNEL_BUILDER(Name("ImageProjectiveTransform")    \
+                              .Device(DEVICE_GPU)             \
+                              .TypeConstraint<TYPE>("dtype"), \
                           ImageProjectiveTransform<GPUDevice, TYPE>)
 
 TF_CALL_uint8(REGISTER);
diff --git a/tensorflow/contrib/image/kernels/image_ops.h b/tensorflow/contrib/image/kernels/image_ops.h
index 2320329b92..ad50133061 100644
--- a/tensorflow/contrib/image/kernels/image_ops.h
+++ b/tensorflow/contrib/image/kernels/image_ops.h
@@ -161,7 +161,7 @@ struct FillProjectiveTransform {
   void operator()(const Device& device, OutputType* output,
                   const InputType& images,
                   const TransformsType& transform) const {
-    output->device(device) = output->generate(
+    output->device(device) = images.generate(
         ProjectiveGenerator<Device, T>(images, transform, interpolation_));
   }
 };
diff --git a/tensorflow/contrib/image/ops/image_ops.cc b/tensorflow/contrib/image/ops/image_ops.cc
index fb62507174..ebdcaea7ab 100644
--- a/tensorflow/contrib/image/ops/image_ops.cc
+++ b/tensorflow/contrib/image/ops/image_ops.cc
@@ -19,56 +19,9 @@ limitations under the License.
 
 namespace tensorflow {
 
-using shape_inference::DimensionHandle;
 using shape_inference::InferenceContext;
 using shape_inference::ShapeHandle;
 
-namespace {
-
-// Sets output[0] to shape [batch_dim,height,width,channel_dim], where
-// height and width come from the size_tensor.
-Status SetOutputToSizedImage(InferenceContext* c, DimensionHandle batch_dim,
-                             int size_input_idx, DimensionHandle channel_dim) {
-  // Verify shape of size input.
-  ShapeHandle size;
-  TF_RETURN_IF_ERROR(c->WithRank(c->input(size_input_idx), 1, &size));
-  DimensionHandle unused;
-  TF_RETURN_IF_ERROR(c->WithValue(c->Dim(size, 0), 2, &unused));
-
-  // Get size values from the size tensor.
-  const Tensor* size_tensor = c->input_tensor(size_input_idx);
-  DimensionHandle width;
-  DimensionHandle height;
-  if (size_tensor == nullptr) {
-    width = c->UnknownDim();
-    height = c->UnknownDim();
-  } else {
-    // TODO(petewarden) - Remove once we have constant evaluation in C++ only.
-    if (size_tensor->dtype() != DT_INT32) {
-      return errors::InvalidArgument(
-          "Bad size input type for SetOutputToSizedImage: Expected DT_INT32 "
-          "but got ",
-          DataTypeString(size_tensor->dtype()), " for input #", size_input_idx,
-          " in ", c->DebugString());
-    }
-    auto vec = size_tensor->vec<int32>();
-    height = c->MakeDim(vec(0));
-    width = c->MakeDim(vec(1));
-  }
-  c->set_output(0, c->MakeShape({batch_dim, height, width, channel_dim}));
-  return Status::OK();
-}
-
-// TODO(qyu): Move this to core/framework/common_shape_fns.h
-Status ResizeShapeFn(InferenceContext* c) {
-  ShapeHandle input;
-  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 4, &input));
-  return SetOutputToSizedImage(c, c->Dim(input, 0), 2 /* size_input_idx */,
-                               c->Dim(input, 3));
-}
-
-}  // namespace
-
 // TODO(ringwalt): Add a "fill_mode" argument with "constant", "mirror", etc.
 // TODO(ringwalt): Add a "fill_constant" argument for constant mode (default 0).
 // TODO(ringwalt): Add an "output_shape" argument. This is sufficient to
@@ -76,11 +29,13 @@ Status ResizeShapeFn(InferenceContext* c) {
 REGISTER_OP("ImageProjectiveTransform")
     .Input("images: dtype")
     .Input("transforms: float32")
-    .Input("output_shape: int32")
     .Attr("dtype: {uint8, int32, int64, float32, float64}")
     .Attr("interpolation: string")
     .Output("transformed_images: dtype")
-    .SetShapeFn(ResizeShapeFn)
+    .SetShapeFn([](InferenceContext* c) {
+      c->set_output(0, c->input(0));
+      return Status::OK();
+    })
     .Doc(R"doc(
 Applies the given transform to each of the images.
 
@@ -94,7 +49,7 @@ If one row of `transforms` is `[a0, a1, a2, b0, b1, b2, c0, c1]`, then it maps
 the *output* point `(x, y)` to a transformed *input* point
 `(x', y') = ((a0 x + a1 y + a2) / k, (b0 x + b1 y + b2) / k)`, where
 `k = c0 x + c1 y + 1`. If the transformed point lays outside of the input
-image, the output pixel is set to 0.
+image, the output pixel is set to 0. The output is the same size as the input,
 
 images: 4D `Tensor`, input image(s) in NHWC format.
 transforms: 2D `Tensor`, projective transform(s) to apply to the image(s).
diff --git a/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py b/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py
index c0151d320f..b50177ae56 100644
--- a/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py
+++ b/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py
@@ -195,40 +195,10 @@ class ImageOpsTest(test_util.TensorFlowTestCase):
           x_init_value=test_image)
       self.assertLess(left_err, 1e-10)
 
-  def _test_grad_different_shape(self, input_shape, output_shape):
-    with self.test_session():
-      test_image_shape = input_shape
-      test_image = np.random.randn(*test_image_shape)
-      test_image_tensor = constant_op.constant(
-          test_image, shape=test_image_shape)
-      test_transform = image_ops.angles_to_projective_transforms(
-          np.pi / 2, 4, 4)
-
-      if len(output_shape) == 2:
-        resize_shape = output_shape
-      elif len(output_shape) == 3:
-        resize_shape = output_shape[0:2]
-      elif len(output_shape) == 4:
-        resize_shape = output_shape[1:3]
-      output = image_ops.transform(
-          images=test_image_tensor,
-          transforms=test_transform,
-          output_shape=resize_shape)
-      left_err = gradient_checker.compute_gradient_error(
-          test_image_tensor,
-          test_image_shape,
-          output,
-          output_shape,
-          x_init_value=test_image)
-      self.assertLess(left_err, 1e-10)
-
   def test_grad(self):
     self._test_grad([16, 16])
     self._test_grad([4, 12, 12])
     self._test_grad([3, 4, 12, 12])
-    self._test_grad_different_shape([16, 16], [8, 8])
-    self._test_grad_different_shape([4, 12, 3], [8, 24, 3])
-    self._test_grad_different_shape([3, 4, 12, 3], [3, 8, 24, 3])
 
 
 class BipartiteMatchTest(test_util.TensorFlowTestCase):
diff --git a/tensorflow/contrib/image/python/ops/image_ops.py b/tensorflow/contrib/image/python/ops/image_ops.py
index 192571ced8..cd984c8054 100644
--- a/tensorflow/contrib/image/python/ops/image_ops.py
+++ b/tensorflow/contrib/image/python/ops/image_ops.py
@@ -23,7 +23,6 @@ from tensorflow.python.framework import common_shapes
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import linalg_ops
@@ -213,11 +212,7 @@ def translations_to_projective_transforms(translations, name=None):
         axis=1)
 
 
-def transform(images,
-              transforms,
-              interpolation="NEAREST",
-              output_shape=None,
-              name=None):
+def transform(images, transforms, interpolation="NEAREST", name=None):
   """Applies the given transform(s) to the image(s).
 
   Args:
@@ -234,10 +229,6 @@ def transform(images,
        the transform mapping input points to output points. Note that gradients
        are not backpropagated into transformation parameters.
     interpolation: Interpolation mode. Supported values: "NEAREST", "BILINEAR".
-    output_shape: Output dimesion after the transform, [height, width].
-       If None, output is the same size as input image.
-
-    name: The name of the op.
 
   Returns:
     Image(s) with the same type and shape as `images`, with the given
@@ -246,7 +237,6 @@ def transform(images,
 
   Raises:
     TypeError: If `image` is an invalid type.
-    ValueError: If output shape is not 1-D int32 Tensor.
   """
   with ops.name_scope(name, "transform"):
     image_or_images = ops.convert_to_tensor(images, name="images")
@@ -265,17 +255,6 @@ def transform(images,
     else:
       raise TypeError("Images should have rank between 2 and 4.")
 
-    if output_shape is None:
-      output_shape = tensor_util.constant_value(
-          array_ops.shape(images)[1:3]) or array_ops.shape(images)[1:3]
-
-    output_shape = ops.convert_to_tensor(
-        output_shape, dtypes.int32, name="output_shape")
-
-    if not output_shape.get_shape().is_compatible_with([2]):
-      raise ValueError("output_shape must be a 1-D Tensor of 2 elements: "
-                       "new_height, new_width")
-
     if len(transform_or_transforms.get_shape()) == 1:
       transforms = transform_or_transforms[None]
     elif transform_or_transforms.get_shape().ndims is None:
@@ -285,12 +264,8 @@ def transform(images,
       transforms = transform_or_transforms
     else:
       raise TypeError("Transforms should have rank 1 or 2.")
-
     output = gen_image_ops.image_projective_transform(
-        images,
-        output_shape=output_shape,
-        transforms=transforms,
-        interpolation=interpolation.upper())
+        images, transforms, interpolation=interpolation.upper())
     if len(image_or_images.get_shape()) == 2:
       return output[0, :, :, 0]
     elif len(image_or_images.get_shape()) == 3:
@@ -400,6 +375,14 @@ def _image_projective_transform_grad(op, grad):
 
   if image_or_images.dtype.base_dtype not in _IMAGE_DTYPES:
     raise TypeError("Invalid dtype %s." % image_or_images.dtype)
+  if len(image_or_images.get_shape()) == 2:
+    images = image_or_images[None, :, :, None]
+  elif len(image_or_images.get_shape()) == 3:
+    images = image_or_images[None, :, :, :]
+  elif len(image_or_images.get_shape()) == 4:
+    images = image_or_images
+  else:
+    raise TypeError("Images should have rank between 2 and 4")
   if len(transform_or_transforms.get_shape()) == 1:
     transforms = transform_or_transforms[None]
   elif len(transform_or_transforms.get_shape()) == 2:
@@ -412,11 +395,13 @@ def _image_projective_transform_grad(op, grad):
   inverse = linalg_ops.matrix_inverse(transforms)
   transforms = matrices_to_flat_transforms(inverse)
   output = gen_image_ops.image_projective_transform(
-      images=grad,
-      transforms=transforms,
-      output_shape=array_ops.shape(image_or_images)[1:3],
-      interpolation=interpolation)
-  return [output, None, None]
+      grad, transforms, interpolation=interpolation)
+  if len(image_or_images.get_shape()) == 2:
+    return [output[0, :, :, 0], None]
+  elif len(image_or_images.get_shape()) == 3:
+    return [output[0, :, :, :], None]
+  else:
+    return [output, None]
 
 
 def bipartite_match(distance_mat,
-- 
GitLab


From a6a862e90d1b336570ab67816ca14e191f5acb32 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Tue, 8 May 2018 08:07:08 -0700
Subject: [PATCH 359/395] [TF:XLA] Fix NaN in StatelessRandomNormal if the
 underlying uniform distribution returned -1.

PiperOrigin-RevId: 195819645
---
 tensorflow/compiler/tests/stateless_random_ops_test.py   | 9 +++++++++
 .../compiler/tf2xla/kernels/stateless_random_ops.cc      | 3 ++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/tensorflow/compiler/tests/stateless_random_ops_test.py b/tensorflow/compiler/tests/stateless_random_ops_test.py
index 4336ebdbd1..b6f8390a45 100644
--- a/tensorflow/compiler/tests/stateless_random_ops_test.py
+++ b/tensorflow/compiler/tests/stateless_random_ops_test.py
@@ -86,6 +86,15 @@ class StatelessRandomOpsTest(XLATestCase):
         # seed were not fixed.
         self.assertTrue(self._chi_squared(y, 10) < 16.92)
 
+  def testRandomNormalIsFinite(self):
+    with self.test_session() as sess, self.test_scope():
+      for dtype in self._random_types():
+        seed_t = array_ops.placeholder(dtypes.int32, shape=[2])
+        x = stateless.stateless_random_uniform(
+            shape=[10000], seed=seed_t, dtype=dtype)
+        y = sess.run(x, {seed_t: [0x12345678, 0xabcdef12]})
+        self.assertTrue(np.all(np.isfinite(y)))
+
   def _normal_cdf(self, x):
     """Cumulative distribution function for a standard normal distribution."""
     return 0.5 + 0.5 * np.vectorize(math.erf)(x / math.sqrt(2))
diff --git a/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc b/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc
index 6340c22518..a99d4ddc7c 100644
--- a/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc
@@ -255,7 +255,8 @@ class StatelessRandomNormalOp : public XlaOpKernel {
                                         seed_shape.DebugString()));
     xla::XlaOp seed = ctx->Input(1);
     xla::XlaBuilder* builder = ctx->builder();
-    auto uniform = RandomUniform(builder, seed, shape, -1.0, 1.0);
+    auto uniform =
+        RandomUniform(builder, seed, shape, std::nextafter(-1.0f, 0.0f), 1.0);
     // Convert uniform distribution to normal distribution by computing
     // sqrt(2) * erfinv(x)
     auto normal = builder->Mul(builder->ConstantR0<float>(std::sqrt(2.0)),
-- 
GitLab


From 4a6e6632eb866a2910396c6bc78d601b5b9b550e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 08:57:45 -0700
Subject: [PATCH 360/395] Update comment clarifying continuous eval behavior.

PiperOrigin-RevId: 195826025
---
 tensorflow/contrib/learn/python/learn/experiment.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/tensorflow/contrib/learn/python/learn/experiment.py b/tensorflow/contrib/learn/python/learn/experiment.py
index 3744abd860..dfc6a393d0 100644
--- a/tensorflow/contrib/learn/python/learn/experiment.py
+++ b/tensorflow/contrib/learn/python/learn/experiment.py
@@ -468,10 +468,15 @@ class Experiment(object):
         on which that evaluation was based.
         At the beginning of evaluation, the passed `eval_results` will be None
         so it's expected that the predicate function handles that gracefully.
-        When `predicate_fn` is not specified, continuous eval will run in an
-        infinite loop (if `train_steps` is None). or exit once global step
-        reaches `train_steps`.
-
+        Continuous eval behavior under different conditions:
+          * When `predicate_fn` is specified:
+            + if `train_steps` is None, run until `predicate_fn` returns False.
+            + if `train_steps` is specified, run until either global step
+              reaches `train_steps` or `predicate_fn` returns False.
+          * When `predicate_fn` is not specified:
+            + if `train_steps` is None, run in an infinite loop.
+            + if `train_steps` is specified, run until global step reaches
+              `train_steps`.
       export: Whether to export from this step. Default is 'True'.
 
     Raises:
-- 
GitLab


From bd606508ebb0e1dbb3215c3ad1d0a41da3507766 Mon Sep 17 00:00:00 2001
From: Shanqing Cai <cais@google.com>
Date: Tue, 8 May 2018 09:04:17 -0700
Subject: [PATCH 361/395] Minor formatting tweaks to distribute.py and
 simple_tfkeras_example.py

PiperOrigin-RevId: 195827029
---
 .../python/examples/simple_tfkeras_example.py | 33 ++++++++++++-------
 tensorflow/python/training/distribute.py      | 16 ++++-----
 2 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/tensorflow/contrib/distribute/python/examples/simple_tfkeras_example.py b/tensorflow/contrib/distribute/python/examples/simple_tfkeras_example.py
index b87224251c..2b05884b9b 100644
--- a/tensorflow/contrib/distribute/python/examples/simple_tfkeras_example.py
+++ b/tensorflow/contrib/distribute/python/examples/simple_tfkeras_example.py
@@ -12,11 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""An example tf.keras model that is trained using MirroredStrategy."""
+"""An example of training tf.keras Model using MirroredStrategy."""
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-from sys import argv
+
+import sys
+
 import numpy as np
 import tensorflow as tf
 
@@ -33,30 +35,37 @@ def input_fn():
 
 def main(args):
   if len(args) < 2:
-    print('You must specify  model_dir for checkpoints such as'
-          ' /tmp/tfkeras_example./')
+    print('You must specify model_dir for checkpoints such as'
+          ' /tmp/tfkeras_example/.')
     return
 
-  print('Using %s to store checkpoints.' % args[1])
-
-  strategy = tf.contrib.distribute.MirroredStrategy(
-      ['/device:GPU:0', '/device:GPU:1'])
-  config = tf.estimator.RunConfig(train_distribute=strategy)
-  optimizer = tf.train.GradientDescentOptimizer(0.2)
+  model_dir = args[1]
+  print('Using %s to store checkpoints.' % model_dir)
 
+  # Define tf.keras Model.
   model = tf.keras.Sequential()
   model.add(tf.keras.layers.Dense(16, activation='relu', input_shape=(10,)))
   model.add(tf.keras.layers.Dense(1, activation='sigmoid'))
 
+  # Compile tf.keras Model.
+  optimizer = tf.train.GradientDescentOptimizer(0.2)
   model.compile(loss='binary_crossentropy', optimizer=optimizer)
   model.summary()
   tf.keras.backend.set_learning_phase(True)
+
+  # Define a DistributionStrategy and convert the tf.keras Model to a
+  # tf.Estimator that utilizes the DistributionStrategy.
+  strategy = tf.contrib.distribute.MirroredStrategy(
+      ['/device:GPU:0', '/device:GPU:1'])
+  config = tf.estimator.RunConfig(train_distribute=strategy)
   keras_estimator = tf.keras.estimator.model_to_estimator(
-      keras_model=model, config=config, model_dir=args[1])
+      keras_model=model, config=config, model_dir=model_dir)
 
+  # Train and evaluate the tf.Estimator.
   keras_estimator.train(input_fn=input_fn, steps=10)
   eval_result = keras_estimator.evaluate(input_fn=input_fn)
   print('Eval result: {}'.format(eval_result))
 
+
 if __name__ == '__main__':
-  tf.app.run(argv=argv)
+  tf.app.run(argv=sys.argv)
diff --git a/tensorflow/python/training/distribute.py b/tensorflow/python/training/distribute.py
index b60f87c05f..6d05a2ee29 100644
--- a/tensorflow/python/training/distribute.py
+++ b/tensorflow/python/training/distribute.py
@@ -357,14 +357,14 @@ class DistributionStrategy(object):
     on different slices of the input data. This is in contrast to
     _model parallelism_ where we divide up a single copy of a model
     across multiple devices.
-    Note: for now we only support data parallelism at this time, but
+    Note: we only support data parallelism for now, but
     hope to add support for model parallelism in the future.
   * A _tower_ is one copy of the model, running on one slice of the
     input data.
-  * _Synchronous_, or more commonly _sync_, training is when the
+  * _Synchronous_, or more commonly _sync_, training is where the
     updates from each tower are aggregated together before updating
     the model variables. This is in contrast to _asynchronous_, or
-    _async_ training where each tower updates the model variables
+    _async_ training, where each tower updates the model variables
     independently.
   * Furthermore you might run your computation on multiple devices
     on one machine (or "host"), or on multiple machines/hosts.
@@ -386,11 +386,11 @@ class DistributionStrategy(object):
   * Reductions and Allreduce: A _reduction_ is some method of
     aggregating multiple values into one value, like "sum" or
     "mean". If doing sync training, we will perform a reduction on the
-    gradients to a parameter from each tower before applying the
+    gradients to a parameter from all towers before applying the
     update. Allreduce is an algorithm for performing a reduction on
     values from multiple devices and making the result available on
     all of those devices.
-  * In the future we will have support for TensorFlows' partitioned
+  * In the future we will have support for TensorFlow's partitioned
     variables, where a single variable is split across multiple
     devices.
 
@@ -419,9 +419,9 @@ class DistributionStrategy(object):
     `tower_fn` can use the `get_tower_context()` API to get enhanced
     behavior in this case.
 
-    You can also create an initializable iterator instead of one shot iterator.
-    In that case, you will need to ensure that you initialize the iterator
-    before calling get_next.
+    You can also create an initializable iterator instead of a one-shot
+    iterator. In that case, you will need to ensure that you initialize the
+    iterator before calling get_next.
     ```
     iterator = my_distribution.distribute_dataset(
         dataset).make_initializable_iterator())
-- 
GitLab


From 77bb984c23aa7ec347c981c31f650598c9624304 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 09:46:45 -0700
Subject: [PATCH 362/395] Free ANeuralNetworksCompilation object in
 NNAPIDelegate destructor

PiperOrigin-RevId: 195832807
---
 tensorflow/contrib/lite/nnapi_delegate.cc | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensorflow/contrib/lite/nnapi_delegate.cc b/tensorflow/contrib/lite/nnapi_delegate.cc
index 6a231dc6bc..eb451397bd 100644
--- a/tensorflow/contrib/lite/nnapi_delegate.cc
+++ b/tensorflow/contrib/lite/nnapi_delegate.cc
@@ -61,6 +61,10 @@ NNAPIAllocation::~NNAPIAllocation() {
 }
 
 NNAPIDelegate::~NNAPIDelegate() {
+  if (nn_compiled_model_) {
+    ANeuralNetworksCompilation_free(nn_compiled_model_);
+    nn_compiled_model_ = nullptr;
+  }
   if (nn_model_) {
     ANeuralNetworksModel_free(nn_model_);
     nn_model_ = nullptr;
-- 
GitLab


From 074d2901e2f6b9807394f300e5ccbc65defcf161 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 10:25:30 -0700
Subject: [PATCH 363/395] Add cost model of depthwiseConv2dNative. Tensorflow
 computes depthwise separable convolutions as depthwiseConv2dNative followed
 by 1x1 Conv2D

PiperOrigin-RevId: 195838887
---
 .../grappler/costs/op_level_cost_estimator.cc | 68 +++++++++++++++----
 .../costs/op_level_cost_estimator_test.cc     | 26 +++++++
 2 files changed, 81 insertions(+), 13 deletions(-)

diff --git a/tensorflow/core/grappler/costs/op_level_cost_estimator.cc b/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
index 199b69452f..2542fa2d67 100644
--- a/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
+++ b/tensorflow/core/grappler/costs/op_level_cost_estimator.cc
@@ -32,6 +32,11 @@ constexpr char kConv2d[] = "Conv2D";
 constexpr char kConv2dBackpropFilter[] = "Conv2DBackpropFilter";
 constexpr char kConv2dBackpropInput[] = "Conv2DBackpropInput";
 constexpr char kFusedConv2dBiasActivation[] = "FusedConv2DBiasActivation";
+constexpr char kDepthwiseConv2dNative[] = "DepthwiseConv2dNative";
+constexpr char kDepthwiseConv2dNativeBackpropFilter[] =
+    "DepthwiseConv2dNativeBackpropFilter";
+constexpr char kDepthwiseConv2dNativeBackpropInput[] =
+    "DepthwiseConv2dNativeBackpropInput";
 constexpr char kMatMul[] = "MatMul";
 constexpr char kSparseMatMul[] = "SparseMatMul";
 constexpr char kPlaceholder[] = "Placeholder";
@@ -201,6 +206,14 @@ OpLevelCostEstimator::OpLevelCostEstimator() {
        wrap(&OpLevelCostEstimator::PredictConv2DBackpropInput)},
       {kFusedConv2dBiasActivation,
        wrap(&OpLevelCostEstimator::PredictFusedConv2DBiasActivation)},
+      // reuse Conv2D for DepthwiseConv2dNative because the caculation is the
+      // same although the actual meaning of the parameters are different. See
+      // comments in PredictConv2D and related functions
+      {kDepthwiseConv2dNative, wrap(&OpLevelCostEstimator::PredictConv2D)},
+      {kDepthwiseConv2dNativeBackpropFilter,
+       wrap(&OpLevelCostEstimator::PredictConv2DBackpropFilter)},
+      {kDepthwiseConv2dNativeBackpropInput,
+       wrap(&OpLevelCostEstimator::PredictConv2DBackpropInput)},
       {kMatMul, wrap(&OpLevelCostEstimator::PredictMatMul)},
       {kSparseMatMul, wrap(&OpLevelCostEstimator::PredictMatMul)},
       {kBatchMatMul, wrap(&OpLevelCostEstimator::PredictBatchMatMul)},
@@ -539,18 +552,30 @@ OpLevelCostEstimator::ConvolutionDimensionsFromInputs(
 int64 OpLevelCostEstimator::CountConv2DOperations(
     const OpInfo& op_features, ConvolutionDimensions* conv_info,
     bool* found_unknown_shapes) const {
-  if (op_features.op() != kConv2d) {
-    LOG(ERROR) << "Invalid Operation";
-    return 0;
-  }
+  DCHECK(op_features.op() == kConv2d ||
+         op_features.op() == kDepthwiseConv2dNative)
+      << "Invalid Operation: not Conv2D nor DepthwiseConv2dNative";
+
   ConvolutionDimensions conv_dims = ConvolutionDimensionsFromInputs(
       op_features.inputs(0).shape(), op_features.inputs(1).shape(), op_features,
       found_unknown_shapes);
 
+  //  in DepthwiseConv2dNative conv_dims.oz is actually the channel depth
+  //  multiplier; The effective output channel depth oz_effective is
+  //  conv_dims.iz * conv_dims.oz. thus # ops = N x H x W x oz_effective x 2RS.
+  //  Compare to Conv2D where # ops =  N x H x W x iz x oz x 2RS,
+  //  oz = oz_effective,  then Conv2D_ops / Depthwise_conv2d_native_ops = iz.
   int64 ops = conv_dims.batch;
   ops *= conv_dims.ox * conv_dims.oy;
   ops *= conv_dims.kx * conv_dims.ky;
-  ops *= conv_dims.iz * conv_dims.oz;
+  if (op_features.op() == kConv2d) {
+    ops *= conv_dims.iz * conv_dims.oz;
+  } else {
+    // To ensure output tensor dims to be correct for DepthwiseConv2DNative,
+    // although ops are the same as Conv2D.
+    conv_dims.oz *= conv_dims.iz;
+    ops *= conv_dims.oz;
+  }
   ops *= kOpsPerMac;
 
   if (conv_info != nullptr) {
@@ -797,7 +822,10 @@ int64 OpLevelCostEstimator::CountConv2DBackpropInputOperations(
     bool* found_unknown_shapes) const {
   int64 ops = 0;
 
-  DCHECK_EQ(kConv2dBackpropInput, op_features.op());
+  DCHECK(op_features.op() == kConv2dBackpropInput ||
+         op_features.op() == kDepthwiseConv2dNativeBackpropInput)
+      << "Invalid Operation: not kConv2dBackpropInput nor"
+         "kDepthwiseConv2dNativeBackpropInput";
 
   if (op_features.inputs_size() < 2) {
     *found_unknown_shapes = true;
@@ -830,10 +858,15 @@ int64 OpLevelCostEstimator::CountConv2DBackpropInputOperations(
   ops = conv_dims.batch;
   ops *= conv_dims.ox * conv_dims.oy;
   ops *= conv_dims.kx * conv_dims.ky;
-  ops *= conv_dims.iz * conv_dims.oz;
-  ops *= kOpsPerMac;
+  if (op_features.op() == kConv2dBackpropInput) {
+    ops *= conv_dims.iz * conv_dims.oz;
+  } else {
+    // conv_dims always use forward path definition regardless
+    conv_dims.oz *= conv_dims.iz;
+    ops *= conv_dims.oz;
+  }
 
-  VLOG(1) << "Operations for Conv2DBackpropInput " << ops;
+  VLOG(1) << "Operations for" << op_features.op() << "  " << ops;
 
   if (returned_conv_dims != nullptr) {
     *returned_conv_dims = conv_dims;
@@ -845,7 +878,11 @@ int64 OpLevelCostEstimator::CountConv2DBackpropFilterOperations(
     const OpInfo& op_features, ConvolutionDimensions* returned_conv_dims,
     bool* found_unknown_shapes) const {
   int64 ops = 0;
-  DCHECK_EQ(kConv2dBackpropFilter, op_features.op());
+
+  DCHECK(op_features.op() == kConv2dBackpropFilter ||
+         op_features.op() == kDepthwiseConv2dNativeBackpropFilter)
+      << "Invalid Operation: not kConv2dBackpropFilter nor"
+         "kDepthwiseConv2dNativeBackpropFilter";
 
   TensorShapeProto filter_shape;
   bool shape_found = false;
@@ -877,10 +914,15 @@ int64 OpLevelCostEstimator::CountConv2DBackpropFilterOperations(
   ops = conv_dims.batch;
   ops *= conv_dims.ox * conv_dims.oy;
   ops *= conv_dims.kx * conv_dims.ky;
-  ops *= conv_dims.iz * conv_dims.oz;
-  ops *= kOpsPerMac;
+  if (op_features.op() == kConv2dBackpropFilter) {
+    ops *= conv_dims.iz * conv_dims.oz;
+  } else {
+    // conv_dims always use forward path definition regardless
+    conv_dims.oz *= conv_dims.iz;
+    ops *= conv_dims.oz;
+  }
 
-  VLOG(1) << "Operations for Conv2DBackpropFilter" << ops;
+  VLOG(1) << "Operations for" << op_features.op() << "  " << ops;
 
   if (returned_conv_dims != nullptr) {
     *returned_conv_dims = conv_dims;
diff --git a/tensorflow/core/grappler/costs/op_level_cost_estimator_test.cc b/tensorflow/core/grappler/costs/op_level_cost_estimator_test.cc
index 13ea43bed6..b2c021b73a 100644
--- a/tensorflow/core/grappler/costs/op_level_cost_estimator_test.cc
+++ b/tensorflow/core/grappler/costs/op_level_cost_estimator_test.cc
@@ -128,6 +128,23 @@ OpContext DescribeConvolution(int batch, int ix, int iy, int iz1, int iz2,
   return op_context;
 }
 
+// Describe DepthwiseConvolution constructs an OpContext for a
+// DepthwiseConv2dNative applied to an input
+// tensor with shape (batch, ix, iy, iz1) and a kernel tensor with shape
+// (kx, ky, iz2, cm). cm is channel multiplier
+
+OpContext DescribeDepthwiseConv2dNative(int batch, int ix, int iy, int iz1,
+                                        int iz2, int kx, int ky, int cm) {
+  OpContext op_context;
+  SetCpuDevice(&op_context.op_info);
+  op_context.op_info.set_op("DepthwiseConv2dNative");
+
+  DescribeTensor4D(batch, ix, iy, iz1, op_context.op_info.add_inputs());
+  DescribeTensor4D(kx, ky, iz2, cm, op_context.op_info.add_inputs());
+
+  return op_context;
+}
+
 // DescribeFusedConv2DBiasActivation constructs an OpContext for a
 // FusedConv2DBiasActivation applied to a convolution input tensor with shape
 // (batch, ix, iy, iz1), a kernel tensor with shape (kx, ky, iz2, oz), a
@@ -505,6 +522,15 @@ TEST_F(OpLevelCostEstimatorTest, Conv2DExecutionTime) {
   EXPECT_FALSE(cost.inaccurate);
 }
 
+TEST_F(OpLevelCostEstimatorTest, DepthwiseConv2dNativeExecutionTime) {
+  auto cost =
+      PredictCosts(DescribeDepthwiseConv2dNative(16, 19, 19, 48, 48, 5, 5, 3));
+  EXPECT_EQ(Costs::Duration(112340), cost.memory_time);
+  EXPECT_EQ(Costs::Duration(4158720), cost.compute_time);
+  EXPECT_EQ(Costs::Duration(4271060), cost.execution_time);
+  EXPECT_FALSE(cost.inaccurate);
+}
+
 TEST_F(OpLevelCostEstimatorTest, DummyExecutionTime) {
   auto cost = PredictCosts(DescribeBinaryOp("Dummy", 1000, 1));
   EXPECT_EQ(Costs::Duration(2000), cost.memory_time);
-- 
GitLab


From 83aa3239b45175fff56e85b07a68caf1e182b455 Mon Sep 17 00:00:00 2001
From: Akshay Agrawal <akshayka@google.com>
Date: Tue, 8 May 2018 11:07:45 -0700
Subject: [PATCH 364/395] When building functions, capture tensors in
 `internal_convert_to_tensor`.

This change is motivated by the fact that, when eager execution is disabled, library functions assume that tensors returned from `internal_convert_to_tensor` are in fact `Tensor`s and not `EagerTensor`s.

PiperOrigin-RevId: 195846039
---
 tensorflow/python/eager/function.py      | 16 +++++++++-------
 tensorflow/python/eager/function_test.py | 15 +++++++++++++++
 tensorflow/python/framework/function.py  |  9 +++++----
 tensorflow/python/framework/ops.py       | 17 +++++++++++++----
 4 files changed, 42 insertions(+), 15 deletions(-)

diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py
index 741bd2ac9c..89257bb20a 100644
--- a/tensorflow/python/eager/function.py
+++ b/tensorflow/python/eager/function.py
@@ -102,13 +102,15 @@ class CapturingGraph(ops.Graph):
   def clear_resource_control_flow_state(self):
     self._last_op_using_resource_tensor = {}
 
-  def maybe_capture_tensor(self, tensor):
+  def capture(self, tensor, name=None):
     if isinstance(tensor, ops.EagerTensor):
-      return capture_value(
-          self.captures, tensor, tensor.dtype, str(ops.uid()))
+      if name is None:
+        name = str(ops.uid())
+      return capture_value(self.captures, tensor, tensor.dtype, name)
     if tensor.graph is not self:
-      return capture_value(
-          self.captures, tensor, tensor.dtype, tensor.op.name)
+      if name is None:
+        name = tensor.op.name
+      return capture_value(self.captures, tensor, tensor.dtype, name)
     return tensor
 
   def create_op(
@@ -126,7 +128,7 @@ class CapturingGraph(ops.Graph):
     # forward the resources such as Identity and Switch can cause serialization
     # to fail.
     for i, inp in enumerate(inputs):
-      inputs[i] = self.maybe_capture_tensor(inp)
+      inputs[i] = self.capture(inp)
     return super(CapturingGraph, self).create_op(
         op_type, inputs, dtypes, input_types, name, attrs, op_def,
         compute_shapes, compute_device)
@@ -598,7 +600,7 @@ def _defun_internal(name, func, args, kwds):
       # call to convert_to_tensor, so we manually capture all such tensors.
       outputs_list = _flatten(func_outputs)
       func_def_outputs = [
-          tmp_graph.maybe_capture_tensor(x) for x in outputs_list
+          tmp_graph.capture(x) for x in outputs_list
           if x is not None
       ]
 
diff --git a/tensorflow/python/eager/function_test.py b/tensorflow/python/eager/function_test.py
index 185f6d981c..f53d6c2608 100644
--- a/tensorflow/python/eager/function_test.py
+++ b/tensorflow/python/eager/function_test.py
@@ -771,6 +771,21 @@ class AutomaticControlDependenciesTest(test.TestCase):
       self.assertAllEqual(val.eval(feed_dict={p: False}), 10.0)
       self.assertAllEqual(val.eval(feed_dict={p: True}), 20.0)
 
+  def testDefunWhileLoopWithCapturedLoopVars(self):
+    n = 3
+    x = constant_op.constant(list(range(n)))
+
+    @function.defun
+    def loop():
+      c = lambda i, x: i < n
+      b = lambda i, x: (i + 1, x + 1)
+      i, out = control_flow_ops.while_loop(c, b, (0, x))
+      return i, out
+
+    i, out = loop()
+    self.assertEqual(int(i), 3)
+    self.assertAllEqual(out, [3, 4, 5])
+
   def testDecorator(self):
     with context.graph_mode(), self.test_session():
       v = resource_variable_ops.ResourceVariable(1.0)
diff --git a/tensorflow/python/framework/function.py b/tensorflow/python/framework/function.py
index e7f9e590af..f82e94b1a3 100644
--- a/tensorflow/python/framework/function.py
+++ b/tensorflow/python/framework/function.py
@@ -696,7 +696,7 @@ class _FuncGraph(ops.Graph):
     return super(_FuncGraph, self).create_op(op_type, inputs, data_types,
                                              **kwargs)
 
-  def capture(self, tensor):
+  def capture(self, tensor, name=None):
     """Adds the given tensor to this graph and returns the captured tensor."""
     if tensor in self._captured:
       # Captured already.
@@ -704,15 +704,16 @@ class _FuncGraph(ops.Graph):
     elif self._capture_by_value:
       return self._add_tensor_and_parents(tensor)
     else:
-      return self._capture_tensor_as_extra_input(tensor)
+      return self._capture_tensor_as_extra_input(tensor, name)
 
-  def _capture_tensor_as_extra_input(self, tensor):
+  def _capture_tensor_as_extra_input(self, tensor, name=None):
     # Substitute with a placeholder.
     self.extra_inputs.append(tensor)
     # Hoist the new input placeholder out of any control flow context
     # we're currently in.
     with ops.control_dependencies(None):
-      ph = array_ops.placeholder(tensor.dtype, shape=tensor.get_shape())
+      ph = array_ops.placeholder(
+          tensor.dtype, shape=tensor.get_shape(), name=name)
     # pylint: disable=protected-access
     if ops._USE_C_SHAPES:
       handle_data = c_api.GetResourceHandleShapeAndType(tensor.graph._c_graph,
diff --git a/tensorflow/python/framework/ops.py b/tensorflow/python/framework/ops.py
index dd9acdd9eb..bf27647d27 100644
--- a/tensorflow/python/framework/ops.py
+++ b/tensorflow/python/framework/ops.py
@@ -1057,13 +1057,19 @@ def internal_convert_to_tensor(value,
 
   """
   if ctx is None: ctx = context.context()
-  if ctx.executing_eagerly():
-    # Fast path for EagerTensors that don't need any conversion.
-    if isinstance(value, EagerTensor):
+  if isinstance(value, EagerTensor):
+    if ctx.executing_eagerly():
+      # Fast path for EagerTensors that don't need any conversion.
       # Note that we don't check that value's dtype matches the dtype
       # argument.  We expect that the C runtime will do that checking
       # when we execute the kernel.
       return value
+    else:
+      graph = get_default_graph()
+      if not graph.building_function:
+        raise RuntimeError("Attempting to capture an EagerTensor without "
+                           "building a function.")
+      return graph.capture(value, name=name)
 
   if dtype is not None:
     dtype = dtypes.as_dtype(dtype)
@@ -1251,7 +1257,10 @@ def internal_convert_to_tensor_or_indexed_slices(value,
   Raises:
     ValueError: If `dtype` does not match the element type of `value`.
   """
-  if isinstance(value, _TensorLike):
+  if isinstance(value, EagerTensor) and not context.executing_eagerly():
+    return internal_convert_to_tensor(
+        value, dtype=dtype, name=name, as_ref=as_ref)
+  elif isinstance(value, _TensorLike):
     if dtype and not dtypes.as_dtype(dtype).is_compatible_with(value.dtype):
       raise ValueError(
           "Tensor conversion requested dtype %s for Tensor with dtype %s: %r" %
-- 
GitLab


From f0a506f67fe316c3adb282b58b7087e11d7c493f Mon Sep 17 00:00:00 2001
From: Andrew Selle <aselle@google.com>
Date: Tue, 8 May 2018 11:10:23 -0700
Subject: [PATCH 365/395] Fix Raspberry Pi build by making PNG not try to use
 Neon (by autodetect).

This involves patching to override the png neon option. In the future
it might be worth enabling PNG optimization.

PiperOrigin-RevId: 195846513
---
 tensorflow/workspace.bzl      |  1 +
 third_party/png_fix_rpi.patch | 16 ++++++++++++++++
 2 files changed, 17 insertions(+)
 create mode 100644 third_party/png_fix_rpi.patch

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
index 8f499976de..01d424f20b 100644
--- a/tensorflow/workspace.bzl
+++ b/tensorflow/workspace.bzl
@@ -228,6 +228,7 @@ def tf_workspace(path_prefix="", tf_repo_name=""):
       sha256 = "e45ce5f68b1d80e2cb9a2b601605b374bdf51e1798ef1c2c2bd62131dfcf9eef",
       strip_prefix = "libpng-1.6.34",
       build_file = clean_dep("//third_party:png.BUILD"),
+      patch_file = clean_dep("//third_party:png_fix_rpi.patch"),
   )
 
   tf_http_archive(
diff --git a/third_party/png_fix_rpi.patch b/third_party/png_fix_rpi.patch
new file mode 100644
index 0000000000..80da7b3c06
--- /dev/null
+++ b/third_party/png_fix_rpi.patch
@@ -0,0 +1,16 @@
+diff -r -u /tmp/libpng-1.6.34/scripts/pnglibconf.h.prebuilt ./scripts/pnglibconf.h.prebuilt
+--- /tmp/libpng-1.6.34/scripts/pnglibconf.h.prebuilt	2017-09-29 01:42:33.000000000 -0700
++++ ./scripts/pnglibconf.h.prebuilt	2018-05-01 09:51:24.719318242 -0700
+@@ -20,6 +20,12 @@
+ #define PNG_ALIGNED_MEMORY_SUPPORTED
+ /*#undef PNG_ARM_NEON_API_SUPPORTED*/
+ /*#undef PNG_ARM_NEON_CHECK_SUPPORTED*/
++
++/* Workaround not having a great build file by forcing
++ * png filter optimization to be disabled on arm */
++#define PNG_ARM_NEON_OPT 0
++
++
+ /*#undef PNG_POWERPC_VSX_API_SUPPORTED*/
+ /*#undef PNG_POWERPC_VSX_CHECK_SUPPORTED*/
+ #define PNG_BENIGN_ERRORS_SUPPORTED
-- 
GitLab


From 4ca46dd6dda433e622e4a382123f9a81487aeef5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 11:12:07 -0700
Subject: [PATCH 366/395] Increase size of test
 //third_party/tensorflow/python:saver_large_variable_test from "small" to
 "medium" to prevent flaky timeouts.

PiperOrigin-RevId: 195846802
---
 tensorflow/python/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 4057e37681..a865e8ca75 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -4135,7 +4135,7 @@ cuda_py_test(
 
 py_test(
     name = "saver_large_variable_test",
-    size = "small",
+    size = "medium",
     srcs = ["training/saver_large_variable_test.py"],
     srcs_version = "PY2AND3",
     tags = [
-- 
GitLab


From b62573f37b1040311b520d55715492df32cac0cf Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 11:15:53 -0700
Subject: [PATCH 367/395] Add affinity binding functionality and documentation
 to OVIC benchmarker.

PiperOrigin-RevId: 195847378
---
 tensorflow/contrib/lite/java/ovic/README.md   |  58 +++++++++-
 .../demo/app/OvicBenchmarkerActivity.java     | 100 ++++++++++++++++--
 2 files changed, 149 insertions(+), 9 deletions(-)

diff --git a/tensorflow/contrib/lite/java/ovic/README.md b/tensorflow/contrib/lite/java/ovic/README.md
index 373a50854c..77799b3569 100644
--- a/tensorflow/contrib/lite/java/ovic/README.md
+++ b/tensorflow/contrib/lite/java/ovic/README.md
@@ -6,7 +6,7 @@ This folder contains building code for track one of the [Low Power ImageNet Reco
 
 Follow the steps [here](https://www.tensorflow.org/mobile/tflite/demo_android) to install Tensorflow, Bazel, and the Android NDK and SDK.
 
-## To test the benchmarker:
+## Test the benchmarker:
 
 The testing utilities helps the developers (you) to make sure that your submissions in TfLite format will be processed as expected in the competition's benchmarking system.
 
@@ -80,3 +80,59 @@ Change `TEST_IMAGE_PATH` to `my_test_image.jpg`. Change either `FLOAT_MODEL_PATH
 Now you can run the bazel tests to catch any runtime issues with the submission.
 
 Note: Please make sure that your submission passes the test. If a submission fails to pass the test it will not be processed by the submission server.
+
+## Measure on-device latency
+
+We provide two ways to measure the on-device latency of your submission. The first is through our competition server, which is reliable and repeatable, but is limited to a few trials per day. The second is through the benchmarker Apk, which requires a device and may not be as accurate as the server, but has a fast turn-around and no access limitations. We recommend that the participants use the benchmarker apk for early development, and reserve the competition server for evaluating promising submissions.
+
+### Running the benchmarker app
+
+Make sure that you have followed instructions in [Test your submissions](#test-your-submissions) to add your model to the testdata folder and to the corresponding build rules.
+
+Modify `tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarkerActivity.java`:
+
+* Add your model to the benchmarker apk by changing `MODEL_PATH` and `TEST_IMAGE_PATH` below to your submission and test image.
+
+```
+  private static final String TEST_IMAGE_PATH = "my_test_image.jpg";
+  private static final String MODEL_PATH = "my_model.lite";
+```
+
+* Adjust the benchmark parameters when needed:
+
+You can chnage the length of each experiment, and the processor affinity below. `BIG_CORE_MASK` is an integer whose binary encoding represents the set of used cores. This number is phone-specific. For example, Pixel 2 has 8 cores: the 4 little cores are represented by the 4 less significant bits, and the 4 big cores by the 4 more significant bits. Therefore a mask value of 16, or in binary `00010000`, represents using only the first big core. The mask 32, or in binary `00100000` uses the second big core and should deliver identical results as the mask 16 because the big cores are interchangeable.
+
+```
+  /** Wall time for each benchmarking experiment. */
+  private static final double WALL_TIME = 3000;
+  /** Maximum number of iterations in each benchmarking experiment. */
+  private static final int MAX_ITERATIONS = 100;
+  /** Mask for binding to a single big core. Pixel 1 (4), Pixel 2 (16). */
+  private static final int BIG_CORE_MASK = 16;
+```
+
+Note: You'll need ROOT access to the phone to change processor affinity.
+
+* Build and install the app.
+
+```
+bazel build -c opt --cxxopt=--std=c++11 --cxxopt=-Wno-all //tensorflow/contrib/lite/java/ovic/demo/app:ovic_benchmarker_binary
+adb install -r bazel-bin/tensorflow/contrib/lite/java/ovic/demo/app/ovic_benchmarker_binary.apk
+```
+
+Start the app and click the `Start` button in dark green. The button should turn bright green, signaling that the experiment is running. The benchmarking results will be displayed after about the `WALL_TIME` you specified above. For example:
+
+```
+my_model.lite: Average latency=158.6ms after 20 runs.
+```
+
+### Sample latencies
+
+Note: the benchmarking results can be quite different depending on the background processes running on the phone. A few things that help stabilize the app's readings are placing the phone on a cooling plate, restarting the phone, and shutting down internet access.
+
+| Model                | Pixel 1 latency (ms)  | Pixel 2 latency (ms) |
+| -------------------- |:---------------------:| --------------------:|
+|  float_model.lite    | 120                   | 155                  |
+| quantized_model.lite | 85                    | 74                   |
+|  low_res_model.lite  | 4.2                   | 4.0                  |
+
diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarkerActivity.java b/tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarkerActivity.java
index a871b869b0..59457c308a 100644
--- a/tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarkerActivity.java
+++ b/tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarkerActivity.java
@@ -20,10 +20,15 @@ import android.content.res.AssetManager;
 import android.graphics.Bitmap;
 import android.graphics.BitmapFactory;
 import android.os.Bundle;
+import android.os.Process;
+import android.os.SystemClock;
 import android.util.Log;
 import android.view.View;
 import android.widget.TextView;
+import java.io.BufferedReader;
+import java.io.File;
 import java.io.FileInputStream;
+import java.io.FileReader;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.MappedByteBuffer;
@@ -50,6 +55,10 @@ public class OvicBenchmarkerActivity extends Activity {
   private static final double WALL_TIME = 3000;
   /** Maximum number of iterations in each benchmarking experiment. */
   private static final int MAX_ITERATIONS = 100;
+  /** Mask for binding to a single big core. Pixel 1 (4), Pixel 2 (16). */
+  private static final int BIG_CORE_MASK = 16;
+  /** Amount of time in milliseconds to wait for affinity to set. */
+  private static final int WAIT_TIME_FOR_AFFINITY = 1000;
 
   /* The model to be benchmarked. */
   private MappedByteBuffer model = null;
@@ -123,6 +132,13 @@ public class OvicBenchmarkerActivity extends Activity {
       Log.e(TAG, "Can't initialize benchmarker.", e);
       throw e;
     }
+    String displayText = "";
+    try {
+      setProcessorAffinity(BIG_CORE_MASK);
+    } catch (IOException e) {
+      Log.e(TAG, e.getMessage());
+      displayText = e.getMessage() + "\n";
+    }
     Log.i(TAG, "Successfully initialized benchmarker.");
     int testIter = 0;
     Boolean iterSuccess = false;
@@ -147,17 +163,85 @@ public class OvicBenchmarkerActivity extends Activity {
 
     if (textView != null) {
       if (testIter > 0) {
-        textView
-            .setText(
-                MODEL_PATH
-                    + ": Average latency="
-                    + df2.format(totalLatency / testIter)
-                    + "ms after "
-                    + testIter
-                    + " runs.");
+        textView.setText(
+            displayText
+                + MODEL_PATH
+                + ": Average latency="
+                + df2.format(totalLatency / testIter)
+                + "ms after "
+                + testIter
+                + " runs.");
       } else {
         textView.setText("Benchmarker failed to run on more than one images.");
       }
     }
   }
+
+  private static void setProcessorAffinity(int mask) throws IOException {
+    int myPid = Process.myPid();
+    Log.i(TAG, String.format("Setting processor affinity to 0x%02x", mask));
+
+    String command = String.format("taskset -a -p %x %d", mask, myPid);
+    try {
+      Runtime.getRuntime().exec(command).waitFor();
+    } catch (InterruptedException e) {
+      throw new IOException("Interrupted: " + e);
+    }
+
+    // Make sure set took effect - try for a second to confirm the change took.  If not then fail.
+    long startTimeMs = SystemClock.elapsedRealtime();
+    while (true) {
+      int readBackMask = readCpusAllowedMask();
+      if (readBackMask == mask) {
+        Log.i(TAG, String.format("Successfully set affinity to 0x%02x", mask));
+        return;
+      }
+      if (SystemClock.elapsedRealtime() > startTimeMs + WAIT_TIME_FOR_AFFINITY) {
+        throw new IOException(
+            String.format(
+                "Core-binding failed: affinity set to 0x%02x but read back as 0x%02x\n"
+                    + "please root device.",
+                mask, readBackMask));
+      }
+
+      try {
+        Thread.sleep(50);
+      } catch (InterruptedException e) {
+        // Ignore sleep interrupted, will sleep again and compare is final cross-check.
+      }
+    }
+  }
+
+  public static int readCpusAllowedMask() throws IOException {
+    // Determine how many CPUs there are total
+    final String pathname = "/proc/self/status";
+    final String resultPrefix = "Cpus_allowed:";
+    File file = new File(pathname);
+    String line = "<NO LINE READ>";
+    String allowedCPU = "";
+    Integer allowedMask = null;
+    BufferedReader bufReader = null;
+    try {
+      bufReader = new BufferedReader(new FileReader(file));
+      while ((line = bufReader.readLine()) != null) {
+        if (line.startsWith(resultPrefix)) {
+          allowedMask = Integer.valueOf(line.substring(resultPrefix.length()).trim(), 16);
+          allowedCPU = bufReader.readLine();
+          break;
+        }
+      }
+    } catch (RuntimeException e) {
+      throw new IOException(
+          "Invalid number in " + pathname + " line: \"" + line + "\": " + e.getMessage());
+    } finally {
+      if (bufReader != null) {
+        bufReader.close();
+      }
+    }
+    if (allowedMask == null) {
+      throw new IOException(pathname + " missing " + resultPrefix + " line");
+    }
+    Log.i(TAG, allowedCPU);
+    return allowedMask;
+  }
 }
-- 
GitLab


From 26749309690949cf355fd51f17e818b7450d3f7f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 11:19:46 -0700
Subject: [PATCH 368/395] Change visibility of hlo_proto.

PiperOrigin-RevId: 195848035
---
 tensorflow/compiler/xla/service/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD
index ec67e19b23..aa3a6261e0 100644
--- a/tensorflow/compiler/xla/service/BUILD
+++ b/tensorflow/compiler/xla/service/BUILD
@@ -26,6 +26,7 @@ xla_proto_library(
 xla_proto_library(
     name = "hlo_proto",
     srcs = ["hlo.proto"],
+    visibility = ["//visibility:public"],
     deps = ["//tensorflow/compiler/xla:xla_data_proto"],
 )
 
-- 
GitLab


From 211e3a20016cd1dd29883d57576eecd477a3dcac Mon Sep 17 00:00:00 2001
From: Ilya Biryukov <ibiryukov@google.com>
Date: Tue, 8 May 2018 11:25:50 -0700
Subject: [PATCH 369/395] Update version of downloadable clang toolchain

PiperOrigin-RevId: 195849091
---
 third_party/clang_toolchain/download_clang.bzl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/third_party/clang_toolchain/download_clang.bzl b/third_party/clang_toolchain/download_clang.bzl
index 54d383d7d7..cfd8bfe98d 100644
--- a/third_party/clang_toolchain/download_clang.bzl
+++ b/third_party/clang_toolchain/download_clang.bzl
@@ -35,18 +35,18 @@ def download_clang(repo_ctx, out_folder):
 
   # Latest CLANG_REVISION and CLANG_SUB_REVISION of the Chromiums's release
   # can be found in https://chromium.googlesource.com/chromium/src/tools/clang/+/master/scripts/update.py
-  CLANG_REVISION = '321529'
+  CLANG_REVISION = '330570'
   CLANG_SUB_REVISION = 2
 
   package_version = '%s-%s' % (CLANG_REVISION, CLANG_SUB_REVISION)
 
   checksums = {
       'Linux_x64':
-          '76d4eb1ad011e3127c4a9de9b9f5d4ac624b5a9395c4d7395c9e0a487b13daf6',
+          '2108e172e05d4904c3c46125a33ab4a1175b36ec2a2226619a243e1d8f397e97',
       'Mac':
-          '4b2a7a65ac1ee892b318c723eec8771f514bb306f346aa8216bb0006f19d87b7',
+          '481b5c6909f0ea250216061bd45e9c982b4befff65cbfca2ee1090c21a109eac',
       'Win':
-          'eba51bb8f84af41a85903113666bd21c22709010c39c4cb19dc20cf1ed14581b',
+          '8f04a3ac99d463d4179eb2f68a13575408c3dddc62887a1e441c77123e35e301',
   }
 
   platform_folder = _get_platform_folder(repo_ctx.os.name)
-- 
GitLab


From 4f7a0bc8c11827dde6986ad29e9fd21c48597367 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 11:45:53 -0700
Subject: [PATCH 370/395] Fix docstring for flush() method

PiperOrigin-RevId: 195852402
---
 tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py b/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py
index 1b184d296b..50cc00afdc 100644
--- a/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py
+++ b/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py
@@ -187,7 +187,7 @@ class QuantileAccumulator(saver.BaseSaverBuilder.SaveableObject):
       stamp_token: Expected current token.
       next_stamp_token: Next value for the token.
     Returns:
-      A list of quantiles or approximate boundaries.
+      The flush operation.
     """
     return gen_quantile_ops.quantile_accumulator_flush(
         quantile_accumulator_handle=self._quantile_accumulator_handle,
-- 
GitLab


From 59bffb7051231c7e0f8020892db8c3d584c555f4 Mon Sep 17 00:00:00 2001
From: Alina Sbirlea <asbirlea@google.com>
Date: Tue, 8 May 2018 11:54:03 -0700
Subject: [PATCH 371/395] Re-land: Optimize dot(DynamicSlice(ConstA),
 ConstantB) by memoizing dot(ConstA, ConstB)

Make transformation when ConstA and ConstB are 2D, and DynamicSlice is slicing a full row, column respectively.
Handle:
dot(DynamicSlice(Index, ConstA), ConstB) => DynamicSlice(Index, dot*(ConstA, ConstB));
and
dot(ConstA, DynamicSlice(Index, ConstB)) => DynamicSlice(Index, dot*(ConstA, ConstB));

Reason to roll forward: Previous issue of getting out of memory errors when generating LLVM constants was resolved by CSE-ing constants before allocation.

PiperOrigin-RevId: 195853680
---
 .../xla/service/algebraic_simplifier.cc       | 141 ++++++++++
 .../xla/service/algebraic_simplifier_test.cc  | 203 +++++++++++++++
 .../compiler/xla/tests/dot_operation_test.cc  | 245 ++++++++++++++++++
 3 files changed, 589 insertions(+)

diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.cc b/tensorflow/compiler/xla/service/algebraic_simplifier.cc
index 8e785de68c..4ec79a0244 100644
--- a/tensorflow/compiler/xla/service/algebraic_simplifier.cc
+++ b/tensorflow/compiler/xla/service/algebraic_simplifier.cc
@@ -291,6 +291,8 @@ class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault {
       const Shape& dot_shape, HloInstruction* lhs, int64 lhs_contracting_dim,
       HloInstruction* rhs, int64 rhs_contracting_dim, bool swapped);
 
+  StatusOr<HloInstruction*> OptimizeDotOfGather(HloInstruction* dot);
+
   // Current HloComputation instance the AlgebraicSimplifierVisitor is
   // traversing.
   HloComputation* computation_;
@@ -912,6 +914,134 @@ StatusOr<HloInstruction*> AlgebraicSimplifierVisitor::OptimizeDotOfConcatHelper(
   return add_result;
 }
 
+StatusOr<HloInstruction*> AlgebraicSimplifierVisitor::OptimizeDotOfGather(
+    HloInstruction* dot) {
+  const DotDimensionNumbers& dnums = dot->dot_dimension_numbers();
+  if (dnums.lhs_contracting_dimensions_size() != 1 ||
+      dnums.rhs_contracting_dimensions_size() != 1 ||
+      dnums.lhs_batch_dimensions_size() != 0 ||
+      dnums.rhs_batch_dimensions_size() != 0 ||
+      dot->shape().dimensions_size() != 2) {  // dot output 2D
+    VLOG(10) << "DotOfGather: Can only optimize 2D, non-batch dot operations.";
+    return nullptr;
+  }
+
+  // Optimize either dot(DS(ctA), ctB)) or dot(ctB, DS(ctA)).
+  // Currently a Gather is a DynamicSlice.
+  auto is_dynamic_slice_constant_combination =
+      [](HloInstruction* a, HloInstruction* b, int a_contracting_dimension) {
+        // First operand is a DynamicSlice(Constant).
+        if (a->opcode() != HloOpcode::kDynamicSlice) {
+          return false;
+        }
+        auto* dynamic_slice_op = a->operand(0);
+        if (dynamic_slice_op->opcode() != HloOpcode::kConstant) {
+          return false;
+        }
+        // Second operand is a Constant.
+        if (b->opcode() != HloOpcode::kConstant) {
+          return false;
+        }
+        // The DynamicSlice output is a vector.
+        const Shape& dynamic_slice_shape = a->shape();
+        if (dynamic_slice_shape.dimensions(1 - a_contracting_dimension) != 1) {
+          return false;
+        }
+        // Constant size is the same before and after slice in the contracting
+        // dimension, otherwise we either must precompute for all possible slice
+        // indices or dot is invalid.
+        const Shape& dynamic_slice_op_shape = dynamic_slice_op->shape();
+        if (dynamic_slice_op_shape.dimensions(a_contracting_dimension) !=
+            dynamic_slice_shape.dimensions(a_contracting_dimension)) {
+          return false;
+        }
+        return true;
+      };
+
+  HloInstruction* lhs = dot->mutable_operand(0);
+  HloInstruction* rhs = dot->mutable_operand(1);
+  int lhs_contracting_dimension = dnums.lhs_contracting_dimensions(0);
+  int rhs_contracting_dimension = dnums.rhs_contracting_dimensions(0);
+
+  if (!is_dynamic_slice_constant_combination(
+          lhs, rhs, /*a_contracting_dimension=*/lhs_contracting_dimension) &&
+      !is_dynamic_slice_constant_combination(
+          rhs, lhs, /*a_contracting_dimension=*/rhs_contracting_dimension)) {
+    VLOG(10) << "DotOfGather: Can only optimize dot(DS(ctA), ctB)) or "
+                "dot(ctB, DS(ctA)), where the two constants have equal "
+                "contracting dimensions.";
+    return nullptr;
+  }
+
+  // LHS is DynamicSlice:
+  // input: dot(DS(ctA), ctB))
+  // where DS(ctA) = DS({M x K}, {start, 0}, {1, K}) and ctB = {K x N}.
+  // => input dimensions: dot({1 x K}, {K x N}) => {1 x N}.
+  // output: DS(dot(ctA, ctB))
+  // => output dimensions: DS ({M x N}, {start, 0}, {1, N}) => {1 x N}.
+
+  // RHS is DynamicSlice:
+  // input: dot(ctA, DS(ctB))
+  // where ctA = {M x K} and DS(ctB) = DS({K x N}, {0, start}, {K, 1}).
+  // => input dimensions: dot({M x K}, {K x 1}) => {M x 1}.
+  // output: DS(dot(ctA, ctB))
+  // => output dimensions: DS ({M x N}, {0, start}, {M, 1}) => {M x 1}.
+
+  bool lhs_is_dynamic_slice = lhs->opcode() == HloOpcode::kDynamicSlice;
+
+  // ctA:
+  HloInstruction* left_operand =
+      lhs_is_dynamic_slice ? lhs->mutable_operand(0) : lhs;
+  // ctB:
+  HloInstruction* right_operand =
+      lhs_is_dynamic_slice ? rhs : rhs->mutable_operand(0);
+  // Build ctA x ctB.
+  const int m = left_operand->shape().dimensions(1 - lhs_contracting_dimension);
+  const int n =
+      right_operand->shape().dimensions(1 - rhs_contracting_dimension);
+  auto memoized_shape = ShapeUtil::MakeShape(F32, {m, n});
+  auto* memoized_inst = computation_->AddInstruction(HloInstruction::CreateDot(
+      memoized_shape, left_operand, right_operand, dnums));
+  // Get pair {start, 0} or {0, start}.
+  HloInstruction* original_start_indices =
+      lhs_is_dynamic_slice ? lhs->mutable_operand(1) : rhs->mutable_operand(1);
+  // Position of start:
+  int index_of_non_zero_start = lhs_is_dynamic_slice
+                                    ? 1 - lhs_contracting_dimension
+                                    : 1 - rhs_contracting_dimension;
+  // Position of zero:
+  int index_of_zero_start = 1 - index_of_non_zero_start;
+
+  // Slice out start and 0 components and reorder if necessary.
+  auto indices_type = original_start_indices->shape().element_type();
+  Shape s_shape = ShapeUtil::MakeShape(indices_type, {1});
+  Shape d_shape = ShapeUtil::MakeShape(indices_type, {2});
+  HloInstruction* non_zero_start =
+      computation_->AddInstruction(HloInstruction::CreateSlice(
+          s_shape, original_start_indices, {index_of_non_zero_start},
+          {index_of_non_zero_start + 1}, {1}));
+  HloInstruction* zero_start =
+      computation_->AddInstruction(HloInstruction::CreateSlice(
+          s_shape, original_start_indices, {index_of_zero_start},
+          {index_of_zero_start + 1}, {1}));
+  HloInstruction* new_start_indices =
+      lhs_is_dynamic_slice
+          ? computation_->AddInstruction(HloInstruction::CreateConcatenate(
+                d_shape, {non_zero_start, zero_start}, 0))
+          : computation_->AddInstruction(HloInstruction::CreateConcatenate(
+                d_shape, {zero_start, non_zero_start}, 0));
+
+  // Build DynamicSlice(ctA x ctB).
+  const int new_slice_m = lhs_is_dynamic_slice ? 1 : m;
+  const int new_slice_n = lhs_is_dynamic_slice ? n : 1;
+  auto* memoized_lookup =
+      computation_->AddInstruction(HloInstruction::CreateDynamicSlice(
+          dot->shape(), memoized_inst, new_start_indices,
+          {new_slice_m, new_slice_n}));
+
+  return memoized_lookup;
+}
+
 Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot) {
   HloInstruction *lhs, *rhs;
   CHECK(Match(dot, m::Dot(m::Op(&lhs), m::Op(&rhs))));
@@ -941,6 +1071,17 @@ Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot) {
     return ReplaceInstruction(dot, dot_of_concat_optimized);
   }
 
+  // Simplify dot(ConstA, Gather(Index, ConstB)) to:
+  // Gather(Index, dot*(ConstA, ConstB)), where dot* is an appropriately
+  // batched version of dot.
+  TF_ASSIGN_OR_RETURN(HloInstruction * dot_of_gather_optimized,
+                      OptimizeDotOfGather(dot));
+  if (dot_of_gather_optimized) {
+    VLOG(10) << "Replaced dot(constA, gather(i, constB)) with "
+                "gather(i, dot*(constA, constB))";
+    return ReplaceInstruction(dot, dot_of_gather_optimized);
+  }
+
   if (enable_dot_strength_reduction_ && !is_layout_sensitive_) {
     TF_ASSIGN_OR_RETURN(bool did_strength_reduction,
                         HandleDotStrengthReduction(dot));
diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc
index d0c99bf818..4e082877c7 100644
--- a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc
+++ b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc
@@ -2963,5 +2963,208 @@ TEST_F(AlgebraicSimplifierTest, DynamicUpdateSliceZeroUpdate) {
 INSTANTIATE_TEST_CASE_P(DotOfConcatSimplificationTestInstantiation,
                         DotOfConcatSimplificationTest,
                         ::testing::ValuesIn(kDotOfConcatTestSpecs));
+
+struct DotOfGatherTestSpec {
+  int64 m;
+  int64 k;
+  int64 n;
+  int s;      // start index for dynamic slice on the non-contracting dimension
+  int64 lcd;  // left contracting dimension
+  int64 rcd;  // right contracting dimension
+  bool neg;   // is negative testcase
+};
+
+class DotOfGatherSimplificationTest
+    : public HloVerifiedTestBase,
+      public ::testing::WithParamInterface<DotOfGatherTestSpec> {};
+
+// input: dot(DS(ctA), ctB))
+// where DS(ctA) = DS({M x K}, {s, 0}, {1, K}) and ctB = {K x N}.
+// => input dimensions: dot({1 x K}, {K x N}) => {1 x N}.
+// output: DS(dot(ctA, ctB))
+// => output dimensions: DS ({M x N}, {s, 0}, {1, N}) => {1 x N}.
+TEST_P(DotOfGatherSimplificationTest, ConstantRHS) {
+  HloComputation::Builder builder(TestName());
+
+  DotOfGatherTestSpec spec = GetParam();
+
+  ASSERT_LE(spec.s, spec.m);
+
+  // For negative tests, increase k of the dynamic slice argument to prevent the
+  // optimization (constants ctA, ctB must have equal contracting dimensions).
+  int64 k_increase = spec.neg ? 5 : 0;
+  int64 lhs_rows = (spec.lcd == 0) ? (spec.k + k_increase) : spec.m;
+  int64 lhs_cols = (spec.lcd == 0) ? spec.m : (spec.k + k_increase);
+  Shape lhs_shape = ShapeUtil::MakeShape(F32, {lhs_rows, lhs_cols});
+  auto* lhs = builder.AddInstruction(
+      HloInstruction::CreateConstant(Literal::CreateR2F32Linspace(
+          /*from=*/10.0, /*to=*/10000.0, /*rows=*/lhs_rows,
+          /*cols=*/lhs_cols)));
+
+  int32 start_row = (spec.lcd == 0) ? 0 : spec.s;
+  int32 start_col = (spec.lcd == 0) ? spec.s : 0;
+  const auto start_indices =
+      builder.AddInstruction(HloInstruction::CreateConstant(
+          Literal::CreateR1<int32>({start_row, start_col})));
+  int64 slice_row_size = (spec.lcd == 0) ? spec.k : 1;
+  int64 slice_col_size = (spec.lcd == 0) ? 1 : spec.k;
+  Shape ds_shape = ShapeUtil::MakeShape(F32, {slice_row_size, slice_col_size});
+  auto* ds = builder.AddInstruction(HloInstruction::CreateDynamicSlice(
+      ds_shape, lhs, start_indices, {slice_row_size, slice_col_size}));
+
+  int64 rhs_rows = (spec.rcd == 0) ? spec.k : spec.n;
+  int64 rhs_cols = (spec.rcd == 0) ? spec.n : spec.k;
+  Shape rhs_shape = ShapeUtil::MakeShape(F32, {rhs_rows, rhs_cols});
+  auto* rhs = builder.AddInstruction(
+      HloInstruction::CreateConstant(Literal::CreateR2F32Linspace(
+          /*from=*/10.0, /*to=*/10000.0, /*rows=*/rhs_rows,
+          /*cols=*/rhs_cols)));
+
+  DotDimensionNumbers dot_dnums;
+  dot_dnums.add_lhs_contracting_dimensions(spec.lcd);
+  dot_dnums.add_rhs_contracting_dimensions(spec.rcd);
+
+  int64 dot_row_size = 1;
+  int64 dot_col_size = spec.n;
+  Shape dot_shape = ShapeUtil::MakeShape(F32, {dot_row_size, dot_col_size});
+  builder.AddInstruction(
+      HloInstruction::CreateDot(dot_shape, ds, rhs, dot_dnums));
+
+  auto computation = module().AddEntryComputation(builder.Build());
+  AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false,
+                                 non_bitcasting_callback());
+  TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(&module()));
+  ASSERT_TRUE(run_successful);
+  EXPECT_TRUE(
+      ShapeUtil::Equal(computation->root_instruction()->shape(), dot_shape));
+
+  if (spec.neg) {
+    EXPECT_NE(computation->root_instruction()->opcode(),
+              HloOpcode::kDynamicSlice);
+  } else {
+    EXPECT_THAT(computation->root_instruction(),
+                op::DynamicSlice(op::Dot(op::Constant(), op::Constant()),
+                                 op::Concatenate()));
+  }
+}
+
+// input: dot(ctA, DS(ctB))
+// where ctA = {M x K} and DS(ctB) = DS({K x N}, {0, s}, {K, 1}).
+// => input dimensions: dot({M x K}, {K x 1}) => {M x 1}.
+// output: DS(dot(ctA, ctB))
+// => output dimensions: DS ({M x N}, {0, s}, {M, 1}) => {M x 1}.
+TEST_P(DotOfGatherSimplificationTest, ConstantLHS) {
+  HloComputation::Builder builder(TestName());
+
+  DotOfGatherTestSpec spec = GetParam();
+
+  ASSERT_LE(spec.s, spec.n);
+
+  int64 lhs_rows = (spec.lcd == 0) ? spec.k : spec.m;
+  int64 lhs_cols = (spec.lcd == 0) ? spec.m : spec.k;
+  Shape lhs_shape = ShapeUtil::MakeShape(F32, {lhs_rows, lhs_cols});
+  auto* lhs = builder.AddInstruction(
+      HloInstruction::CreateConstant(Literal::CreateR2F32Linspace(
+          /*from=*/10.0, /*to=*/10000.0, /*rows=*/lhs_rows,
+          /*cols=*/lhs_cols)));
+
+  // For negative tests increase k of the dynamic slice argument to prevent the
+  // optimization
+  int64 k_increase = spec.neg ? 5 : 0;
+  int64 rhs_rows = (spec.rcd == 0) ? (spec.k + k_increase) : spec.n;
+  int64 rhs_cols = (spec.rcd == 0) ? spec.n : (spec.k + k_increase);
+  Shape rhs_shape = ShapeUtil::MakeShape(F32, {rhs_rows, rhs_cols});
+  auto* rhs = builder.AddInstruction(
+      HloInstruction::CreateConstant(Literal::CreateR2F32Linspace(
+          /*from=*/10.0, /*to=*/10000.0, /*rows=*/rhs_rows,
+          /*cols=*/rhs_cols)));
+
+  int32 start_row = (spec.rcd == 0) ? 0 : spec.s;
+  int32 start_col = (spec.rcd == 0) ? spec.s : 0;
+  const auto start_indices =
+      builder.AddInstruction(HloInstruction::CreateConstant(
+          Literal::CreateR1<int32>({start_row, start_col})));
+  int64 slice_row_size = (spec.rcd == 0) ? spec.k : 1;
+  int64 slice_col_size = (spec.rcd == 0) ? 1 : spec.k;
+  Shape ds_shape = ShapeUtil::MakeShape(F32, {slice_row_size, slice_col_size});
+  auto* ds = builder.AddInstruction(HloInstruction::CreateDynamicSlice(
+      ds_shape, rhs, start_indices, {slice_row_size, slice_col_size}));
+
+  DotDimensionNumbers dot_dnums;
+  dot_dnums.add_lhs_contracting_dimensions(spec.lcd);
+  dot_dnums.add_rhs_contracting_dimensions(spec.rcd);
+
+  int64 dot_row_size = spec.m;
+  int64 dot_col_size = 1;
+  Shape dot_shape = ShapeUtil::MakeShape(F32, {dot_row_size, dot_col_size});
+  builder.AddInstruction(
+      HloInstruction::CreateDot(dot_shape, lhs, ds, dot_dnums));
+
+  auto computation = module().AddEntryComputation(builder.Build());
+  AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false,
+                                 non_bitcasting_callback());
+  TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(&module()));
+  ASSERT_TRUE(run_successful);
+  EXPECT_TRUE(
+      ShapeUtil::Equal(computation->root_instruction()->shape(), dot_shape));
+
+  if (spec.neg) {
+    EXPECT_NE(computation->root_instruction()->opcode(),
+              HloOpcode::kDynamicSlice);
+  } else {
+    EXPECT_THAT(computation->root_instruction(),
+                op::DynamicSlice(op::Dot(op::Constant(), op::Constant()),
+                                 op::Concatenate()));
+  }
+}
+
+std::vector<DotOfGatherTestSpec> DotOfGatherPositiveNegativeTests() {
+  std::vector<DotOfGatherTestSpec> positives = {
+      // "Classical dot", i.e. matrix multiply:
+      {/*m=*/10, /*k=*/10, /*n=*/5, /*s=*/0, /*lcd=*/1, /*rcd=*/0,
+       /*neg=*/false},
+      {/*m=*/20, /*k=*/20, /*n=*/3, /*s=*/2, /*lcd=*/1, /*rcd=*/0,
+       /*neg=*/false},
+      {/*m=*/10, /*k=*/3, /*n=*/10, /*s=*/9, /*lcd=*/1, /*rcd=*/0,
+       /*neg=*/false},
+      // Note: testing for m=1 and n=1 is unnecessary, as this optimizes to
+      // dot(ct, ct) before DotOfGather optimization kicks in.
+      // Contract on rows:
+      {/*m=*/10, /*k=*/10, /*n=*/5, /*s=*/0, /*lcd=*/0, /*rcd=*/0,
+       /*neg=*/false},
+      {/*m=*/20, /*k=*/20, /*n=*/3, /*s=*/2, /*lcd=*/0, /*rcd=*/0,
+       /*neg=*/false},
+      {/*m=*/10, /*k=*/3, /*n=*/10, /*s=*/9, /*lcd=*/0, /*rcd=*/0,
+       /*neg=*/false},
+      // Reverse matrix multiply:
+      {/*m=*/10, /*k=*/10, /*n=*/5, /*s=*/0, /*lcd=*/0, /*rcd=*/1,
+       /*neg=*/false},
+      {/*m=*/20, /*k=*/20, /*n=*/3, /*s=*/2, /*lcd=*/0, /*rcd=*/1,
+       /*neg=*/false},
+      {/*m=*/10, /*k=*/3, /*n=*/10, /*s=*/9, /*lcd=*/0, /*rcd=*/1,
+       /*neg=*/false},
+      // Contract on columns:
+      {/*m=*/10, /*k=*/10, /*n=*/5, /*s=*/0, /*lcd=*/1, /*rcd=*/1,
+       /*neg=*/false},
+      {/*m=*/20, /*k=*/20, /*n=*/3, /*s=*/2, /*lcd=*/1, /*rcd=*/1,
+       /*neg=*/false},
+      {/*m=*/10, /*k=*/3, /*n=*/10, /*s=*/9, /*lcd=*/1, /*rcd=*/1,
+       /*neg=*/false},
+  };
+  std::vector<DotOfGatherTestSpec> all;
+  for (int i = 0; i < positives.size(); i++) {
+    DotOfGatherTestSpec positive_test = positives[i];
+    all.push_back(positive_test);
+    DotOfGatherTestSpec negative_test = positive_test;
+    negative_test.neg = true;
+    all.push_back(negative_test);
+  }
+  return all;
+}
+
+INSTANTIATE_TEST_CASE_P(
+    DotOfGatherSimplificationTestInstantiation, DotOfGatherSimplificationTest,
+    ::testing::ValuesIn(DotOfGatherPositiveNegativeTests()));
+
 }  // namespace
 }  // namespace xla
diff --git a/tensorflow/compiler/xla/tests/dot_operation_test.cc b/tensorflow/compiler/xla/tests/dot_operation_test.cc
index 6b3efba4f8..efa5aed2d1 100644
--- a/tensorflow/compiler/xla/tests/dot_operation_test.cc
+++ b/tensorflow/compiler/xla/tests/dot_operation_test.cc
@@ -798,5 +798,250 @@ XLA_TYPED_TEST(DotOperationTest_F16F32F64,
       this->error_spec_);
 }
 
+TEST_F(DotOperationTest, DotOfGatherOptimizationWithConstRHSClassicMM) {
+  std::unique_ptr<Array2D<float>> constant_lhs_array(new Array2D<float>(
+      {{1.0, 2.0, 3.0, 4.0, 5.0, 6.0}, {6.0, 5.0, 4.0, 3.0, 2.0, 1.0}}));
+  std::unique_ptr<Array2D<float>> constant_rhs_array(
+      new Array2D<float>({{1.0, 2.0, 3.0},
+                          {4.0, 5.0, 6.0},
+                          {7.0, 8.0, 9.0},
+                          {9.0, 8.0, 7.0},
+                          {6.0, 5.0, 4.0},
+                          {3.0, 2.0, 1.0}}));
+  // Dot result to slice from: {{114, 105, 96}, {96, 105, 114}}
+
+  XlaBuilder builder(TestName());
+  auto lhs_constant = builder.ConstantR2FromArray2D(*constant_lhs_array);
+  auto rhs_constant = builder.ConstantR2FromArray2D(*constant_rhs_array);
+  auto start_constant = builder.ConstantR1<int32>({1, 0});
+  auto dynamic_slice =
+      builder.DynamicSlice(lhs_constant, start_constant, {1, 6});
+
+  DotDimensionNumbers dot_dnums;
+  dot_dnums.add_lhs_contracting_dimensions(1);
+  dot_dnums.add_rhs_contracting_dimensions(0);
+  auto result = builder.DotGeneral(dynamic_slice, rhs_constant, dot_dnums);
+
+  Array2D<float> expected({{96.0, 105.0, 114.0}});
+  ComputeAndCompareR2<float>(&builder, expected, {}, error_spec_);
+}
+
+TEST_F(DotOperationTest, DotOfGatherOptimizationWithConstLHSClassicMM) {
+  std::unique_ptr<Array2D<float>> constant_lhs_array(new Array2D<float>(
+      {{1.0, 2.0, 3.0, 4.0, 5.0, 6.0}, {6.0, 5.0, 4.0, 3.0, 2.0, 1.0}}));
+  std::unique_ptr<Array2D<float>> constant_rhs_array(
+      new Array2D<float>({{1.0, 2.0, 3.0},
+                          {4.0, 5.0, 6.0},
+                          {7.0, 8.0, 9.0},
+                          {9.0, 8.0, 7.0},
+                          {6.0, 5.0, 4.0},
+                          {3.0, 2.0, 1.0}}));
+  // Dot result to slice from: {{114, 105, 96}, {96, 105, 114}}
+
+  XlaBuilder builder(TestName());
+  auto lhs_constant = builder.ConstantR2FromArray2D(*constant_lhs_array);
+  auto rhs_constant = builder.ConstantR2FromArray2D(*constant_rhs_array);
+  auto start_constant = builder.ConstantR1<int32>({0, 1});
+  auto dynamic_slice =
+      builder.DynamicSlice(rhs_constant, start_constant, {6, 1});
+
+  DotDimensionNumbers dot_dnums;
+  dot_dnums.add_lhs_contracting_dimensions(1);
+  dot_dnums.add_rhs_contracting_dimensions(0);
+  auto result = builder.DotGeneral(lhs_constant, dynamic_slice, dot_dnums);
+
+  Array2D<float> expected({{105.0}, {105.0}});
+  ComputeAndCompareR2<float>(&builder, expected, {}, error_spec_);
+}
+
+// TODO (b/69062148) Enable when Dot implements general contracting dimensions.
+TEST_F(DotOperationTest,
+       DISABLED_ON_CPU(DISABLED_ON_GPU(DISABLED_ON_INTERPRETER(
+           DotOfGatherOptimizationWithConstRHSReverseMM)))) {
+  std::unique_ptr<Array2D<float>> constant_lhs_array(
+      new Array2D<float>({{1.0, 2.0, 3.0},
+                          {4.0, 5.0, 6.0},
+                          {7.0, 8.0, 9.0},
+                          {9.0, 8.0, 7.0},
+                          {6.0, 5.0, 4.0},
+                          {3.0, 2.0, 1.0}}));
+  std::unique_ptr<Array2D<float>> constant_rhs_array(new Array2D<float>(
+      {{1.0, 2.0, 3.0, 4.0, 5.0, 6.0}, {6.0, 5.0, 4.0, 3.0, 2.0, 1.0}}));
+  // Dot result to slice from: {{114, 96}, {105, 105}, {96, 114}}
+
+  XlaBuilder builder(TestName());
+  auto lhs_constant = builder.ConstantR2FromArray2D(*constant_lhs_array);
+  auto rhs_constant = builder.ConstantR2FromArray2D(*constant_rhs_array);
+  auto start_constant = builder.ConstantR1<int32>({0, 1});
+  auto dynamic_slice =
+      builder.DynamicSlice(lhs_constant, start_constant, {6, 1});
+
+  DotDimensionNumbers dot_dnums;
+  dot_dnums.add_lhs_contracting_dimensions(0);
+  dot_dnums.add_rhs_contracting_dimensions(1);
+  auto result = builder.DotGeneral(dynamic_slice, rhs_constant, dot_dnums);
+
+  Array2D<float> expected({{105.0, 105.0}});
+  ComputeAndCompareR2<float>(&builder, expected, {}, error_spec_);
+}
+
+// TODO (b/69062148) Enable when Dot implements general contracting dimensions.
+TEST_F(DotOperationTest,
+       DISABLED_ON_CPU(DISABLED_ON_GPU(DISABLED_ON_INTERPRETER(
+           DotOfGatherOptimizationWithConstLHSReverseMM)))) {
+  std::unique_ptr<Array2D<float>> constant_lhs_array(
+      new Array2D<float>({{1.0, 2.0, 3.0},
+                          {4.0, 5.0, 6.0},
+                          {7.0, 8.0, 9.0},
+                          {9.0, 8.0, 7.0},
+                          {6.0, 5.0, 4.0},
+                          {3.0, 2.0, 1.0}}));
+  std::unique_ptr<Array2D<float>> constant_rhs_array(new Array2D<float>(
+      {{1.0, 2.0, 3.0, 4.0, 5.0, 6.0}, {6.0, 5.0, 4.0, 3.0, 2.0, 1.0}}));
+  // Dot result to slice from: {{114, 96}, {105, 105}, {96, 114}}
+
+  XlaBuilder builder(TestName());
+  auto lhs_constant = builder.ConstantR2FromArray2D(*constant_lhs_array);
+  auto rhs_constant = builder.ConstantR2FromArray2D(*constant_rhs_array);
+  auto start_constant = builder.ConstantR1<int32>({1, 0});
+  auto dynamic_slice =
+      builder.DynamicSlice(rhs_constant, start_constant, {1, 6});
+
+  DotDimensionNumbers dot_dnums;
+  dot_dnums.add_lhs_contracting_dimensions(0);
+  dot_dnums.add_rhs_contracting_dimensions(1);
+  auto result = builder.DotGeneral(lhs_constant, dynamic_slice, dot_dnums);
+
+  Array2D<float> expected({{96.0}, {105.0}, {114.0}});
+  ComputeAndCompareR2<float>(&builder, expected, {}, error_spec_);
+}
+
+// TODO (b/69062148) Enable when Dot implements general contracting dimensions.
+TEST_F(DotOperationTest,
+       DISABLED_ON_CPU(DISABLED_ON_GPU(
+           DISABLED_ON_INTERPRETER(DotOfGatherOptimizationWithConstRHSRows)))) {
+  std::unique_ptr<Array2D<float>> constant_lhs_array(
+      new Array2D<float>({{1.0, 2.0},
+                          {3.0, 4.0},
+                          {5.0, 6.0},
+                          {6.0, 5.0},
+                          {4.0, 3.0},
+                          {2.0, 1.0}}));
+  std::unique_ptr<Array2D<float>> constant_rhs_array(
+      new Array2D<float>({{1.0, 2.0, 3.0},
+                          {4.0, 5.0, 6.0},
+                          {7.0, 8.0, 9.0},
+                          {9.0, 8.0, 7.0},
+                          {6.0, 5.0, 4.0},
+                          {3.0, 2.0, 1.0}}));
+  // Dot result to slice from: {{132, 129, 126}, {126, 129, 132}}
+
+  XlaBuilder builder(TestName());
+  auto lhs_constant = builder.ConstantR2FromArray2D(*constant_lhs_array);
+  auto rhs_constant = builder.ConstantR2FromArray2D(*constant_rhs_array);
+  auto start_constant = builder.ConstantR1<int32>({0, 1});
+  auto dynamic_slice =
+      builder.DynamicSlice(lhs_constant, start_constant, {6, 1});
+
+  DotDimensionNumbers dot_dnums;
+  dot_dnums.add_lhs_contracting_dimensions(0);
+  dot_dnums.add_rhs_contracting_dimensions(0);
+  auto result = builder.DotGeneral(dynamic_slice, rhs_constant, dot_dnums);
+
+  Array2D<float> expected({{126.0, 129.0, 132.0}});
+  ComputeAndCompareR2<float>(&builder, expected, {}, error_spec_);
+}
+
+// TODO (b/69062148) Enable when Dot implements general contracting dimensions.
+TEST_F(DotOperationTest,
+       DISABLED_ON_CPU(DISABLED_ON_GPU(
+           DISABLED_ON_INTERPRETER(DotOfGatherOptimizationWithConstLHSRows)))) {
+  std::unique_ptr<Array2D<float>> constant_lhs_array(
+      new Array2D<float>({{1.0, 2.0},
+                          {3.0, 4.0},
+                          {5.0, 6.0},
+                          {6.0, 5.0},
+                          {4.0, 3.0},
+                          {2.0, 1.0}}));
+  std::unique_ptr<Array2D<float>> constant_rhs_array(
+      new Array2D<float>({{1.0, 2.0, 3.0},
+                          {4.0, 5.0, 6.0},
+                          {7.0, 8.0, 9.0},
+                          {9.0, 8.0, 7.0},
+                          {6.0, 5.0, 4.0},
+                          {3.0, 2.0, 1.0}}));
+  // Dot result to slice from: {{132, 129, 126}, {126, 129, 132}}
+
+  XlaBuilder builder(TestName());
+  auto lhs_constant = builder.ConstantR2FromArray2D(*constant_lhs_array);
+  auto rhs_constant = builder.ConstantR2FromArray2D(*constant_rhs_array);
+  auto start_constant = builder.ConstantR1<int32>({0, 1});
+  auto dynamic_slice =
+      builder.DynamicSlice(rhs_constant, start_constant, {6, 1});
+
+  DotDimensionNumbers dot_dnums;
+  dot_dnums.add_lhs_contracting_dimensions(0);
+  dot_dnums.add_rhs_contracting_dimensions(0);
+  auto result = builder.DotGeneral(lhs_constant, dynamic_slice, dot_dnums);
+
+  Array2D<float> expected({{129.0}, {129.0}});
+  ComputeAndCompareR2<float>(&builder, expected, {}, error_spec_);
+}
+
+// TODO (b/69062148) Enable when Dot implements general contracting dimensions.
+TEST_F(DotOperationTest,
+       DISABLED_ON_CPU(DISABLED_ON_GPU(
+           DISABLED_ON_INTERPRETER(DotOfGatherOptimizationWithConstRHSCols)))) {
+  std::unique_ptr<Array2D<float>> constant_lhs_array(new Array2D<float>(
+      {{1.0, 2.0, 3.0, 4.0, 5.0, 6.0}, {6.0, 5.0, 4.0, 3.0, 2.0, 1.0}}));
+  std::unique_ptr<Array2D<float>> constant_rhs_array(
+      new Array2D<float>({{1.0, 2.0, 3.0, 4.0, 5.0, 6.0},
+                          {7.0, 8.0, 9.0, 9.0, 8.0, 7.0},
+                          {6.0, 5.0, 4.0, 3.0, 2.0, 1.0}}));
+  // Dot result to slice from: {{91, 168, 56}, {56, 168, 91}}
+
+  XlaBuilder builder(TestName());
+  auto lhs_constant = builder.ConstantR2FromArray2D(*constant_lhs_array);
+  auto rhs_constant = builder.ConstantR2FromArray2D(*constant_rhs_array);
+  auto start_constant = builder.ConstantR1<int32>({1, 0});
+  auto dynamic_slice =
+      builder.DynamicSlice(lhs_constant, start_constant, {1, 6});
+
+  DotDimensionNumbers dot_dnums;
+  dot_dnums.add_lhs_contracting_dimensions(1);
+  dot_dnums.add_rhs_contracting_dimensions(1);
+  auto result = builder.DotGeneral(dynamic_slice, rhs_constant, dot_dnums);
+
+  Array2D<float> expected({{56.0, 168.0, 91.0}});
+  ComputeAndCompareR2<float>(&builder, expected, {}, error_spec_);
+}
+
+// TODO (b/69062148) Enable when Dot implements general contracting dimensions.
+TEST_F(DotOperationTest,
+       DISABLED_ON_CPU(DISABLED_ON_GPU(
+           DISABLED_ON_INTERPRETER(DotOfGatherOptimizationWithConstLHSCols)))) {
+  std::unique_ptr<Array2D<float>> constant_lhs_array(new Array2D<float>(
+      {{1.0, 2.0, 3.0, 4.0, 5.0, 6.0}, {6.0, 5.0, 4.0, 3.0, 2.0, 1.0}}));
+  std::unique_ptr<Array2D<float>> constant_rhs_array(
+      new Array2D<float>({{1.0, 2.0, 3.0, 4.0, 5.0, 6.0},
+                          {7.0, 8.0, 9.0, 9.0, 8.0, 7.0},
+                          {6.0, 5.0, 4.0, 3.0, 2.0, 1.0}}));
+  // Dot result to slice from: {{91, 168, 56}, {56, 168, 91}}
+
+  XlaBuilder builder(TestName());
+  auto lhs_constant = builder.ConstantR2FromArray2D(*constant_lhs_array);
+  auto rhs_constant = builder.ConstantR2FromArray2D(*constant_rhs_array);
+  auto start_constant = builder.ConstantR1<int32>({1, 0});
+  auto dynamic_slice =
+      builder.DynamicSlice(rhs_constant, start_constant, {1, 6});
+
+  DotDimensionNumbers dot_dnums;
+  dot_dnums.add_lhs_contracting_dimensions(1);
+  dot_dnums.add_rhs_contracting_dimensions(1);
+  auto result = builder.DotGeneral(lhs_constant, dynamic_slice, dot_dnums);
+
+  Array2D<float> expected({{168.0}, {168.0}});
+  ComputeAndCompareR2<float>(&builder, expected, {}, error_spec_);
+}
 }  // namespace
 }  // namespace xla
-- 
GitLab


From b15500be31f29850c73804b8694e4f0f01b82305 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 12:04:38 -0700
Subject: [PATCH 372/395] Remove outdated CUDA SDK string (the text is now
 consistent with other version choices, and the '9.0' format is already
 present in the default).

PiperOrigin-RevId: 195855416
---
 configure.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/configure.py b/configure.py
index b745e374a2..7d04d3a14f 100644
--- a/configure.py
+++ b/configure.py
@@ -845,8 +845,8 @@ def reformat_version_sequence(version_str, sequence_count):
 def set_tf_cuda_version(environ_cp):
   """Set CUDA_TOOLKIT_PATH and TF_CUDA_VERSION."""
   ask_cuda_version = (
-      'Please specify the CUDA SDK version you want to use, '
-      'e.g. 7.0. [Leave empty to default to CUDA %s]: ') % _DEFAULT_CUDA_VERSION
+      'Please specify the CUDA SDK version you want to use. '
+      '[Leave empty to default to CUDA %s]: ') % _DEFAULT_CUDA_VERSION
 
   for _ in range(_DEFAULT_PROMPT_ASK_ATTEMPTS):
     # Configure the Cuda SDK version to use.
-- 
GitLab


From 8fcf64732bb43d6df5df99171346e9de6c15e7ed Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 12:10:36 -0700
Subject: [PATCH 373/395] Better wrapping of stream executor's cuDNN API calls.
 Replacing mutex locking and setting the cuDNN stream followed by calling
 wrap::cudnn... with an RAII CudnnHandle object that handles the former two
 operations.

Distinguish three different API types:

A) APIs that don't take a cudnnHandle_t: These are thread-safe APIs that don't enqueue any CUDA work on a stream. They can be called directly without any extra precautions.

B) APIs that take a cudnnHandle_t and perform CUDA work. The CUDA context needs to be acquired and the stream needs to be set beforehand, calls need to be serialized. A CudnnHandle instance guarantees that this work has been performed before calling cuDNN.

C) APIs that do take a cudnnHandle_t, but (presumably, the API makes no guarantees) still don't perform any CUDA work. This is limited to the API to setup RNN descriptors. Calls need to be serialized, but most likely we wouldn't need to acquire the CUDA context or set the stream. We still do though using the legacy default stream, because there are no guarantees.

PiperOrigin-RevId: 195856300
---
 tensorflow/core/platform/default/mutex.h      |    4 +-
 .../stream_executor/cuda/cuda_activation.cc   |    6 +
 .../stream_executor/cuda/cuda_activation.h    |    3 +-
 tensorflow/stream_executor/cuda/cuda_dnn.cc   | 1448 +++++++----------
 tensorflow/stream_executor/cuda/cuda_dnn.h    |   51 +-
 5 files changed, 596 insertions(+), 916 deletions(-)

diff --git a/tensorflow/core/platform/default/mutex.h b/tensorflow/core/platform/default/mutex.h
index a12d92795e..89e57d58a0 100644
--- a/tensorflow/core/platform/default/mutex.h
+++ b/tensorflow/core/platform/default/mutex.h
@@ -77,9 +77,7 @@ class SCOPED_LOCKABLE mutex_lock {
 
   // Manually nulls out the source to prevent double-free.
   // (std::move does not null the source pointer by default.)
-  explicit mutex_lock(mutex_lock&& ml) noexcept : mu_(ml.mu_) {
-    ml.mu_ = nullptr;
-  }
+  mutex_lock(mutex_lock&& ml) noexcept : mu_(ml.mu_) { ml.mu_ = nullptr; }
   ~mutex_lock() UNLOCK_FUNCTION() {
     if (mu_ != nullptr) {
       mu_->unlock();
diff --git a/tensorflow/stream_executor/cuda/cuda_activation.cc b/tensorflow/stream_executor/cuda/cuda_activation.cc
index cf6b9e2c6e..02371c3c3a 100644
--- a/tensorflow/stream_executor/cuda/cuda_activation.cc
+++ b/tensorflow/stream_executor/cuda/cuda_activation.cc
@@ -38,5 +38,11 @@ ScopedActivateExecutorContext::~ScopedActivateExecutorContext() {
   delete static_cast<ScopedActivateContext *>(driver_scoped_activate_context_);
 }
 
+ScopedActivateExecutorContext::ScopedActivateExecutorContext(
+    ScopedActivateExecutorContext &&other)
+    : driver_scoped_activate_context_(other.driver_scoped_activate_context_) {
+  other.driver_scoped_activate_context_ = nullptr;
+}
+
 }  // namespace cuda
 }  // namespace stream_executor
diff --git a/tensorflow/stream_executor/cuda/cuda_activation.h b/tensorflow/stream_executor/cuda/cuda_activation.h
index 04ffaef364..ef9807820f 100644
--- a/tensorflow/stream_executor/cuda/cuda_activation.h
+++ b/tensorflow/stream_executor/cuda/cuda_activation.h
@@ -44,10 +44,11 @@ class ScopedActivateExecutorContext {
   // fatal failure if it is not CUDA inside.
   explicit ScopedActivateExecutorContext(StreamExecutor* stream_exec);
 
+  ScopedActivateExecutorContext(ScopedActivateExecutorContext&& other);
+
   ~ScopedActivateExecutorContext();
 
  private:
-
   // The cuda.h-using datatype that we wrap.
   ScopedActivateContext* driver_scoped_activate_context_;
 
diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.cc b/tensorflow/stream_executor/cuda/cuda_dnn.cc
index 316f4c4f1e..af78efe81d 100644
--- a/tensorflow/stream_executor/cuda/cuda_dnn.cc
+++ b/tensorflow/stream_executor/cuda/cuda_dnn.cc
@@ -46,8 +46,20 @@ limitations under the License.
 #include "cuda/include/cudnn.h"
 // clang-format on
 
+namespace stream_executor {
+namespace cuda {
+
+PLUGIN_REGISTRY_DEFINE_PLUGIN_ID(kCuDnnPlugin);
+
 namespace {
 
+// TODO(csigg): remove dnn namespace qualifier from the RNN code below.
+using ::stream_executor::dnn::BatchDescriptor;
+using ::stream_executor::dnn::ConvolutionDescriptor;
+using ::stream_executor::dnn::FilterDescriptor;
+using ::stream_executor::dnn::NormalizeDescriptor;
+using ::stream_executor::dnn::PoolingDescriptor;
+
 // Converts (via narrowing) a type T value to a type U, and checks that the
 // value has no value change due to the conversion.
 template <typename WideT, typename NarrowT>
@@ -58,20 +70,6 @@ NarrowT CheckedNarrowing(const WideT& wide) {
   return narrow;
 }
 
-}  // namespace
-
-namespace stream_executor {
-
-using dnn::BatchDescriptor;
-using dnn::FilterDescriptor;
-using dnn::ConvolutionDescriptor;
-using dnn::PoolingDescriptor;
-using dnn::NormalizeDescriptor;
-
-namespace cuda {
-
-PLUGIN_REGISTRY_DEFINE_PLUGIN_ID(kCuDnnPlugin);
-
 string ToString(cudnnStatus_t status) {
   switch (status) {
     case CUDNN_STATUS_SUCCESS:
@@ -136,208 +134,82 @@ cudnnDataType_t GetCudnnDataType<Eigen::half>() {
   return CUDNN_DATA_HALF;
 }
 
-namespace wrap {
-
-#define STREAM_EXECUTOR_CUDNN_WRAP(__name)                         \
-  struct WrapperShim__##__name {                                   \
-    template <typename... Args>                                    \
-    cudnnStatus_t operator()(CUDAExecutor* parent, Args... args) { \
-      cuda::ScopedActivateExecutorContext sac{parent};             \
-      cudnnStatus_t retval = ::__name(args...);                    \
-      return retval;                                               \
-    }                                                              \
-  } __name;
-
-#define STREAM_EXECUTOR_CUDNN_WRAP_WITH_CHECKED_STREAM(__name)           \
-  struct WrapperShim__##__name {                                         \
-    template <typename... Args>                                          \
-    cudnnStatus_t operator()(CudnnSupport* dnn, Stream* s, Args... args) \
-        SHARED_LOCKS_REQUIRED(dnn->dnn_handle_mutex_) {                  \
-      CHECK_NOTNULL(s);                                                  \
-      CHECK_EQ(s, dnn->GetCurrentDnnStream())                            \
-          << "Stream is not set correctly!";                             \
-      cuda::ScopedActivateExecutorContext sac{dnn->GetParentExecutor()}; \
-      cudnnStatus_t retval = ::__name(args...);                          \
-      return retval;                                                     \
-    }                                                                    \
-  } __name;
-
-// Handles cudnnSetStream differently in order to add debug information.
-// It stores a reference to 'stream' in 'dnn', and checks that all calls from
-// that dnn instance use the same stream (see
-// STREAM_EXECUTOR_CUDNN_WRAP_WITH_CHECKED_STREAM macro).
-struct WrapperShim__cudnnSetStream {
-  cudnnStatus_t operator()(CudnnSupport* dnn, Stream* stream,
-                           cudnnHandle_t handle)
-      EXCLUSIVE_LOCKS_REQUIRED(dnn->dnn_handle_mutex_) {
-    dnn->SetCurrentDnnStream(stream);
-    cuda::ScopedActivateExecutorContext sac{dnn->GetParentExecutor()};
-    cudnnStatus_t retval = ::cudnnSetStream(handle, AsCUDAStreamValue(stream));
-    return retval;
-  }
-} cudnnSetStream;
-
-// clang-format off
-#define CUDNN_DNN_ROUTINE_EACH(__macro)                   \
-  __macro(cudnnGetConvolutionNdForwardOutputDim)          \
-  __macro(cudnnGetConvolutionForwardAlgorithm)            \
-  __macro(cudnnCreateTensorDescriptor)                    \
-  __macro(cudnnDestroyTensorDescriptor)                   \
-  __macro(cudnnCreateFilterDescriptor)                    \
-  __macro(cudnnSetPoolingNdDescriptor)                    \
-  __macro(cudnnSetLRNDescriptor)                          \
-  __macro(cudnnDestroyFilterDescriptor)                   \
-  __macro(cudnnCreateConvolutionDescriptor)               \
-  __macro(cudnnCreatePoolingDescriptor)                   \
-  __macro(cudnnDestroyPoolingDescriptor)                  \
-  __macro(cudnnCreateLRNDescriptor)                       \
-  __macro(cudnnDestroyLRNDescriptor)                      \
-  __macro(cudnnDestroyConvolutionDescriptor)              \
-  __macro(cudnnCreate)                                    \
-  __macro(cudnnDestroy)                                   \
-  __macro(cudnnGetConvolutionForwardWorkspaceSize)        \
-  __macro(cudnnSetConvolutionNdDescriptor)                \
-  __macro(cudnnSetTensor4dDescriptor)                     \
-  __macro(cudnnSetTensorNdDescriptor)                     \
-  __macro(cudnnSetFilterNdDescriptor)
-
-// clang-format on
-CUDNN_DNN_ROUTINE_EACH(STREAM_EXECUTOR_CUDNN_WRAP)
-#undef CUDNN_DNN_ROUTINE_EACH
-
-// clang-format off
-#define CUDNN_DNN_ROUTINE_EACH_WITH_STREAM(__macro)       \
-  __macro(cudnnBatchNormalizationBackward)                \
-  __macro(cudnnBatchNormalizationForwardInference)        \
-  __macro(cudnnBatchNormalizationForwardTraining)         \
-  __macro(cudnnActivationForward)                         \
-  __macro(cudnnConvolutionForward)                        \
-  __macro(cudnnConvolutionBackwardBias)                   \
-  __macro(cudnnTransformTensor)                           \
-  __macro(cudnnPoolingForward)                            \
-  __macro(cudnnPoolingBackward)                           \
-  __macro(cudnnLRNCrossChannelForward)                    \
-  __macro(cudnnLRNCrossChannelBackward)                   \
-  __macro(cudnnAddTensor)                                 \
-  __macro(cudnnConvolutionBackwardData)                   \
-  __macro(cudnnConvolutionBackwardFilter)
-
-// clang-format on
-CUDNN_DNN_ROUTINE_EACH_WITH_STREAM(
-    STREAM_EXECUTOR_CUDNN_WRAP_WITH_CHECKED_STREAM)
-#undef CUDNN_DNN_ROUTINE_EACH_WITH_STREAM
-
-// APIs available after R3:
-#if CUDNN_VERSION >= 3000
-#define CUDNN_DNN_ROUTINE_EACH_AFTER_R3(__macro)              \
-  __macro(cudnnGetConvolutionBackwardFilterWorkspaceSize)     \
-  __macro(cudnnGetConvolutionBackwardDataAlgorithm)           \
-  __macro(cudnnGetConvolutionBackwardFilterAlgorithm)         \
-  __macro(cudnnGetConvolutionBackwardDataWorkspaceSize)
-CUDNN_DNN_ROUTINE_EACH_AFTER_R3(STREAM_EXECUTOR_CUDNN_WRAP)
-#undef CUDNN_DNN_ROUTINE_EACH_AFTER_R3
-#endif
-
-// APIs in R3 but not in R5
-// clang-format off
-#if CUDNN_VERSION >= 3000 && CUDNN_VERSION < 5000
-#define CUDNN_DNN_ROUTINE_EACH_R3_WITH_STREAM(__macro)        \
-  __macro(cudnnAddTensor_v3)                                  \
-  __macro(cudnnConvolutionBackwardData_v3)                    \
-  __macro(cudnnConvolutionBackwardFilter_v3)
-// clang-format on
-
-CUDNN_DNN_ROUTINE_EACH_R3_WITH_STREAM(
-    STREAM_EXECUTOR_CUDNN_WRAP_WITH_CHECKED_STREAM)
-#undef CUDNN_DNN_ROUTINE_EACH_R3_WITH_STREAM
-#endif
-
-// APIs in R5
-// clang-format off
-#if CUDNN_VERSION >= 5000
-#define CUDNN_DNN_ROUTINE_EACH_R5(__macro)                    \
-  __macro(cudnnCreateActivationDescriptor)                    \
-  __macro(cudnnSetActivationDescriptor)                       \
-  __macro(cudnnGetActivationDescriptor)                       \
-  __macro(cudnnDestroyActivationDescriptor)                   \
-  __macro(cudnnCreateDropoutDescriptor)                       \
-  __macro(cudnnDestroyDropoutDescriptor)                      \
-  __macro(cudnnSetDropoutDescriptor)                          \
-  __macro(cudnnDropoutGetStatesSize)                          \
-  __macro(cudnnCreateRNNDescriptor)                           \
-  __macro(cudnnDestroyRNNDescriptor)                          \
-  __macro(cudnnGetRNNParamsSize)                              \
-  __macro(cudnnGetRNNWorkspaceSize)                           \
-  __macro(cudnnGetRNNTrainingReserveSize)                     \
-  __macro(cudnnGetRNNLinLayerMatrixParams)                    \
-  __macro(cudnnGetRNNLinLayerBiasParams)                      \
-  __macro(cudnnSetRNNDescriptor)                              \
-  __macro(cudnnGetFilterNdDescriptor)
-
-// clang-format on
-CUDNN_DNN_ROUTINE_EACH_R5(STREAM_EXECUTOR_CUDNN_WRAP)
-#undef CUDNN_DNN_ROUTINE_EACH_R5
-
-// clang-format off
-#define CUDNN_DNN_ROUTINE_EACH_R5_WITH_STREAM(__macro)        \
-  __macro(cudnnRNNForwardInference)                           \
-  __macro(cudnnRNNForwardTraining)                            \
-  __macro(cudnnRNNBackwardData)                               \
-  __macro(cudnnRNNBackwardWeights)
+// RAII wrapper for all calls to cuDNN with a cuDNN handle argument.
+//
+// See CudnnAccess::GetHandle() for details.
+class CudnnHandle {
+ public:
+  // Takes ownership of the executor context and the lock to access cuDNN
+  // using handle.
+  CudnnHandle(cuda::ScopedActivateExecutorContext context, mutex_lock lock,
+              cudnnHandle_t handle)
+      : context_(std::move(context)), lock_(std::move(lock)), handle_(handle) {}
 
-// clang-format on
-CUDNN_DNN_ROUTINE_EACH_R5_WITH_STREAM(
-    STREAM_EXECUTOR_CUDNN_WRAP_WITH_CHECKED_STREAM)
-#undef CUDNN_DNN_ROUTINE_EACH_R5_WITH_STREAM
-#endif
+  // Returns cuDNN handle. To be passed directly to cuDNN APIs, don't keep
+  // a copy.
+  cudnnHandle_t handle() const { return handle_; }
 
-// APIs in R6
-// clang-format off
-#if CUDNN_VERSION >= 6000
-#define CUDNN_DNN_ROUTINE_EACH_R6(__macro)                    \
-  __macro(cudnnSetRNNDescriptor_v6)                           \
-  __macro(cudnnCreatePersistentRNNPlan)                       \
-  __macro(cudnnDestroyPersistentRNNPlan)                      \
-  __macro(cudnnSetPersistentRNNPlan)
+ private:
+  cuda::ScopedActivateExecutorContext context_;
+  mutex_lock lock_;
+  cudnnHandle_t handle_;  // Not owned.
+};
 
-// clang-format on
-CUDNN_DNN_ROUTINE_EACH_R6(STREAM_EXECUTOR_CUDNN_WRAP)
-#undef CUDNN_DNN_ROUTINE_EACH_R6
+}  // namespace
 
-// clang-format off
-#define CUDNN_DNN_ROUTINE_EACH_R6_WITH_STREAM(__macro)        \
-  __macro(cudnnConvolutionBiasActivationForward)
+// Wraps a cuDNN handle and provides access to it through CudnnHandle instances,
+// which also locks a mutex, acquires the CUDA context, and sets the stream
+// that cuDNN should use to enqueue any work.
+//
+// Note: CudnnSupport::cudnn_ should be the only instantiation of this class.
+class CudnnAccess {
+ public:
+  // Takes ownership of the handle.
+  explicit CudnnAccess(cudnnHandle_t handle) : handle_(handle) {}
 
-// clang-format on
-CUDNN_DNN_ROUTINE_EACH_R6_WITH_STREAM(
-    STREAM_EXECUTOR_CUDNN_WRAP_WITH_CHECKED_STREAM)
-#undef CUDNN_DNN_ROUTINE_EACH_R6_WITH_STREAM
-#endif
+  ~CudnnAccess() {
+    mutex_lock lock(mutex_);
+    cudnnDestroy(handle_);
+  }
 
-// APIs in R7
-// clang-format off
-#if CUDNN_VERSION >= 7000
-#define CUDNN_DNN_ROUTINE_EACH_R7(__macro)                    \
-  __macro(cudnnSetConvolutionMathType)                        \
-  __macro(cudnnSetRNNMatrixMathType)                          \
-  __macro(cudnnSetConvolutionGroupCount)                      \
-  __macro(cudnnGetConvolutionGroupCount)
+  // Creates a CudnnHandle instance for stream.
+  //
+  // cuDNN API calls using the same handle instance need to be serialized across
+  // threads. This is guaranteed by CudnnHandle instances locking the mutex
+  // owned by this class.
+  //
+  // Most cuDNN APIs taking a handle perform work on a CUDA stream. The
+  // CudnnHandle instance acquires the executor's CUDA context and sets cuDNN to
+  // use the provided stream.
+  //
+  // The stream argument may be null, which translates to the legacy default
+  // stream. See
+  // https://docs.nvidia.com/cuda/cuda-driver-api/stream-sync-behavior.html.
+  // The legacy default stream synchronizes with all other streams and it is
+  // therefore a bad idea (performance wise) to call any cuDNN APIs that
+  // enqueue work in the stream.
+  CudnnHandle GetHandle(CUDAExecutor* executor, Stream* stream) {
+    mutex_lock lock(mutex_);
+    cuda::ScopedActivateExecutorContext context(executor);
+    CUstream cu_stream = stream ? AsCUDAStreamValue(stream) : cudaStreamLegacy;
+    auto status = cudnnSetStream(handle_, cu_stream);
+    CHECK_EQ(status, CUDNN_STATUS_SUCCESS) << "Failed to set cuDNN stream.";
+    using my_mutex_lock = mutex_lock;
+    return CudnnHandle(std::move(context), std::move(lock), handle_);
+  }
 
-// clang-format on
-CUDNN_DNN_ROUTINE_EACH_R7(STREAM_EXECUTOR_CUDNN_WRAP)
-#undef CUDNN_DNN_ROUTINE_EACH_R7
-#endif
+ private:
+  // Guards the enqueueing of cuDNN operations via the handle_ below.
+  mutex mutex_;
 
-}  // namespace wrap
+  // cuDNN library handle.
+  cudnnHandle_t handle_ GUARDED_BY(mutex_);  // Owned.
+};
 
 namespace {
 
 cudnnDataType_t GetRnnComputeType(dnn::DataType data_type);
 
-cudnnHandle_t ToHandle(void* opaque_handle) {
-  return static_cast<cudnnHandle_t>(opaque_handle);
-}
-
 cudnnConvolutionFwdAlgo_t ToConvForwardAlgo(dnn::AlgorithmDesc algorithm) {
   cudnnConvolutionFwdAlgo_t algo =
       cudnnConvolutionFwdAlgo_t(algorithm.algo_id());
@@ -414,7 +286,7 @@ port::Status GetCudnnProperty(libraryPropertyType type, int* value) {
         port::StrCat("cudnnGetProperty failed for type: ", ToString(type),
                      " with status: ", ToString(status));
     LOG(ERROR) << error;
-    return port::Status{port::error::INTERNAL, error};
+    return port::Status(port::error::INTERNAL, error);
   }
   return port::Status::OK();
 }
@@ -453,19 +325,11 @@ port::Status GetLoadedCudnnVersion(CudnnVersion* version) {
 
 }  // namespace
 
-CudnnSupport::CudnnSupport(CUDAExecutor* parent)
-    : parent_(parent), dnn_handle_(nullptr), current_dnn_stream_(nullptr) {}
-
-CudnnSupport::~CudnnSupport() {
-  auto status = wrap::cudnnDestroy(parent_, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(ERROR) << "could not destroy cudnn handle: " << ToString(status);
-  }
-}
+CudnnSupport::CudnnSupport(CUDAExecutor* parent) : parent_(parent) {}
 
 port::Status CudnnSupport::Init() {
-  auto status = wrap::cudnnCreate(
-      parent_, reinterpret_cast<cudnnHandle_t*>(&dnn_handle_));
+  cudnnHandle_t cudnn_handle = nullptr;
+  auto status = cudnnCreate(&cudnn_handle);
   if (status == CUDNN_STATUS_SUCCESS) {
     CudnnVersion source_version(CUDNN_MAJOR, CUDNN_MINOR, CUDNN_PATCHLEVEL);
 
@@ -481,9 +345,10 @@ port::Status CudnnSupport::Init() {
           "from sources, make sure the library loaded at runtime is compatible "
           "with the version specified during compile configuration.");
       LOG(ERROR) << error;
-      return port::Status{port::error::INTERNAL, error};
+      return port::Status(port::error::INTERNAL, error);
     }
 
+    cudnn_.reset(new CudnnAccess(cudnn_handle));
     return port::Status::OK();
   }
 
@@ -507,9 +372,9 @@ port::Status CudnnSupport::Init() {
     }
   }
 
-  return port::Status{port::error::INTERNAL,
+  return port::Status(port::error::INTERNAL,
                       port::StrCat("cudnn library could not create a handle: ",
-                                   ToString(status))};
+                                   ToString(status)));
 }
 
 port::StatusOr<perftools::gputools::dnn::VersionInfo>
@@ -520,14 +385,15 @@ CudnnSupport::GetVersion() {
       version.major_version, version.minor_version, version.patch_level);
 }
 
+namespace {
+
 // Turns a BatchDescriptor structure into a cudnn tensor handle within a scope.
 class ScopedTensorDescriptor {
  public:
-  ScopedTensorDescriptor(CUDAExecutor* parent,
-                         const BatchDescriptor& batch_descriptor,
+  ScopedTensorDescriptor(const BatchDescriptor& batch_descriptor,
                          cudnnDataType_t elem_type)
-      : parent_(parent), handle_(nullptr) {
-    cudnnStatus_t status = wrap::cudnnCreateTensorDescriptor(parent_, &handle_);
+      : handle_(nullptr) {
+    cudnnStatus_t status = cudnnCreateTensorDescriptor(&handle_);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(FATAL) << "could not create cudnn tensor descriptor: "
                  << ToString(status);
@@ -550,8 +416,8 @@ class ScopedTensorDescriptor {
                        &CheckedNarrowing<int64, int>);
         std::transform(dims64.cbegin(), dims64.cend(), dims.begin(),
                        &CheckedNarrowing<int64, int>);
-        status = wrap::cudnnSetTensorNdDescriptor(
-            parent_, handle_, elem_type, nd, dims.data(), strides.data());
+        status = cudnnSetTensorNdDescriptor(handle_, elem_type, nd, dims.data(),
+                                            strides.data());
 
         if (status != CUDNN_STATUS_SUCCESS) {
           LOG(FATAL) << "could not convert BatchDescriptor "
@@ -561,8 +427,8 @@ class ScopedTensorDescriptor {
       } break;
 #if CUDNN_VERSION >= 6000
       case dnn::DataLayout::kBatchDepthYX4: {
-        status = wrap::cudnnSetTensor4dDescriptor(
-            parent_, handle_, CUDNN_TENSOR_NCHW_VECT_C, elem_type,
+        status = cudnnSetTensor4dDescriptor(
+            handle_, CUDNN_TENSOR_NCHW_VECT_C, elem_type,
             batch_descriptor.count(), batch_descriptor.feature_map_count(),
             batch_descriptor.height(), batch_descriptor.width());
         if (status != CUDNN_STATUS_SUCCESS) {
@@ -580,7 +446,7 @@ class ScopedTensorDescriptor {
   }
 
   ~ScopedTensorDescriptor() {
-    cudnnStatus_t status = wrap::cudnnDestroyTensorDescriptor(parent_, handle_);
+    cudnnStatus_t status = cudnnDestroyTensorDescriptor(handle_);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(ERROR) << "could not destroy cudnn tensor descriptor: "
                  << ToString(status);
@@ -590,7 +456,6 @@ class ScopedTensorDescriptor {
   cudnnTensorDescriptor_t handle() const { return handle_; }
 
  private:
-  CUDAExecutor* parent_;            // Parent executor. Not owned.
   cudnnTensorDescriptor_t handle_;  // Owned.
 
   SE_DISALLOW_COPY_AND_ASSIGN(ScopedTensorDescriptor);
@@ -599,12 +464,10 @@ class ScopedTensorDescriptor {
 // Turns a FilterDescriptor structure into a cudnn filter handle within a scope.
 class ScopedFilterDescriptor {
  public:
-  ScopedFilterDescriptor(CUDAExecutor* parent,
-                         const FilterDescriptor& filter_descriptor,
-                         const BatchDescriptor& batch_descriptor,
+  ScopedFilterDescriptor(const FilterDescriptor& filter_descriptor,
                          cudnnDataType_t elem_type)
-      : parent_(parent), handle_(nullptr) {
-    cudnnStatus_t status = wrap::cudnnCreateFilterDescriptor(parent_, &handle_);
+      : handle_(nullptr) {
+    cudnnStatus_t status = cudnnCreateFilterDescriptor(&handle_);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(FATAL) << "could not create cudnn filter descriptor: "
                  << ToString(status);
@@ -638,11 +501,11 @@ class ScopedFilterDescriptor {
     const auto& spatial_dims = filter_descriptor.input_filter_dims();
     std::copy(spatial_dims.begin(), spatial_dims.end(), dims.begin() + 2);
 
-    status = wrap::cudnnSetFilterNdDescriptor(parent_, handle_, elem_type,
+    status = cudnnSetFilterNdDescriptor(handle_, elem_type,
 #if CUDNN_VERSION >= 5000
-                                              format,
+                                        format,
 #endif
-                                              dims.size(), dims.data());
+                                        dims.size(), dims.data());
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(FATAL) << "could not set cudnn filter descriptor: "
                  << ToString(status);
@@ -650,7 +513,7 @@ class ScopedFilterDescriptor {
   }
 
   ~ScopedFilterDescriptor() {
-    cudnnStatus_t status = wrap::cudnnDestroyFilterDescriptor(parent_, handle_);
+    cudnnStatus_t status = cudnnDestroyFilterDescriptor(handle_);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(ERROR) << "could not destroy cudnn filter descriptor: "
                  << ToString(status);
@@ -660,11 +523,7 @@ class ScopedFilterDescriptor {
   cudnnFilterDescriptor_t handle() const { return handle_; }
 
  private:
-  // Parent executor object. Not owned.
-  CUDAExecutor* parent_;
-
-  // cudnn filter descriptor this object creates. Owned.
-  cudnnFilterDescriptor_t handle_;
+  cudnnFilterDescriptor_t handle_;  // Owned.
 
   SE_DISALLOW_COPY_AND_ASSIGN(ScopedFilterDescriptor);
 };
@@ -718,11 +577,10 @@ static bool BatchnormSpatialPersistentEnabled() {
 class ScopedConvolutionDescriptor {
  public:
   ScopedConvolutionDescriptor(
-      CUDAExecutor* parent, const ConvolutionDescriptor& convolution_descriptor,
+      const ConvolutionDescriptor& convolution_descriptor,
       cudnnDataType_t data_type)
-      : parent_(parent), handle_(nullptr) {
-    cudnnStatus_t status =
-        wrap::cudnnCreateConvolutionDescriptor(parent_, &handle_);
+      : handle_(nullptr) {
+    cudnnStatus_t status = cudnnCreateConvolutionDescriptor(&handle_);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(FATAL) << "could not create cudnn convolution descriptor: "
                  << ToString(status);
@@ -748,9 +606,9 @@ class ScopedConvolutionDescriptor {
     std::transform(dilations64.cbegin(), dilations64.cend(), dilations.begin(),
                    &CheckedNarrowing<int64, int>);
 
-    status = wrap::cudnnSetConvolutionNdDescriptor(
-        parent_, handle_, convolution_descriptor.ndims(), padding.data(),
-        strides.data(), dilations.data(),
+    status = cudnnSetConvolutionNdDescriptor(
+        handle_, convolution_descriptor.ndims(), padding.data(), strides.data(),
+        dilations.data(),
         // NOTE(keveman): cuDNN supports convolution and cross correlation.
         // However, almost all the use cases do cross correlation, so just
         // hard coding it here.
@@ -767,8 +625,8 @@ class ScopedConvolutionDescriptor {
 #if CUDNN_MAJOR >= 7
     VLOG(2) << "Requesting grouped convolution: "
             << convolution_descriptor.group_count();
-    status = wrap::cudnnSetConvolutionGroupCount(
-        parent_, handle_, convolution_descriptor.group_count());
+    status = cudnnSetConvolutionGroupCount(
+        handle_, convolution_descriptor.group_count());
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(FATAL) << "could not set cudnn convolution group count: "
                  << ToString(status);
@@ -784,8 +642,7 @@ class ScopedConvolutionDescriptor {
     cudnnMathType_t math_type =
         (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH);
     if (TensorOpMathEnabled()) {
-      cudnnStatus_t status =
-          wrap::cudnnSetConvolutionMathType(parent_, handle_, math_type);
+      cudnnStatus_t status = cudnnSetConvolutionMathType(handle_, math_type);
       if (status != CUDNN_STATUS_SUCCESS) {
         LOG(FATAL) << "could not set cudnn convolution math type: "
                    << ToString(status);
@@ -795,8 +652,7 @@ class ScopedConvolutionDescriptor {
   }
 
   ~ScopedConvolutionDescriptor() {
-    cudnnStatus_t status =
-        wrap::cudnnDestroyConvolutionDescriptor(parent_, handle_);
+    cudnnStatus_t status = cudnnDestroyConvolutionDescriptor(handle_);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(ERROR) << "could not destroy cudnn convolution descriptor: "
                  << ToString(status);
@@ -806,7 +662,6 @@ class ScopedConvolutionDescriptor {
   cudnnConvolutionDescriptor_t handle() const { return handle_; }
 
  private:
-  CUDAExecutor* parent_;                 // Parent executor. Not owned.
   cudnnConvolutionDescriptor_t handle_;  // Owned.
 
   SE_DISALLOW_COPY_AND_ASSIGN(ScopedConvolutionDescriptor);
@@ -816,11 +671,9 @@ class ScopedConvolutionDescriptor {
 // within a scope.
 class ScopedPoolingDescriptor {
  public:
-  ScopedPoolingDescriptor(CUDAExecutor* parent,
-                          const PoolingDescriptor& pooling_descriptor)
-      : parent_(parent), handle_(nullptr) {
-    cudnnStatus_t status =
-        wrap::cudnnCreatePoolingDescriptor(parent_, &handle_);
+  explicit ScopedPoolingDescriptor(const PoolingDescriptor& pooling_descriptor)
+      : handle_(nullptr) {
+    cudnnStatus_t status = cudnnCreatePoolingDescriptor(&handle_);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(FATAL) << "could not create cudnn pooling descriptor: "
                  << ToString(status);
@@ -840,8 +693,8 @@ class ScopedPoolingDescriptor {
     std::transform(shape64.cbegin(), shape64.cend(), shape.begin(),
                    &CheckedNarrowing<int64, int>);
     bool propagate_nans = pooling_descriptor.propagate_nans();
-    status = wrap::cudnnSetPoolingNdDescriptor(
-        parent_, handle_,
+    status = cudnnSetPoolingNdDescriptor(
+        handle_,
         (pooling_descriptor.mode() == dnn::PoolingMode::kMaximum
              ? CUDNN_POOLING_MAX
              : CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING),
@@ -855,8 +708,7 @@ class ScopedPoolingDescriptor {
     }
   }
   ~ScopedPoolingDescriptor() {
-    cudnnStatus_t status =
-        wrap::cudnnDestroyPoolingDescriptor(parent_, handle_);
+    cudnnStatus_t status = cudnnDestroyPoolingDescriptor(handle_);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(ERROR) << "could not destroy cudnn pooling descriptor: "
                  << ToString(status);
@@ -866,7 +718,6 @@ class ScopedPoolingDescriptor {
   cudnnPoolingDescriptor_t handle() const { return handle_; }
 
  private:
-  CUDAExecutor* parent_;             // Parent executor. Not owned.
   cudnnPoolingDescriptor_t handle_;  // Owned.
 
   SE_DISALLOW_COPY_AND_ASSIGN(ScopedPoolingDescriptor);
@@ -875,10 +726,10 @@ class ScopedPoolingDescriptor {
 // Turns a NormalizeDescriptor structure into a cudnn LRN descriptor handle.
 class ScopedNormalizeDescriptor {
  public:
-  ScopedNormalizeDescriptor(CUDAExecutor* parent,
-                            const NormalizeDescriptor& normalize_descriptor)
-      : parent_(parent), handle_(nullptr) {
-    cudnnStatus_t status = wrap::cudnnCreateLRNDescriptor(parent_, &handle_);
+  explicit ScopedNormalizeDescriptor(
+      const NormalizeDescriptor& normalize_descriptor)
+      : handle_(nullptr) {
+    cudnnStatus_t status = cudnnCreateLRNDescriptor(&handle_);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(FATAL) << "could not create cudnn LRN descriptor: "
                  << ToString(status);
@@ -904,15 +755,14 @@ class ScopedNormalizeDescriptor {
 
     double lrnBeta = normalize_descriptor.beta();
     double lrnK = normalize_descriptor.bias();
-    status = wrap::cudnnSetLRNDescriptor(parent_, handle_, lrnN, lrnAlpha,
-                                         lrnBeta, lrnK);
+    status = cudnnSetLRNDescriptor(handle_, lrnN, lrnAlpha, lrnBeta, lrnK);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(FATAL) << "could not set cudnn LRN descriptor: " << ToString(status);
     }
   }
 
   ~ScopedNormalizeDescriptor() {
-    cudnnStatus_t status = wrap::cudnnDestroyLRNDescriptor(parent_, handle_);
+    cudnnStatus_t status = cudnnDestroyLRNDescriptor(handle_);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(ERROR) << "could not destroy cudnn LRN descriptor: "
                  << ToString(status);
@@ -922,7 +772,6 @@ class ScopedNormalizeDescriptor {
   cudnnLRNDescriptor_t handle() const { return handle_; }
 
  private:
-  CUDAExecutor* parent_;         // Parent executor. Not owned.
   cudnnLRNDescriptor_t handle_;  // Owned.
 
   SE_DISALLOW_COPY_AND_ASSIGN(ScopedNormalizeDescriptor);
@@ -933,13 +782,11 @@ class ScopedNormalizeDescriptor {
 // descriptor handle within a scope.
 class ScopedActivationDescriptor {
  public:
-  ScopedActivationDescriptor(CUDAExecutor* parent,
-                             dnn::ActivationMode activation_mode,
+  ScopedActivationDescriptor(dnn::ActivationMode activation_mode,
                              cudnnNanPropagation_t nan_propagation,
                              double value_max)
-      : parent_(parent), handle_(nullptr) {
-    cudnnStatus_t status =
-        wrap::cudnnCreateActivationDescriptor(parent_, &handle_);
+      : handle_(nullptr) {
+    cudnnStatus_t status = cudnnCreateActivationDescriptor(&handle_);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(FATAL) << "could not create cudnn activation descriptor: "
                  << ToString(status);
@@ -970,8 +817,8 @@ class ScopedActivationDescriptor {
                    << static_cast<int>(activation_mode);
     }
 
-    status = wrap::cudnnSetActivationDescriptor(parent_, handle_, mode,
-                                                nan_propagation, relu_ceiling);
+    status = cudnnSetActivationDescriptor(handle_, mode, nan_propagation,
+                                          relu_ceiling);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(FATAL) << "could not set cudnn activation descriptor: "
                  << ToString(status);
@@ -979,8 +826,7 @@ class ScopedActivationDescriptor {
   }
 
   ~ScopedActivationDescriptor() {
-    cudnnStatus_t status =
-        wrap::cudnnDestroyActivationDescriptor(parent_, handle_);
+    cudnnStatus_t status = cudnnDestroyActivationDescriptor(handle_);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(ERROR) << "could not destroy cudnn activation descriptor: "
                  << ToString(status);
@@ -990,14 +836,12 @@ class ScopedActivationDescriptor {
   cudnnActivationDescriptor_t handle() const { return handle_; }
 
  private:
-  CUDAExecutor* parent_;                // Parent executor. Not owned.
   cudnnActivationDescriptor_t handle_;  // Owned.
 
   SE_DISALLOW_COPY_AND_ASSIGN(ScopedActivationDescriptor);
 };
 #endif
 
-namespace {
 cudnnDataType_t ToCudnnDataType(
     dnn::DataType data_type,
     dnn::DataLayout data_layout = dnn::DataLayout::kBatchDepthYX) {
@@ -1072,8 +916,6 @@ class MixinBase : public Base {};
 template <>
 class MixinBase<void> {};
 
-}  // namespace
-
 #if CUDNN_VERSION >= 5000
 
 #define CUDNN_RETURN_IF_FAIL(STATUS, ...)                                \
@@ -1084,6 +926,7 @@ class MixinBase<void> {};
     return;                                                              \
   }
 
+// TODO(csigg): Remove inheritance for code reuse.
 template <typename Base>
 class CudnnDescriptorCommon : public MixinBase<Base> {
  public:
@@ -1097,12 +940,11 @@ class CudnnDescriptorCommon : public MixinBase<Base> {
 
 class CudnnDropoutDescriptor : public CudnnDescriptorCommon<void> {
  public:
-  CudnnDropoutDescriptor(CUDAExecutor* parent, cudnnHandle_t cudnn_handle,
-                         float dropout, uint64 seed,
+  CudnnDropoutDescriptor(const CudnnHandle& cudnn, float dropout, uint64 seed,
                          ScratchAllocator* state_allocator)
-      : parent_(parent), handle_(nullptr) {
+      : handle_(nullptr) {
     cudnnStatus_t status;
-    status = wrap::cudnnCreateDropoutDescriptor(parent_, &handle_);
+    status = cudnnCreateDropoutDescriptor(&handle_);
     CUDNN_RETURN_IF_FAIL(status, "Failed to create dropout descriptor");
 
     if (dropout == 0.f) {
@@ -1112,8 +954,7 @@ class CudnnDropoutDescriptor : public CudnnDescriptorCommon<void> {
     DeviceMemory<uint8> state_memory;
     if (state_allocator) {
       size_t state_sizes_in_bytes = 0;
-      status = wrap::cudnnDropoutGetStatesSize(parent_, cudnn_handle,
-                                               &state_sizes_in_bytes);
+      status = cudnnDropoutGetStatesSize(cudnn.handle(), &state_sizes_in_bytes);
       CUDNN_RETURN_IF_FAIL(status, "Failed to query dropout state sizes");
 
       auto allocated =
@@ -1128,9 +969,9 @@ class CudnnDropoutDescriptor : public CudnnDescriptorCommon<void> {
         return;
       }
     }
-    status = wrap::cudnnSetDropoutDescriptor(parent_, handle_, cudnn_handle,
-                                             dropout, state_memory.opaque(),
-                                             state_memory.size(), seed);
+    status = cudnnSetDropoutDescriptor(handle_, cudnn.handle(), dropout,
+                                       state_memory.opaque(),
+                                       state_memory.size(), seed);
     CUDNN_RETURN_IF_FAIL(
         status, port::StrCat(
                     "Failed to set dropout descriptor with state memory size: ",
@@ -1138,11 +979,9 @@ class CudnnDropoutDescriptor : public CudnnDescriptorCommon<void> {
   }
 
   ~CudnnDropoutDescriptor() {
-    if (handle_) {
-      cudnnStatus_t status =
-          wrap::cudnnDestroyDropoutDescriptor(parent_, handle_);
-      CUDNN_RETURN_IF_FAIL(status, "Failed to destroy Cudnn dropout handle: ");
-    }
+    cudnnStatus_t status = cudnnDestroyDropoutDescriptor(handle_);
+    // TODO(csigg): This is a no-op (error is not reported). Same below.
+    CUDNN_RETURN_IF_FAIL(status, "Failed to destroy Cudnn dropout handle: ");
   }
 
   cudnnDropoutDescriptor_t handle() const {
@@ -1151,8 +990,7 @@ class CudnnDropoutDescriptor : public CudnnDescriptorCommon<void> {
   }
 
  private:
-  CUDAExecutor* parent_;
-  cudnnDropoutDescriptor_t handle_;
+  cudnnDropoutDescriptor_t handle_;  // Owned.
   float dropout_;
   uint64 seed_;
   SE_DISALLOW_COPY_AND_ASSIGN(CudnnDropoutDescriptor);
@@ -1162,10 +1000,10 @@ class CudnnRnnParamsDescriptor : public CudnnDescriptorCommon<void> {
  public:
   typedef dnn::RnnDescriptor::ParamsRegion ParamsRegion;
   typedef dnn::RnnDescriptor::ParamsRegions ParamsRegions;
-  CudnnRnnParamsDescriptor(CUDAExecutor* parent, cudnnHandle_t cudnn_handle,
+  CudnnRnnParamsDescriptor(const CudnnHandle& cudnn,
                            const CudnnRnnDescriptor& rnn_desc);
   ~CudnnRnnParamsDescriptor() {
-    cudnnStatus_t status = wrap::cudnnDestroyFilterDescriptor(parent_, handle_);
+    cudnnStatus_t status = cudnnDestroyFilterDescriptor(handle_);
     CUDNN_RETURN_IF_FAIL(status, "Failed to destroy RNN filter descriptor");
   }
   cudnnFilterDescriptor_t handle() const {
@@ -1184,7 +1022,6 @@ class CudnnRnnParamsDescriptor : public CudnnDescriptorCommon<void> {
 
  private:
   int GetRegionCountPerLayer() const;
-  CUDAExecutor* parent_;
   cudnnFilterDescriptor_t handle_;
   const CudnnRnnDescriptor* rnn_desc_;
   int64 params_size_in_bytes_;
@@ -1193,19 +1030,20 @@ class CudnnRnnParamsDescriptor : public CudnnDescriptorCommon<void> {
   SE_DISALLOW_COPY_AND_ASSIGN(CudnnRnnParamsDescriptor);
 };
 
+}  // namespace
+
 class CudnnRnnDescriptor : public CudnnDescriptorCommon<dnn::RnnDescriptor> {
  public:
-  CudnnRnnDescriptor(CUDAExecutor* parent, cudnnHandle_t cudnn_handle,
-                     int num_layers, int hidden_size, int input_size,
-                     int batch_size, cudnnRNNInputMode_t input_mode,
+  CudnnRnnDescriptor(const CudnnHandle& cudnn, int num_layers, int hidden_size,
+                     int input_size, int batch_size,
+                     cudnnRNNInputMode_t input_mode,
                      cudnnDirectionMode_t direction_mode,
                      cudnnRNNMode_t rnn_mode, cudnnDataType_t data_type,
                      cudnnDataType_t compute_type,
                      const dnn::AlgorithmConfig& algorithm_config,
                      float dropout, uint64 seed,
                      ScratchAllocator* state_allocator)
-      : parent_(parent),
-        rnn_desc_(nullptr),
+      : rnn_desc_(nullptr),
         num_layers_(num_layers),
         hidden_size_(hidden_size),
         input_size_(input_size),
@@ -1220,21 +1058,21 @@ class CudnnRnnDescriptor : public CudnnDescriptorCommon<dnn::RnnDescriptor> {
         compute_type_(compute_type),
         algorithm_config_(algorithm_config) {
     // Create the dropout handle.
-    cudnn_dropout_desc_.reset(new CudnnDropoutDescriptor(
-        parent, cudnn_handle, dropout, seed, state_allocator));
+    cudnn_dropout_desc_.reset(
+        new CudnnDropoutDescriptor(cudnn, dropout, seed, state_allocator));
     if (!cudnn_dropout_desc_->ok()) {
       SetFailure(cudnn_dropout_desc_->Status());
       return;
     }
 
     // Create the RNN handle
-    cudnnStatus_t status = wrap::cudnnCreateRNNDescriptor(parent_, &rnn_desc_);
+    cudnnStatus_t status = cudnnCreateRNNDescriptor(&rnn_desc_);
     CUDNN_RETURN_IF_FAIL(status, "Unable to create RNN descriptor");
 #if CUDNN_VERSION >= 6000
     // TODO: allow the user to choose an algorithm.
     rnn_algo_ = ToCudnnRNNAlgo(algorithm_config_.algorithm());
-    status = wrap::cudnnSetRNNDescriptor_v6(
-        parent, cudnn_handle, /*rnnDesc=*/rnn_desc_, /*hiddenSize=*/hidden_size,
+    status = cudnnSetRNNDescriptor_v6(
+        cudnn.handle(), /*rnnDesc=*/rnn_desc_, /*hiddenSize=*/hidden_size,
         /*numLayers=*/num_layers, /*dropoutDesc=*/dropout_handle(),
         /*inputMode=*/input_mode, /*direction=*/direction_mode,
         /*mode=*/rnn_mode, /*algo=*/rnn_algo_, /*dataType=*/compute_type);
@@ -1246,26 +1084,25 @@ class CudnnRnnDescriptor : public CudnnDescriptorCommon<dnn::RnnDescriptor> {
 
     if (rnn_algo_ == CUDNN_RNN_ALGO_PERSIST_DYNAMIC) {
       CHECK_GE(batch_size_, 0);
-      status = wrap::cudnnCreatePersistentRNNPlan(
-          parent, rnn_desc_, batch_size_, data_type_, &rnn_plan_);
+      status = cudnnCreatePersistentRNNPlan(rnn_desc_, batch_size_, data_type_,
+                                            &rnn_plan_);
       CUDNN_RETURN_IF_FAIL(status, "Unable to create persistent RNN plan.");
-      status = wrap::cudnnSetPersistentRNNPlan(parent, rnn_desc_, rnn_plan_);
+      status = cudnnSetPersistentRNNPlan(rnn_desc_, rnn_plan_);
       CUDNN_RETURN_IF_FAIL(status, "Unable to update persistent RNN plan.");
     }
 #else
     CHECK(algorithm_config_.is_default())
         << "Non-default algorithm not supported for CUDA version < 6.0";
-    status = wrap::cudnnSetRNNDescriptor(
-        parent, rnn_desc_ /*rnnDesc*/, hidden_size /*hiddenSize*/,
-        num_layers /*numLayers*/, dropout_handle() /*dropoutDesc*/,
-        input_mode /*inputMode*/, direction_mode /*direction*/,
-        rnn_mode /*mode*/, compute_type /*dataType*/);
+    status = cudnnSetRNNDescriptor(
+        /*rnnDesc=*/rnn_desc_, /*hiddenSize=*/hidden_size,
+        /*numLayers=*/num_layers, /*dropoutDesc=*/dropout_handle(),
+        /*inputMode=*/input_mode, /*direction=*/direction_mode,
+        /*mode=*/rnn_mode, /*dataType=*/compute_type);
     CUDNN_RETURN_IF_FAIL(status, "Unable to update RNN descriptor");
 #endif
 
     // Create the params handle.
-    cudnn_params_desc_.reset(
-        new CudnnRnnParamsDescriptor(parent, cudnn_handle, *this));
+    cudnn_params_desc_.reset(new CudnnRnnParamsDescriptor(cudnn, *this));
     if (!cudnn_params_desc_->ok()) {
       SetFailure(cudnn_params_desc_->Status());
       return;
@@ -1277,11 +1114,11 @@ class CudnnRnnDescriptor : public CudnnDescriptorCommon<dnn::RnnDescriptor> {
       cudnnStatus_t status;
 #if CUDNN_VERSION >= 6000
       if (rnn_algo_ == CUDNN_RNN_ALGO_PERSIST_DYNAMIC && rnn_plan_) {
-        status = wrap::cudnnDestroyPersistentRNNPlan(parent_, rnn_plan_);
+        status = cudnnDestroyPersistentRNNPlan(rnn_plan_);
         CUDNN_RETURN_IF_FAIL(status, "Unable to destroy persistent RNN plan.");
       }
 #endif
-      status = wrap::cudnnDestroyRNNDescriptor(parent_, rnn_desc_);
+      status = cudnnDestroyRNNDescriptor(rnn_desc_);
       CUDNN_RETURN_IF_FAIL(status, "Unable to destroy RNN descriptor");
     }
   }
@@ -1290,11 +1127,9 @@ class CudnnRnnDescriptor : public CudnnDescriptorCommon<dnn::RnnDescriptor> {
     cudnnMathType_t math_type =
         (use_tensor_op_math ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH);
     if (RnnTensorOpMathEnabled()) {
-      cudnnStatus_t status =
-          wrap::cudnnSetRNNMatrixMathType(parent_, rnn_desc_, math_type);
+      cudnnStatus_t status = cudnnSetRNNMatrixMathType(rnn_desc_, math_type);
       if (status != CUDNN_STATUS_SUCCESS) {
-        LOG(FATAL) << "could not set cudnn RNN math type: "
-                   << ToString(status);
+        LOG(FATAL) << "could not set cudnn RNN math type: " << ToString(status);
       }
     }
 #endif
@@ -1336,7 +1171,6 @@ class CudnnRnnDescriptor : public CudnnDescriptorCommon<dnn::RnnDescriptor> {
   }
 
  private:
-  CUDAExecutor* parent_;
   cudnnRNNDescriptor_t rnn_desc_;
   int num_layers_;
   int hidden_size_;
@@ -1359,30 +1193,28 @@ class CudnnRnnDescriptor : public CudnnDescriptorCommon<dnn::RnnDescriptor> {
   SE_DISALLOW_COPY_AND_ASSIGN(CudnnRnnDescriptor);
 };
 
+namespace {
+
 CudnnRnnParamsDescriptor::CudnnRnnParamsDescriptor(
-    CUDAExecutor* parent, cudnnHandle_t cudnn_handle,
-    const CudnnRnnDescriptor& rnn_desc)
-    : parent_(parent),
-      handle_(nullptr),
-      rnn_desc_(&rnn_desc),
-      params_size_in_bytes_(0) {
+    const CudnnHandle& cudnn, const CudnnRnnDescriptor& rnn_desc)
+    : handle_(nullptr), rnn_desc_(&rnn_desc), params_size_in_bytes_(0) {
   cudnnTensorDescriptor_t input_desc = nullptr;
   {
     // Query the params size.
-    auto status = wrap::cudnnCreateTensorDescriptor(parent, &input_desc);
+    auto status = cudnnCreateTensorDescriptor(&input_desc);
     CUDNN_RETURN_IF_FAIL(status, "Cudnn fails to create tensor descriptor");
     int dims[] = {1, rnn_desc.input_size(), 1};
     int strides[] = {dims[1] * dims[2], dims[2], 1};
-    status = wrap::cudnnSetTensorNdDescriptor(
-        parent, input_desc /*tensorDesc*/, rnn_desc.data_type() /*dataType*/,
-        sizeof(dims) / sizeof(dims[0]) /*nbDims*/, dims /*dimA*/,
-        strides /*strideA*/);
+    status = cudnnSetTensorNdDescriptor(
+        /*tensorDesc=*/input_desc, rnn_desc.data_type() /*dataType*/,
+        sizeof(dims) / sizeof(dims[0]) /*nbDims*/, /*dimA=*/dims,
+        /*strideA=*/strides);
     CUDNN_RETURN_IF_FAIL(status, "Cudnn fails to set tensor descriptor");
 
     size_t params_size = 0;
-    status = wrap::cudnnGetRNNParamsSize(
-        parent, cudnn_handle /*handle*/, rnn_desc.handle() /*rnnDesc*/,
-        input_desc /*xDesc*/, &params_size /*sizeInBytes*/,
+    status = cudnnGetRNNParamsSize(
+        cudnn.handle() /*handle*/, rnn_desc.handle() /*rnnDesc*/,
+        /*xDesc=*/input_desc, /*sizeInBytes=*/&params_size,
         rnn_desc.data_type() /*dataType*/);
     CUDNN_RETURN_IF_FAIL(status, "Cudnn fails to get RNN parameter size");
     params_size_in_bytes_ = static_cast<int64>(params_size);
@@ -1390,13 +1222,13 @@ CudnnRnnParamsDescriptor::CudnnRnnParamsDescriptor(
 
   {
     // Create the params descriptor.
-    auto status = wrap::cudnnCreateFilterDescriptor(parent, &handle_);
+    auto status = cudnnCreateFilterDescriptor(&handle_);
     CUDNN_RETURN_IF_FAIL(status, "Cudnn fails to create RNN filter descriptor");
     int dims[] = {static_cast<int>(params_size_in_bytes_), 1, 1};
-    status = wrap::cudnnSetFilterNdDescriptor(
-        parent, handle_ /*filterDesc*/, rnn_desc.data_type() /*dataType*/,
-        CUDNN_TENSOR_NCHW /*format*/, sizeof(dims) / sizeof(dims[0]) /*nbDims*/,
-        dims /*filterDimA*/);
+    status = cudnnSetFilterNdDescriptor(
+        /*filterDesc=*/handle_, rnn_desc.data_type() /*dataType*/,
+        /*format=*/CUDNN_TENSOR_NCHW, sizeof(dims) / sizeof(dims[0]) /*nbDims*/,
+        /*filterDimA=*/dims);
     CUDNN_RETURN_IF_FAIL(status, "Cudnn fails to update RNN filter descriptor");
   }
 
@@ -1404,8 +1236,7 @@ CudnnRnnParamsDescriptor::CudnnRnnParamsDescriptor(
     // Create the weights and biases into the params buffer
     int region_count_per_layer = GetRegionCountPerLayer();
     cudnnFilterDescriptor_t region_desc_handle = nullptr;
-    auto status =
-        wrap::cudnnCreateFilterDescriptor(parent, &region_desc_handle);
+    auto status = cudnnCreateFilterDescriptor(&region_desc_handle);
     CUDNN_RETURN_IF_FAIL(status, "Cudnn fails to create filter descriptor");
     const int layer_count = rnn_desc.direction_mode() == CUDNN_UNIDIRECTIONAL
                                 ? rnn_desc.num_layers()
@@ -1415,21 +1246,21 @@ CudnnRnnParamsDescriptor::CudnnRnnParamsDescriptor(
         for (int type = 0; type < 2; type++) {
           void* offset = nullptr;
           if (type == 0) {
-            status = wrap::cudnnGetRNNLinLayerMatrixParams(
-                parent, cudnn_handle /*handle*/, rnn_desc.handle() /*rnnDesc*/,
-                layer /*layer*/, input_desc /*xDesc*/, handle_ /*wDesc*/,
-                nullptr /*w*/, region /*linLayerID*/,
-                region_desc_handle /*linLayerMatDesc*/,
-                &offset /*linLayerMat*/);
+            status = cudnnGetRNNLinLayerMatrixParams(
+                cudnn.handle() /*handle*/, rnn_desc.handle() /*rnnDesc*/,
+                /*layer=*/layer, /*xDesc=*/input_desc, /*wDesc=*/handle_,
+                /*w=*/nullptr, /*linLayerID=*/region,
+                /*linLayerMatDesc=*/region_desc_handle,
+                /*linLayerMat=*/&offset);
             CUDNN_RETURN_IF_FAIL(
                 status, "Cudnn fails to call cudnnGetRNNLinLayerMatrixParams");
           } else {
-            status = wrap::cudnnGetRNNLinLayerBiasParams(
-                parent, cudnn_handle /*rnnDesc*/, rnn_desc.handle() /*rnnDesc*/,
-                layer /*layer*/, input_desc /*xDesc*/, handle_ /*wDesc*/,
-                nullptr /*w*/, region /*linLayerID*/,
-                region_desc_handle /*linLayerBiasDesc*/,
-                &offset /*linLayerBias*/);
+            status = cudnnGetRNNLinLayerBiasParams(
+                cudnn.handle() /*rnnDesc*/, rnn_desc.handle() /*rnnDesc*/,
+                /*layer=*/layer, /*xDesc=*/input_desc, /*wDesc=*/handle_,
+                /*w=*/nullptr, /*linLayerID=*/region,
+                /*linLayerBiasDesc=*/region_desc_handle,
+                /*linLayerBias=*/&offset);
             CUDNN_RETURN_IF_FAIL(
                 status, "Cudnn fails to call cudnnGetRNNLinLayerBiasParams");
           }
@@ -1437,15 +1268,15 @@ CudnnRnnParamsDescriptor::CudnnRnnParamsDescriptor(
           cudnnDataType_t data_type;
           cudnnTensorFormat_t tensor_format;
           int n_dims;
-          status = wrap::cudnnGetFilterNdDescriptor(
-              parent, region_desc_handle /*filterDesc*/,
+          status = cudnnGetFilterNdDescriptor(
+              /*filterDesc=*/region_desc_handle,
               sizeof(dims) / sizeof(dims[0]) /*nbDimsRequested*/,
-              &data_type /*dataType*/, &tensor_format /*format*/,
-              &n_dims /*nbDims*/, dims /*filterDimA*/);
+              /*dataType=*/&data_type, /*format=*/&tensor_format,
+              /*nbDims=*/&n_dims, /*filterDimA=*/dims);
           CUDNN_RETURN_IF_FAIL(status, "Cudnn fails to get filter description");
           int64 size = dims[0] * dims[1] * dims[2] *
                        CudnnDataTypeToByteSize(rnn_desc.data_type());
-          auto region = ParamsRegion{reinterpret_cast<int64>(offset), size};
+          ParamsRegion region = {reinterpret_cast<int64>(offset), size};
           if (type == 0) {
             weights_.push_back(region);
           } else {
@@ -1454,13 +1285,13 @@ CudnnRnnParamsDescriptor::CudnnRnnParamsDescriptor(
         }
       }
     }
-    status = wrap::cudnnDestroyFilterDescriptor(parent, region_desc_handle);
+    status = cudnnDestroyFilterDescriptor(region_desc_handle);
     CUDNN_RETURN_IF_FAIL(status, "Cudnn fails to destroy filter descriptor");
   }
 
   {
     // Release the dummy input tensor descriptor.
-    auto status = wrap::cudnnDestroyTensorDescriptor(parent, input_desc);
+    auto status = cudnnDestroyTensorDescriptor(input_desc);
     CUDNN_RETURN_IF_FAIL(status, "Cudnn fails to destroy tensor descriptor");
   }
 }
@@ -1480,6 +1311,8 @@ int CudnnRnnParamsDescriptor::GetRegionCountPerLayer() const {
   }
 }
 
+}  // namespace
+
 class CudnnRnnSequenceTensorDescriptor
     : public CudnnDescriptorCommon<dnn::RnnSequenceTensorDescriptor> {
  public:
@@ -1499,14 +1332,14 @@ class CudnnRnnSequenceTensorDescriptor
       SetFailure(port::Status(port::error::UNKNOWN, error_msg));
       return;
     }
-    cudnnStatus_t status = wrap::cudnnCreateTensorDescriptor(parent, &handle);
+    cudnnStatus_t status = cudnnCreateTensorDescriptor(&handle);
     CUDNN_RETURN_IF_FAIL(status, "Failed to create tensor descriptor");
     int dims[] = {batch_size, data_size, 1};
     int strides[] = {dims[1] * dims[2], dims[2], 1};
-    status = wrap::cudnnSetTensorNdDescriptor(
-        parent, handle /*tensorDesc*/, data_type /*dataType*/,
-        sizeof(dims) / sizeof(dims[0]) /*nbDims*/, dims /*dimA*/,
-        strides /*strideA*/);
+    status = cudnnSetTensorNdDescriptor(
+        /*tensorDesc=*/handle, /*dataType=*/data_type,
+        sizeof(dims) / sizeof(dims[0]) /*nbDims*/, /*dimA=*/dims,
+        /*strideA=*/strides);
     CUDNN_RETURN_IF_FAIL(status, "Failed to update tensor descriptor");
     // Replicate handle across the number of steps.
     handles_.assign(seq_length, handle);
@@ -1514,8 +1347,7 @@ class CudnnRnnSequenceTensorDescriptor
 
   ~CudnnRnnSequenceTensorDescriptor() override {
     // Only the first one needs to be destroyed. All others are the same.
-    cudnnStatus_t status =
-        wrap::cudnnDestroyTensorDescriptor(parent_, handles_[0]);
+    cudnnStatus_t status = cudnnDestroyTensorDescriptor(handles_[0]);
     CUDNN_RETURN_IF_FAIL(status,
                          "Failed to destroy sequence tensor descriptor");
   }
@@ -1552,21 +1384,20 @@ class CudnnRnnStateTensorDescriptor
         batch_size_(batch_size),
         data_size_(data_size),
         data_type_(data_type) {
-    cudnnStatus_t status = wrap::cudnnCreateTensorDescriptor(parent, &handle_);
+    cudnnStatus_t status = cudnnCreateTensorDescriptor(&handle_);
     CUDNN_RETURN_IF_FAIL(status, "Failed to create tensor descriptor");
     int dims[] = {num_layers, batch_size, data_size};
     int strides[] = {dims[1] * dims[2], dims[2], 1};
-    status = wrap::cudnnSetTensorNdDescriptor(
-        parent, handle_ /*tensorDesc*/, data_type /*dataType*/,
-        sizeof(dims) / sizeof(dims[0]) /*nbDims*/, dims /*dimA*/,
-        strides /*strideA*/);
+    status = cudnnSetTensorNdDescriptor(
+        /*tensorDesc=*/handle_, /*dataType=*/data_type,
+        sizeof(dims) / sizeof(dims[0]) /*nbDims*/, /*dimA=*/dims,
+        /*strideA=*/strides);
     CUDNN_RETURN_IF_FAIL(status, "Failed to update tensor descriptor");
   }
 
   ~CudnnRnnStateTensorDescriptor() override {
     if (!handle_) {
-      cudnnStatus_t status =
-          wrap::cudnnDestroyTensorDescriptor(parent_, handle_);
+      cudnnStatus_t status = cudnnDestroyTensorDescriptor(handle_);
       CUDNN_RETURN_IF_FAIL(status, "Unable to destroy RNN state tensor");
     }
   }
@@ -1661,13 +1492,13 @@ bool ExtractAndCheckRnnForward(
   return true;
 }
 
-bool CheckRNNParameterSize(CUDAExecutor* parent, cudnnHandle_t cudnn_handle,
+bool CheckRNNParameterSize(const CudnnHandle& cudnn,
                            const CudnnRnnDescriptor& rnn_desc,
                            const CudnnRnnSequenceTensorDescriptor& input_desc) {
   size_t params_size_in_bytes = 0;
-  cudnnStatus_t status = wrap::cudnnGetRNNParamsSize(
-      parent, cudnn_handle /*handle*/, rnn_desc.handle() /*rnnDesc*/,
-      input_desc.handles()[0] /*xDesc*/, &params_size_in_bytes /*sizeInBytes*/,
+  cudnnStatus_t status = cudnnGetRNNParamsSize(
+      /*handle=*/cudnn.handle(), rnn_desc.handle() /*rnnDesc*/,
+      input_desc.handles()[0] /*xDesc*/, /*sizeInBytes=*/&params_size_in_bytes,
       rnn_desc.data_type() /*dataType*/);
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "Unable to check RNN param size: " << ToString(status);
@@ -1677,18 +1508,17 @@ bool CheckRNNParameterSize(CUDAExecutor* parent, cudnnHandle_t cudnn_handle,
          rnn_desc.ParamsSizeInBytes();
 }
 
-bool CreateRnnWorkspace(Stream* stream, CUDAExecutor* parent,
-                        cudnnHandle_t cudnn_handle,
+bool CreateRnnWorkspace(Stream* stream, const CudnnHandle& cudnn,
                         const CudnnRnnDescriptor& rnn_desc,
                         const CudnnRnnSequenceTensorDescriptor& input_desc,
                         ScratchAllocator* workspace_allocator,
                         DeviceMemory<uint8>* workspace) {
   // Query the workspace size.
   size_t workspace_size_in_bytes = 0;
-  cudnnStatus_t status = wrap::cudnnGetRNNWorkspaceSize(
-      parent, cudnn_handle /*handle*/, rnn_desc.handle() /*rnnDesc*/,
-      input_desc.seq_length() /*seqLength*/, input_desc.handles() /*xDesc*/,
-      &workspace_size_in_bytes /*sizeInBytes*/);
+  cudnnStatus_t status = cudnnGetRNNWorkspaceSize(
+      /*handle=*/cudnn.handle(), /*rnnDesc=*/rnn_desc.handle(),
+      /*seqLength=*/input_desc.seq_length(), /*xDesc=*/input_desc.handles(),
+      /*sizeInBytes=*/&workspace_size_in_bytes);
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "Unable to query workspace size: " << ToString(status);
     return false;
@@ -1740,25 +1570,18 @@ bool CudnnSupport::DoRnnForwardImpl(
     return false;
   }
 
-  // check params size
-  mutex_lock lock{dnn_handle_mutex_};
-  auto set_stream_status =
-      wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (set_stream_status != CUDNN_STATUS_SUCCESS) {
-    LOG(FATAL) << "failed to set stream for cudnn handle: "
-               << ToString(set_stream_status);
-  }
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
 
-  if (!CheckRNNParameterSize(parent_, ToHandle(dnn_handle_), rnn_desc,
-                             input_desc)) {
+  // check params size
+  if (!CheckRNNParameterSize(cudnn, rnn_desc, input_desc)) {
     LOG(ERROR) << "Invalid parameters";
     return false;
   }
 
   // create the workspace
   DeviceMemory<uint8> workspace;
-  if (!CreateRnnWorkspace(stream, parent_, ToHandle(dnn_handle_), rnn_desc,
-                          input_desc, workspace_allocator, &workspace)) {
+  if (!CreateRnnWorkspace(stream, cudnn, rnn_desc, input_desc,
+                          workspace_allocator, &workspace)) {
     LOG(ERROR) << "Unable to create rnn workspace";
     return false;
   }
@@ -1768,11 +1591,10 @@ bool CudnnSupport::DoRnnForwardImpl(
   DeviceMemory<uint8> reserve_space;
   if (is_training) {
     size_t reserve_space_size_in_bytes = 0;
-    cudnnStatus_t status = wrap::cudnnGetRNNTrainingReserveSize(
-        parent_, ToHandle(dnn_handle_) /*handle*/,
-        rnn_desc.handle() /*rnnDesc*/, model_dims.seq_length /*seqLength*/,
-        input_desc.handles() /*xDesc*/,
-        &reserve_space_size_in_bytes /*sizeInBytes*/);
+    cudnnStatus_t status = cudnnGetRNNTrainingReserveSize(
+        cudnn.handle() /*handle*/, rnn_desc.handle() /*rnnDesc*/,
+        /*seqLength=*/model_dims.seq_length, input_desc.handles() /*xDesc*/,
+        /*sizeInBytes=*/&reserve_space_size_in_bytes);
     if (status != CUDNN_STATUS_SUCCESS) {
       LOG(ERROR) << "Unable to query reserve space size: " << ToString(status);
       return false;
@@ -1807,30 +1629,28 @@ bool CudnnSupport::DoRnnForwardImpl(
   // make the forward call
   cudnnStatus_t status;
   if (!is_training) {
-    status = wrap::cudnnRNNForwardInference(
-        this, stream, ToHandle(dnn_handle_) /*handle*/,
-        rnn_desc.handle() /*rnnDesc*/, model_dims.seq_length /*seqLength*/,
-        input_desc.handles() /*xDesc*/, input_data.opaque() /*x*/,
-        input_h_desc.handle() /*hxDesc*/, input_h_data.opaque() /*hx*/,
-        input_c_desc.handle() /*cxDesc*/, input_c_data.opaque() /*cx*/,
-        rnn_desc.params_handle() /*wDesc*/, params.opaque() /*w*/,
-        output_desc.handles() /*yDesc*/, output_data->opaque() /*y*/,
-        output_h_desc.handle() /*hyDesc*/, output_h_data->opaque() /*hy*/,
-        output_c_desc.handle() /*cyDesc*/, output_c_data->opaque() /*cy*/,
-        workspace.opaque() /*workspace*/,
+    status = cudnnRNNForwardInference(
+        cudnn.handle() /*handle*/, rnn_desc.handle() /*rnnDesc*/,
+        model_dims.seq_length /*seqLength*/, input_desc.handles() /*xDesc*/,
+        input_data.opaque() /*x*/, input_h_desc.handle() /*hxDesc*/,
+        input_h_data.opaque() /*hx*/, input_c_desc.handle() /*cxDesc*/,
+        input_c_data.opaque() /*cx*/, rnn_desc.params_handle() /*wDesc*/,
+        params.opaque() /*w*/, output_desc.handles() /*yDesc*/,
+        output_data->opaque() /*y*/, output_h_desc.handle() /*hyDesc*/,
+        output_h_data->opaque() /*hy*/, output_c_desc.handle() /*cyDesc*/,
+        output_c_data->opaque() /*cy*/, workspace.opaque() /*workspace*/,
         workspace.size() /*workSpaceSizeInBytes*/);
   } else {
-    status = wrap::cudnnRNNForwardTraining(
-        this, stream, ToHandle(dnn_handle_) /*handle*/,
-        rnn_desc.handle() /*rnnDesc*/, model_dims.seq_length /*seqLength*/,
-        input_desc.handles() /*xDesc*/, input_data.opaque() /*x*/,
-        input_h_desc.handle() /*hxDesc*/, input_h_data.opaque() /*hx*/,
-        input_c_desc.handle() /*cxDesc*/, input_c_data.opaque() /*cx*/,
-        rnn_desc.params_handle() /*wDesc*/, params.opaque() /*w*/,
-        output_desc.handles() /*yDesc*/, output_data->opaque() /*y*/,
-        output_h_desc.handle() /*hyDesc*/, output_h_data->opaque() /*hy*/,
-        output_c_desc.handle() /*cyDesc*/, output_c_data->opaque() /*cy*/,
-        workspace.opaque() /*workspace*/,
+    status = cudnnRNNForwardTraining(
+        cudnn.handle() /*handle*/, rnn_desc.handle() /*rnnDesc*/,
+        model_dims.seq_length /*seqLength*/, input_desc.handles() /*xDesc*/,
+        input_data.opaque() /*x*/, input_h_desc.handle() /*hxDesc*/,
+        input_h_data.opaque() /*hx*/, input_c_desc.handle() /*cxDesc*/,
+        input_c_data.opaque() /*cx*/, rnn_desc.params_handle() /*wDesc*/,
+        params.opaque() /*w*/, output_desc.handles() /*yDesc*/,
+        output_data->opaque() /*y*/, output_h_desc.handle() /*hyDesc*/,
+        output_h_data->opaque() /*hy*/, output_c_desc.handle() /*cyDesc*/,
+        output_c_data->opaque() /*cy*/, workspace.opaque() /*workspace*/,
         workspace.size() /*workSpaceSizeInBytes*/,
         reserve_space.opaque() /*reserveSpace*/,
         reserve_space.size() /*reserveSpaceSizeInBytes*/);
@@ -1896,25 +1716,18 @@ bool CudnnSupport::DoRnnBackwardImpl(
     return false;
   }
 
-  // check params size
-  mutex_lock lock{dnn_handle_mutex_};
-  auto set_stream_status =
-      wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (set_stream_status != CUDNN_STATUS_SUCCESS) {
-    LOG(FATAL) << "failed to set stream for cudnn handle: "
-               << ToString(set_stream_status);
-  }
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
 
-  if (!CheckRNNParameterSize(parent_, ToHandle(dnn_handle_), rnn_desc,
-                             input_desc)) {
+  // check params size
+  if (!CheckRNNParameterSize(cudnn, rnn_desc, input_desc)) {
     LOG(ERROR) << "Invalid parameters";
     return false;
   }
 
   // create the workspace
   DeviceMemory<uint8> workspace;
-  if (!CreateRnnWorkspace(stream, parent_, ToHandle(dnn_handle_), rnn_desc,
-                          input_desc, workspace_allocator, &workspace)) {
+  if (!CreateRnnWorkspace(stream, cudnn, rnn_desc, input_desc,
+                          workspace_allocator, &workspace)) {
     LOG(ERROR) << "Unable to create rnn workspace";
     return false;
   }
@@ -1934,12 +1747,11 @@ bool CudnnSupport::DoRnnBackwardImpl(
     }
   }
   // make the backward data call
-  cudnnStatus_t status = wrap::cudnnRNNBackwardData(
-      this, stream, ToHandle(dnn_handle_) /*handle*/,
-      rnn_desc.handle() /*rnnDesc*/, model_dims.seq_length /*seqLength*/,
-      output_desc.handles() /*yDesc*/, output_data.opaque() /*y*/,
-      output_desc.handles() /*dyDesc*/, output_backprop_data.opaque() /*dy*/,
-      output_h_desc.handle() /*dhyDesc*/,
+  cudnnStatus_t status = cudnnRNNBackwardData(
+      cudnn.handle() /*handle*/, rnn_desc.handle() /*rnnDesc*/,
+      model_dims.seq_length /*seqLength*/, output_desc.handles() /*yDesc*/,
+      output_data.opaque() /*y*/, output_desc.handles() /*dyDesc*/,
+      output_backprop_data.opaque() /*dy*/, output_h_desc.handle() /*dhyDesc*/,
       output_h_backprop_data.opaque() /*dhy*/,
       output_c_desc.handle() /*dcyDesc*/,
       output_c_backprop_data.opaque() /*dcy*/,
@@ -1967,13 +1779,12 @@ bool CudnnSupport::DoRnnBackwardImpl(
     // Clear the dw to zeros.
     stream->ThenMemZero(params_backprop_data, params_backprop_data->size());
     // make the backward weight call
-    status = wrap::cudnnRNNBackwardWeights(
-        this, stream, ToHandle(dnn_handle_) /*handle*/,
-        rnn_desc.handle() /*rnnDesc*/, model_dims.seq_length /*seqLength*/,
-        input_desc.handles() /*xDesc*/, input_data.opaque() /*x*/,
-        input_h_desc.handle() /*hxDesc*/, input_h_data.opaque() /*hx*/,
-        output_desc.handles() /*yDesc*/, output_data.opaque() /*y*/,
-        workspace.opaque() /*workspace*/,
+    status = cudnnRNNBackwardWeights(
+        cudnn.handle() /*handle*/, rnn_desc.handle() /*rnnDesc*/,
+        model_dims.seq_length /*seqLength*/, input_desc.handles() /*xDesc*/,
+        input_data.opaque() /*x*/, input_h_desc.handle() /*hxDesc*/,
+        input_h_data.opaque() /*hx*/, output_desc.handles() /*yDesc*/,
+        output_data.opaque() /*y*/, workspace.opaque() /*workspace*/,
         workspace.size() /*workSpaceSizeInBytes*/,
         rnn_desc.params_handle() /*dwDesc*/,
         params_backprop_data->opaque() /*dw*/,
@@ -2011,13 +1822,15 @@ CudnnSupport::createRnnDescriptor(
     const dnn::AlgorithmConfig& algorithm_config, float dropout, uint64 seed,
     ScratchAllocator* state_allocator) {
 #if CUDNN_VERSION >= 5000
-  mutex_lock lock{dnn_handle_mutex_};
+  // Setting up a cudnnRNNDescriptor requires a cuDNN handle, but because it's
+  // not enqueueing anything into a stream, we pass in the null stream.
+  auto cudnn = cudnn_->GetHandle(parent_, /*stream=*/nullptr);
   std::unique_ptr<CudnnRnnDescriptor> rnn_desc(new CudnnRnnDescriptor(
-      parent_, ToHandle(dnn_handle_), num_layers, hidden_size, input_size,
-      batch_size, ToCudnnRnnInputMode(input_mode),
-      ToCudnnRnnDirectionMode(direction_mode), ToCudnnRnnMode(rnn_mode),
-      ToCudnnDataType(data_type), GetRnnComputeType(data_type),
-      algorithm_config, dropout, seed, state_allocator));
+      cudnn, num_layers, hidden_size, input_size, batch_size,
+      ToCudnnRnnInputMode(input_mode), ToCudnnRnnDirectionMode(direction_mode),
+      ToCudnnRnnMode(rnn_mode), ToCudnnDataType(data_type),
+      GetRnnComputeType(data_type), algorithm_config, dropout, seed,
+      state_allocator));
   if (!rnn_desc->ok()) {
     return rnn_desc->Status();
   }
@@ -2028,7 +1841,7 @@ CudnnSupport::createRnnDescriptor(
       port::StrCat("createRnnDescriptor needs at least Cudnn 5.0 to work. ",
                    "Current Cudnn version: ", CUDNN_VERSION, ". ");
   LOG(ERROR) << error_msg;
-  return port::Status{port::error::UNIMPLEMENTED, error_msg};
+  return port::Status(port::error::UNIMPLEMENTED, error_msg);
 #endif  // CUDNN_VERSION
 }
 
@@ -2051,7 +1864,7 @@ CudnnSupport::createRnnSequenceTensorDescriptor(int seq_length, int batch_size,
       "createRnnSequenceTensorDescriptor needs at least Cudnn 5.0 to work. ",
       "Current Cudnn version: ", CUDNN_VERSION, ". ");
   LOG(ERROR) << error_msg;
-  return port::Status{port::error::UNIMPLEMENTED, error_msg};
+  return port::Status(port::error::UNIMPLEMENTED, error_msg);
 #endif  // CUDNN_VERSION
 }
 
@@ -2073,7 +1886,7 @@ CudnnSupport::createRnnStateTensorDescriptor(int num_layer, int batch_size,
       "createRnnStateTensorDescriptor needs at least Cudnn 5.0 to work. ",
       "Current Cudnn version: ", CUDNN_VERSION, ". ");
   LOG(ERROR) << error_msg;
-  return port::Status{port::error::UNIMPLEMENTED, error_msg};
+  return port::Status(port::error::UNIMPLEMENTED, error_msg);
 #endif  // CUDNN_VERSION
 }
 
@@ -2375,35 +2188,26 @@ bool CudnnSupport::DoRnnBackward(
 namespace {
 
 inline cudnnConvolutionFwdAlgo_t GetCudnnConvolutionForwardAlgo(
-    Stream* stream, CUDAExecutor* parent, void* dnn_handle,
-    const ScopedTensorDescriptor& input_nd,
+    const CudnnHandle& cudnn, const ScopedTensorDescriptor& input_nd,
     const ScopedFilterDescriptor& filter,
     const ScopedConvolutionDescriptor& conv,
     const ScopedTensorDescriptor& output_nd, bool specify_workspace_limit,
-    ScratchAllocator* scratch_allocator) {
+    size_t memory_limit_bytes) {
   cudnnConvolutionFwdPreference_t preference =
       specify_workspace_limit ? CUDNN_CONVOLUTION_FWD_SPECIFY_WORKSPACE_LIMIT
                               : CUDNN_CONVOLUTION_FWD_NO_WORKSPACE;
-  auto memory_limit_bytes =
-      scratch_allocator == nullptr
-          ? 0
-          : scratch_allocator->GetMemoryLimitInBytes(stream);
-  if (memory_limit_bytes < 0) {
-    memory_limit_bytes = 0;
-  }
 
   cudnnConvolutionFwdAlgo_t algo_to_use;
-  auto status = wrap::cudnnGetConvolutionForwardAlgorithm(
-      parent, ToHandle(dnn_handle), input_nd.handle(), filter.handle(),
-      conv.handle(), output_nd.handle(), preference, memory_limit_bytes,
-      &algo_to_use);
+  auto status = cudnnGetConvolutionForwardAlgorithm(
+      cudnn.handle(), input_nd.handle(), filter.handle(), conv.handle(),
+      output_nd.handle(), preference, memory_limit_bytes, &algo_to_use);
   CHECK_EQ(status, CUDNN_STATUS_SUCCESS)
       << "Unable to find a suitable algorithm for doing forward convolution";
   return algo_to_use;
 }
 
 dnn::AlgorithmDesc GetCudnnConvolutionForwardAlgorithm(
-    Stream* stream, CUDAExecutor* parent, void* dnn_handle,
+    Stream* stream, const CudnnHandle& cudnn,
     const dnn::AlgorithmConfig& algorithm_config, bool is_profiling,
     const ScopedTensorDescriptor& input_nd,
     const ScopedFilterDescriptor& filter,
@@ -2414,19 +2218,29 @@ dnn::AlgorithmDesc GetCudnnConvolutionForwardAlgorithm(
   bool use_tensor_ops;
   if (algorithm_config.algorithm().is_default()) {
     use_tensor_ops = true;
+
+    auto memory_limit_bytes =
+        scratch_allocator == nullptr
+            ? 0
+            : scratch_allocator->GetMemoryLimitInBytes(stream);
+    if (memory_limit_bytes < 0) {
+      memory_limit_bytes = 0;
+    }
+
     algo = GetCudnnConvolutionForwardAlgo(
-        stream, parent, dnn_handle, input_nd, filter, conv, output_nd,
+        cudnn, input_nd, filter, conv, output_nd,
         /*specify_workspace_limit=*/scratch_allocator != nullptr,
-        scratch_allocator);
+        memory_limit_bytes);
   } else {
     use_tensor_ops = algorithm_config.algorithm().tensor_ops_enabled();
     algo = ToConvForwardAlgo(algorithm_config.algorithm());
   }
   size_t size_in_bytes;
-  auto status = wrap::cudnnGetConvolutionForwardWorkspaceSize(
-      parent, ToHandle(dnn_handle), /*srcDesc=*/input_nd.handle(),
-      /*filterDesc=*/filter.handle(), /*convDesc=*/conv.handle(),
-      /*destDesc=*/output_nd.handle(), /*algo=*/algo,
+  auto status = cudnnGetConvolutionForwardWorkspaceSize(
+      cudnn.handle(),
+      /*xDesc=*/input_nd.handle(),
+      /*wDesc=*/filter.handle(), /*convDesc=*/conv.handle(),
+      /*yDesc=*/output_nd.handle(), /*algo=*/algo,
       /*sizeInBytes=*/&size_in_bytes);
   int64 size_in_bytes_int64 = size_in_bytes;
   if (TF_PREDICT_FALSE(status != CUDNN_STATUS_SUCCESS)) {
@@ -2466,8 +2280,8 @@ dnn::AlgorithmDesc GetCudnnConvolutionForwardAlgorithm(
       if (algorithm_config.algorithm_no_scratch().is_default()) {
         use_tensor_ops = true;
         algo = GetCudnnConvolutionForwardAlgo(
-            stream, parent, dnn_handle, input_nd, filter, conv, output_nd,
-            /*specify_workspace_limit=*/false, nullptr);
+            cudnn, input_nd, filter, conv, output_nd,
+            /*specify_workspace_limit=*/false, 0);
       } else {
         use_tensor_ops = algorithm_config.algorithm().tensor_ops_enabled();
         algo = ToConvForwardAlgo(algorithm_config.algorithm_no_scratch());
@@ -2596,11 +2410,12 @@ cudnnDataType_t GetRnnComputeType(dnn::DataType data_type) {
       LOG(FATAL) << "Invalid RNN data type: " << static_cast<int>(data_type);
   }
 }
+
 }  // namespace
 
 template <class T>
 bool CudnnSupport::DoConvolveImpl(
-    Stream* stream, const BatchDescriptor& batch_descriptor,
+    Stream* stream, const BatchDescriptor& input_descriptor,
     const DeviceMemory<T>& input_data,
     const FilterDescriptor& filter_descriptor,
     const DeviceMemory<T>& filter_data,
@@ -2610,18 +2425,13 @@ bool CudnnSupport::DoConvolveImpl(
     const dnn::AlgorithmConfig& algorithm_config,
     dnn::ProfileResult* output_profile_result) {
   cudnnDataType_t cudnn_type = GetCudnnDataType<T>();
-  ScopedTensorDescriptor input_nd{parent_, batch_descriptor, cudnn_type};
-  ScopedTensorDescriptor output_nd{parent_, output_descriptor, cudnn_type};
-  ScopedFilterDescriptor filter{parent_, filter_descriptor, batch_descriptor,
-                                cudnn_type};
-  ScopedConvolutionDescriptor conv{parent_, convolution_descriptor,
-                                   GetConvComputeType<T>()};
-
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(FATAL) << "failed to set stream for cudnn handle: " << ToString(status);
-  }
+  ScopedTensorDescriptor input_nd(input_descriptor, cudnn_type);
+  ScopedTensorDescriptor output_nd(output_descriptor, cudnn_type);
+  ScopedFilterDescriptor filter(filter_descriptor, cudnn_type);
+  ScopedConvolutionDescriptor conv(convolution_descriptor,
+                                   GetConvComputeType<T>());
+
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
   // Alpha is the scaling factor for input.
   float falpha = 1.0;
   double dalpha = 1.0;
@@ -2642,42 +2452,41 @@ bool CudnnSupport::DoConvolveImpl(
   //   GetCudnnConvolutionForwardAlgorithm().
   if (algorithm_config.algorithm().is_default()) {
     // With the default algorithm, use Cudnn's heuristics.
-    auto get_algorithm =
-        [&](bool specify_limit) SHARED_LOCKS_REQUIRED(dnn_handle_mutex_) {
-          cudnnConvolutionFwdPreference_t preference =
-              specify_limit ? CUDNN_CONVOLUTION_FWD_SPECIFY_WORKSPACE_LIMIT
-                            : CUDNN_CONVOLUTION_FWD_NO_WORKSPACE;
-
-          auto memory_limit_bytes =
-              scratch_allocator == nullptr
-                  ? 0
-                  : scratch_allocator->GetMemoryLimitInBytes(stream);
-          if (memory_limit_bytes < 0) {
-            memory_limit_bytes = 0;
-          }
+    auto get_algorithm = [&](bool specify_limit) {
+      cudnnConvolutionFwdPreference_t preference =
+          specify_limit ? CUDNN_CONVOLUTION_FWD_SPECIFY_WORKSPACE_LIMIT
+                        : CUDNN_CONVOLUTION_FWD_NO_WORKSPACE;
 
-          cudnnConvolutionFwdAlgo_t algo_to_use;
-          status = wrap::cudnnGetConvolutionForwardAlgorithm(
-              parent_, ToHandle(dnn_handle_), input_nd.handle(),
-              filter.handle(), conv.handle(), output_nd.handle(),
-              /*preference=*/preference,
-              /*memoryLimitInBytes=*/memory_limit_bytes,
-              /*algo=*/&algo_to_use);
-          CHECK_EQ(status, CUDNN_STATUS_SUCCESS)
-              << "Unable to find a suitable "
-                 "algorithm for doing forward "
-                 "convolution";
-          return algo_to_use;
-        };
+      auto memory_limit_bytes =
+          scratch_allocator == nullptr
+              ? 0
+              : scratch_allocator->GetMemoryLimitInBytes(stream);
+      if (memory_limit_bytes < 0) {
+        memory_limit_bytes = 0;
+      }
+
+      cudnnConvolutionFwdAlgo_t algo_to_use;
+      auto status = cudnnGetConvolutionForwardAlgorithm(
+          cudnn.handle(), input_nd.handle(), filter.handle(), conv.handle(),
+          output_nd.handle(),
+          /*preference=*/preference,
+          /*memoryLimitInBytes=*/memory_limit_bytes,
+          /*algo=*/&algo_to_use);
+      CHECK_EQ(status, CUDNN_STATUS_SUCCESS) << "Unable to find a suitable "
+                                                "algorithm for doing forward "
+                                                "convolution";
+      return algo_to_use;
+    };
 
     algo = get_algorithm(/*specify_limit=*/scratch_allocator != nullptr);
     use_tensor_ops = true;
     if (scratch_allocator != nullptr) {
       size_t size_in_bytes;
-      status = wrap::cudnnGetConvolutionForwardWorkspaceSize(
-          parent_, ToHandle(dnn_handle_), /*srcDesc=*/input_nd.handle(),
-          /*filterDesc=*/filter.handle(), /*convDesc=*/conv.handle(),
-          /*destDesc=*/output_nd.handle(), /*algo=*/algo,
+      auto status = cudnnGetConvolutionForwardWorkspaceSize(
+          cudnn.handle(),
+          /*xDesc=*/input_nd.handle(),
+          /*wDesc=*/filter.handle(), /*convDesc=*/conv.handle(),
+          /*yDesc=*/output_nd.handle(), /*algo=*/algo,
           /*sizeInBytes=*/&size_in_bytes);
       int64 size_in_bytes_int64 = size_in_bytes;
       if (status == CUDNN_STATUS_SUCCESS && size_in_bytes_int64 != 0) {
@@ -2709,10 +2518,11 @@ bool CudnnSupport::DoConvolveImpl(
     use_tensor_ops = algotype.tensor_ops_enabled();
     conv.set_use_tensor_op_math(use_tensor_ops);
     size_t size_in_bytes;
-    status = wrap::cudnnGetConvolutionForwardWorkspaceSize(
-        parent_, ToHandle(dnn_handle_), /*srcDesc=*/input_nd.handle(),
-        /*filterDesc=*/filter.handle(), /*convDesc=*/conv.handle(),
-        /*destDesc=*/output_nd.handle(), /*algo=*/algo,
+    auto status = cudnnGetConvolutionForwardWorkspaceSize(
+        cudnn.handle(),
+        /*xDesc=*/input_nd.handle(),
+        /*wDesc=*/filter.handle(), /*convDesc=*/conv.handle(),
+        /*yDesc=*/output_nd.handle(), /*algo=*/algo,
         /*sizeInBytes=*/&size_in_bytes);
     if (status != CUDNN_STATUS_SUCCESS) {
       if (is_profiling) {
@@ -2767,8 +2577,8 @@ bool CudnnSupport::DoConvolveImpl(
       return false;
     }
   }
-  status = wrap::cudnnConvolutionForward(
-      this, stream, ToHandle(dnn_handle_),
+  auto status = cudnnConvolutionForward(
+      cudnn.handle(),
       /*alpha=*/alpha, /*srcDesc=*/input_nd.handle(),
       /*srcData=*/input_data.opaque(), /*filterDesc=*/filter.handle(),
       /*filterData=*/filter_data.opaque(), /*convDesc=*/conv.handle(),
@@ -2822,30 +2632,22 @@ bool CudnnSupport::DoFusedConvolveImpl(
                 "supported for cuDNN version >= 6";
   return false;
 #else
-  ScopedTensorDescriptor conv_input_nd{
-      parent_, conv_input_descriptor,
-      static_cast<cudnnDataType_t>(cudnn_data_type)};
-  ScopedTensorDescriptor output_nd{
-      parent_, output_descriptor,
-      static_cast<cudnnDataType_t>(cudnn_data_type)};
-  ScopedFilterDescriptor filter{parent_, filter_descriptor,
-                                conv_input_descriptor,
-                                static_cast<cudnnDataType_t>(cudnn_data_type)};
-  ScopedTensorDescriptor bias_nd{parent_, bias_descriptor, CUDNN_DATA_FLOAT};
-  ScopedConvolutionDescriptor conv{
-      parent_, convolution_descriptor,
-      static_cast<cudnnDataType_t>(cudnn_compute_type)};
-
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  CHECK(status == CUDNN_STATUS_SUCCESS)
-      << "failed to set stream for cudnn handle: " << ToString(status);
-
+  ScopedTensorDescriptor conv_input_nd(
+      conv_input_descriptor, static_cast<cudnnDataType_t>(cudnn_data_type));
+  ScopedTensorDescriptor output_nd(
+      output_descriptor, static_cast<cudnnDataType_t>(cudnn_data_type));
+  ScopedFilterDescriptor filter(filter_descriptor,
+                                static_cast<cudnnDataType_t>(cudnn_data_type));
+  ScopedTensorDescriptor bias_nd(bias_descriptor, CUDNN_DATA_FLOAT);
+  ScopedConvolutionDescriptor conv(
+      convolution_descriptor, static_cast<cudnnDataType_t>(cudnn_compute_type));
+
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
   const bool is_profiling = output_profile_result != nullptr;
   DeviceMemory<uint8> scratch;
   dnn::AlgorithmDesc algotype = GetCudnnConvolutionForwardAlgorithm(
-      stream, parent_, dnn_handle_, algorithm_config, is_profiling,
-      conv_input_nd, filter, conv, output_nd, scratch_allocator, &scratch);
+      stream, cudnn, algorithm_config, is_profiling, conv_input_nd, filter,
+      conv, output_nd, scratch_allocator, &scratch);
   if (algotype.is_default()) {
     if (!is_profiling) {
       LOG(ERROR) << "No suitable algorithm found";
@@ -2879,9 +2681,8 @@ bool CudnnSupport::DoFusedConvolveImpl(
   // activation descriptor. Note that this will change the nan propagation
   // behavior from separate conv, bias, and relu (which by default is
   // CUDNN_PROPAGATE_NAN.
-  ScopedActivationDescriptor activation_desc{parent_, activation_mode,
-                                             CUDNN_NOT_PROPAGATE_NAN,
-                                             output_descriptor.value_max()};
+  ScopedActivationDescriptor activation_desc(
+      activation_mode, CUDNN_NOT_PROPAGATE_NAN, output_descriptor.value_max());
   auto side_input_data_ptr = (side_input_scale == 0) ? output_data->opaque()
                                                      : side_input_data.opaque();
 
@@ -2902,8 +2703,9 @@ bool CudnnSupport::DoFusedConvolveImpl(
           << "\noutput_nd.handle() = " << output_nd.handle()
           << "\noutput_data->opaque() = " << output_data->opaque();
 
-  status = wrap::cudnnConvolutionBiasActivationForward(
-      this, stream, ToHandle(dnn_handle_), /*alpha1=*/&conv_input_scale,
+  auto status = cudnnConvolutionBiasActivationForward(
+      cudnn.handle(),
+      /*alpha1=*/&conv_input_scale,
       /*srcDesc=*/conv_input_nd.handle(), /*srcData=*/conv_input_data.opaque(),
       /*filterDesc=*/filter.handle(), /*filterData=*/filter_data.opaque(),
       /*convDesc=*/conv.handle(), algo, /*workSpace=*/scratch.opaque(),
@@ -3107,17 +2909,9 @@ bool CudnnSupport::DoBatchNormalizationForwardImpl(
     DeviceMemory<U>* saved_mean, DeviceMemory<U>* saved_inv_var,
     bool is_training, std::function<const DeviceMemory<U>&()> var_to_inv_var,
     std::function<void()> inv_var_to_var) {
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status);
-    return false;
-  }
-
-  ScopedTensorDescriptor x_descriptor{parent_, x_desc,
-                                      ToCudnnDataType(input_data_type)};
-  ScopedTensorDescriptor scale_offset_descriptor{
-      parent_, scale_offset_desc, ToCudnnDataType(scale_data_type)};
+  ScopedTensorDescriptor x_descriptor(x_desc, ToCudnnDataType(input_data_type));
+  ScopedTensorDescriptor scale_offset_descriptor(
+      scale_offset_desc, ToCudnnDataType(scale_data_type));
   cudnnBatchNormMode_t mode = CUDNN_BATCHNORM_SPATIAL;
 #if CUDNN_VERSION >= 7000
   if (BatchnormSpatialPersistentEnabled() && is_training) {
@@ -3126,7 +2920,9 @@ bool CudnnSupport::DoBatchNormalizationForwardImpl(
 #endif
   float one = 1.0;
   float zero = 0.0;
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
 
+  auto status = CUDNN_STATUS_SUCCESS;
   if (is_training) {
     CHECK_EQ(batch_mean->is_null(), batch_var->is_null())
         << "batch_mean and batch_var must both be null or both be non-null";
@@ -3143,11 +2939,11 @@ bool CudnnSupport::DoBatchNormalizationForwardImpl(
       batch_var_opaque = nullptr;
     }
 
-    status = wrap::cudnnBatchNormalizationForwardTraining(
-        this, stream, ToHandle(dnn_handle_), mode, &one, &zero,
-        x_descriptor.handle(), x.opaque(), x_descriptor.handle(), y->opaque(),
-        scale_offset_descriptor.handle(), scale.opaque(), offset.opaque(), 1.0,
-        batch_mean_opaque, batch_var_opaque, epsilon, saved_mean->opaque(),
+    status = cudnnBatchNormalizationForwardTraining(
+        cudnn.handle(), mode, &one, &zero, x_descriptor.handle(), x.opaque(),
+        x_descriptor.handle(), y->opaque(), scale_offset_descriptor.handle(),
+        scale.opaque(), offset.opaque(), 1.0, batch_mean_opaque,
+        batch_var_opaque, epsilon, saved_mean->opaque(),
         saved_inv_var->opaque());
 #if CUDNN_VERSION < 5000
     CHECK(inv_var_to_var);
@@ -3160,11 +2956,11 @@ bool CudnnSupport::DoBatchNormalizationForwardImpl(
 #else
     const void* maybe_inv_var = estimated_variance.opaque();
 #endif
-    status = wrap::cudnnBatchNormalizationForwardInference(
-        this, stream, ToHandle(dnn_handle_), mode, &one, &zero,
-        x_descriptor.handle(), x.opaque(), x_descriptor.handle(), y->opaque(),
-        scale_offset_descriptor.handle(), scale.opaque(), offset.opaque(),
-        estimated_mean.opaque(), maybe_inv_var, epsilon);
+    status = cudnnBatchNormalizationForwardInference(
+        cudnn.handle(), mode, &one, &zero, x_descriptor.handle(), x.opaque(),
+        x_descriptor.handle(), y->opaque(), scale_offset_descriptor.handle(),
+        scale.opaque(), offset.opaque(), estimated_mean.opaque(), maybe_inv_var,
+        epsilon);
   }
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "failed to enqueue forward batch normalization on stream: "
@@ -3211,18 +3007,10 @@ bool CudnnSupport::DoBatchNormalizationBackwardImpl(
     const dnn::BatchDescriptor& scale_offset_desc, const double epsilon,
     DeviceMemory<T>* x_backprop, DeviceMemory<U>* scale_backprop,
     DeviceMemory<U>* offset_backprop) {
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status);
-    return false;
-  }
-
-  ScopedTensorDescriptor x_descriptor{
-      parent_, x_desc, static_cast<cudnnDataType_t>(cudnn_input_type)};
-  ScopedTensorDescriptor scale_offset_descriptor{
-      parent_, scale_offset_desc,
-      static_cast<cudnnDataType_t>(cudnn_scale_type)};
+  ScopedTensorDescriptor x_descriptor(
+      x_desc, static_cast<cudnnDataType_t>(cudnn_input_type));
+  ScopedTensorDescriptor scale_offset_descriptor(
+      scale_offset_desc, static_cast<cudnnDataType_t>(cudnn_scale_type));
   cudnnBatchNormMode_t mode = CUDNN_BATCHNORM_SPATIAL;
 #if CUDNN_VERSION >= 7000
   if (BatchnormSpatialPersistentEnabled()) {
@@ -3232,10 +3020,12 @@ bool CudnnSupport::DoBatchNormalizationBackwardImpl(
   float one = 1.0;
   float zero = 0.0;
 
-  status = wrap::cudnnBatchNormalizationBackward(
-      this, stream, ToHandle(dnn_handle_), mode, &one, &zero, &one, &zero,
-      x_descriptor.handle(), x.opaque(), x_descriptor.handle(),
-      y_backprop.opaque(), x_descriptor.handle(), x_backprop->opaque(),
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+
+  auto status = cudnnBatchNormalizationBackward(
+      cudnn.handle(), mode, &one, &zero, &one, &zero, x_descriptor.handle(),
+      x.opaque(), x_descriptor.handle(), y_backprop.opaque(),
+      x_descriptor.handle(), x_backprop->opaque(),
       scale_offset_descriptor.handle(), scale.opaque(),
       scale_backprop->opaque(), offset_backprop->opaque(), epsilon,
       mean.opaque(), inv_var.opaque());
@@ -3398,11 +3188,21 @@ bool CudnnSupport::DoFusedConvolve(
 #endif
 }
 
-template<class T>
-DeviceMemory<T> CudnnSupport::MaybeTransformLayout(
-    Stream* stream,
-    BatchDescriptor* output_descriptor,
-    DeviceMemory<T> backward_output_data,
+namespace {
+// NOTE(keveman): Temporary data layout transformation until cuDNN supports
+// kBatchYXDepth for backward pass. This function allocates temporary memory,
+// lays out the source data into the temporary but in the kBatchDepthXY
+// layout, and returns the temporary memory. The caller is responsible for
+// deallocating the temporary. Since the allocation is done using Stream's
+// AllocateTemporaryMemory, a later BlockHostUntilDone could be used for
+// deallocation.
+//
+// transform_scratch is populated with a legitimate temporary allocation iff
+// the original output data needs to be transformed.
+template <class T>
+DeviceMemory<T> MaybeTransformLayout(
+    Stream* stream, const CudnnHandle& cudnn,
+    BatchDescriptor* output_descriptor, DeviceMemory<T> backward_output_data,
     std::unique_ptr<TemporaryDeviceMemory<T>>* transform_scratch) {
   if (output_descriptor->layout() == dnn::DataLayout::kBatchDepthYX) {
     return backward_output_data;
@@ -3415,15 +3215,14 @@ DeviceMemory<T> CudnnSupport::MaybeTransformLayout(
   transformed_output_descriptor.CloneFrom(*output_descriptor);
   transformed_output_descriptor.set_layout(dnn::DataLayout::kBatchDepthYX);
   cudnnDataType_t cudnn_type = GetCudnnDataType<T>();
-  ScopedTensorDescriptor orig_out_back_nd{parent_, *output_descriptor,
-                                          cudnn_type};
-  ScopedTensorDescriptor transformed_out_back_nd{
-      parent_, transformed_output_descriptor, cudnn_type};
+  ScopedTensorDescriptor orig_out_back_nd(*output_descriptor, cudnn_type);
+  ScopedTensorDescriptor transformed_out_back_nd(transformed_output_descriptor,
+                                                 cudnn_type);
 
   float alpha = 1.0f;
   float beta = 0.0f;
-  auto status = wrap::cudnnTransformTensor(
-      this, stream, ToHandle(dnn_handle_), &alpha, orig_out_back_nd.handle(),
+  auto status = cudnnTransformTensor(
+      cudnn.handle(), &alpha, orig_out_back_nd.handle(),
       backward_output_data.opaque(), &beta, transformed_out_back_nd.handle(),
       (*transform_scratch)->mutable_device_memory()->opaque());
 
@@ -3433,6 +3232,7 @@ DeviceMemory<T> CudnnSupport::MaybeTransformLayout(
   output_descriptor->set_layout(dnn::DataLayout::kBatchDepthYX);
   return (*transform_scratch)->device_memory();
 }
+}  // namespace
 
 bool CudnnSupport::DoTransformTensor(Stream* stream,
                                      const dnn::BatchDescriptor& input_desc,
@@ -3441,21 +3241,15 @@ bool CudnnSupport::DoTransformTensor(Stream* stream,
                                      const dnn::BatchDescriptor& output_desc,
                                      dnn::DataType output_type, float scale,
                                      DeviceMemoryBase* output_data) {
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(FATAL) << "failed to set stream for cudnn handle: " << ToString(status);
-  }
-
   float beta = 0.0f;
   ScopedTensorDescriptor input_tensor_desc(
-      parent_, input_desc, ToCudnnDataType(input_type, input_desc.layout()));
+      input_desc, ToCudnnDataType(input_type, input_desc.layout()));
   ScopedTensorDescriptor output_tensor_desc(
-      parent_, output_desc, ToCudnnDataType(output_type, output_desc.layout()));
-  status = wrap::cudnnTransformTensor(
-      this, stream, ToHandle(dnn_handle_), &scale, input_tensor_desc.handle(),
-      input_data.opaque(), &beta, output_tensor_desc.handle(),
-      output_data->opaque());
+      output_desc, ToCudnnDataType(output_type, output_desc.layout()));
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+  auto status = cudnnTransformTensor(
+      cudnn.handle(), &scale, input_tensor_desc.handle(), input_data.opaque(),
+      &beta, output_tensor_desc.handle(), output_data->opaque());
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "Could not transform a tensor with layout "
                << input_desc.ToString() << " and data type "
@@ -3469,8 +3263,7 @@ bool CudnnSupport::DoTransformTensor(Stream* stream,
 
 template <class T>
 bool CudnnSupport::DoConvolveBackwardDataImpl(
-    Stream* stream,
-    const FilterDescriptor& filter_descriptor,
+    Stream* stream, const FilterDescriptor& filter_descriptor,
     const DeviceMemory<T>& filter_data,
     const BatchDescriptor& output_descriptor_in,
     DeviceMemory<T> backward_output_data,
@@ -3479,12 +3272,6 @@ bool CudnnSupport::DoConvolveBackwardDataImpl(
     DeviceMemory<T>* backward_input_data, ScratchAllocator* scratch_allocator,
     const dnn::AlgorithmConfig& algorithm_config,
     dnn::ProfileResult* output_profile_result) {
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(FATAL) << "failed to set stream for cudnn handle: " << ToString(status);
-  }
-
   cudnnDataType_t cudnn_type = GetCudnnDataType<T>();
   // Alpha is the scaling factor for input.
   float falpha = 1.0;
@@ -3497,19 +3284,21 @@ bool CudnnSupport::DoConvolveBackwardDataImpl(
   void* beta = cudnn_type == CUDNN_DATA_DOUBLE ? static_cast<void*>(&dbeta)
                                                : static_cast<void*>(&fbeta);
 
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+
   // TBD(keveman): remove once cuDNN supports kBatchYXDepth for backward pass.
   BatchDescriptor output_descriptor;
   output_descriptor.CloneFrom(output_descriptor_in);
   std::unique_ptr<TemporaryDeviceMemory<T>> transform_scratch;
-  backward_output_data = MaybeTransformLayout(
-      stream, &output_descriptor, backward_output_data, &transform_scratch);
+  backward_output_data =
+      MaybeTransformLayout(stream, cudnn, &output_descriptor,
+                           backward_output_data, &transform_scratch);
 
-  ScopedTensorDescriptor out_back_nd{parent_, output_descriptor, cudnn_type};
-  ScopedTensorDescriptor in_back_nd{parent_, input_descriptor, cudnn_type};
-  ScopedFilterDescriptor filter{parent_, filter_descriptor, input_descriptor,
-                                cudnn_type};
-  ScopedConvolutionDescriptor conv{parent_, convolution_descriptor,
-                                   GetConvComputeType<T>()};
+  ScopedTensorDescriptor out_back_nd(output_descriptor, cudnn_type);
+  ScopedTensorDescriptor in_back_nd(input_descriptor, cudnn_type);
+  ScopedFilterDescriptor filter(filter_descriptor, cudnn_type);
+  ScopedConvolutionDescriptor conv(convolution_descriptor,
+                                   GetConvComputeType<T>());
 
   const bool is_profiling = output_profile_result != nullptr;
   cudnnConvolutionBwdDataAlgo_t algo;
@@ -3517,8 +3306,8 @@ bool CudnnSupport::DoConvolveBackwardDataImpl(
 
   if (algorithm_config.algorithm().is_default()) {
     // With the default algorithm, use Cudnn's heuristics.
-    auto get_algorithm = [&](bool specify_limit) SHARED_LOCKS_REQUIRED(
-        dnn_handle_mutex_) -> cudnnConvolutionBwdDataAlgo_t {
+    auto get_algorithm =
+        [&](bool specify_limit) -> cudnnConvolutionBwdDataAlgo_t {
       cudnnConvolutionBwdDataPreference_t preference =
           specify_limit ? CUDNN_CONVOLUTION_BWD_DATA_SPECIFY_WORKSPACE_LIMIT
                         : CUDNN_CONVOLUTION_BWD_DATA_NO_WORKSPACE;
@@ -3531,8 +3320,8 @@ bool CudnnSupport::DoConvolveBackwardDataImpl(
         memory_limit_bytes = 0;
       }
       cudnnConvolutionBwdDataAlgo_t algo_to_use;
-      cudnnStatus_t status = wrap::cudnnGetConvolutionBackwardDataAlgorithm(
-          parent_, ToHandle(dnn_handle_),
+      cudnnStatus_t status = cudnnGetConvolutionBackwardDataAlgorithm(
+          cudnn.handle(),
           /*filterDesc=*/filter.handle(),
           /*diffDesc=*/out_back_nd.handle(),
           /*convDesc=*/conv.handle(),
@@ -3550,8 +3339,8 @@ bool CudnnSupport::DoConvolveBackwardDataImpl(
 
     if (scratch_allocator != nullptr) {
       size_t size_in_bytes;
-      status = wrap::cudnnGetConvolutionBackwardDataWorkspaceSize(
-          parent_, ToHandle(dnn_handle_),
+      auto status = cudnnGetConvolutionBackwardDataWorkspaceSize(
+          cudnn.handle(),
           /*filterDesc=*/filter.handle(),
           /*diffDesc=*/out_back_nd.handle(),
           /*convDesc=*/conv.handle(),
@@ -3587,8 +3376,8 @@ bool CudnnSupport::DoConvolveBackwardDataImpl(
     algo = ToConvBackwardDataAlgo(algotype);
     conv.set_use_tensor_op_math(algotype.tensor_ops_enabled());
     size_t size_in_bytes;
-    status = wrap::cudnnGetConvolutionBackwardDataWorkspaceSize(
-        parent_, ToHandle(dnn_handle_),
+    auto status = cudnnGetConvolutionBackwardDataWorkspaceSize(
+        cudnn.handle(),
         /*filterDesc=*/filter.handle(),
         /*diffDesc=*/out_back_nd.handle(),
         /*convDesc=*/conv.handle(),
@@ -3645,23 +3434,24 @@ bool CudnnSupport::DoConvolveBackwardDataImpl(
   }
 
 #if CUDNN_VERSION >= 5000
-  status = wrap::cudnnConvolutionBackwardData(
+  auto status =
+      cudnnConvolutionBackwardData(cudnn.handle(),
 #else
-  status = wrap::cudnnConvolutionBackwardData_v3(
+  auto status =
+      cudnnConvolutionBackwardData_v3(cudnn.handle(),
 #endif
-      this, stream, ToHandle(dnn_handle_),
-      /*alpha=*/alpha,
-      /*filterDesc=*/filter.handle(),
-      /*filterData=*/filter_data.opaque(),
-      /*diffDesc=*/out_back_nd.handle(),
-      /*diffData=*/backward_output_data.opaque(),
-      /*convDesc=*/conv.handle(),
-      /*algo=*/algo,
-      /*workSpace=*/scratch.opaque(),
-      /*workSpaceSizeInBytes=*/scratch.size(),
-      /*beta=*/beta,
-      /*gradDesc=*/in_back_nd.handle(),
-      /*gradData=*/backward_input_data->opaque());
+                                   /*alpha=*/alpha,
+                                   /*wDesc=*/filter.handle(),
+                                   /*w=*/filter_data.opaque(),
+                                   /*dyDesc=*/out_back_nd.handle(),
+                                   /*dy=*/backward_output_data.opaque(),
+                                   /*convDesc=*/conv.handle(),
+                                   /*algo=*/algo,
+                                   /*workSpace=*/scratch.opaque(),
+                                   /*workSpaceSizeInBytes=*/scratch.size(),
+                                   /*beta=*/beta,
+                                   /*dxDesc=*/in_back_nd.handle(),
+                                   /*dx=*/backward_input_data->opaque());
   if (is_profiling) {
     timer->Stop(AsCUDAStream(stream));
     if (status == CUDNN_STATUS_SUCCESS) {
@@ -3749,12 +3539,6 @@ bool CudnnSupport::DoConvolveBackwardFilterImpl(
     DeviceMemory<T>* backward_filter_data, ScratchAllocator* scratch_allocator,
     const dnn::AlgorithmConfig& algorithm_config,
     dnn::ProfileResult* output_profile_result) {
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(FATAL) << "failed to set stream for cudnn handle: " << ToString(status);
-  }
-
   cudnnDataType_t cudnn_type = GetCudnnDataType<T>();
   // Alpha is the scaling factor for input.
   float falpha = 1.0;
@@ -3767,19 +3551,21 @@ bool CudnnSupport::DoConvolveBackwardFilterImpl(
   void* beta = cudnn_type == CUDNN_DATA_DOUBLE ? static_cast<void*>(&dbeta)
                                                : static_cast<void*>(&fbeta);
 
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+
   // TBD(keveman): remove once cuDNN supports kBatchYXDepth for backward pass.
   BatchDescriptor output_descriptor;
   output_descriptor.CloneFrom(output_descriptor_in);
   std::unique_ptr<TemporaryDeviceMemory<T>> transform_scratch;
-  backward_output_data = MaybeTransformLayout(
-      stream, &output_descriptor, backward_output_data, &transform_scratch);
+  backward_output_data =
+      MaybeTransformLayout(stream, cudnn, &output_descriptor,
+                           backward_output_data, &transform_scratch);
 
-  ScopedTensorDescriptor out_back_nd{parent_, output_descriptor, cudnn_type};
-  ScopedTensorDescriptor input_nd{parent_, input_descriptor, cudnn_type};
-  ScopedFilterDescriptor filter{parent_, filter_descriptor, input_descriptor,
-                                cudnn_type};
-  ScopedConvolutionDescriptor conv{parent_, convolution_descriptor,
-                                   GetConvComputeType<T>()};
+  ScopedTensorDescriptor out_back_nd(output_descriptor, cudnn_type);
+  ScopedTensorDescriptor input_nd(input_descriptor, cudnn_type);
+  ScopedFilterDescriptor filter(filter_descriptor, cudnn_type);
+  ScopedConvolutionDescriptor conv(convolution_descriptor,
+                                   GetConvComputeType<T>());
 
   const bool is_profiling = output_profile_result != nullptr;
   cudnnConvolutionBwdFilterAlgo_t algo;
@@ -3791,8 +3577,7 @@ bool CudnnSupport::DoConvolveBackwardFilterImpl(
     // Lambda that retrieves the algorithm.
     // specify_limit will occur when we have a scratch allocator and it succeeds
     // in allocating; otherwise, we'll fall back to the "no workspace" version.
-    auto get_algorithm = [&](bool specify_limit) SHARED_LOCKS_REQUIRED(
-        dnn_handle_mutex_) {
+    auto get_algorithm = [&](bool specify_limit) {
       cudnnConvolutionBwdFilterPreference_t preference =
           specify_limit ? CUDNN_CONVOLUTION_BWD_FILTER_SPECIFY_WORKSPACE_LIMIT
                         : CUDNN_CONVOLUTION_BWD_FILTER_NO_WORKSPACE;
@@ -3806,8 +3591,8 @@ bool CudnnSupport::DoConvolveBackwardFilterImpl(
       }
 
       cudnnConvolutionBwdFilterAlgo_t algo_to_use;
-      cudnnStatus_t status = wrap::cudnnGetConvolutionBackwardFilterAlgorithm(
-          parent_, ToHandle(dnn_handle_),
+      cudnnStatus_t status = cudnnGetConvolutionBackwardFilterAlgorithm(
+          cudnn.handle(),
           /*srcDesc=*/input_nd.handle(),
           /*diffDesc=*/out_back_nd.handle(),
           /*convDesc=*/conv.handle(),
@@ -3825,9 +3610,10 @@ bool CudnnSupport::DoConvolveBackwardFilterImpl(
 
     if (scratch_allocator != nullptr) {
       size_t size_in_bytes;
-      status = wrap::cudnnGetConvolutionBackwardFilterWorkspaceSize(
-          parent_, ToHandle(dnn_handle_), /*srcDesc=*/input_nd.handle(),
-          /*diffDesc=*/out_back_nd.handle(), /*convDesc=*/conv.handle(),
+      auto status = cudnnGetConvolutionBackwardFilterWorkspaceSize(
+          cudnn.handle(),
+          /*xDesc=*/input_nd.handle(),
+          /*dyDesc=*/out_back_nd.handle(), /*convDesc=*/conv.handle(),
           /*gradDesc=*/filter.handle(), /*algo=*/algo,
           /*sizeInBytes=*/&size_in_bytes);
       int64 size_in_bytes_int64 = size_in_bytes;
@@ -3860,9 +3646,10 @@ bool CudnnSupport::DoConvolveBackwardFilterImpl(
     conv.set_use_tensor_op_math(algotype.tensor_ops_enabled());
 
     size_t size_in_bytes;
-    status = wrap::cudnnGetConvolutionBackwardFilterWorkspaceSize(
-        parent_, ToHandle(dnn_handle_), /*srcDesc=*/input_nd.handle(),
-        /*diffDesc=*/out_back_nd.handle(), /*convDesc=*/conv.handle(),
+    auto status = cudnnGetConvolutionBackwardFilterWorkspaceSize(
+        cudnn.handle(),
+        /*xDesc=*/input_nd.handle(),
+        /*dyDesc=*/out_back_nd.handle(), /*convDesc=*/conv.handle(),
         /*gradDesc=*/filter.handle(), /*algo=*/algo,
         /*sizeInBytes=*/&size_in_bytes);
     if (status != CUDNN_STATUS_SUCCESS) {
@@ -3916,11 +3703,13 @@ bool CudnnSupport::DoConvolveBackwardFilterImpl(
   }
 
 #if CUDNN_VERSION >= 5000
-  status = wrap::cudnnConvolutionBackwardFilter(
+  auto status = cudnnConvolutionBackwardFilter(
+      cudnn.handle(),
 #else
-  status = wrap::cudnnConvolutionBackwardFilter_v3(
+  auto status = cudnnConvolutionBackwardFilter_v3(
+      cudnn.handle(),
 #endif
-      this, stream, ToHandle(dnn_handle_), /*alpha=*/alpha,
+      /*alpha=*/alpha,
       /*srcDesc=*/input_nd.handle(),
       /*srcData=*/input_data.opaque(),
       /*diffDesc=*/out_back_nd.handle(),
@@ -4015,25 +3804,19 @@ bool CudnnSupport::DoConvolveBackwardBiasImpl(
     const DeviceMemory<T>& input_data,
     const dnn::BatchDescriptor& bias_descriptor,
     DeviceMemory<T>* backward_bias_data) {
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(FATAL) << "failed to set stream for cudnn handle: " << ToString(status);
-  }
-
   cudnnDataType_t cudnn_type = GetCudnnDataType<T>();
-  ScopedTensorDescriptor input_nd{parent_, input_descriptor, cudnn_type};
-  ScopedTensorDescriptor bias_nd{parent_, bias_descriptor, cudnn_type};
+  ScopedTensorDescriptor input_nd(input_descriptor, cudnn_type);
+  ScopedTensorDescriptor bias_nd(bias_descriptor, cudnn_type);
 
   // Alpha is the scaling factor for input.
   float alpha = 1.0;
   // Beta is the scaling factor for output.
   float beta = 0.0;
 
-  status = wrap::cudnnConvolutionBackwardBias(
-      this, stream, ToHandle(dnn_handle_), &alpha, input_nd.handle(),
-      input_data.opaque(), &beta, bias_nd.handle(),
-      backward_bias_data->opaque());
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+  auto status = cudnnConvolutionBackwardBias(
+      cudnn.handle(), &alpha, input_nd.handle(), input_data.opaque(), &beta,
+      bias_nd.handle(), backward_bias_data->opaque());
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "failed to enqueue backward convolution on stream: "
                << ToString(status);
@@ -4209,8 +3992,7 @@ bool CudnnSupport::DoBiasAdd(Stream* stream,
                              const DeviceMemory<float>& biases,
                              const dnn::BatchDescriptor& dimensions,
                              DeviceMemory<float>* output_data) {
-  ScopedTensorDescriptor input_descriptor{parent_, dimensions,
-                                          CUDNN_DATA_FLOAT};
+  ScopedTensorDescriptor input_descriptor(dimensions, CUDNN_DATA_FLOAT);
 
   BatchDescriptor bias_dimensions;
   bias_dimensions.set_count(1)
@@ -4218,8 +4000,7 @@ bool CudnnSupport::DoBiasAdd(Stream* stream,
       .set_height(1)
       .set_width(1)
       .set_layout(dnn::DataLayout::kBatchYXDepth);
-  ScopedTensorDescriptor bias_descriptor{parent_, bias_dimensions,
-                                         CUDNN_DATA_FLOAT};
+  ScopedTensorDescriptor bias_descriptor(bias_dimensions, CUDNN_DATA_FLOAT);
 
   // cudnnAddTensor after R3 is in-place, so we need to copy input_data to
   // output_data before doing the addition, unless the input and
@@ -4235,23 +4016,18 @@ bool CudnnSupport::DoBiasAdd(Stream* stream,
     }
   }
 
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status);
-    return false;
-  }
-
   const float alpha = 1.0f;
   const float beta = 1.0f;
 
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+
 #if CUDNN_VERSION >= 5000
-  status = wrap::cudnnAddTensor(
+  auto status = cudnnAddTensor(
 #else
-  status = wrap::cudnnAddTensor_v3(
+  auto status = cudnnAddTensor_v3(
 #endif
-      this, stream, ToHandle(dnn_handle_), &alpha, bias_descriptor.handle(),
-      biases.opaque(), &beta, input_descriptor.handle(), output_data->opaque());
+      cudnn.handle(), &alpha, bias_descriptor.handle(), biases.opaque(), &beta,
+      input_descriptor.handle(), output_data->opaque());
 
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "stream " << stream << " could not enqueue bias addition.";
@@ -4267,16 +4043,9 @@ bool CudnnSupport::DoActivate(Stream* stream,
                               const DeviceMemory<float>& input_data,
                               DeviceMemory<float>* output_data,
                               uint64 options) {
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status);
-    return false;
-  }
-
 #if CUDNN_VERSION >= 5000
-  ScopedActivationDescriptor activation_desc{
-      parent_, activation_mode, CUDNN_PROPAGATE_NAN, dimensions.value_max()};
+  ScopedActivationDescriptor activation_desc(
+      activation_mode, CUDNN_PROPAGATE_NAN, dimensions.value_max());
 #else
   cudnnActivationMode_t mode;
   switch (activation_mode) {
@@ -4306,20 +4075,22 @@ bool CudnnSupport::DoActivate(Stream* stream,
   }
 #endif
 
-  ScopedTensorDescriptor input_nd{parent_, dimensions, CUDNN_DATA_FLOAT};
+  ScopedTensorDescriptor input_nd(dimensions, CUDNN_DATA_FLOAT);
   // Alpha is the input scaling factor.
   float alpha = 1.0;
   // Beta is the output scaling factor.
   float beta = 0.0;
-  status = wrap::cudnnActivationForward(
-      this, stream, ToHandle(dnn_handle_),
+
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+  auto status =
+      cudnnActivationForward(cudnn.handle(),
 #if CUDNN_VERSION >= 5000
-      activation_desc.handle(),
+                             activation_desc.handle(),
 #else
-      mode,
+                             mode,
 #endif
-      &alpha, input_nd.handle(), input_data.opaque(), &beta, input_nd.handle(),
-      output_data->opaque());
+                             &alpha, input_nd.handle(), input_data.opaque(),
+                             &beta, input_nd.handle(), output_data->opaque());
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "stream " << stream
                << " could not enqueue activation: " << ToString(status);
@@ -4335,26 +4106,19 @@ bool CudnnSupport::DoPoolForward(
     const DeviceMemory<double>& input_data,
     const dnn::BatchDescriptor& output_dimensions,
     DeviceMemory<double>* output_data) {
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status);
-    return false;
-  }
-
   // Alpha is the scaling factor for input.
   double alpha = 1.0;
   // Beta is the scaling factor for output.
   double beta = 0.0;
 
-  ScopedTensorDescriptor src_desc{parent_, input_dimensions, CUDNN_DATA_DOUBLE};
-  ScopedTensorDescriptor dest_desc{parent_, output_dimensions,
-                                   CUDNN_DATA_DOUBLE};
-  ScopedPoolingDescriptor pooling_desc{parent_, pooling_dimensions};
-  status = wrap::cudnnPoolingForward(
-      this, stream, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha,
-      src_desc.handle(), input_data.opaque(), &beta, dest_desc.handle(),
-      output_data->opaque());
+  ScopedTensorDescriptor src_desc(input_dimensions, CUDNN_DATA_DOUBLE);
+  ScopedTensorDescriptor dest_desc(output_dimensions, CUDNN_DATA_DOUBLE);
+  ScopedPoolingDescriptor pooling_desc(pooling_dimensions);
+
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+  auto status = cudnnPoolingForward(
+      cudnn.handle(), pooling_desc.handle(), &alpha, src_desc.handle(),
+      input_data.opaque(), &beta, dest_desc.handle(), output_data->opaque());
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "failed to enqueue forward pooling on stream: "
                << ToString(status);
@@ -4369,26 +4133,19 @@ bool CudnnSupport::DoPoolForward(
     const DeviceMemory<float>& input_data,
     const dnn::BatchDescriptor& output_dimensions,
     DeviceMemory<float>* output_data) {
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status);
-    return false;
-  }
-
   // Alpha is the scaling factor for input.
   float alpha = 1.0;
   // Beta is the scaling factor for output.
   float beta = 0.0;
 
-  ScopedTensorDescriptor src_desc{parent_, input_dimensions, CUDNN_DATA_FLOAT};
-  ScopedTensorDescriptor dest_desc{parent_, output_dimensions,
-                                   CUDNN_DATA_FLOAT};
-  ScopedPoolingDescriptor pooling_desc{parent_, pooling_dimensions};
-  status = wrap::cudnnPoolingForward(
-      this, stream, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha,
-      src_desc.handle(), input_data.opaque(), &beta, dest_desc.handle(),
-      output_data->opaque());
+  ScopedTensorDescriptor src_desc(input_dimensions, CUDNN_DATA_FLOAT);
+  ScopedTensorDescriptor dest_desc(output_dimensions, CUDNN_DATA_FLOAT);
+  ScopedPoolingDescriptor pooling_desc(pooling_dimensions);
+
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+  auto status = cudnnPoolingForward(
+      cudnn.handle(), pooling_desc.handle(), &alpha, src_desc.handle(),
+      input_data.opaque(), &beta, dest_desc.handle(), output_data->opaque());
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "failed to enqueue forward pooling on stream: "
                << ToString(status);
@@ -4403,25 +4160,18 @@ bool CudnnSupport::DoPoolForward(
     const DeviceMemory<Eigen::half>& input_data,
     const dnn::BatchDescriptor& output_dimensions,
     DeviceMemory<Eigen::half>* output_data) {
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status);
-    return false;
-  }
-
   // Alpha is the scaling factor for input.
   float alpha = 1.0;
   // Beta is the scaling factor for output.
   float beta = 0.0;
 
-  ScopedTensorDescriptor src_desc{parent_, input_dimensions, CUDNN_DATA_HALF};
-  ScopedTensorDescriptor dest_desc{parent_, output_dimensions, CUDNN_DATA_HALF};
-  ScopedPoolingDescriptor pooling_desc{parent_, pooling_dimensions};
-  status = wrap::cudnnPoolingForward(
-      this, stream, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha,
-      src_desc.handle(), input_data.opaque(), &beta, dest_desc.handle(),
-      output_data->opaque());
+  ScopedTensorDescriptor src_desc(input_dimensions, CUDNN_DATA_HALF);
+  ScopedTensorDescriptor dest_desc(output_dimensions, CUDNN_DATA_HALF);
+  ScopedPoolingDescriptor pooling_desc(pooling_dimensions);
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+  auto status = cudnnPoolingForward(
+      cudnn.handle(), pooling_desc.handle(), &alpha, src_desc.handle(),
+      input_data.opaque(), &beta, dest_desc.handle(), output_data->opaque());
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "failed to enqueue forward pooling on stream: "
                << ToString(status);
@@ -4438,27 +4188,21 @@ bool CudnnSupport::DoPoolBackward(
     const DeviceMemory<double>& output_data,
     const DeviceMemory<double>& input_diff_data,
     DeviceMemory<double>* output_diff_data) {
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status);
-    return false;
-  }
-
   // Alpha is the scaling factor for input.
   double alpha = 1.0;
   // Beta is the scaling factor for output.
   double beta = 0.0;
 
-  ScopedTensorDescriptor src_desc{parent_, input_dimensions, CUDNN_DATA_DOUBLE};
-  ScopedTensorDescriptor dest_desc{parent_, output_dimensions,
-                                   CUDNN_DATA_DOUBLE};
-  ScopedPoolingDescriptor pooling_desc{parent_, pooling_dimensions};
-  status = wrap::cudnnPoolingBackward(
-      this, stream, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha,
-      dest_desc.handle(), output_data.opaque(), dest_desc.handle(),
-      input_diff_data.opaque(), src_desc.handle(), input_data.opaque(), &beta,
-      src_desc.handle(), output_diff_data->opaque());
+  ScopedTensorDescriptor src_desc(input_dimensions, CUDNN_DATA_DOUBLE);
+  ScopedTensorDescriptor dest_desc(output_dimensions, CUDNN_DATA_DOUBLE);
+  ScopedPoolingDescriptor pooling_desc(pooling_dimensions);
+
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+  auto status = cudnnPoolingBackward(
+      cudnn.handle(), pooling_desc.handle(), &alpha, dest_desc.handle(),
+      output_data.opaque(), dest_desc.handle(), input_diff_data.opaque(),
+      src_desc.handle(), input_data.opaque(), &beta, src_desc.handle(),
+      output_diff_data->opaque());
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "failed to enqueue backward pooling on stream: "
                << ToString(status);
@@ -4475,27 +4219,21 @@ bool CudnnSupport::DoPoolBackward(
     const DeviceMemory<float>& output_data,
     const DeviceMemory<float>& input_diff_data,
     DeviceMemory<float>* output_diff_data) {
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status);
-    return false;
-  }
-
   // Alpha is the scaling factor for input.
   float alpha = 1.0;
   // Beta is the scaling factor for output.
   float beta = 0.0;
 
-  ScopedTensorDescriptor src_desc{parent_, input_dimensions, CUDNN_DATA_FLOAT};
-  ScopedTensorDescriptor dest_desc{parent_, output_dimensions,
-                                   CUDNN_DATA_FLOAT};
-  ScopedPoolingDescriptor pooling_desc{parent_, pooling_dimensions};
-  status = wrap::cudnnPoolingBackward(
-      this, stream, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha,
-      dest_desc.handle(), output_data.opaque(), dest_desc.handle(),
-      input_diff_data.opaque(), src_desc.handle(), input_data.opaque(), &beta,
-      src_desc.handle(), output_diff_data->opaque());
+  ScopedTensorDescriptor src_desc(input_dimensions, CUDNN_DATA_FLOAT);
+  ScopedTensorDescriptor dest_desc(output_dimensions, CUDNN_DATA_FLOAT);
+  ScopedPoolingDescriptor pooling_desc(pooling_dimensions);
+
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+  auto status = cudnnPoolingBackward(
+      cudnn.handle(), pooling_desc.handle(), &alpha, dest_desc.handle(),
+      output_data.opaque(), dest_desc.handle(), input_diff_data.opaque(),
+      src_desc.handle(), input_data.opaque(), &beta, src_desc.handle(),
+      output_diff_data->opaque());
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "failed to enqueue backward pooling on stream: "
                << ToString(status);
@@ -4512,26 +4250,21 @@ bool CudnnSupport::DoPoolBackward(
     const DeviceMemory<Eigen::half>& output_data,
     const DeviceMemory<Eigen::half>& input_diff_data,
     DeviceMemory<Eigen::half>* output_diff_data) {
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status);
-    return false;
-  }
-
   // Alpha is the scaling factor for input.
   float alpha = 1.0;
   // Beta is the scaling factor for output.
   float beta = 0.0;
 
-  ScopedTensorDescriptor src_desc{parent_, input_dimensions, CUDNN_DATA_HALF};
-  ScopedTensorDescriptor dest_desc{parent_, output_dimensions, CUDNN_DATA_HALF};
-  ScopedPoolingDescriptor pooling_desc{parent_, pooling_dimensions};
-  status = wrap::cudnnPoolingBackward(
-      this, stream, ToHandle(dnn_handle_), pooling_desc.handle(), &alpha,
-      dest_desc.handle(), output_data.opaque(), dest_desc.handle(),
-      input_diff_data.opaque(), src_desc.handle(), input_data.opaque(), &beta,
-      src_desc.handle(), output_diff_data->opaque());
+  ScopedTensorDescriptor src_desc(input_dimensions, CUDNN_DATA_HALF);
+  ScopedTensorDescriptor dest_desc(output_dimensions, CUDNN_DATA_HALF);
+  ScopedPoolingDescriptor pooling_desc(pooling_dimensions);
+
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+  auto status = cudnnPoolingBackward(
+      cudnn.handle(), pooling_desc.handle(), &alpha, dest_desc.handle(),
+      output_data.opaque(), dest_desc.handle(), input_diff_data.opaque(),
+      src_desc.handle(), input_data.opaque(), &beta, src_desc.handle(),
+      output_diff_data->opaque());
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "failed to enqueue backward pooling on stream: "
                << ToString(status);
@@ -4553,7 +4286,7 @@ bool CudnnSupport::DoNormalizeWithDimensions(
     const DeviceMemory<float>& input_data, DeviceMemory<float>* output_data) {
   // Check for unsupported modes.
   if (normalize_descriptor.wrap_around()) {
-    LOG(ERROR) << "CUDA LRN does not support wrap-around mode";
+    LOG(ERROR) << "CUDA LRN does not support cudnn-around mode";
     return false;
   }
   if (normalize_descriptor.segment_size()) {
@@ -4561,26 +4294,21 @@ bool CudnnSupport::DoNormalizeWithDimensions(
     return false;
   }
 
-  // Launch the normalization.
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status);
-    return false;
-  }
-
-  ScopedTensorDescriptor dims{parent_, dimensions, CUDNN_DATA_FLOAT};
-  ScopedNormalizeDescriptor normalize{parent_, normalize_descriptor};
+  ScopedTensorDescriptor dims(dimensions, CUDNN_DATA_FLOAT);
+  ScopedNormalizeDescriptor normalize(normalize_descriptor);
 
   // Alpha is the scaling factor for input.
   float alpha = 1.0f;
   // Beta is the scaling factor for output.
   float beta = 0.0f;
 
-  status = wrap::cudnnLRNCrossChannelForward(
-      this, stream, ToHandle(dnn_handle_), normalize.handle(),
-      CUDNN_LRN_CROSS_CHANNEL_DIM1, &alpha, dims.handle(), input_data.opaque(),
-      &beta, dims.handle(), output_data->opaque());
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+
+  // Launch the normalization.
+  auto status = cudnnLRNCrossChannelForward(
+      cudnn.handle(), normalize.handle(), CUDNN_LRN_CROSS_CHANNEL_DIM1, &alpha,
+      dims.handle(), input_data.opaque(), &beta, dims.handle(),
+      output_data->opaque());
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "failed to run cudnnLRNCrossChannelForward";
     return false;
@@ -4596,7 +4324,7 @@ bool CudnnSupport::DoNormalizeBackwardWithDimensions(
     DeviceMemory<float>* raw_variable_gradient) {
   // Check for unsupported modes.
   if (normalize_descriptor.wrap_around()) {
-    LOG(ERROR) << "CUDA LRN does not support wrap-around mode";
+    LOG(ERROR) << "CUDA LRN does not support cudnn-around mode";
     return false;
   }
   if (normalize_descriptor.segment_size()) {
@@ -4604,23 +4332,16 @@ bool CudnnSupport::DoNormalizeBackwardWithDimensions(
     return false;
   }
 
-  mutex_lock lock{dnn_handle_mutex_};
-  auto status = wrap::cudnnSetStream(this, stream, ToHandle(dnn_handle_));
-  if (status != CUDNN_STATUS_SUCCESS) {
-    LOG(ERROR) << "failed to set stream for cudnn handle: " << ToString(status);
-    return false;
-  }
-
-  ScopedTensorDescriptor dims{parent_, dimensions, CUDNN_DATA_FLOAT};
-  ScopedNormalizeDescriptor normalize{parent_, normalize_descriptor};
+  ScopedTensorDescriptor dims(dimensions, CUDNN_DATA_FLOAT);
+  ScopedNormalizeDescriptor normalize(normalize_descriptor);
 
   float alpha = 1.0f;
   float beta = 0.0f;
 
-  status = wrap::cudnnLRNCrossChannelBackward(
-      this, stream, ToHandle(dnn_handle_), normalize.handle(),
-      CUDNN_LRN_CROSS_CHANNEL_DIM1, &alpha, dims.handle(),
-      normalized_data.opaque(), dims.handle(),
+  auto cudnn = cudnn_->GetHandle(parent_, stream);
+  auto status = cudnnLRNCrossChannelBackward(
+      cudnn.handle(), normalize.handle(), CUDNN_LRN_CROSS_CHANNEL_DIM1, &alpha,
+      dims.handle(), normalized_data.opaque(), dims.handle(),
       normalized_variable_gradient.opaque(), dims.handle(), raw_data.opaque(),
       &beta, dims.handle(), raw_variable_gradient->opaque());
   if (status != CUDNN_STATUS_SUCCESS) {
@@ -4736,17 +4457,14 @@ bool CudnnSupport::DeriveOutputBatchDescriptor(
     const FilterDescriptor& filter_descriptor,
     const dnn::ConvolutionDescriptor& convolution_descriptor,
     dnn::BatchDescriptor* output_batch_descriptor) {
-  ScopedTensorDescriptor input_nd{parent_, batch_descriptor, CUDNN_DATA_FLOAT};
-  ScopedFilterDescriptor filter{parent_, filter_descriptor, batch_descriptor,
-                                CUDNN_DATA_FLOAT};
-  ScopedConvolutionDescriptor conv{parent_, convolution_descriptor,
-                                   CUDNN_DATA_FLOAT};
+  ScopedTensorDescriptor input_nd(batch_descriptor, CUDNN_DATA_FLOAT);
+  ScopedFilterDescriptor filter(filter_descriptor, CUDNN_DATA_FLOAT);
+  ScopedConvolutionDescriptor conv(convolution_descriptor, CUDNN_DATA_FLOAT);
 
   int dn = batch_descriptor.ndims() + 2;
   std::vector<int> dims(dn);  // in BDYX
-  auto status = wrap::cudnnGetConvolutionNdForwardOutputDim(
-      parent_, conv.handle(), input_nd.handle(), filter.handle(), dn,
-      dims.data());
+  auto status = cudnnGetConvolutionNdForwardOutputDim(
+      conv.handle(), input_nd.handle(), filter.handle(), dn, dims.data());
   if (status != CUDNN_STATUS_SUCCESS) {
     LOG(ERROR) << "could not get output tensor for convolution: "
                << ToString(status);
diff --git a/tensorflow/stream_executor/cuda/cuda_dnn.h b/tensorflow/stream_executor/cuda/cuda_dnn.h
index 8a0458bc80..e2de3c62d8 100644
--- a/tensorflow/stream_executor/cuda/cuda_dnn.h
+++ b/tensorflow/stream_executor/cuda/cuda_dnn.h
@@ -19,6 +19,7 @@ limitations under the License.
 #ifndef TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_DNN_H_
 #define TENSORFLOW_STREAM_EXECUTOR_CUDA_CUDA_DNN_H_
 
+#include "tensorflow/stream_executor/cuda/cuda_activation.h"
 #include "tensorflow/stream_executor/dnn.h"
 #include "tensorflow/stream_executor/lib/status.h"
 #include "tensorflow/stream_executor/platform/mutex.h"
@@ -42,7 +43,6 @@ extern const PluginId kCuDnnPlugin;
 class CudnnSupport : public dnn::DnnSupport {
  public:
   explicit CudnnSupport(CUDAExecutor* parent);
-  ~CudnnSupport() override;
 
   port::Status Init() override;
   port::StatusOr<perftools::gputools::dnn::VersionInfo> GetVersion() override;
@@ -624,54 +624,11 @@ class CudnnSupport : public dnn::DnnSupport {
                          dnn::DataType output_type, float scale,
                          DeviceMemoryBase* output_data) override;
 
-  const Stream* GetCurrentDnnStream() const
-      SHARED_LOCKS_REQUIRED(dnn_handle_mutex_) {
-    return current_dnn_stream_;
-  }
-
-  void SetCurrentDnnStream(Stream* stream)
-      EXCLUSIVE_LOCKS_REQUIRED(dnn_handle_mutex_) {
-    current_dnn_stream_ = stream;
-  }
-
-  CUDAExecutor* GetParentExecutor() { return parent_; }
-
-  // Guards the enqueueing of DNN operations via the dnn_handle_ below, and
-  // access to current_dnn_stream_.
-  //
-  // This is a public member because we need to add thread safty annotations in
-  // the cudnn wrapper functions in the cc file, which need to access this
-  // mutex (the annotations require C++ permission checks).
-  mutex dnn_handle_mutex_;
-
  private:
   CUDAExecutor* parent_;  // Parent executor object. Not owned.
 
-  // cudnn library handle. cudnnHandle_t type is not present in this header to
-  // prevent third-party library header inclusions from leaking outside the
-  // single cuda_dnn translation unit.
-  void* dnn_handle_ GUARDED_BY(dnn_handle_mutex_);
-
-  // The current cudnn stream that is set by SetCurrentDnnStream().
-  Stream* current_dnn_stream_ GUARDED_BY(dnn_handle_mutex_);
-
-  // NOTE(keveman): Temporary data layout transformation until cuDNN supports
-  // kBatchYXDepth for backward pass. This function allocates temporary memory,
-  // lays out the source data into the temporary but in the kBatchDepthXY
-  // layout, and returns the temporary memory. The caller is responsible for
-  // deallocating the temporary. Since the allocation is done using Stream's
-  // AllocateTemporaryMemory, a later BlockHostUntilDone could be used for
-  // deallocation.
-  //
-  // transform_scratch is populated with a legitimate temporary allocation iff
-  // the original output data needs to be transformed.
-  template<class T>
-  DeviceMemory<T> MaybeTransformLayout(
-      Stream* stream,
-      dnn::BatchDescriptor* output_descriptor,
-      DeviceMemory<T> backward_output_data,
-      std::unique_ptr<TemporaryDeviceMemory<T>>* transform_scratch)
-      EXCLUSIVE_LOCKS_REQUIRED(dnn_handle_mutex_);
+  // Provides access to the cuDNN handle.
+  std::unique_ptr<class CudnnAccess> cudnn_;
 
   template <class T, class U>
   bool DoBatchNormalizationForwardImpl(
@@ -700,7 +657,7 @@ class CudnnSupport : public dnn::DnnSupport {
 
   template <class T>
   bool DoConvolveImpl(Stream* stream,
-                      const dnn::BatchDescriptor& batch_descriptor,
+                      const dnn::BatchDescriptor& input_descriptor,
                       const DeviceMemory<T>& input_data,
                       const dnn::FilterDescriptor& filter_descriptor,
                       const DeviceMemory<T>& filter_data,
-- 
GitLab


From 62c50e197e25c661048fe90fdd177a87eda47376 Mon Sep 17 00:00:00 2001
From: Allen Lavoie <allenl@google.com>
Date: Tue, 8 May 2018 14:00:30 -0700
Subject: [PATCH 374/395] Avoid string formatting in assert_same_float_dtype
 unless there's an error

Especially helpful when executing eagerly

PiperOrigin-RevId: 195871887
---
 tensorflow/python/ops/check_ops.py | 30 +++++++++++++++++++++++-------
 1 file changed, 23 insertions(+), 7 deletions(-)

diff --git a/tensorflow/python/ops/check_ops.py b/tensorflow/python/ops/check_ops.py
index 306055d202..cabc1e724c 100644
--- a/tensorflow/python/ops/check_ops.py
+++ b/tensorflow/python/ops/check_ops.py
@@ -1169,19 +1169,35 @@ def _assert_same_base_type(items, expected_type=None):
   Raises:
     ValueError: If any types do not match.
   """
-  original_item_str = None
+  original_expected_type = expected_type
+  mismatch = False
   for item in items:
     if item is not None:
       item_type = item.dtype.base_dtype
       if not expected_type:
         expected_type = item_type
-        original_item_str = item.name if hasattr(item, 'name') else str(item)
       elif expected_type != item_type:
-        raise ValueError('%s, type=%s, must be of the same type (%s)%s.' % (
-            item.name if hasattr(item, 'name') else str(item),
-            item_type, expected_type,
-            (' as %s' % original_item_str) if original_item_str else ''))
-  return expected_type
+        mismatch = True
+        break
+  if mismatch:
+    # Loop back through and build up an informative error message (this is very
+    # slow, so we don't do it unless we found an error above).
+    expected_type = original_expected_type
+    original_item_str = None
+    for item in items:
+      if item is not None:
+        item_type = item.dtype.base_dtype
+        if not expected_type:
+          expected_type = item_type
+          original_item_str = item.name if hasattr(item, 'name') else str(item)
+        elif expected_type != item_type:
+          raise ValueError('%s, type=%s, must be of the same type (%s)%s.' % (
+              item.name if hasattr(item, 'name') else str(item),
+              item_type, expected_type,
+              (' as %s' % original_item_str) if original_item_str else ''))
+    return expected_type  # Should be unreachable
+  else:
+    return expected_type
 
 
 @tf_export('assert_same_float_dtype')
-- 
GitLab


From 1d94ed775417bad963a91cd6831a51e7538d797b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 14:00:48 -0700
Subject: [PATCH 375/395] Increase size of test
 tensorflow/contrib/layers:rev_block_lib_test to medium to avoid flaky
 timeouts.

PiperOrigin-RevId: 195871947
---
 tensorflow/contrib/layers/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/layers/BUILD b/tensorflow/contrib/layers/BUILD
index d5b3b279a1..7355a403ae 100644
--- a/tensorflow/contrib/layers/BUILD
+++ b/tensorflow/contrib/layers/BUILD
@@ -381,7 +381,7 @@ py_test(
 
 py_test(
     name = "rev_block_lib_test",
-    size = "small",
+    size = "medium",
     srcs = ["python/layers/rev_block_lib_test.py"],
     srcs_version = "PY2AND3",
     deps = [
-- 
GitLab


From d3f3fb5b5f2db18f890838b29cac94ba88335f0a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 14:41:48 -0700
Subject: [PATCH 376/395] Increase shard count of
 tensorflow/contrib/distributions:mixture_test to avoid flaky timeouts in asan
 mode

PiperOrigin-RevId: 195878809
---
 tensorflow/contrib/distributions/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/contrib/distributions/BUILD b/tensorflow/contrib/distributions/BUILD
index 47f2ebca77..8021ec6141 100644
--- a/tensorflow/contrib/distributions/BUILD
+++ b/tensorflow/contrib/distributions/BUILD
@@ -372,6 +372,7 @@ cuda_py_test(
         "//tensorflow/python:random_ops",
         "//tensorflow/python:variables",
     ],
+    shard_count = 4,
 )
 
 cuda_py_test(
-- 
GitLab


From f58effe44dea9e8c7bf092c6779cd430994f7a72 Mon Sep 17 00:00:00 2001
From: Alexandre Passos <apassos@google.com>
Date: Tue, 8 May 2018 14:42:35 -0700
Subject: [PATCH 377/395] Do not differentiage integers in the eager API.

This is similar to the change made in:
https://github.com/tensorflow/tensorflow/commit/f63750645826df65b05cad505546a86f0e347674
for backpropagation during graph construction via tf.gradients()

PiperOrigin-RevId: 195878952
---
 tensorflow/c/eager/tape.h                   | 36 +++++++++---
 tensorflow/contrib/eager/python/tfe_test.py |  6 +-
 tensorflow/python/eager/backprop.py         |  5 ++
 tensorflow/python/eager/backprop_test.py    | 10 +++-
 tensorflow/python/eager/pywrap_tensor.cc    |  6 ++
 tensorflow/python/eager/pywrap_tensor.h     |  1 +
 tensorflow/python/eager/pywrap_tfe_src.cc   | 62 ++++++++++++++++++---
 7 files changed, 107 insertions(+), 19 deletions(-)

diff --git a/tensorflow/c/eager/tape.h b/tensorflow/c/eager/tape.h
index 8026076b9e..e9ed3395c4 100644
--- a/tensorflow/c/eager/tape.h
+++ b/tensorflow/c/eager/tape.h
@@ -130,13 +130,15 @@ class GradientTape {
     }
   }
 
-  bool ShouldRecord(gtl::ArraySlice<int64> tensor_ids);
+  bool ShouldRecord(gtl::ArraySlice<int64> tensor_ids,
+                    gtl::ArraySlice<tensorflow::DataType> dtypes);
 
   void Watch(int64 tensor_id);
 
   void RecordOperation(const string& op_type,
                        gtl::ArraySlice<TapeTensor> output_tensors,
                        gtl::ArraySlice<int64> input_tensor_id,
+                       gtl::ArraySlice<tensorflow::DataType> input_dtypes,
                        BackwardFunction* backward_function,
                        const std::function<void()>& backward_function_deleter);
 
@@ -170,12 +172,30 @@ class GradientTape {
 
 // Template instantiations here
 
+inline bool IsDtypeTrainable(DataType dtype) {
+  switch (dtype) {
+    case DT_HALF:
+    case DT_BFLOAT16:
+    case DT_FLOAT:
+    case DT_DOUBLE:
+    case DT_COMPLEX64:
+    case DT_COMPLEX128:
+    case DT_RESOURCE:
+    case DT_VARIANT:
+      return true;
+    default:
+      return false;
+  }
+}
+
 template <typename Gradient, typename BackwardFunction>
 bool GradientTape<Gradient, BackwardFunction>::ShouldRecord(
-    gtl::ArraySlice<int64> tensor_ids) {
-  for (int64 i : tensor_ids) {
-    if (tensor_tape_.find(i) != tensor_tape_.end()) {
-      return true;
+    gtl::ArraySlice<int64> tensor_ids,
+    gtl::ArraySlice<tensorflow::DataType> dtypes) {
+  CHECK_EQ(tensor_ids.size(), dtypes.size());
+  for (int i = 0; i < tensor_ids.size(); ++i) {
+    if (tensor_tape_.find(tensor_ids[i]) != tensor_tape_.end()) {
+      return IsDtypeTrainable(dtypes[i]);
     }
   }
   return false;
@@ -189,9 +209,11 @@ void GradientTape<Gradient, BackwardFunction>::Watch(int64 tensor_id) {
 template <typename Gradient, typename BackwardFunction>
 void GradientTape<Gradient, BackwardFunction>::RecordOperation(
     const string& op_type, gtl::ArraySlice<TapeTensor> output_tensors,
-    gtl::ArraySlice<int64> input_tensor_id, BackwardFunction* backward_function,
+    gtl::ArraySlice<int64> input_tensor_id,
+    gtl::ArraySlice<tensorflow::DataType> input_dtypes,
+    BackwardFunction* backward_function,
     const std::function<void()>& backward_function_deleter) {
-  if (!ShouldRecord(input_tensor_id)) {
+  if (!ShouldRecord(input_tensor_id, input_dtypes)) {
     backward_function_deleter();
     return;
   }
diff --git a/tensorflow/contrib/eager/python/tfe_test.py b/tensorflow/contrib/eager/python/tfe_test.py
index e80ccbb74d..db50b33af2 100644
--- a/tensorflow/contrib/eager/python/tfe_test.py
+++ b/tensorflow/contrib/eager/python/tfe_test.py
@@ -57,7 +57,7 @@ class TFETest(test_util.TensorFlowTestCase):
       return math_ops.multiply(x, x)
 
     grad = tfe.gradients_function(square)
-    self.assertEquals([6], [x.numpy() for x in grad(3)])
+    self.assertEquals([6], [x.numpy() for x in grad(3.)])
 
   def testGradOfGrad(self):
 
@@ -66,7 +66,7 @@ class TFETest(test_util.TensorFlowTestCase):
 
     grad = tfe.gradients_function(square)
     gradgrad = tfe.gradients_function(lambda x: grad(x)[0])
-    self.assertEquals([2], [x.numpy() for x in gradgrad(3)])
+    self.assertEquals([2], [x.numpy() for x in gradgrad(3.)])
 
   def testCustomGrad(self):
 
@@ -80,7 +80,7 @@ class TFETest(test_util.TensorFlowTestCase):
       return y, grad_fn
 
     grad = tfe.gradients_function(f)
-    self.assertEquals([12], [x.numpy() for x in grad(3)])
+    self.assertEquals([12], [x.numpy() for x in grad(3.)])
 
   def testGPU(self):
     if tfe.num_gpus() <= 0:
diff --git a/tensorflow/python/eager/backprop.py b/tensorflow/python/eager/backprop.py
index d04b004451..967c128280 100644
--- a/tensorflow/python/eager/backprop.py
+++ b/tensorflow/python/eager/backprop.py
@@ -358,6 +358,8 @@ def gradients_function(f, params=None):
   assert y_grad.numpy() == (2 ** 3) - 2 * 2 * 3
   ```
 
+  Note that only tensors with real or complex dtypes are differentiable.
+
   Args:
    f: function to be differentiated. If `f` returns a scalar, this scalar will
      be differentiated. If `f` returns a tensor or list of tensors, by default
@@ -700,6 +702,9 @@ class GradientTape(object):
   dz_dx = g.gradient(z, x)  # 108.0 (4*x^3 at x = 3)
   dy_dx = g.gradient(y, x)  # 6.0
   del g  # Drop the reference to the tape
+  ```
+
+  Note that only tensors with real or complex dtypes are differentiable.
   """
 
   def __init__(self, persistent=False):
diff --git a/tensorflow/python/eager/backprop_test.py b/tensorflow/python/eager/backprop_test.py
index 8d9959fe20..be674487f1 100644
--- a/tensorflow/python/eager/backprop_test.py
+++ b/tensorflow/python/eager/backprop_test.py
@@ -124,6 +124,14 @@ class BackpropTest(test.TestCase):
     grad_fn = backprop.gradients_function(f)
     self.assertAllEqual(2., grad_fn(1., dy=2.)[0])
 
+  def testGradientInteger(self):
+
+    def f(x):
+      return x + x
+
+    int_tensor = constant_op.constant(1)
+    self.assertEqual(backprop.gradients_function(f)(int_tensor)[0], None)
+
   def testErrors(self):
 
     @custom_gradient.custom_gradient
@@ -753,7 +761,7 @@ class BackpropTest(test.TestCase):
       return result, grad
 
     x = resource_variable_ops.ResourceVariable(
-        initial_value=3, name='X.' + self.id())
+        initial_value=3., name='X.' + self.id())
 
     def f():
       return my_square(x)
diff --git a/tensorflow/python/eager/pywrap_tensor.cc b/tensorflow/python/eager/pywrap_tensor.cc
index b5b4e394e3..b3aadd55ce 100644
--- a/tensorflow/python/eager/pywrap_tensor.cc
+++ b/tensorflow/python/eager/pywrap_tensor.cc
@@ -650,6 +650,12 @@ tensorflow::int64 EagerTensor_id(const PyObject* tensor) {
   return reinterpret_cast<const EagerTensor*>(tensor)->id;
 }
 
+tensorflow::DataType EagerTensor_dtype(const PyObject* tensor) {
+  CHECK(EagerTensor_CheckExact(tensor));
+  return static_cast<tensorflow::DataType>(TFE_TensorHandleDataType(
+      reinterpret_cast<const EagerTensor*>(tensor)->handle));
+}
+
 PyObject* TFE_Py_InitEagerTensor(PyObject* base_class) {
   if (!PyType_Check(base_class)) {
     PyErr_SetString(
diff --git a/tensorflow/python/eager/pywrap_tensor.h b/tensorflow/python/eager/pywrap_tensor.h
index 63ab1ed84d..88982b0c85 100644
--- a/tensorflow/python/eager/pywrap_tensor.h
+++ b/tensorflow/python/eager/pywrap_tensor.h
@@ -21,6 +21,7 @@ limitations under the License.
 
 bool EagerTensor_CheckExact(const PyObject* o);
 tensorflow::int64 EagerTensor_id(const PyObject* tensor);
+tensorflow::DataType EagerTensor_dtype(const PyObject* tensor);
 
 namespace tensorflow {
 TFE_TensorHandle* ConvertToEagerTensor(PyObject* value, PyObject* dtype);
diff --git a/tensorflow/python/eager/pywrap_tfe_src.cc b/tensorflow/python/eager/pywrap_tfe_src.cc
index 4ecba1a46b..48a5b21dc7 100644
--- a/tensorflow/python/eager/pywrap_tfe_src.cc
+++ b/tensorflow/python/eager/pywrap_tfe_src.cc
@@ -843,6 +843,24 @@ static tensorflow::int64 FastTensorId(PyObject* tensor) {
   return id;
 }
 
+static tensorflow::DataType FastTensorDtype(PyObject* tensor) {
+  if (EagerTensor_CheckExact(tensor)) {
+    return EagerTensor_dtype(tensor);
+  }
+  PyObject* dtype_field = PyObject_GetAttrString(tensor, "dtype");
+  if (dtype_field == nullptr) {
+    return tensorflow::DT_INVALID;
+  }
+  PyObject* enum_field = PyObject_GetAttrString(dtype_field, "_type_enum");
+  Py_DECREF(dtype_field);
+  if (dtype_field == nullptr) {
+    return tensorflow::DT_INVALID;
+  }
+  tensorflow::int64 id = MakeInt(enum_field);
+  Py_DECREF(enum_field);
+  return static_cast<tensorflow::DataType>(id);
+}
+
 class GradientTape
     : public tensorflow::eager::GradientTape<PyObject, PyObject> {
  public:
@@ -1053,15 +1071,18 @@ PyObject* TFE_Py_TapeSetShouldRecord(PyObject* tensors) {
   // TODO(apassos) consider not building a list and changing the API to check
   // each tensor individually.
   std::vector<tensorflow::int64> tensor_ids;
+  std::vector<tensorflow::DataType> dtypes;
   tensor_ids.reserve(len);
+  dtypes.reserve(len);
   for (int i = 0; i < len; ++i) {
     PyObject* item = PySequence_Fast_GET_ITEM(seq, i);
     tensor_ids.push_back(FastTensorId(item));
+    dtypes.push_back(FastTensorDtype(item));
   }
   Py_DECREF(seq);
   auto tape_set = *tape_set_ptr;
   for (TFE_Py_Tape* tape : tape_set) {
-    if (tape->tape->ShouldRecord(tensor_ids)) {
+    if (tape->tape->ShouldRecord(tensor_ids, dtypes)) {
       Py_RETURN_TRUE;
     }
   }
@@ -1169,9 +1190,27 @@ PyObject* TFE_Py_TapeWatchedVariables(PyObject* tape) {
 }
 
 namespace {
-void TapeSetRecordOperation(PyObject* op_type, PyObject* output_tensors,
-                            const std::vector<tensorflow::int64>& input_ids,
-                            PyObject* backward_function) {
+std::vector<tensorflow::DataType> MakeTensorDtypeList(PyObject* tensors) {
+  PyObject* seq = PySequence_Fast(tensors, "expected a sequence");
+  if (seq == nullptr) {
+    return {};
+  }
+  int len = PySequence_Fast_GET_SIZE(seq);
+  std::vector<tensorflow::DataType> list;
+  list.reserve(len);
+  for (int i = 0; i < len; ++i) {
+    PyObject* tensor = PySequence_Fast_GET_ITEM(seq, i);
+    list.push_back(FastTensorDtype(tensor));
+  }
+  Py_DECREF(seq);
+  return list;
+}
+
+void TapeSetRecordOperation(
+    PyObject* op_type, PyObject* output_tensors,
+    const std::vector<tensorflow::int64>& input_ids,
+    const std::vector<tensorflow::DataType>& input_dtypes,
+    PyObject* backward_function) {
   std::vector<tensorflow::eager::TapeTensor> output_info;
   PyObject* seq = PySequence_Fast(output_tensors,
                                   "expected a sequence of integer tensor ids");
@@ -1206,7 +1245,7 @@ void TapeSetRecordOperation(PyObject* op_type, PyObject* output_tensors,
   for (TFE_Py_Tape* tape : SafeTapeSet()) {
     Py_INCREF(backward_function);
     tape->tape->RecordOperation(
-        op_type_str, output_info, input_ids, backward_function,
+        op_type_str, output_info, input_ids, input_dtypes, backward_function,
         [backward_function]() { Py_DECREF(backward_function); });
   }
 }
@@ -1221,7 +1260,11 @@ void TFE_Py_TapeSetRecordOperation(PyObject* op_type, PyObject* output_tensors,
   std::vector<tensorflow::int64> input_ids = MakeTensorIDList(input_tensors);
   if (PyErr_Occurred()) return;
 
-  TapeSetRecordOperation(op_type, output_tensors, input_ids, backward_function);
+  std::vector<tensorflow::DataType> input_dtypes =
+      MakeTensorDtypeList(input_tensors);
+  if (PyErr_Occurred()) return;
+  TapeSetRecordOperation(op_type, output_tensors, input_ids, input_dtypes,
+                         backward_function);
 }
 
 void TFE_Py_TapeSetDeleteTrace(tensorflow::int64 tensor_id) {
@@ -1710,10 +1753,12 @@ PyObject* RecordGradient(PyObject* op_name, PyObject* inputs, PyObject* attrs,
                          PyObject* results, PyObject* name) {
   std::vector<tensorflow::int64> input_ids = MakeTensorIDList(inputs);
   if (PyErr_Occurred()) return nullptr;
+  std::vector<tensorflow::DataType> input_dtypes = MakeTensorDtypeList(inputs);
+  if (PyErr_Occurred()) return nullptr;
 
   bool should_record = false;
   for (TFE_Py_Tape* tape : SafeTapeSet()) {
-    if (tape->tape->ShouldRecord(input_ids)) {
+    if (tape->tape->ShouldRecord(input_ids, input_dtypes)) {
       should_record = true;
       break;
     }
@@ -1744,7 +1789,8 @@ PyObject* RecordGradient(PyObject* op_name, PyObject* inputs, PyObject* attrs,
   Py_DECREF(callback_args);
   if (backward_function == nullptr) return nullptr;
 
-  TapeSetRecordOperation(op_name, results, input_ids, backward_function);
+  TapeSetRecordOperation(op_name, results, input_ids, input_dtypes,
+                         backward_function);
 
   Py_DECREF(backward_function);
 
-- 
GitLab


From 96fa17d853149f9bdf33c09b89abdd8c6521044d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 14:45:01 -0700
Subject: [PATCH 378/395] Increase shard_count of
 tensorflow/python/estimator:estimator_test to avoid flaky asan timeouts

PiperOrigin-RevId: 195879364
---
 tensorflow/python/estimator/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/estimator/BUILD b/tensorflow/python/estimator/BUILD
index b25cc7aa26..2d9a084bc6 100644
--- a/tensorflow/python/estimator/BUILD
+++ b/tensorflow/python/estimator/BUILD
@@ -489,6 +489,7 @@ py_library(
 py_test(
     name = "estimator_test",
     srcs = ["estimator_test.py"],
+    shard_count = 4,
     srcs_version = "PY2AND3",
     tags = ["notsan"],  # b/67510291
     deps = [
-- 
GitLab


From 34f6241fd822b15c66085dbd1cbec092196d0225 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 14:50:24 -0700
Subject: [PATCH 379/395] Add missing ":haswell" match to list of platform
 selectors.

PiperOrigin-RevId: 195880275
---
 tensorflow/contrib/lite/kernels/internal/BUILD | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorflow/contrib/lite/kernels/internal/BUILD b/tensorflow/contrib/lite/kernels/internal/BUILD
index 7ec4782f96..54188217d9 100644
--- a/tensorflow/contrib/lite/kernels/internal/BUILD
+++ b/tensorflow/contrib/lite/kernels/internal/BUILD
@@ -387,6 +387,9 @@ cc_library(
         ":armv7a": [
             ":neon_tensor_utils",
         ],
+        ":haswell": [
+            ":neon_tensor_utils",
+        ],
         ":ios_armv7": [
             ":neon_tensor_utils",
         ],
-- 
GitLab


From 71387153307a7df94bcdc5307de95e6e228a95a9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 15:26:44 -0700
Subject: [PATCH 380/395] Increase shard count of
 tensorflow/python/keras:lstm_test to avoid flaky timeouts

PiperOrigin-RevId: 195886372
---
 tensorflow/python/keras/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/keras/BUILD b/tensorflow/python/keras/BUILD
index 37b24841bd..77db07b86b 100755
--- a/tensorflow/python/keras/BUILD
+++ b/tensorflow/python/keras/BUILD
@@ -604,6 +604,7 @@ py_test(
     name = "lstm_test",
     size = "medium",
     srcs = ["_impl/keras/layers/lstm_test.py"],
+    shard_count = 4,
     srcs_version = "PY2AND3",
     tags = [
         "noasan",  # times out b/63678675
-- 
GitLab


From 24d9492f07e8cba89ae94cf01a1bcae22fcf438b Mon Sep 17 00:00:00 2001
From: jjsjann123 <alex.jann2012@gmail.com>
Date: Tue, 8 May 2018 16:30:08 -0700
Subject: [PATCH 381/395] [tftrt update] (#19135)

* [tftrt update]
  code cleaning, removed some boilerplate code

* addressing comments
---
 .../contrib/tensorrt/convert/convert_nodes.cc | 396 ++++++------------
 1 file changed, 130 insertions(+), 266 deletions(-)

diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
index 3767596f8c..be559d30e0 100644
--- a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
+++ b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc
@@ -346,11 +346,10 @@ void ReorderCKtoKC(const TRT_ShapedWeights& iweights,
       break;
     }
     case tensorflow::DataType::DT_HALF: {
-      Reorder2(
-          {k, c}, static_cast<Eigen::half const*>(iweights.GetValues()),
-          istrides,
-          static_cast<Eigen::half*>(const_cast<void*>(oweights->GetValues())),
-          ostrides);
+      Reorder2({k, c}, static_cast<Eigen::half const*>(iweights.GetValues()),
+               istrides, static_cast<Eigen::half*>(
+                             const_cast<void*>(oweights->GetValues())),
+               ostrides);
       break;
     }
     default:
@@ -1159,9 +1158,9 @@ tensorflow::Status BinaryTensorOpTensor(
   CHECK_EQ_TYPE(tensor_r->getType(), dtype);
   auto op_pair = ops.find(node_def.op());
   if (op_pair == ops.end())
-    return tensorflow::errors::Unimplemented(
-        "binary op: " + node_def.op() +
-        " not supported at: " + node_def.name());
+    return tensorflow::errors::Unimplemented("binary op: " + node_def.op() +
+                                             " not supported at: " +
+                                             node_def.name());
 
   nvinfer1::IElementWiseLayer* layer = ctx.network()->addElementWise(
       *const_cast<nvinfer1::ITensor*>(tensor_l),
@@ -2214,64 +2213,63 @@ tensorflow::Status ConvertCalibrationNodeToEngineNode(
   return tensorflow::Status::OK();
 }
 
-tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) {
-  // Visit nodes in reverse topological order and construct the TRT network.
-
-  // Toposort
+tensorflow::Status ReverseTopologicalSort(
+    const tensorrt::convert::SubGraphParams& s,
+    std::list<tensorflow::Node*>* order) {
   std::vector<tensorflow::Node*> order_vec;
   tensorflow::GetPostOrder(s.graph, &order_vec);
   // Select just the subgraph
-  std::list<tensorflow::Node*> order;
   for (tensorflow::Node* node : order_vec) {
     if (s.subgraph_node_ids.count(node->id())) {
-      order.push_front(node);  // we want topological order to construct the
+      // We want topological order to contstruct the
       // network layer by layer
+      order->push_front(node);
     }
   }
-  // topological order is needed to build TRT network
-  static int static_id = 0;
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status SetInputList(
+    const tensorrt::convert::SubGraphParams& s,
+    tensorflow::NodeDefBuilder* op_builder,
+    const std::vector<string>* input_names,
+    std::vector<tensorflow::DataType>* input_dtypes) {
+  std::vector<tensorflow::NodeDefBuilder::NodeOut> income_edges;
+  VLOG(2) << "input edge size: " << input_names->size();
+  for (size_t i = 0; i < input_names->size(); ++i) {
+    VLOG(2) << "input edges: " << i << " " << input_names->at(i);
+    int output_idx = s.input_inds.at(i).second;
+    // we wired up the input here already, it is redundant to do it again in
+    //  ConvertSubGraphToTensorRT(convert_graph.cc)
+    auto incoming_edge = tensorflow::NodeDefBuilder::NodeOut(
+        input_names->at(i), output_idx, input_dtypes->at(i));
+    income_edges.push_back(incoming_edge);
+  }
+  tensorflow::gtl::ArraySlice<tensorflow::NodeDefBuilder::NodeOut> input_list(
+      income_edges);
+  op_builder->Input(input_list);
+  return tensorflow::Status::OK();
+}
+
+string SubgraphNameScopeGenerator(const std::list<tensorflow::Node*>* order) {
   string subgraph_name_scope;
-  if (!order.empty()) {
-    subgraph_name_scope = order.front()->name();
+  if (!order->empty()) {
+    subgraph_name_scope = order->front()->name();
   }
-  for (const tensorflow::Node* node : order) {
+  for (const tensorflow::Node* node : *order) {
     subgraph_name_scope = GetCommonNameScope(subgraph_name_scope, node->name());
   }
   // TODO(sami,ben,jie): proper naming!
-  string calib_op_name =
-      StrCat(subgraph_name_scope, "my_trt_calib_op_", static_id);
-  string engine_name = StrCat(subgraph_name_scope, "my_trt_op", static_id);
-  static_id++;
-  auto trt_rmgr = tensorflow::tensorrt::TRTResourceManager::instance();
-  auto op_rmgr = trt_rmgr->getManager("TRTCalibOps");
-  auto op_res = new tensorflow::tensorrt::TRTCalibrationResource();
-  TF_CHECK_OK(op_rmgr->Create(calib_op_name, calib_op_name, op_res));
-  op_res->logger_ = new tensorflow::tensorrt::Logger();
-  cudaSetDevice(s.cuda_gpu_id_);
-  op_res->builder_ = nvinfer1::createInferBuilder(*(op_res->logger_));
-  op_res->allocator_ = s.allocator_;
-#if NV_TENSORRT_MAJOR > 3
-  op_res->builder_->setGpuAllocator(s.allocator_.get());
-#endif
-  if (!op_res->builder_) {
-    return tensorflow::errors::Internal(
-        "failed to create TensorRT builder object");
-  }
-
-  op_res->network_ = op_res->builder_->createNetwork();
-  if (!op_res->network_) {
-    return tensorflow::errors::Internal(
-        "failed to create TensorRT network object");
-  }
-
-  // Build the network
-  auto weight_rmgr = trt_rmgr->getManager("WeightStore");
-  auto ws = new tensorflow::tensorrt::TRTWeightStore();
-  TF_CHECK_OK(weight_rmgr->Create(calib_op_name, calib_op_name, ws));
-  Converter converter(op_res->network_, ws, s.precision_mode == FP16MODE);
+  return subgraph_name_scope;
+}
 
-  std::vector<string> input_names;
-  std::vector<tensorflow::DataType> input_dtypes;
+tensorflow::Status ConvertSubgraph(
+    Converter& converter, tensorrt::convert::SubGraphParams& s,
+    std::list<tensorflow::Node*>* order, std::vector<string>* input_names,
+    std::vector<tensorflow::DataType>* input_dtypes,
+    std::vector<string>* output_names,
+    std::vector<tensorflow::DataType>* output_dtypes,
+    const string& engine_name) {
   for (const std::pair<int, int>& input : s.input_inds) {
     VLOG(2) << "parsing input. Node id= " << input.first;
     int node_id = input.first;
@@ -2314,19 +2312,18 @@ tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) {
 
     auto op_info = op_info_vec.at(shape_inference_output_idx);
     tensorflow::DataType tf_dtype = op_info.dtype();
-    input_dtypes.push_back(tf_dtype);
+    input_dtypes->push_back(tf_dtype);
 
     nvinfer1::DataType dtype(nvinfer1::DataType::kFLOAT);
     auto type_status = ConvertDType(tf_dtype, &dtype);
     if (type_status != tensorflow::Status::OK()) {
-      LOG(WARNING) << "Data type conversion for input '" << node_name
-                   << "' failed";
+      LOG(WARNING) << "Type conversion failed for " << node_name;
       return type_status;
     }
 
-    VLOG(2) << "accessing output index of: " << output_idx
+    VLOG(2) << "Accessing output index of: " << output_idx
             << ", at node: " << node_name
-            << "with output entry from shape_map: " << op_info_vec.size();
+            << " with output entry from shape_map: " << op_info_vec.size();
     // TODO(ben,jie): update TRT input format/dimension
     nvinfer1::DimsCHW input_dim_pseudo_chw;
     for (int i = 0; i < 3; i++) input_dim_pseudo_chw.d[i] = 1;
@@ -2352,33 +2349,29 @@ tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) {
       input_tensor_name = StrCat(node_name, ":", output_idx);
     }
 
-    input_names.push_back(input_tensor_name);
+    input_names->push_back(input_tensor_name);
     nvinfer1::ITensor* input_tensor = converter.network()->addInput(
         input_tensor_name.c_str(), dtype, input_dim_pseudo_chw);
 
     if (!input_tensor)
       return tensorflow::errors::InvalidArgument(
           "Failed to create Input layer");
-    VLOG(2) << "input tensor name :" << input_tensor_name;
+    VLOG(2) << "Input tensor name :" << input_tensor_name;
 
     if (!converter.insert_input_tensor(input_tensor_name, input_tensor))
       return tensorflow::errors::AlreadyExists(
-          "output tensor already exists for op: " + input_tensor_name);
+          "Output tensor already exists for op: " + input_tensor_name);
   }
 
-  VLOG(2) << "finished sorting";
-
-  for (const tensorflow::Node* node : order) {
+  for (const tensorflow::Node* node : *order) {
     const tensorflow::NodeDef& node_def = node->def();
-    VLOG(2) << "converting node: " << node_def.name() << " , " << node_def.op();
+    VLOG(2) << "Converting node: " << node_def.name() << " , " << node_def.op();
     TF_RETURN_IF_ERROR(converter.convert_node(node_def));
   }
 
-  VLOG(2) << "finished conversion";
+  VLOG(2) << "Finished conversion";
 
   // Gather output metadata
-  std::vector<string> output_names;
-  std::vector<tensorflow::DataType> output_dtypes;
   int trt_engine_op_output_idx = 0;
   for (const std::pair<int, int>& output : s.output_inds) {
     int node_id = output.first;
@@ -2393,14 +2386,13 @@ tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) {
              : StrCat(engine_name, ":", trt_engine_op_output_idx),
          {output_idx, tensor_name}});
     trt_engine_op_output_idx++;
-    if (output_idx != 0) {
-      tensor_name = StrCat(tensor_name, ":", output_idx);
-    }
-    VLOG(1) << "output tensor name: " << tensor_name;
-    output_names.push_back(tensor_name);
+    if (output_idx != 0)
+      tensorflow::strings::StrAppend(&tensor_name, ":", output_idx);
+    VLOG(2) << "Output tensor name: " << tensor_name;
+    output_names->push_back(tensor_name);
     auto tensor_or_weights = converter.get_tensor(tensor_name);
     if (!tensor_or_weights.is_tensor()) {
-      return tensorflow::errors::InvalidArgument("Output node'" + tensor_name +
+      return tensorflow::errors::InvalidArgument("Output node '" + tensor_name +
                                                  "' is weights not tensor");
     }
     nvinfer1::ITensor* tensor = tensor_or_weights.tensor();
@@ -2410,12 +2402,65 @@ tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) {
     }
     converter.network()->markOutput(*tensor);
     tensorflow::DataType tf_dtype = node->output_type(output_idx);
-    output_dtypes.push_back(tf_dtype);
+    output_dtypes->push_back(tf_dtype);
     nvinfer1::DataType trt_dtype = nvinfer1::DataType::kFLOAT;
     TF_RETURN_IF_ERROR(ConvertDType(tf_dtype, &trt_dtype));
     tensor->setType(trt_dtype);
   }
 
+  return tensorflow::Status::OK();
+}
+
+tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) {
+  // Visit nodes in reverse topological order and construct the TRT network.
+  // Toposort
+  std::list<tensorflow::Node*> order;
+  TF_RETURN_IF_ERROR(ReverseTopologicalSort(s, &order));
+
+  static int static_id = 0;
+  string subgraph_name_scope = SubgraphNameScopeGenerator(&order);
+  // TODO(sami,ben,jie): proper naming!
+  string calib_op_name =
+      StrCat(subgraph_name_scope, "my_trt_calib_op_", static_id);
+  string engine_name = StrCat(subgraph_name_scope, "my_trt_op", static_id);
+  static_id++;
+
+  auto trt_rmgr = tensorflow::tensorrt::TRTResourceManager::instance();
+  auto op_rmgr = trt_rmgr->getManager("TRTCalibOps");
+  auto op_res = new tensorflow::tensorrt::TRTCalibrationResource();
+  TF_CHECK_OK(op_rmgr->Create(calib_op_name, calib_op_name, op_res));
+  op_res->logger_ = new tensorflow::tensorrt::Logger();
+  cudaSetDevice(s.cuda_gpu_id_);
+  op_res->builder_ = nvinfer1::createInferBuilder(*(op_res->logger_));
+  op_res->allocator_ = s.allocator_;
+#if NV_TENSORRT_MAJOR > 3
+  op_res->builder_->setGpuAllocator(s.allocator_.get());
+#endif
+  if (!op_res->builder_) {
+    return tensorflow::errors::Internal(
+        "failed to create TensorRT builder object");
+  }
+
+  op_res->network_ = op_res->builder_->createNetwork();
+  if (!op_res->network_) {
+    return tensorflow::errors::Internal(
+        "failed to create TensorRT network object");
+  }
+
+  // Build the network
+  auto weight_rmgr = trt_rmgr->getManager("WeightStore");
+  auto ws = new tensorflow::tensorrt::TRTWeightStore();
+  TF_CHECK_OK(weight_rmgr->Create(calib_op_name, calib_op_name, ws));
+  Converter converter(op_res->network_, ws, s.precision_mode == FP16MODE);
+
+  std::vector<string> input_names;
+  std::vector<tensorflow::DataType> input_dtypes;
+  std::vector<string> output_names;
+  std::vector<tensorflow::DataType> output_dtypes;
+  TF_RETURN_IF_ERROR(ConvertSubgraph(converter, s, &order, &input_names,
+                                     &input_dtypes, &output_names,
+                                     &output_dtypes, engine_name));
+
   VLOG(2) << "Finished processing outputs";
 
   // Build the engine
@@ -2427,21 +2472,8 @@ tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) {
   // Build the TRT op
   // TODO(sami,ben,jie): proper naming!
   tensorflow::NodeDefBuilder op_builder(calib_op_name, "TRTCalibOp");
-  std::vector<tensorflow::NodeDefBuilder::NodeOut> income_edges;
-  for (size_t i = 0; i < input_names.size(); ++i) {
-    int output_idx = s.input_inds.at(i).second;
-    // we wired up the input here already, it is redundant to do it again in
-    //  ConvertSubGraphToTensorRT(convert_graph.cc)
-    auto incoming_edge = tensorflow::NodeDefBuilder::NodeOut(
-        input_names.at(i), output_idx, input_dtypes.at(i));
-    VLOG(1) << calib_op_name << " input " << i << " = " << input_names.at(i)
-            << ":" << output_idx
-            << " dType= " << tensorflow::DataTypeString(input_dtypes.at(i));
-    income_edges.push_back(incoming_edge);
-  }
-  tensorflow::gtl::ArraySlice<tensorflow::NodeDefBuilder::NodeOut> input_list(
-      income_edges);
-  op_builder.Input(input_list);
+  SetInputList(s, &op_builder, &input_names, &input_dtypes);
+
   std::vector<string> segment_names;
   segment_names.reserve(s.subgraph_node_ids.size());
   for (int i : s.subgraph_node_ids) {
@@ -2465,20 +2497,12 @@ tensorflow::Status InjectCalibrationNode(tensorrt::convert::SubGraphParams& s) {
 tensorflow::Status ConvertSubGraphToTensorRTNodeDef(
     tensorrt::convert::SubGraphParams& s) {
   // Visit nodes in reverse topological order and construct the TRT network.
-
-  // Toposort
-  std::vector<tensorflow::Node*> order_vec;
-  tensorflow::GetPostOrder(s.graph, &order_vec);
-  // Select just the subgraph
   std::list<tensorflow::Node*> order;
-  for (tensorflow::Node* node : order_vec) {
-    if (s.subgraph_node_ids.count(node->id())) {
-      // We want topological order to contstruct the
-      // network layer by layer
-      order.push_front(node);
-    }
-  }
-  // Topological order is needed to build TRT network
+  TF_RETURN_IF_ERROR(ReverseTopologicalSort(s, &order));
+
+  static int static_id = 0;
+  string subgraph_name_scope = SubgraphNameScopeGenerator(&order);
+  string engine_name = StrCat(subgraph_name_scope, "my_trt_op", static_id++);
 
   tensorflow::tensorrt::Logger trt_logger;
   cudaSetDevice(s.cuda_gpu_id_);
@@ -2496,17 +2520,6 @@ tensorflow::Status ConvertSubGraphToTensorRTNodeDef(
         "Failed to create TensorRT network object");
   }
 
-  string subgraph_name_scope;
-  if (!order.empty()) {
-    subgraph_name_scope = order.front()->name();
-  }
-  for (const tensorflow::Node* node : order) {
-    subgraph_name_scope = GetCommonNameScope(subgraph_name_scope, node->name());
-  }
-  static int static_id = 0;
-  // TODO(sami,ben,jie): proper naming!
-  string engine_name = StrCat(subgraph_name_scope, "my_trt_op");
-  engine_name = StrCat(engine_name, static_id++);
   auto trt_rmgr = tensorflow::tensorrt::TRTResourceManager::instance();
   auto weight_rmgr = trt_rmgr->getManager("WeightStore");
   auto ws = new tensorflow::tensorrt::TRTWeightStore();
@@ -2517,147 +2530,11 @@ tensorflow::Status ConvertSubGraphToTensorRTNodeDef(
 
   std::vector<string> input_names;
   std::vector<tensorflow::DataType> input_dtypes;
-  for (const std::pair<int, int>& input : s.input_inds) {
-    VLOG(2) << "parsing input. Node id= " << input.first;
-    int node_id = input.first;
-    int output_idx = input.second;
-    tensorflow::Node* node = s.graph.FindNodeId(node_id);
-    auto node_name = node->name();
-    // input_names should use the node name in the graph
-    // here it should be the input tensor name -> matching the binding
-    // insert original node name without port
-    auto tensor_name = node_name;
-    if (output_idx != 0) {
-      tensor_name = StrCat(tensor_name, ":", output_idx);
-    }
-
-    VLOG(2) << "input name: " << node_name << " tensor_name: " << tensor_name
-            << " idx: " << output_idx;
-
-    auto shape_inference_node_name = node_name;
-    auto shape_inference_output_idx = output_idx;
-    // rewire the shape inference to original node in the graph
-    if (s.output_edge_map->count(tensor_name)) {
-      shape_inference_node_name = s.output_edge_map->at(tensor_name).second;
-      shape_inference_output_idx = s.output_edge_map->at(tensor_name).first;
-    }
-    if (shape_inference_output_idx < 0) continue;
-    VLOG(2) << "shapeinference name: " << shape_inference_node_name
-            << " idx: " << shape_inference_output_idx;
-
-    if (!s.graph_properties.HasOutputProperties(shape_inference_node_name))
-      return tensorflow::errors::Internal("failed to find input node: " +
-                                          shape_inference_node_name);
-
-    auto op_info_vec =
-        s.graph_properties.GetOutputProperties(shape_inference_node_name);
-    if (static_cast<int>(op_info_vec.size()) <= shape_inference_output_idx)
-      return tensorflow::errors::Internal(
-          "accessing output index of: ", shape_inference_output_idx,
-          ", at node: ", shape_inference_node_name,
-          " with output entry from shape_map: ", op_info_vec.size());
-
-    auto op_info = op_info_vec.at(shape_inference_output_idx);
-    tensorflow::DataType tf_dtype = op_info.dtype();
-    input_dtypes.push_back(tf_dtype);
-
-    nvinfer1::DataType dtype(nvinfer1::DataType::kFLOAT);
-    auto type_status = ConvertDType(tf_dtype, &dtype);
-    if (type_status != tensorflow::Status::OK()) {
-      LOG(WARNING) << "Type conversion failed for " << node_name;
-      return type_status;
-    }
-
-    VLOG(2) << "Accessing output index of: " << output_idx
-            << ", at node: " << node_name
-            << " with output entry from shape_map: " << op_info_vec.size();
-    // TODO(ben,jie): update TRT input format/dimension
-    nvinfer1::DimsCHW input_dim_pseudo_chw;
-    for (int i = 0; i < 3; i++) input_dim_pseudo_chw.d[i] = 1;
-
-    // TODO(jie): TRT 3.x only support 4 dimensional input tensor.
-    //            update the code once TRT 4.0 comes out.
-    if (op_info.shape().dim_size() != 4) {
-      string err_str = "Require 4 dimensional input.";
-      StrAppend(&err_str, " Got ", op_info.shape().dim_size(), " ",
-                shape_inference_node_name);
-      return tensorflow::errors::Unimplemented(err_str);
-    }
-
-    for (int i = 1; i < op_info.shape().dim_size(); i++) {
-      VLOG(2) << "dimension: " << i
-              << " , size: " << op_info.shape().dim(i).size();
-      input_dim_pseudo_chw.d[i - 1] = op_info.shape().dim(i).size();
-    }
-
-    // TODO(ben,jie): proper way to restore input tensor name?
-    auto input_tensor_name = node_name;
-    if (output_idx != 0) {
-      input_tensor_name = StrCat(node_name, ":", output_idx);
-    }
-
-    input_names.push_back(input_tensor_name);
-    nvinfer1::ITensor* input_tensor = converter.network()->addInput(
-        input_tensor_name.c_str(), dtype, input_dim_pseudo_chw);
-
-    if (!input_tensor)
-      return tensorflow::errors::InvalidArgument(
-          "Failed to create Input layer");
-    VLOG(2) << "Input tensor name :" << input_tensor_name;
-
-    if (!converter.insert_input_tensor(input_tensor_name, input_tensor))
-      return tensorflow::errors::AlreadyExists(
-          "Output tensor already exists for op: " + input_tensor_name);
-  }
-
-  VLOG(2) << "Finished sorting";
-
-  for (const tensorflow::Node* node : order) {
-    const tensorflow::NodeDef& node_def = node->def();
-    VLOG(2) << "Converting node: " << node_def.name() << " , " << node_def.op();
-    TF_RETURN_IF_ERROR(converter.convert_node(node_def));
-  }
-
-  VLOG(2) << "Finished conversion";
-
-  // Gather output metadata
   std::vector<string> output_names;
   std::vector<tensorflow::DataType> output_dtypes;
-  int trt_engine_op_output_idx = 0;
-  for (const std::pair<int, int>& output : s.output_inds) {
-    int node_id = output.first;
-    int output_idx = output.second;
-    tensorflow::Node* node = s.graph.FindNodeId(node_id);
-    string op_name = node->name();
-    string tensor_name = op_name;
-
-    s.output_edge_map->insert(
-        {trt_engine_op_output_idx == 0
-             ? engine_name
-             : StrCat(engine_name, ":", trt_engine_op_output_idx),
-         {output_idx, tensor_name}});
-    trt_engine_op_output_idx++;
-    if (output_idx != 0)
-      tensorflow::strings::StrAppend(&tensor_name, ":", output_idx);
-    VLOG(2) << "Output tensor name: " << tensor_name;
-    output_names.push_back(tensor_name);
-    auto tensor_or_weights = converter.get_tensor(tensor_name);
-    if (!tensor_or_weights.is_tensor()) {
-      return tensorflow::errors::InvalidArgument("Output node '" + tensor_name +
-                                                 "' is weights not tensor");
-    }
-    nvinfer1::ITensor* tensor = tensor_or_weights.tensor();
-    if (!tensor) {
-      return tensorflow::errors::NotFound("Output tensor not found: " +
-                                          tensor_name);
-    }
-    converter.network()->markOutput(*tensor);
-    tensorflow::DataType tf_dtype = node->output_type(output_idx);
-    output_dtypes.push_back(tf_dtype);
-    nvinfer1::DataType trt_dtype = nvinfer1::DataType::kFLOAT;
-    TF_RETURN_IF_ERROR(ConvertDType(tf_dtype, &trt_dtype));
-    tensor->setType(trt_dtype);
-  }
+  TF_RETURN_IF_ERROR(ConvertSubgraph(converter, s, &order, &input_names,
+                                     &input_dtypes, &output_names,
+                                     &output_dtypes, engine_name));
 
   VLOG(2) << "Finished output";
 
@@ -2693,20 +2570,7 @@ tensorflow::Status ConvertSubGraphToTensorRTNodeDef(
 
   // Build the TRT op
   tensorflow::NodeDefBuilder op_builder(engine_name, "TRTEngineOp");
-  std::vector<tensorflow::NodeDefBuilder::NodeOut> income_edges;
-  VLOG(2) << "input edge size: " << input_names.size();
-  for (size_t i = 0; i < input_names.size(); ++i) {
-    VLOG(2) << "input edges: " << i << " " << input_names.at(i);
-    int output_idx = s.input_inds.at(i).second;
-    // we wired up the input here already, it is redundant to do it again in
-    //  ConvertSubGraphToTensorRT(convert_graph.cc)
-    auto incoming_edge = tensorflow::NodeDefBuilder::NodeOut(
-        input_names.at(i), output_idx, input_dtypes.at(i));
-    income_edges.push_back(incoming_edge);
-  }
-  tensorflow::gtl::ArraySlice<tensorflow::NodeDefBuilder::NodeOut> input_list(
-      income_edges);
-  op_builder.Input(input_list);
+  SetInputList(s, &op_builder, &input_names, &input_dtypes);
 
   VLOG(0) << "Finished op preparation";
 
-- 
GitLab


From c0fb9413914d983cad2ea6bb4997033a1f0dd722 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lgeiger@users.noreply.github.com>
Date: Wed, 9 May 2018 01:31:39 +0200
Subject: [PATCH 382/395] [tfgan] Allow to add custom eval metrics to
 GANEstimator (#19133)

---
 .../estimator/python/gan_estimator_impl.py    |  7 ++++-
 .../estimator/python/gan_estimator_test.py    |  9 +++++++
 .../gan/python/estimator/python/head_impl.py  | 27 ++++++++++++++-----
 .../gan/python/estimator/python/head_test.py  |  7 ++++-
 4 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_impl.py b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_impl.py
index e3fc6bf0f0..4092b32004 100644
--- a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_impl.py
+++ b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_impl.py
@@ -112,6 +112,7 @@ class GANEstimator(estimator.Estimator):
                generator_optimizer=None,
                discriminator_optimizer=None,
                get_hooks_fn=None,
+               get_eval_metric_ops_fn=None,
                add_summaries=None,
                use_loss_summaries=True,
                config=None):
@@ -146,6 +147,9 @@ class GANEstimator(estimator.Estimator):
         list of hooks. These hooks are run on the generator and discriminator
         train ops, and can be used to implement the GAN training scheme.
         Defaults to `train.get_sequential_train_hooks()`.
+      get_eval_metric_ops_fn: A function that takes a `GANModel`, and returns a
+        dict of metric results keyed by name. The output of this function is
+        passed into `tf.estimator.EstimatorSpec` during evaluation.
       add_summaries: `None`, a single `SummaryType`, or a list of `SummaryType`.
       use_loss_summaries: If `True`, add loss summaries. If `False`, does not.
         If `None`, uses defaults.
@@ -160,7 +164,8 @@ class GANEstimator(estimator.Estimator):
               else discriminator_optimizer)
       gan_head = head_lib.gan_head(
           generator_loss_fn, discriminator_loss_fn, gopt, dopt,
-          use_loss_summaries, get_hooks_fn=get_hooks_fn)
+          use_loss_summaries, get_hooks_fn=get_hooks_fn,
+          get_eval_metric_ops_fn=get_eval_metric_ops_fn)
       return _gan_model_fn(
           features, labels, mode, generator_fn, discriminator_fn, gan_head,
           add_summaries)
diff --git a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py
index 6bbd173f86..955482599b 100644
--- a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py
+++ b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py
@@ -38,6 +38,7 @@ from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import metrics as metrics_lib
 from tensorflow.python.ops import parsing_ops
 from tensorflow.python.platform import test
 from tensorflow.python.summary.writer import writer_cache
@@ -194,6 +195,12 @@ class GANEstimatorIntegrationTest(test.TestCase):
       lr = learning_rate_decay.exponential_decay(1.0, gstep, 10, 0.9)
       return training.GradientDescentOptimizer(lr)
 
+    def get_metrics(gan_model):
+      return {
+          'mse_custom_metric': metrics_lib.mean_squared_error(
+              gan_model.real_data, gan_model.generated_data)
+      }
+
     gopt = make_opt if lr_decay else training.GradientDescentOptimizer(1.0)
     dopt = make_opt if lr_decay else training.GradientDescentOptimizer(1.0)
     est = estimator.GANEstimator(
@@ -203,6 +210,7 @@ class GANEstimatorIntegrationTest(test.TestCase):
         discriminator_loss_fn=losses.wasserstein_discriminator_loss,
         generator_optimizer=gopt,
         discriminator_optimizer=dopt,
+        get_eval_metric_ops_fn=get_metrics,
         model_dir=self._model_dir)
 
     # TRAIN
@@ -215,6 +223,7 @@ class GANEstimatorIntegrationTest(test.TestCase):
     self.assertIn('loss', six.iterkeys(scores))
     self.assertEqual(scores['discriminator_loss'] + scores['generator_loss'],
                      scores['loss'])
+    self.assertIn('mse_custom_metric', six.iterkeys(scores))
 
     # PREDICT
     predictions = np.array([x for x in est.predict(predict_input_fn)])
diff --git a/tensorflow/contrib/gan/python/estimator/python/head_impl.py b/tensorflow/contrib/gan/python/estimator/python/head_impl.py
index d174cb3bb2..ff903a78cc 100644
--- a/tensorflow/contrib/gan/python/estimator/python/head_impl.py
+++ b/tensorflow/contrib/gan/python/estimator/python/head_impl.py
@@ -39,7 +39,7 @@ def _summary_key(head_name, val):
 def gan_head(generator_loss_fn, discriminator_loss_fn, generator_optimizer,
              discriminator_optimizer, use_loss_summaries=True,
              get_hooks_fn=tfgan_train.get_sequential_train_hooks(),
-             name=None):
+             get_eval_metric_ops_fn=None, name=None):
   """Creates a `GANHead`.
 
   Args:
@@ -51,9 +51,12 @@ def gan_head(generator_loss_fn, discriminator_loss_fn, generator_optimizer,
     discriminator_optimizer: Same as `generator_optimizer`, but for the
       discriminator updates.
     use_loss_summaries: If `True`, add loss summaries. If `False`, does not.
-        If `None`, uses defaults.
-    get_hooks_fn: A function that takes a GANTrainOps tuple and returns a list
-        of hooks.
+      If `None`, uses defaults.
+    get_hooks_fn: A function that takes a `GANTrainOps` tuple and returns a
+      list of hooks.
+    get_eval_metric_ops_fn: A function that takes a `GANModel`, and returns a
+      dict of metric results keyed by name. The output of this function is
+      passed into `tf.estimator.EstimatorSpec` during evaluation.
     name: name of the head. If provided, summary and metrics keys will be
       suffixed by `"/" + name`.
 
@@ -66,6 +69,7 @@ def gan_head(generator_loss_fn, discriminator_loss_fn, generator_optimizer,
                  discriminator_optimizer=discriminator_optimizer,
                  use_loss_summaries=use_loss_summaries,
                  get_hooks_fn=get_hooks_fn,
+                 get_eval_metric_ops_fn=get_eval_metric_ops_fn,
                  name=name)
 
 
@@ -76,6 +80,7 @@ class GANHead(head._Head):  # pylint: disable=protected-access
                generator_optimizer, discriminator_optimizer,
                use_loss_summaries=True,
                get_hooks_fn=None,
+               get_eval_metric_ops_fn=None,
                name=None):
     """`Head` for GAN training.
 
@@ -89,8 +94,11 @@ class GANHead(head._Head):  # pylint: disable=protected-access
         discriminator updates.
       use_loss_summaries: If `True`, add loss summaries. If `False`, does not.
         If `None`, uses defaults.
-      get_hooks_fn: A function that takes a GANTrainOps tuple and returns a list
-        of hooks. Defaults to `train.get_sequential_train_hooks()`
+      get_hooks_fn: A function that takes a `GANTrainOps` tuple and returns a
+        list of hooks. Defaults to `train.get_sequential_train_hooks()`
+      get_eval_metric_ops_fn: A function that takes a `GANModel`, and returns a
+        dict of metric results keyed by name. The output of this function is
+        passed into `tf.estimator.EstimatorSpec` during evaluation.
       name: name of the head. If provided, summary and metrics keys will be
         suffixed by `"/" + name`.
     """
@@ -108,6 +116,7 @@ class GANHead(head._Head):  # pylint: disable=protected-access
     self._generator_optimizer = generator_optimizer
     self._discriminator_optimizer = discriminator_optimizer
     self._get_hooks_fn = get_hooks_fn
+    self._get_eval_metric_ops_fn = get_eval_metric_ops_fn
     self._name = name
 
   @property
@@ -187,6 +196,12 @@ class GANHead(head._Head):  # pylint: disable=protected-access
               _summary_key(self._name, 'discriminator_loss'):
                   metrics_lib.mean(gan_loss.discriminator_loss)
           }
+          if self._get_eval_metric_ops_fn is not None:
+            custom_eval_metric_ops = self._get_eval_metric_ops_fn(gan_model)
+            if not isinstance(custom_eval_metric_ops, dict):
+              raise TypeError('get_eval_metric_ops_fn must return a dict, '
+                              'received: {}'.format(custom_eval_metric_ops))
+            eval_metric_ops.update(custom_eval_metric_ops)
         return model_fn_lib.EstimatorSpec(
             mode=model_fn_lib.ModeKeys.EVAL,
             predictions=gan_model.generated_data,
diff --git a/tensorflow/contrib/gan/python/estimator/python/head_test.py b/tensorflow/contrib/gan/python/estimator/python/head_test.py
index 8168f005cd..6587f1fc60 100644
--- a/tensorflow/contrib/gan/python/estimator/python/head_test.py
+++ b/tensorflow/contrib/gan/python/estimator/python/head_test.py
@@ -62,9 +62,14 @@ class GANHeadTest(test.TestCase):
         generator_loss_fn=dummy_loss,
         discriminator_loss_fn=dummy_loss,
         generator_optimizer=training.GradientDescentOptimizer(1.0),
-        discriminator_optimizer=training.GradientDescentOptimizer(1.0))
+        discriminator_optimizer=training.GradientDescentOptimizer(1.0),
+        get_eval_metric_ops_fn=self.get_metrics)
     self.assertTrue(isinstance(self.gan_head, head.GANHead))
 
+  def get_metrics(self, gan_model):
+    self.assertTrue(isinstance(gan_model, tfgan_tuples.GANModel))
+    return {}
+
   def _test_modes_helper(self, mode):
     self.gan_head.create_estimator_spec(
         features=None,
-- 
GitLab


From 8039c947c3a2e0f3d780d0a1458bd40c6acd2145 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 15:33:37 -0700
Subject: [PATCH 383/395] Increase size of
 tensorflow/contrib/distributions:batch_reshape_test to medium to avoid flaky
 timeouts

PiperOrigin-RevId: 195887374
---
 tensorflow/contrib/distributions/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/contrib/distributions/BUILD b/tensorflow/contrib/distributions/BUILD
index 8021ec6141..a1d56066b4 100644
--- a/tensorflow/contrib/distributions/BUILD
+++ b/tensorflow/contrib/distributions/BUILD
@@ -460,7 +460,7 @@ cuda_py_test(
 
 cuda_py_test(
     name = "batch_reshape_test",
-    size = "small",
+    size = "medium",
     srcs = ["python/kernel_tests/batch_reshape_test.py"],
     additional_deps = [
         ":distributions_py",
-- 
GitLab


From d4d97591d036bed4ddedc48d66b55500a31b4ab5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 15:34:40 -0700
Subject: [PATCH 384/395] Increase shard count of
 tensorflow/contrib/learn:state_saving_rnn_estimator_test to avoid flaky
 timeouts

PiperOrigin-RevId: 195887546
---
 tensorflow/contrib/learn/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/contrib/learn/BUILD b/tensorflow/contrib/learn/BUILD
index 3b053cd4c6..4a360711f8 100644
--- a/tensorflow/contrib/learn/BUILD
+++ b/tensorflow/contrib/learn/BUILD
@@ -485,6 +485,7 @@ py_test(
     name = "state_saving_rnn_estimator_test",
     size = "medium",
     srcs = ["python/learn/estimators/state_saving_rnn_estimator_test.py"],
+    shard_count = 4,
     srcs_version = "PY2AND3",
     tags = ["noasan"],
     deps = [
-- 
GitLab


From 241e828794162436d1eb08c42e072249388f171f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 15:43:34 -0700
Subject: [PATCH 385/395] Add test to test suite.

PiperOrigin-RevId: 195888932
---
 tensorflow/contrib/lite/kernels/internal/BUILD                | 4 ++++
 .../contrib/lite/kernels/internal/quantization_util_test.cc   | 4 ++++
 2 files changed, 8 insertions(+)

diff --git a/tensorflow/contrib/lite/kernels/internal/BUILD b/tensorflow/contrib/lite/kernels/internal/BUILD
index 54188217d9..d8340d426a 100644
--- a/tensorflow/contrib/lite/kernels/internal/BUILD
+++ b/tensorflow/contrib/lite/kernels/internal/BUILD
@@ -5,6 +5,7 @@ package(default_visibility = [
 licenses(["notice"])  # Apache 2.0
 
 load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts")
+load("//tensorflow/contrib/lite:special_rules.bzl", "tflite_portable_test_suite")
 
 tflite_deps_intel = [
     "@arm_neon_2_x86_sse",
@@ -428,6 +429,7 @@ cc_test(
         "//conditions:default": [],
     }),
     linkstatic = 1,
+    tags = ["tflite_not_portable_ios"],
     deps = [
         ":tensor_utils",
         "//tensorflow/contrib/lite:builtin_op_data",
@@ -462,3 +464,5 @@ cc_test(
 )
 
 exports_files(["optimized/eigen_tensor_reduced_instantiations_oss.h"])
+
+tflite_portable_test_suite()
diff --git a/tensorflow/contrib/lite/kernels/internal/quantization_util_test.cc b/tensorflow/contrib/lite/kernels/internal/quantization_util_test.cc
index 3e9a3c29ee..2d74b3d384 100644
--- a/tensorflow/contrib/lite/kernels/internal/quantization_util_test.cc
+++ b/tensorflow/contrib/lite/kernels/internal/quantization_util_test.cc
@@ -167,6 +167,7 @@ TEST(QuantizationUtilTest, ChooseQuantizationParamsZeroPointOnMinBoundary) {
   EXPECT_EQ(qp.zero_point, 0);
 }
 
+#ifdef GTEST_HAS_DEATH_TEST
 TEST(QuantizationUtilTest, ChooseQuantizationParamsZeroNotInRange) {
   // Assumption is that zero is within the range.
   EXPECT_DEATH(ChooseQuantizationParams<uint8>(10.0, 30.0), "");
@@ -176,6 +177,7 @@ TEST(QuantizationUtilTest, ChooseQuantizationParamsEmptyRangePositive) {
   // Assumption is that zero is within the range.
   EXPECT_DEATH(ChooseQuantizationParams<uint8>(30.0, 30.0), "");
 }
+#endif  // GTEST_HAS_DEATH_TEST
 
 TEST(QuantizationUtilTest, ChooseQuantizationParamsEmptyRangeZero) {
   QuantizationParams qp = ChooseQuantizationParams<uint8>(0.0, 0.0);
@@ -189,6 +191,7 @@ TEST(QuantizationUtilTest, ChooseQuantizationParamsZeroPointOnMaxBoundary) {
   EXPECT_EQ(qp.zero_point, 255);
 }
 
+#ifdef GTEST_HAS_DEATH_TEST
 TEST(QuantizationUtilTest, ChooseQuantizationParamsInvalidRange) {
   EXPECT_DEATH(ChooseQuantizationParams<uint8>(10.0, -30.0), "");
 }
@@ -261,6 +264,7 @@ TEST(QuantizationUtilTest, PreprocessSoftmaxScaling) {
   EXPECT_THAT(quantize(2.0, 16.0, 5), Pair(2147483647, 31));
   EXPECT_THAT(quantize(2.0, 8.0, 5), Pair(1073741824, 31));
 }
+#endif  // GTEST_HAS_DEATH_TEST
 
 TEST(QuantizationUtilTest, CalculateInputRadius) {
   EXPECT_EQ(CalculateInputRadius(4, 27), 15);
-- 
GitLab


From 0028bf843d8846bd16b25bf5447b1649fde10fb7 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 16:16:57 -0700
Subject: [PATCH 386/395] add test for pruning useless function lib in graph.

PiperOrigin-RevId: 195893756
---
 .../optimizers/function_optimizer_test.cc     | 25 +++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/tensorflow/core/grappler/optimizers/function_optimizer_test.cc b/tensorflow/core/grappler/optimizers/function_optimizer_test.cc
index a2dbab3ded..0aaf57e947 100644
--- a/tensorflow/core/grappler/optimizers/function_optimizer_test.cc
+++ b/tensorflow/core/grappler/optimizers/function_optimizer_test.cc
@@ -835,5 +835,30 @@ TEST_F(FunctionOptimizerTest, SpecializeFunction_OncePerUniqueContext) {
   test::ExpectTensorEqual<float>(tensors_expected[5], tensors[5]);
 }
 
+TEST_F(FunctionOptimizerTest, PruningUselessLibraryFunctions) {
+  using test::function::NDef;
+  FunctionOptimizer optimizer(RewriterConfig::DEFAULT);
+  DisableFunctionSpecialization(&optimizer);
+  auto func = test::function::XTimesTwo();
+  (*func.mutable_attr())["_noinline"].set_b(true);
+  GrapplerItem item;
+  item.graph = test::function::GDef(
+      {NDef("x", "Placeholder", {}, {{"dtype", DT_FLOAT}}, "/device:CPU:0"),
+       NDef("y", "XTimesTwo", {"x"}, {{"T", DT_FLOAT}}, "/device:CPU:0"),
+       NDef("z", "Identity", {"y"}, {{"T", DT_FLOAT}}, "/device:CPU:0")},
+      // FunctionLib
+      {
+          func,
+          test::function::XTimesTwoInt32(),
+          test::function::XTimes16(),
+      });
+  GraphDef output;
+  Status status = optimizer.Optimize(nullptr, item, &output);
+  TF_EXPECT_OK(status);
+
+  EXPECT_EQ(output.library().function().size(), 1);
+  EXPECT_EQ(output.library().function(0).signature().name(), "XTimesTwo");
+}
+
 }  // namespace grappler
 }  // namespace tensorflow
-- 
GitLab


From bbebae04db61e137e4013a031f429543422ae373 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 16:20:02 -0700
Subject: [PATCH 387/395] Only use integer values for event_ndims.

event_ndims have the semantics of being an integer. However, other code paths (such as const_value)
can return back numpy wrapped arrays, which can mess with how values are cached. Instead extract
everything as an integer.

PiperOrigin-RevId: 195894216
---
 .../kernel_tests/bijectors/chain_test.py      | 10 ++++
 .../python/ops/bijectors/chain.py             | 44 ++++++++---------
 .../kernel_tests/distributions/util_test.py   | 26 ++++++++++
 .../python/ops/distributions/bijector_impl.py | 49 ++++++++++---------
 tensorflow/python/ops/distributions/util.py   | 24 +++++++++
 5 files changed, 106 insertions(+), 47 deletions(-)

diff --git a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/chain_test.py b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/chain_test.py
index ca20442c39..dc45114b1c 100644
--- a/tensorflow/contrib/distributions/python/kernel_tests/bijectors/chain_test.py
+++ b/tensorflow/contrib/distributions/python/kernel_tests/bijectors/chain_test.py
@@ -26,6 +26,7 @@ from tensorflow.contrib.distributions.python.ops.bijectors.exp import Exp
 from tensorflow.contrib.distributions.python.ops.bijectors.softmax_centered import SoftmaxCentered
 from tensorflow.contrib.distributions.python.ops.bijectors.softplus import Softplus
 from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
 from tensorflow.python.ops.distributions import bijector
 from tensorflow.python.ops.distributions.bijector_test_util import assert_scalar_congruency
 from tensorflow.python.platform import test
@@ -188,6 +189,15 @@ class ChainBijectorTest(test.TestCase):
         -np.log(6, dtype=np.float32) - np.sum(x),
         self.evaluate(chain.inverse_log_det_jacobian(y, event_ndims=1)))
 
+  def testChainIldjWithPlaceholder(self):
+    chain = Chain((Exp(), Exp()))
+    samples = array_ops.placeholder(
+        dtype=np.float32, shape=[None, 10], name="samples")
+    ildj = chain.inverse_log_det_jacobian(samples, event_ndims=0)
+    self.assertTrue(ildj is not None)
+    with self.test_session():
+      ildj.eval({samples: np.zeros([2, 10], np.float32)})
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/chain.py b/tensorflow/contrib/distributions/python/ops/bijectors/chain.py
index 85ad23e413..b158a51bb0 100644
--- a/tensorflow/contrib/distributions/python/ops/bijectors/chain.py
+++ b/tensorflow/contrib/distributions/python/ops/bijectors/chain.py
@@ -20,10 +20,9 @@ from __future__ import print_function
 
 import itertools
 
-from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor_util
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
 from tensorflow.python.ops.distributions import bijector
 
 
@@ -36,15 +35,6 @@ def _use_static_shape(input_tensor, ndims):
   return input_tensor.shape.is_fully_defined() and isinstance(ndims, int)
 
 
-def _maybe_get_event_ndims_statically(event_ndims):
-  static_event_ndims = (event_ndims if isinstance(event_ndims, int)
-                        else tensor_util.constant_value(event_ndims))
-  if static_event_ndims is not None:
-    return static_event_ndims
-
-  return event_ndims
-
-
 def _compute_min_event_ndims(bijector_list, compute_forward=True):
   """Computes the min_event_ndims associated with the give list of bijectors.
 
@@ -238,13 +228,13 @@ class Chain(bijector.Bijector):
     return y
 
   def _inverse_log_det_jacobian(self, y, **kwargs):
-    ildj = constant_op.constant(
-        0., dtype=y.dtype.base_dtype, name="inverse_log_det_jacobian")
+    y = ops.convert_to_tensor(y, name="y")
+    ildj = math_ops.cast(0., dtype=y.dtype.base_dtype)
 
     if not self.bijectors:
       return ildj
 
-    event_ndims = _maybe_get_event_ndims_statically(
+    event_ndims = self._maybe_get_event_ndims_statically(
         self.inverse_min_event_ndims)
 
     if _use_static_shape(y, event_ndims):
@@ -258,11 +248,12 @@ class Chain(bijector.Bijector):
 
       if _use_static_shape(y, event_ndims):
         event_shape = b.inverse_event_shape(event_shape)
-        event_ndims = _maybe_get_event_ndims_statically(event_shape.ndims)
+        event_ndims = self._maybe_get_event_ndims_statically(
+            event_shape.ndims)
       else:
         event_shape = b.inverse_event_shape_tensor(event_shape)
-        event_ndims = _maybe_get_event_ndims_statically(
-            array_ops.rank(event_shape))
+        event_ndims = self._maybe_get_event_ndims_statically(
+            array_ops.size(event_shape))
       y = b.inverse(y, **kwargs.get(b.name, {}))
     return ildj
 
@@ -274,13 +265,12 @@ class Chain(bijector.Bijector):
   def _forward_log_det_jacobian(self, x, **kwargs):
     x = ops.convert_to_tensor(x, name="x")
 
-    fldj = constant_op.constant(
-        0., dtype=x.dtype, name="inverse_log_det_jacobian")
+    fldj = math_ops.cast(0., dtype=x.dtype.base_dtype)
 
     if not self.bijectors:
       return fldj
 
-    event_ndims = _maybe_get_event_ndims_statically(
+    event_ndims = self._maybe_get_event_ndims_statically(
         self.forward_min_event_ndims)
 
     if _use_static_shape(x, event_ndims):
@@ -293,13 +283,21 @@ class Chain(bijector.Bijector):
           x, event_ndims=event_ndims, **kwargs.get(b.name, {}))
       if _use_static_shape(x, event_ndims):
         event_shape = b.forward_event_shape(event_shape)
-        event_ndims = _maybe_get_event_ndims_statically(event_shape.ndims)
+        event_ndims = self._maybe_get_event_ndims_statically(event_shape.ndims)
       else:
         event_shape = b.forward_event_shape_tensor(event_shape)
-        event_ndims = _maybe_get_event_ndims_statically(
-            array_ops.rank(event_shape))
+        event_ndims = self._maybe_get_event_ndims_statically(
+            array_ops.size(event_shape))
 
       x = b.forward(x, **kwargs.get(b.name, {}))
 
     return fldj
 
+  def _maybe_get_event_ndims_statically(self, event_ndims):
+    event_ndims_ = super(Chain, self)._maybe_get_event_ndims_statically(
+        event_ndims)
+    if event_ndims_ is None:
+      return event_ndims
+    return event_ndims_
+
+
diff --git a/tensorflow/python/kernel_tests/distributions/util_test.py b/tensorflow/python/kernel_tests/distributions/util_test.py
index f54f146e0a..b9fe197679 100644
--- a/tensorflow/python/kernel_tests/distributions/util_test.py
+++ b/tensorflow/python/kernel_tests/distributions/util_test.py
@@ -147,6 +147,32 @@ class AssertCloseTest(test.TestCase):
           array_ops.identity(w).eval(feed_dict=feed_dict)
 
 
+class MaybeGetStaticTest(test.TestCase):
+
+  def testGetStaticInt(self):
+    x = 2
+    self.assertEqual(x, du.maybe_get_static_value(x))
+    self.assertAllClose(
+        np.array(2.), du.maybe_get_static_value(x, dtype=np.float64))
+
+  def testGetStaticNumpyArray(self):
+    x = np.array(2, dtype=np.int32)
+    self.assertEqual(x, du.maybe_get_static_value(x))
+    self.assertAllClose(
+        np.array(2.), du.maybe_get_static_value(x, dtype=np.float64))
+
+  def testGetStaticConstant(self):
+    x = constant_op.constant(2, dtype=dtypes.int32)
+    self.assertEqual(np.array(2, dtype=np.int32), du.maybe_get_static_value(x))
+    self.assertAllClose(
+        np.array(2.), du.maybe_get_static_value(x, dtype=np.float64))
+
+  def testGetStaticPlaceholder(self):
+    x = array_ops.placeholder(dtype=dtypes.int32, shape=[1])
+    self.assertEqual(None, du.maybe_get_static_value(x))
+    self.assertEqual(None, du.maybe_get_static_value(x, dtype=np.float64))
+
+
 @test_util.with_c_api
 class GetLogitsAndProbsTest(test.TestCase):
 
diff --git a/tensorflow/python/ops/distributions/bijector_impl.py b/tensorflow/python/ops/distributions/bijector_impl.py
index 36eee5ce78..caceadf53a 100644
--- a/tensorflow/python/ops/distributions/bijector_impl.py
+++ b/tensorflow/python/ops/distributions/bijector_impl.py
@@ -33,6 +33,7 @@ from tensorflow.python.framework import tensor_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import check_ops
 from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.distributions import util as distribution_util
 
 
 __all__ = [
@@ -527,8 +528,6 @@ class Bijector(object):
       ValueError:  If a member of `graph_parents` is not a `Tensor`.
     """
     self._graph_parents = graph_parents or []
-    forward_min_event_ndims = get_static_value(forward_min_event_ndims)
-    inverse_min_event_ndims = get_static_value(inverse_min_event_ndims)
 
     if forward_min_event_ndims is None and inverse_min_event_ndims is None:
       raise ValueError("Must specify at least one of `forward_min_event_ndims` "
@@ -538,12 +537,23 @@ class Bijector(object):
     elif forward_min_event_ndims is None:
       forward_min_event_ndims = inverse_min_event_ndims
 
+    if not isinstance(forward_min_event_ndims, int):
+      raise TypeError("Expected forward_min_event_ndims to be of "
+                      "type int, got {}".format(
+                          type(forward_min_event_ndims).__name__))
+
+    if not isinstance(inverse_min_event_ndims, int):
+      raise TypeError("Expected inverse_min_event_ndims to be of "
+                      "type int, got {}".format(
+                          type(inverse_min_event_ndims).__name__))
+
     if forward_min_event_ndims < 0:
       raise ValueError("forward_min_event_ndims must be a non-negative "
                        "integer.")
     if inverse_min_event_ndims < 0:
       raise ValueError("inverse_min_event_ndims must be a non-negative "
                        "integer.")
+
     self._forward_min_event_ndims = forward_min_event_ndims
     self._inverse_min_event_ndims = inverse_min_event_ndims
     self._is_constant_jacobian = is_constant_jacobian
@@ -994,7 +1004,6 @@ class Bijector(object):
   def _reduce_jacobian_det_over_event(
       self, y, ildj, min_event_ndims, event_ndims):
     """Reduce jacobian over event_ndims - min_event_ndims."""
-    assert_static(min_event_ndims)
 
     if not self.is_constant_jacobian:
       return math_ops.reduce_sum(
@@ -1012,7 +1021,7 @@ class Bijector(object):
         axis=self._get_event_reduce_dims(min_event_ndims, event_ndims))
     # The multiplication by ones can change the inferred static shape so we try
     # to recover as much as possible.
-    event_ndims_ = get_static_value(event_ndims)
+    event_ndims_ = self._maybe_get_event_ndims_statically(event_ndims)
     if (event_ndims_ is not None and
         y.shape.ndims is not None and
         ildj.shape.ndims is not None):
@@ -1027,8 +1036,7 @@ class Bijector(object):
 
   def _get_event_reduce_dims(self, min_event_ndims, event_ndims):
     """Compute the reduction dimensions given event_ndims."""
-    assert_static(min_event_ndims)
-    event_ndims_ = get_static_value(event_ndims, np.int32)
+    event_ndims_ = self._maybe_get_event_ndims_statically(event_ndims)
 
     if event_ndims_ is not None:
       return [-index for index in range(1, event_ndims_ - min_event_ndims + 1)]
@@ -1038,8 +1046,7 @@ class Bijector(object):
 
   def _check_valid_event_ndims(self, min_event_ndims, event_ndims):
     """Check whether event_ndims is atleast min_event_ndims."""
-    assert_static(min_event_ndims)
-    event_ndims_ = get_static_value(event_ndims, np.int32)
+    event_ndims_ = self._maybe_get_event_ndims_statically(event_ndims)
     assertions = []
     if event_ndims_ is not None:
       if min_event_ndims > event_ndims_:
@@ -1051,21 +1058,15 @@ class Bijector(object):
           check_ops.assert_greater_equal(event_ndims, min_event_ndims)]
     return assertions
 
+  def _maybe_get_event_ndims_statically(self, event_ndims):
+    """Helper which returns tries to return an integer static value."""
+    event_ndims_ = distribution_util.maybe_get_static_value(event_ndims)
 
-def get_static_value(x, dtype=None):
-  """Helper which returns static value; casting when dtype is preferred."""
-  if x is None:
-    return x
-  try:
-    x_ = tensor_util.constant_value(x)
-  except TypeError:
-    x_ = x
-  if x_ is None or dtype is None:
-    return x_
-  return np.array(x_, dtype)
-
+    if isinstance(event_ndims_, np.ndarray):
+      if (event_ndims_.dtype not in (np.int32, np.int64) or
+          len(event_ndims_.shape)):
+        raise ValueError("Expected a scalar integer, got {}".format(
+            event_ndims_))
+      event_ndims_ = event_ndims_.tolist()
 
-def assert_static(x):
-  """Helper which asserts that input arg is known statically."""
-  if x is None or type(x) != type(get_static_value(x)):  # pylint: disable=unidiomatic-typecheck
-    raise TypeError("Input must be known statically.")
+    return event_ndims_
diff --git a/tensorflow/python/ops/distributions/util.py b/tensorflow/python/ops/distributions/util.py
index 2e067eab45..3afa85fda0 100644
--- a/tensorflow/python/ops/distributions/util.py
+++ b/tensorflow/python/ops/distributions/util.py
@@ -162,6 +162,30 @@ def same_dynamic_shape(a, b):
       lambda: constant_op.constant(False))
 
 
+def maybe_get_static_value(x, dtype=None):
+  """Helper which tries to return a static value.
+
+  Given `x`, extract it's value statically, optionally casting to a specific
+  dtype. If this is not possible, None is returned.
+
+  Args:
+    x: `Tensor` for which to extract a value statically.
+    dtype: Optional dtype to cast to.
+
+  Returns:
+    Statically inferred value if possible, otherwise None.
+  """
+  if x is None:
+    return x
+  try:
+    x_ = tensor_util.constant_value(x)
+  except TypeError:
+    x_ = x
+  if x_ is None or dtype is None:
+    return x_
+  return np.array(x_, dtype)
+
+
 def get_logits_and_probs(logits=None,
                          probs=None,
                          multidimensional=False,
-- 
GitLab


From 79b773a4395caf7f0b17ce9ac84a1f34dd277bb9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 8 May 2018 16:23:27 -0700
Subject: [PATCH 388/395] Set size of
 tensorflow/python/keras:normalization_test to medium to avoid flaky timeouts

PiperOrigin-RevId: 195894737
---
 tensorflow/python/keras/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/keras/BUILD b/tensorflow/python/keras/BUILD
index 77db07b86b..523eb67935 100755
--- a/tensorflow/python/keras/BUILD
+++ b/tensorflow/python/keras/BUILD
@@ -563,7 +563,7 @@ py_test(
 
 py_test(
     name = "normalization_test",
-    size = "small",
+    size = "medium",
     srcs = ["_impl/keras/layers/normalization_test.py"],
     srcs_version = "PY2AND3",
     tags = ["notsan"],
-- 
GitLab


From 14d5f219f33b1ab8e0a67b84d97204d046adb91f Mon Sep 17 00:00:00 2001
From: Igor Ganichev <iga@google.com>
Date: Tue, 8 May 2018 16:43:54 -0700
Subject: [PATCH 389/395] Make eager functions runable on TPU

PiperOrigin-RevId: 195897321
---
 tensorflow/compiler/jit/BUILD                 |  24 ++
 .../compiler/jit/create_xla_launch_op.cc      | 207 ++++++++++++++----
 .../compiler/jit/create_xla_launch_op.h       |  35 +++
 .../compiler/jit/create_xla_launch_op_test.cc | 145 ++++++++++++
 .../compiler/jit/kernels/xla_launch_op.cc     |  90 ++++++--
 .../compiler/jit/kernels/xla_launch_op.h      |  51 +++--
 .../compiler/jit/xla_compile_on_demand_op.cc  |   3 +-
 tensorflow/compiler/jit/xla_launch_util.cc    |  18 +-
 tensorflow/compiler/jit/xla_launch_util.h     |  15 +-
 tensorflow/compiler/tests/BUILD               |   4 +
 tensorflow/compiler/tests/eager_test.py       | 112 +++++++++-
 .../python/examples/resnet50/resnet50_test.py |  55 +++--
 tensorflow/python/eager/function.py           | 127 +++++++----
 13 files changed, 722 insertions(+), 164 deletions(-)
 create mode 100644 tensorflow/compiler/jit/create_xla_launch_op.h
 create mode 100644 tensorflow/compiler/jit/create_xla_launch_op_test.cc

diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD
index 07136d6a74..a6b3ce394c 100644
--- a/tensorflow/compiler/jit/BUILD
+++ b/tensorflow/compiler/jit/BUILD
@@ -261,6 +261,7 @@ cc_library(
     name = "create_xla_launch_op",
     srcs = [
         "create_xla_launch_op.cc",
+        "create_xla_launch_op.h",
     ],
     deps = [
         ":common",
@@ -270,6 +271,29 @@ cc_library(
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/memory",
+    ],
+    alwayslink = 1,
+)
+
+tf_cc_test(
+    name = "create_xla_launch_op_test",
+    srcs = [
+        "create_xla_launch_op.h",
+        "create_xla_launch_op_test.cc",
+    ],
+    deps = [
+        ":create_xla_launch_op",
+        "//tensorflow/core:core_cpu_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:session_options",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+        "@com_google_absl//absl/memory",
     ],
 )
 
diff --git a/tensorflow/compiler/jit/create_xla_launch_op.cc b/tensorflow/compiler/jit/create_xla_launch_op.cc
index 18d901323f..f35e916eb9 100644
--- a/tensorflow/compiler/jit/create_xla_launch_op.cc
+++ b/tensorflow/compiler/jit/create_xla_launch_op.cc
@@ -12,7 +12,9 @@ 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.
 ==============================================================================*/
+#include "tensorflow/compiler/jit/create_xla_launch_op.h"
 
+#include "absl/memory/memory.h"
 #include "tensorflow/compiler/jit/defs.h"
 #include "tensorflow/compiler/jit/kernels/xla_launch_op.h"
 #include "tensorflow/compiler/jit/mark_for_compilation_pass.h"
@@ -25,78 +27,189 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
-// Givens a NodeDef 'ndef' and the function library runtime 'flr', if
-// 'ndef' is a call to a compilable function defined in 'flr', returns OK
-// and fills in 'kernel' with a XlaLaunchOp kernel which computes the
-// node. Otherwise, returns a non-OK.
+// Utility which searches for values in a sorted list by scanning over it once.
+// No matter how many times ScanForValue is called, the list is scanned at most
+// once. However, if a call to ScanForValue skips over a value, that value is
+// not revisited in future calls to ScanForValue, so callers must take
+// care to order their calls.
 //
-// This routine is here so that FunctionLibraryRuntime can jit a
-// specific function call as requested.
-Status CreateXlaLaunchOp(FunctionLibraryRuntime* flr, const NodeDef& ndef,
-                         std::unique_ptr<OpKernel>* kernel) {
-  bool xla_compile = false;
-  if (!flr->GetFunctionLibraryDefinition()
-           ->GetAttr(ndef, kXlaCompileAttr, &xla_compile)
-           .ok() ||
-      !xla_compile) {
-    // Not marked as _XlaCompile=true.
-    return errors::InvalidArgument("No ", kXlaCompileAttr, " for ", ndef.op());
+// Useful for merging multiple sorted lists in O(n) time.
+class SinglePassSearch {
+ public:
+  // Creates a SinglePassSearch object that can be used to search in `values`.
+  // Does not take ownership of `values`. `values` must outlive this.
+  // `values` must be sorted.
+  explicit SinglePassSearch(const std::vector<int>* values)
+      : current_index_(0), values_(values) {}
+
+  // Scans forward in the vector looking for "value", updating the internal
+  // position in to the vector.
+  // Returns true iff the vector contains the given value at or after current
+  // position.
+  // Not thread-safe.
+  bool ScanForValue(int value) {
+    while (current_index_ < values_->size() &&
+           (*values_)[current_index_] <= value) {
+      if ((*values_)[current_index_] == value) {
+        current_index_++;
+        return true;
+      }
+      current_index_++;
+    }
+    return false;
   }
-  // Make sure that kernels have been registered on the JIT device.
-  XlaOpRegistry::RegisterCompilationKernels();
-  if (!IsCompilable(flr, ndef)) {
-    // ndef is calling a function that XLA can't compile.
-    return errors::InvalidArgument("Not compilable: ", ndef.ShortDebugString());
+
+ private:
+  int current_index_;
+  const std::vector<int>* values_;
+};
+
+Status CompilationRequested(const FunctionLibraryRuntime& flr,
+                            const NodeDef& node_def) {
+  bool xla_compile = false;
+  // Check if op is marked _XlaCompile=true.
+  Status status = flr.GetFunctionLibraryDefinition()->GetAttr(
+      node_def, kXlaCompileAttr, &xla_compile);
+  if (!status.ok() || !xla_compile) {
+    if (VLOG_IS_ON(3)) {
+      if (!status.ok()) {
+        VLOG(3) << "No " << kXlaCompileAttr << " attr defined for "
+                << node_def.op() << ". status=" << status.ToString();
+      } else {
+        VLOG(3) << node_def.op() << " is explicitly marked not to be compiled";
+      }
+    }
+    return Status(error::INVALID_ARGUMENT, "");
   }
+  return Status::OK();
+}
+
+// Given a FunctionLibraryRuntime and a NodeDef calling a function in the
+// runtime, returns this function's body in `fbody` as well as the indices
+// of its constant and resource arguments.
+// `fbody` is owned by `flr`.
+// `constant_arg_indices` and `resource_arg_indices` should be empty vector.
+// They are sorted in ascending order on this function's return.
+Status GetBodyAndConstantsAndResources(FunctionLibraryRuntime* flr,
+                                       const NodeDef& node_def,
+                                       const FunctionBody** fbody,
+                                       std::vector<int>* constant_arg_indices,
+                                       std::vector<int>* resource_arg_indices) {
   FunctionLibraryRuntime::Handle handle;
-  // If ndef is not instantiable, e.g., the function does not exist,
+  // If node_def is not instantiable, e.g., the function does not exist,
   // simply bail out.
   TF_RETURN_IF_ERROR(
-      flr->Instantiate(ndef.op(), AttrSlice(&ndef.attr()), &handle));
-  const FunctionBody* fbody = flr->GetFunctionBody(handle);
-  CHECK(fbody);  // Can't be nullptr since we just instantiated it.
-  std::vector<bool> const_args(fbody->arg_types.size());
+      flr->Instantiate(node_def.op(), AttrSlice(&node_def.attr()), &handle));
+  *fbody = flr->GetFunctionBody(handle);
+  CHECK(*fbody);  // Can't be nullptr since we just instantiated it.
+  const DataTypeVector& arg_types = (*fbody)->arg_types;
+  std::vector<bool> const_args(arg_types.size());
   // If we can't analyze the const args. Bail out.
-  TF_RETURN_IF_ERROR(BackwardsConstAnalysis(*(fbody->graph), &const_args));
+  TF_RETURN_IF_ERROR(BackwardsConstAnalysis(*((*fbody)->graph), &const_args));
 
   for (int i = 0; i < const_args.size(); ++i) {
     if (const_args[i]) {
-      // There is a const arg. Bail out.
-      return errors::InvalidArgument("Const arg: ", i, " in ",
-                                     DebugString(fbody->fdef));
+      constant_arg_indices->push_back(i);
+    }
+  }
+
+  // There can be hundreds of resource variables. Reserve the space for them.
+  // We don't reserve for constants above as they are usually few.
+  resource_arg_indices->reserve(arg_types.size());
+  for (int i = 0; i < arg_types.size(); ++i) {
+    if (arg_types[i] == DT_RESOURCE) {
+      resource_arg_indices->push_back(i);
     }
   }
 
-  NodeDef launch_def;
-  launch_def.set_name(ndef.name());
-  launch_def.set_op("_XlaLaunch");
-  launch_def.set_device(flr->device()->name());
-  AddNodeAttr("Tconstants", DataTypeVector{}, &launch_def);
-  AddNodeAttr("Nresources", 0, &launch_def);
-  AddNodeAttr("Targs", fbody->arg_types, &launch_def);
-  AddNodeAttr("Tresults", fbody->ret_types, &launch_def);
-  NameAttrList func;
-  func.set_name(ndef.op());
-  *(func.mutable_attr()) = ndef.attr();
-  AddNodeAttr("function", func, &launch_def);
-
-  // TODO(b/32387911): Handles the host memory types across function
-  // calls properly. For now, we assume all inputs and outputs are on
-  // the device memory.
+  return Status::OK();
+}
+
+}  // namespace
+
+Status CreateXlaLaunchOp(FunctionLibraryRuntime* flr, const NodeDef& node_def,
+                         std::unique_ptr<OpKernel>* kernel) {
+  TF_RETURN_IF_ERROR(CompilationRequested(*flr, node_def));
+
+  VLOG(3) << "Creating XlaLaunchOp for " << node_def.DebugString();
+
+  // Make sure that kernels have been registered on the JIT device.
+  XlaOpRegistry::RegisterCompilationKernels();
+  if (!IsCompilable(flr, node_def)) {
+    // node_def is calling a function that XLA can't compile.
+    return errors::InvalidArgument("Not compilable: ",
+                                   node_def.ShortDebugString());
+  }
+
+  // Get function body, constant args, and resource args.
+  const FunctionBody* fbody = nullptr;
+  std::vector<int> constant_arg_indices;
+  std::vector<int> resource_arg_indices;
+  TF_RETURN_IF_ERROR(GetBodyAndConstantsAndResources(
+      flr, node_def, &fbody, &constant_arg_indices, &resource_arg_indices));
+
+  // Set input and output memory types.
   MemoryTypeVector input_memory_types(fbody->arg_types.size(), DEVICE_MEMORY);
+  // These indices are used only for optimization purposes. They allow us
+  // to loop over constant_arg_indices and resource_arg_indices only once
+  // while iterating over all the function arguments checking if it is a
+  // resource or a constant.
+  // The reason we optimized this code is because functions can have a lot of
+  // captured arguments. For example, the backward pass of ResNet50 takes in all
+  // 214 variables and a similar number of activations.
+  SinglePassSearch constants_search(&constant_arg_indices);
+  SinglePassSearch resources_search(&resource_arg_indices);
+  for (int i = 0; i < fbody->arg_types.size(); ++i) {
+    if (resources_search.ScanForValue(i) || constants_search.ScanForValue(i)) {
+      // Compile-time constants and resource handles are expected to be in
+      // host memory.
+      input_memory_types[i] = HOST_MEMORY;
+    }
+  }
+  // One might wonder, about the case where a compile-time constant argument
+  // (which must be in host memory) is also used as an input into an op,
+  // e.g. Add, that expects its inputs in device memory. Here is how it
+  // works now.
+  // First, what do we mean by "op expects an input in XYZ memory"?
+  // There are two types of "ops" here: the tf2xla kernel and the HLO
+  // computation it builds. The tf2xla kernel needs to retrieve the actual
+  // numeric value of the compile-time constant tensors, so it really expects
+  // them to be on in host memory. However, for other inputs, it refers to them
+  // using xla::ComputationDataHandle, which is just a symbolic handle that
+  // xla::ComputationBuilder assigns. How does this handle gets assigned for
+  // constant arguments? Even constant arguments get an _Arg node in the graph
+  // instatiated for Function compilation. The tf2xla kernel for constant _Arg
+  // nodes takes the constant value, converts it to XlaLiteral, and feeds it
+  // to xla::ComputationBuilder.ConstantLiteral, which returns the handle. This
+  // constant XlaLiteral is included in the HLO graph, and subsequently, in
+  // the actual executable, which is copied to the device before being
+  // executed. Thus, when this executable runs, the constant is available in
+  // device memory.
+
+  // XlaLaunch kernel keeps all outputs (including constants, which it copies),
+  // in device memory
   MemoryTypeVector output_memory_types(fbody->ret_types.size(), DEVICE_MEMORY);
 
+  // Create the kernel.
+  NameAttrList function;
+  function.set_name(node_def.op());
+  *(function.mutable_attr()) = node_def.attr();
+
   Device* dev = flr->device();
   Status s;
   OpKernelConstruction construction(
       DeviceType(dev->device_type()), dev,
-      dev->GetAllocator(AllocatorAttributes()), &launch_def,
+      dev->GetAllocator(AllocatorAttributes()), &node_def,
       &fbody->fdef.signature(), flr, fbody->arg_types, input_memory_types,
       fbody->ret_types, output_memory_types, flr->graph_def_version(), &s);
-  kernel->reset(new XlaLocalLaunchOp(&construction));
+
+  *kernel = absl::make_unique<XlaLocalLaunchBase>(
+      &construction, constant_arg_indices, resource_arg_indices, function);
   return s;
 }
 
+namespace {
+
 bool RegisterLaunchOpCreator() {
   RegisterDefaultCustomKernelCreator(CreateXlaLaunchOp);
   return true;
diff --git a/tensorflow/compiler/jit/create_xla_launch_op.h b/tensorflow/compiler/jit/create_xla_launch_op.h
new file mode 100644
index 0000000000..98a22e3515
--- /dev/null
+++ b/tensorflow/compiler/jit/create_xla_launch_op.h
@@ -0,0 +1,35 @@
+/* Copyright 2018 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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_JIT_CREATE_XLA_LAUNCH_OP_H_
+#define TENSORFLOW_COMPILER_JIT_CREATE_XLA_LAUNCH_OP_H_
+
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+
+namespace tensorflow {
+
+class FunctionLibraryRuntime;
+class OpKernel;
+
+// Given a NodeDef 'node_def' and the function library runtime 'flr', if
+// 'node_def' is a call to a compilable function defined in 'flr', returns OK
+// and fills in 'kernel' with a XlaLaunchOp kernel which computes the
+// node. Otherwise, returns a non-OK.
+Status CreateXlaLaunchOp(FunctionLibraryRuntime* flr, const NodeDef& node_def,
+                         std::unique_ptr<OpKernel>* kernel);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_JIT_CREATE_XLA_LAUNCH_OP_H_
diff --git a/tensorflow/compiler/jit/create_xla_launch_op_test.cc b/tensorflow/compiler/jit/create_xla_launch_op_test.cc
new file mode 100644
index 0000000000..bcd5e75c7e
--- /dev/null
+++ b/tensorflow/compiler/jit/create_xla_launch_op_test.cc
@@ -0,0 +1,145 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/jit/create_xla_launch_op.h"
+
+#include "absl/memory/memory.h"
+#include "tensorflow/core/common_runtime/device_factory.h"
+#include "tensorflow/core/common_runtime/function.h"
+#include "tensorflow/core/framework/function_testlib.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/public/session_options.h"
+#include "tensorflow/core/public/version.h"
+
+namespace tensorflow {
+
+NodeDef ToNodeDef(const string& text) {
+  NodeDef node_def;
+  EXPECT_TRUE(protobuf::TextFormat::MergeFromString(text, &node_def));
+  return node_def;
+}
+
+// Create a FunctionDef that takes one resource and one regular param
+FunctionDef XTimesY() {
+  return FunctionDefHelper::Define(
+      // Name
+      "XTimesY",
+      // Args
+      {"x: float", "y: resource"},
+      // Return values
+      {"z: float"},
+      // Attr def
+      {},
+      // Nodes
+      {
+          {{"y0"}, "ReadVariableOp", {"y"}, {{"dtype", DT_FLOAT}}},
+          {{"z"}, "Mul", {"x", "y0"}, {{"T", DT_FLOAT}}},
+      });
+}
+
+class CreateXlaLaunchOpTest : public ::testing::Test {
+ protected:
+  void Init(const std::vector<FunctionDef>& flib) {
+    SessionOptions options;
+    auto* device_count = options.config.mutable_device_count();
+    device_count->insert({"CPU", 1});
+    TF_CHECK_OK(DeviceFactory::AddDevices(
+        options, "/job:localhost/replica:0/task:0", &devices_));
+
+    FunctionDefLibrary proto;
+    for (const auto& fdef : flib) {
+      *(proto.add_function()) = fdef;
+    }
+    lib_def_ = absl::make_unique<FunctionLibraryDefinition>(
+        OpRegistry::Global(), proto);
+    OptimizerOptions opts;
+    device_mgr_ = absl::make_unique<DeviceMgr>(devices_);
+    pflr_ = absl::make_unique<ProcessFunctionLibraryRuntime>(
+        device_mgr_.get(), Env::Default(), TF_GRAPH_DEF_VERSION, lib_def_.get(),
+        opts, /*default_thread_pool=*/nullptr, /*cluster_flr=*/nullptr);
+    flr_ = pflr_->GetFLR("/job:localhost/replica:0/task:0/cpu:0");
+  }
+
+  FunctionLibraryRuntime* flr_;
+  std::vector<Device*> devices_;
+  std::unique_ptr<DeviceMgr> device_mgr_;
+  std::unique_ptr<FunctionLibraryDefinition> lib_def_;
+  std::unique_ptr<ProcessFunctionLibraryRuntime> pflr_;
+
+  std::unique_ptr<OpKernel> kernel_;
+};
+
+AttrValue BoolAttr(bool b) {
+  AttrValue v;
+  v.set_b(b);
+  return v;
+}
+
+TEST_F(CreateXlaLaunchOpTest, OneFloatOneResourceArgument) {
+  FunctionDef fdef = XTimesY();
+  (*fdef.mutable_attr())["_XlaCompile"] = BoolAttr(true);
+  Init({fdef});
+
+  Status status = CreateXlaLaunchOp(
+      flr_, ToNodeDef(R"pb(
+        name: 'XTimesY' op: 'XTimesY' input: 'a' input: 'b'
+      )pb"), &kernel_);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+
+  EXPECT_EQ("XTimesY", kernel_->name());
+  EXPECT_EQ("XTimesY", kernel_->type_string());
+
+  EXPECT_EQ(2, kernel_->num_inputs());
+  EXPECT_EQ(DT_FLOAT, kernel_->input_type(0));
+  EXPECT_EQ(DT_RESOURCE, kernel_->input_type(1));
+  EXPECT_EQ(DEVICE_MEMORY, kernel_->input_memory_types()[0]);
+  EXPECT_EQ(HOST_MEMORY, kernel_->input_memory_types()[1]);
+
+  EXPECT_EQ(1, kernel_->num_outputs());
+  EXPECT_EQ(DT_FLOAT, kernel_->output_type(0));
+  EXPECT_EQ(DEVICE_MEMORY, kernel_->output_memory_types()[0]);
+}
+
+TEST_F(CreateXlaLaunchOpTest, FailsIfXlaCompileAttrNotSet) {
+  FunctionDef fdef = XTimesY();
+  Init({fdef});
+
+  Status status = CreateXlaLaunchOp(flr_, ToNodeDef(R"proto(
+                                      name: 'XTimesY'
+                                      op: 'XTimesY'
+                                      input: 'a'
+                                      input: 'b'
+                                    )proto"), &kernel_);
+  EXPECT_TRUE(errors::IsInvalidArgument(status)) << status.ToString();
+}
+
+TEST_F(CreateXlaLaunchOpTest, FailsIfXlaCompileAttrIsSetToFalse) {
+  FunctionDef fdef = XTimesY();
+  (*fdef.mutable_attr())["_XlaCompile"] = BoolAttr(false);
+  Init({fdef});
+
+  Status status = CreateXlaLaunchOp(flr_, ToNodeDef(R"proto(
+                                      name: 'XTimesY'
+                                      op: 'XTimesY'
+                                      input: 'a'
+                                      input: 'b'
+                                    )proto"), &kernel_);
+  EXPECT_TRUE(errors::IsInvalidArgument(status)) << status.ToString();
+}
+
+}  // namespace tensorflow
diff --git a/tensorflow/compiler/jit/kernels/xla_launch_op.cc b/tensorflow/compiler/jit/kernels/xla_launch_op.cc
index 049d170fa4..86a9fd3b8e 100644
--- a/tensorflow/compiler/jit/kernels/xla_launch_op.cc
+++ b/tensorflow/compiler/jit/kernels/xla_launch_op.cc
@@ -39,15 +39,15 @@ limitations under the License.
 
 namespace tensorflow {
 
-XlaLocalLaunchOp::XlaLocalLaunchOp(OpKernelConstruction* ctx)
-    : OpKernel(ctx), device_type_(ctx->device_type()) {
-  const NameAttrList* func;
-  OP_REQUIRES_OK(ctx, ctx->GetAttr("function", &func));
-  function_ = *func;
-  DataTypeVector constant_types;
-  OP_REQUIRES_OK(ctx, ctx->GetAttr("Tconstants", &constant_types));
-  num_constant_args_ = constant_types.size();
-  OP_REQUIRES_OK(ctx, ctx->GetAttr("Nresources", &num_resource_args_));
+XlaLocalLaunchBase::XlaLocalLaunchBase(OpKernelConstruction* ctx,
+                                       const std::vector<int>& constants,
+                                       const std::vector<int>& resources,
+                                       const NameAttrList& function)
+    : OpKernel(ctx),
+      constants_(constants),
+      resources_(resources),
+      device_type_(ctx->device_type()),
+      function_(function) {
   if (device_type_ == DeviceType(DEVICE_CPU)) {
     platform_id_ = se::host::kHostPlatformId;
   } else if (device_type_ == DeviceType(DEVICE_GPU)) {
@@ -57,8 +57,8 @@ XlaLocalLaunchOp::XlaLocalLaunchOp(OpKernelConstruction* ctx)
   }
 }
 
-Status XlaLocalLaunchOp::BuildCompilationCache(OpKernelContext* ctx,
-                                               XlaCompilationCache** cache) {
+Status XlaLocalLaunchBase::BuildCompilationCache(OpKernelContext* ctx,
+                                                 XlaCompilationCache** cache) {
   const XlaDevice::Metadata* metadata;
   Status s = XlaDevice::GetMetadata(ctx, &metadata);
   if (s.ok()) {
@@ -90,8 +90,8 @@ Status XlaLocalLaunchOp::BuildCompilationCache(OpKernelContext* ctx,
   return Status::OK();
 }
 
-void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) {
-  VLOG(1) << "XlaLocalLaunchOp::Compute "
+void XlaLocalLaunchBase::Compute(OpKernelContext* ctx) {
+  VLOG(1) << "XlaLocalLaunchOpBase::Compute "
           << Canonicalize(function_.name(), AttrSlice(&function_.attr()));
   // We store information about the JIT-compiled XLA computation
   // in the ResourceMgr.
@@ -124,7 +124,7 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) {
   }
 
   std::map<int, OptionalTensor> variables =
-      SnapshotResourceVariables(ctx, num_resource_args_);
+      SnapshotResourceVariables(ctx, resources_);
 
   xla::LocalClient* client = static_cast<xla::LocalClient*>(cache->client());
 
@@ -161,7 +161,7 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) {
   xla::LocalExecutable* executable;
 
   std::map<int, Tensor> constant_args;
-  for (int i = 0; i < num_constant_args_; ++i) {
+  for (int i : constants_) {
     constant_args.insert({i, ctx->input(i)});
   }
   OP_REQUIRES_OK(ctx, cache->Compile(options, function_, constant_args,
@@ -170,8 +170,8 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) {
 
   VLOG(1) << "Executing XLA Computation...";
 
-  XlaComputationLaunchContext launch_context(
-      num_resource_args_, client, xla_allocator, allocate_xla_tensors);
+  XlaComputationLaunchContext launch_context(client, xla_allocator,
+                                             allocate_xla_tensors);
   launch_context.PopulateInputs(ctx, kernel, variables);
 
   // Execute the computation.
@@ -194,6 +194,62 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) {
   VLOG(1) << "Done";
 }
 
+namespace {
+
+// OP_REQUIRES_OK_RETURN is the same as OP_REQUIRES_OK except that
+// in error case, it returns RET instead of void.
+#define OP_REQUIRES_OK_RETURN(CTX, RET, ...)                \
+  do {                                                      \
+    ::tensorflow::Status _s(__VA_ARGS__);                   \
+    if (!TF_PREDICT_TRUE(_s.ok())) {                        \
+      (CTX)->CtxFailureWithWarning(__FILE__, __LINE__, _s); \
+      return RET;                                           \
+    }                                                       \
+  } while (0)
+
+// Helper static functions to construct parameters for
+// XlaLocalLaunchBase constructor from OpKernelConstruction.
+std::vector<int> ConstantsVector(OpKernelConstruction* ctx) {
+  DataTypeVector constant_types;
+  OP_REQUIRES_OK_RETURN(ctx, std::vector<int>(),
+                        ctx->GetAttr("Tconstants", &constant_types));
+  std::vector<int> constants(constant_types.size());
+  std::iota(constants.begin(), constants.end(), 0);
+  return constants;
+}
+
+std::vector<int> ResourcesVector(OpKernelConstruction* ctx) {
+  DataTypeVector constant_types;
+  OP_REQUIRES_OK_RETURN(ctx, std::vector<int>(),
+                        ctx->GetAttr("Tconstants", &constant_types));
+
+  DataTypeVector arg_types;
+  OP_REQUIRES_OK_RETURN(ctx, std::vector<int>(),
+                        ctx->GetAttr("Targs", &arg_types));
+
+  int num_resources;
+  OP_REQUIRES_OK_RETURN(ctx, std::vector<int>(),
+                        ctx->GetAttr("Nresources", &num_resources));
+
+  std::vector<int> resources(num_resources);
+  std::iota(resources.begin(), resources.end(),
+            constant_types.size() + arg_types.size());
+  return resources;
+}
+
+NameAttrList FunctionAttr(OpKernelConstruction* ctx) {
+  const NameAttrList* func;
+  OP_REQUIRES_OK_RETURN(ctx, NameAttrList(), ctx->GetAttr("function", &func));
+  return *func;
+}
+
+#undef OP_REQUIRES_OK_RETURN
+}  // namespace
+
+XlaLocalLaunchOp::XlaLocalLaunchOp(OpKernelConstruction* ctx)
+    : XlaLocalLaunchBase(ctx, ConstantsVector(ctx), ResourcesVector(ctx),
+                         FunctionAttr(ctx)) {}
+
 XlaLocalLaunchOp::~XlaLocalLaunchOp() {
   VLOG(1) << "XlaLocalLaunchOp destroyed";
 }
diff --git a/tensorflow/compiler/jit/kernels/xla_launch_op.h b/tensorflow/compiler/jit/kernels/xla_launch_op.h
index 8f8e646f0f..8dfc4b382d 100644
--- a/tensorflow/compiler/jit/kernels/xla_launch_op.h
+++ b/tensorflow/compiler/jit/kernels/xla_launch_op.h
@@ -26,6 +26,41 @@ limitations under the License.
 
 namespace tensorflow {
 
+// XlaLocalLaunchBase is almost the same as XlaLocalLaunchOp.
+// The only difference is that it does not require arguments to follow
+// the "constants, then regular args, then resources" order.
+// It takes vectors of constant and resource arguments explicitly.
+// It does not have corresponding OpDef because it is never present
+// in the GraphDef.
+// Currently, it is used by eager runtime. FunctionLibraryRuntime creates
+// this kernel when asked to create a kernel for an XLA-compiled function.
+class XlaLocalLaunchBase : public OpKernel {
+ public:
+  XlaLocalLaunchBase(OpKernelConstruction* ctx,
+                     const std::vector<int>& constants,
+                     const std::vector<int>& resources,
+                     const NameAttrList& function);
+  XlaLocalLaunchBase(const XlaLocalLaunchBase&) = delete;
+  XlaLocalLaunchBase& operator=(const XlaLocalLaunchBase&) = delete;
+  ~XlaLocalLaunchBase() override = default;
+
+  void Compute(OpKernelContext* ctx) override;
+
+ protected:
+  // Builds a XlaCompilationCache class suitable for the current device.
+  Status BuildCompilationCache(OpKernelContext* ctx,
+                               XlaCompilationCache** cache);
+
+  // Indexes of compile-time constant inputs
+  std::vector<int> constants_;
+  // Indexes of resource inputs
+  std::vector<int> resources_;
+
+  DeviceType device_type_;
+  NameAttrList function_;
+  se::Platform::Id platform_id_;
+};
+
 // XlaLocalLaunchOp is used to replace a region of the TensorFlow graph
 // which will be compiled and executed using XLA.  The XlaLocalLaunchOp is
 // responsible for handling interactions with the TensorFlow executor.
@@ -35,26 +70,12 @@ namespace tensorflow {
 // XlaLocalLaunchOp uses xla::LocalClient::Compile() and
 // xla::LocalExecutable::Run(), and passes arguments into/out of XLA in device
 // memory.
-class XlaLocalLaunchOp : public OpKernel {
+class XlaLocalLaunchOp : public XlaLocalLaunchBase {
  public:
   explicit XlaLocalLaunchOp(OpKernelConstruction* ctx);
   ~XlaLocalLaunchOp() override;
 
-  void Compute(OpKernelContext* ctx) override;
-
  private:
-  // Builds a XlaCompilationCache class suitable for the current device.
-  Status BuildCompilationCache(OpKernelContext* ctx,
-                               XlaCompilationCache** compiler);
-
-  DeviceType device_type_;
-  NameAttrList function_;
-  int num_constant_args_;
-  // Number of resource variable arguments.
-  int num_resource_args_;
-
-  se::Platform::Id platform_id_;
-
   TF_DISALLOW_COPY_AND_ASSIGN(XlaLocalLaunchOp);
 };
 
diff --git a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc
index 60458f6f33..6b83cf67ff 100644
--- a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc
+++ b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc
@@ -48,13 +48,12 @@ Status XlaCompileOnDemandOp::Run(OpKernelContext* ctx,
                                  const XlaCompiler::CompilationResult* result,
                                  xla::LocalExecutable* executable) {
   std::map<int, OptionalTensor> variables = GetVariables(ctx);
-  int64 num_resource_args = variables.size();
 
   xla::LocalClient* client = metadata.client();
 
   // Builds an XLA allocator for the device.
   XlaComputationLaunchContext launch_context(
-      num_resource_args, client, client->backend().memory_allocator(), true);
+      client, client->backend().memory_allocator(), true);
 
   launch_context.PopulateInputs(ctx, result, variables);
 
diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc
index 33e53612b9..0223f97a03 100644
--- a/tensorflow/compiler/jit/xla_launch_util.cc
+++ b/tensorflow/compiler/jit/xla_launch_util.cc
@@ -38,14 +38,13 @@ using xla::ScopedShapedBuffer;
 using xla::ShapedBuffer;
 }  // anonymous namespace
 
-std::map<int, OptionalTensor> SnapshotResourceVariables(OpKernelContext* ctx,
-                                                        int num_variables) {
+std::map<int, OptionalTensor> SnapshotResourceVariables(
+    OpKernelContext* ctx, const std::vector<int>& variables) {
   std::map<int, OptionalTensor> snapshot;
-  int first_variable = ctx->num_inputs() - num_variables;
-  for (int i = 0; i < num_variables; ++i) {
+  for (int i : variables) {
     Var* variable = nullptr;
-    ResourceHandle handle = HandleFromInput(ctx, first_variable + i);
-    OptionalTensor& tensor = snapshot[first_variable + i];
+    ResourceHandle handle = HandleFromInput(ctx, i);
+    OptionalTensor& tensor = snapshot[i];
     if (LookupResource(ctx, handle, &variable).ok()) {
       tf_shared_lock lock(*variable->mu());
       tensor.name = handle.name();
@@ -112,10 +111,9 @@ ScopedShapedBuffer ExtractSubShapedBuffer(
 using internal::ExtractSubShapedBuffer;
 
 XlaComputationLaunchContext::XlaComputationLaunchContext(
-    int64 num_resource_args, xla::LocalClient* client,
-    xla::DeviceMemoryAllocator* xla_allocator, bool allocate_xla_tensors)
-    : num_resource_args_(num_resource_args),
-      client_(client),
+    xla::LocalClient* client, xla::DeviceMemoryAllocator* xla_allocator,
+    bool allocate_xla_tensors)
+    : client_(client),
       xla_allocator_(xla_allocator),
       allocate_xla_tensors_(allocate_xla_tensors) {}
 
diff --git a/tensorflow/compiler/jit/xla_launch_util.h b/tensorflow/compiler/jit/xla_launch_util.h
index 38291b0bd4..a2431253f8 100644
--- a/tensorflow/compiler/jit/xla_launch_util.h
+++ b/tensorflow/compiler/jit/xla_launch_util.h
@@ -31,15 +31,17 @@ limitations under the License.
 namespace tensorflow {
 class XlaAllocator;
 
-// Takes a snapshot of the values of resource variable arguments, which are
-// the last `num_variables` arguments. We snapshot tensors that back
+// Takes a snapshot of the values of resource variable arguments, whose
+// indices are specified in `variables` argument. We snapshot tensors that back
 // resource variables since concurrent updates may modify the shape, and it is
 // important that the shapes used for compilation match the true shapes of the
 // buffers.
 //
-// Returns a map of TensorFlow argument index to resource variable.
-std::map<int, OptionalTensor> SnapshotResourceVariables(OpKernelContext* ctx,
-                                                        int num_variables);
+// Returns a map of TensorFlow argument index to resource variable. If a
+// resource variable is not initialized, the corresponding OptionalTensor
+// will have its `present` field set to false.
+std::map<int, OptionalTensor> SnapshotResourceVariables(
+    OpKernelContext* ctx, const std::vector<int>& variables);
 
 // Adapter class that wraps a Tensorflow allocator as an XLA allocator.
 // Assumes that the Tensorflow allocator permits asynchronous deallocation:
@@ -72,7 +74,7 @@ class XlaComputationLaunchContext {
   // Create a new launch context. 'allocate_xla_tensors' is true if allocated
   // output tensors and variables are always XlaTensors. If false they are
   // assumed to be "normal" device pointers.
-  XlaComputationLaunchContext(int64 num_resource_args, xla::LocalClient* client,
+  XlaComputationLaunchContext(xla::LocalClient* client,
                               xla::DeviceMemoryAllocator* xla_allocator,
                               bool allocate_xla_tensors);
 
@@ -92,7 +94,6 @@ class XlaComputationLaunchContext {
   const std::vector<xla::ShapedBuffer*>& arguments() const { return arg_ptrs_; }
 
  private:
-  int64 num_resource_args_;
   xla::LocalClient* client_;
   xla::DeviceMemoryAllocator* xla_allocator_;
   bool allocate_xla_tensors_;
diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index aaea83ae9c..9791792f29 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -327,7 +327,11 @@ tf_xla_py_test(
         ":xla_test",
         "//tensorflow/python:array_ops",
         "//tensorflow/python:framework_for_generated_wrappers",
+        "//tensorflow/python:layers",
+        "//tensorflow/python:math_ops",
+        "//tensorflow/python:nn",
         "//tensorflow/python:platform_test",
+        "//tensorflow/python/eager:function",
     ],
 )
 
diff --git a/tensorflow/compiler/tests/eager_test.py b/tensorflow/compiler/tests/eager_test.py
index bdd0185dfe..5ab1585f8c 100644
--- a/tensorflow/compiler/tests/eager_test.py
+++ b/tensorflow/compiler/tests/eager_test.py
@@ -24,10 +24,16 @@ from tensorflow.compiler.tests.xla_test import XLATestCase
 from tensorflow.core.protobuf import config_pb2
 from tensorflow.python.eager import backprop
 from tensorflow.python.eager import context
+from tensorflow.python.eager import function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
+from tensorflow.python.layers import convolutional
+from tensorflow.python.layers import pooling
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
 from tensorflow.python.ops import resource_variable_ops
 from tensorflow.python.platform import googletest
 
@@ -43,7 +49,7 @@ class EagerTest(XLATestCase):
 
   def testExecuteListOutputLen0(self):
     with self.test_scope():
-      empty = constant_op.constant([], dtype=dtypes.int32)
+      empty = constant_op.constant([], dtype=dtypes.float32)
       result = array_ops.unstack(empty, 0)
       self.assertTrue(isinstance(result, list))
       self.assertEqual(0, len(result))
@@ -51,7 +57,7 @@ class EagerTest(XLATestCase):
   def testExecuteListOutputLen1(self):
     with self.test_scope():
       split_dim = constant_op.constant(1)
-      value = constant_op.constant([[0, 1, 2], [3, 4, 5]])
+      value = constant_op.constant([[0., 1., 2.], [3., 4., 5.]])
       result = array_ops.split(value, 1, axis=split_dim)
       self.assertTrue(isinstance(result, list))
       self.assertEqual(1, len(result))
@@ -60,7 +66,7 @@ class EagerTest(XLATestCase):
   def testExecuteListOutputLen3(self):
     with self.test_scope():
       split_dim = constant_op.constant(1)
-      value = constant_op.constant([[0, 1, 2], [3, 4, 5]])
+      value = constant_op.constant([[0., 1., 2.], [3., 4., 5.]])
       result = array_ops.split(value, 3, axis=split_dim)
       self.assertTrue(isinstance(result, list))
       self.assertEqual(3, len(result))
@@ -131,7 +137,105 @@ class EagerTest(XLATestCase):
     self.assertEqual(2., grads[0][0].numpy())
 
 
-if __name__ == "__main__":
+class EagerFunctionTest(XLATestCase):
+
+  def testBasic(self):
+    with self.test_scope():
+      matmul = function.defun(math_ops.matmul, compiled=True)
+      t = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
+      sq = matmul(t, t, transpose_a=True)
+      self.assertAllEqual(sq.numpy().reshape(-1), [10, 14, 14, 20])
+
+  def testConv(self):
+    if 'GPU' in self.device:
+      # TODO(b/32333178)
+      self.skipTest('Current implementation of RandomStandardNormal kernel '
+                    'is very slow on GPU, and has been blacklisted.')
+    with self.test_scope():
+      data_format = 'channels_last'
+      conv = convolutional.Conv2D(
+          filters=1, kernel_size=2, padding='VALID',
+          data_format=data_format, activation=nn_ops.relu,
+          kernel_initializer=init_ops.ones_initializer(),
+          bias_initializer=init_ops.zeros_initializer())
+      pool = pooling.MaxPooling2D(2, 2, data_format=data_format)
+
+      def model(x):
+        x = conv(x)
+        return pool(x)
+      model = function.defun(model, compiled=True)
+
+      x = array_ops.ones([1, 4, 4, 1])
+      y = model(x)
+      self.assertAllEqual(y.numpy(), [[[[4.]]]])
+
+  def testReadVariable(self):
+    with self.test_scope():
+      v = resource_variable_ops.ResourceVariable(1.0)
+
+      @function.defun(compiled=True)
+      def f():
+        return v.read_value()
+
+      var = f()
+      self.assertEqual(1.0, var.numpy())
+
+  def testUpdateVariable(self):
+    with self.test_scope():
+      v = resource_variable_ops.ResourceVariable(1.0)
+
+      def f(v):
+        v.assign_add(1.0)
+        return v
+
+      f = function.defun(f, compiled=True)
+
+      var = f(v)
+      self.assertEqual(2.0, var.numpy())
+
+  def testAllArgumentKinds(self):
+    """Test a complex function that takes different argument kinds.
+
+    tf2xla machinery that translates, compiles, and runs defuns
+    classifies arguments into: compile-time constants, regular tensors,
+    and resources. This test creates a function with a mix of all these
+    kinds. Moreover, the order of function arguments is intentionally mixed up.
+
+    This also tests the case when the same argument is a compile-time constant
+    as well as used in an operation that normally expects its inputs to be
+    in device memory - addition in this case.
+    """
+    with self.test_scope():
+      def foo(c1, r1, v1, c2, v2, r2):
+        # c1 and c2 are compile-time constants
+        # r1 and r2 are regular tensors
+        # v1 and v2 are resource variables
+        a = c1 + r1
+        b = math_ops.cast(c2, dtypes.float32) + v2
+        c = array_ops.slice(v1, c1, c2)
+        d = r2 * v2
+        return a, b, c, d
+
+      foo = function.defun(foo, compiled=True)
+
+      c1 = [0, 0]
+      c2 = array_ops.ones([2], dtype=dtypes.int32)
+
+      r1 = array_ops.ones([2])
+      r2 = [[2., 2.], [3., 3.]]
+
+      v1 = resource_variable_ops.ResourceVariable([[1., 2.], [3., 4.]])
+      v2 = resource_variable_ops.ResourceVariable([[10., 20.], [30., 40.]])
+
+      a, b, c, d = foo(c1, r1, v1, c2, v2, r2)
+
+      self.assertAllEqual([1, 1], a.numpy())
+      self.assertAllEqual([[11., 21.], [31., 41.]], b.numpy())
+      self.assertAllEqual([[1.]], c.numpy())
+      self.assertAllEqual([[20., 40.], [90., 120.]], d.numpy())
+
+
+if __name__ == '__main__':
   ops.enable_eager_execution(
       config=config_pb2.ConfigProto(log_device_placement=True))
   googletest.main()
diff --git a/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py b/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py
index 8517a3bf7b..b8f352d5f5 100644
--- a/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py
+++ b/tensorflow/contrib/eager/python/examples/resnet50/resnet50_test.py
@@ -36,9 +36,7 @@ def device_and_data_format():
                                                               'channels_last')
 
 
-def random_batch(batch_size, device_and_format=None):
-  _, data_format = device_and_format or device_and_data_format()
-
+def random_batch(batch_size, data_format):
   shape = (3, 224, 224) if data_format == 'channels_first' else (224, 224, 3)
   shape = (batch_size,) + shape
 
@@ -70,7 +68,7 @@ class ResNet50Test(tf.test.TestCase):
     if defun:
       model.call = tfe.defun(model.call)
     with tf.device(device), tfe.execution_mode(execution_mode):
-      images, _ = random_batch(2)
+      images, _ = random_batch(2, data_format)
       output = model(images, training=False)
       tfe.async_wait()
     self.assertEqual((2, 1000), output.shape)
@@ -91,7 +89,7 @@ class ResNet50Test(tf.test.TestCase):
     device, data_format = device_and_data_format()
     model = resnet50.ResNet50(data_format, include_top=False)
     with tf.device(device):
-      images, _ = random_batch(2)
+      images, _ = random_batch(2, data_format)
       output = model(images, training=False)
     output_shape = ((2, 2048, 1, 1)
                     if data_format == 'channels_first' else (2, 1, 1, 2048))
@@ -101,7 +99,7 @@ class ResNet50Test(tf.test.TestCase):
     device, data_format = device_and_data_format()
     model = resnet50.ResNet50(data_format, include_top=False, pooling='avg')
     with tf.device(device):
-      images, _ = random_batch(2)
+      images, _ = random_batch(2, data_format)
       output = model(images, training=False)
     self.assertEqual((2, 2048), output.shape)
 
@@ -115,7 +113,7 @@ class ResNet50Test(tf.test.TestCase):
         name='t0').as_default(), tf.contrib.summary.always_record_summaries():
       with tf.device(device), tfe.execution_mode(execution_mode):
         optimizer = tf.train.GradientDescentOptimizer(0.1)
-        images, labels = random_batch(2)
+        images, labels = random_batch(2, data_format)
         train_one_step(model, images, labels, optimizer)
         self.assertEqual(320, len(model.variables))
         tfe.async_wait()
@@ -134,7 +132,7 @@ class ResNet50Test(tf.test.TestCase):
     model = resnet50.ResNet50(data_format)
     optimizer = tf.train.GradientDescentOptimizer(0.1)
     with tf.device(device):
-      images, labels = random_batch(2)
+      images, labels = random_batch(2, data_format)
       gc.disable()
       # Warm up. Note that this first run does create significant amounts of
       # garbage to be collected. The hope is that this is a build-only effect,
@@ -202,18 +200,18 @@ class ResNet50Benchmarks(tf.test.Benchmark):
     # which forces a sync. This is a roundabout way, yes.
     tf.constant(1.).cpu()
 
-  def _benchmark_eager_apply(self, label, defun=False, execution_mode=None,
-                             device_and_format=None):
+  def _benchmark_eager_apply(self, label, device_and_format, defun=False,
+                             execution_mode=None, compiled=False):
     with tfe.execution_mode(execution_mode):
-      device, data_format = device_and_format or device_and_data_format()
+      device, data_format = device_and_format
       model = resnet50.ResNet50(data_format)
       if defun:
-        model.call = tfe.defun(model.call)
+        model.call = tfe.defun(model.call, compiled=compiled)
       batch_size = 64
       num_burn = 5
       num_iters = 30
       with tf.device(device):
-        images, _ = random_batch(batch_size, device_and_format)
+        images, _ = random_batch(batch_size, data_format)
         for _ in xrange(num_burn):
           model(images, training=False).cpu()
         if execution_mode:
@@ -227,30 +225,34 @@ class ResNet50Benchmarks(tf.test.Benchmark):
         self._report(label, start, num_iters, device, batch_size, data_format)
 
   def benchmark_eager_apply_sync(self):
-    self._benchmark_eager_apply('eager_apply', defun=False)
+    self._benchmark_eager_apply('eager_apply', device_and_data_format(),
+                                defun=False)
 
   def benchmark_eager_apply_async(self):
     self._benchmark_eager_apply(
-        'eager_apply_async', defun=False, execution_mode=tfe.ASYNC)
+        'eager_apply_async', device_and_data_format(), defun=False,
+        execution_mode=tfe.ASYNC)
 
   def benchmark_eager_apply_with_defun(self):
-    self._benchmark_eager_apply('eager_apply_with_defun', defun=True)
+    self._benchmark_eager_apply('eager_apply_with_defun',
+                                device_and_data_format(), defun=True)
 
   def _benchmark_eager_train(self,
                              label,
                              make_iterator,
+                             device_and_format,
                              defun=False,
                              execution_mode=None,
-                             device_and_format=None):
+                             compiled=False):
     with tfe.execution_mode(execution_mode):
-      device, data_format = device_and_format or device_and_data_format()
+      device, data_format = device_and_format
       for batch_size in self._train_batch_sizes():
-        (images, labels) = random_batch(batch_size, device_and_format)
+        (images, labels) = random_batch(batch_size, data_format)
         num_burn = 3
         num_iters = 10
         model = resnet50.ResNet50(data_format)
         if defun:
-          model.call = tfe.defun(model.call)
+          model.call = tfe.defun(model.call, compiled=compiled)
         optimizer = tf.train.GradientDescentOptimizer(0.1)
 
         with tf.device(device):
@@ -273,18 +275,21 @@ class ResNet50Benchmarks(tf.test.Benchmark):
           self._report(label, start, num_iters, device, batch_size, data_format)
 
   def benchmark_eager_train_sync(self):
-    self._benchmark_eager_train('eager_train', MockIterator, defun=False)
+    self._benchmark_eager_train('eager_train', MockIterator,
+                                device_and_data_format(), defun=False)
 
   def benchmark_eager_train_async(self):
     self._benchmark_eager_train(
         'eager_train_async',
         MockIterator,
+        device_and_data_format(),
         defun=False,
         execution_mode=tfe.ASYNC)
 
   def benchmark_eager_train_with_defun(self):
     self._benchmark_eager_train(
-        'eager_train_with_defun', MockIterator, defun=True)
+        'eager_train_with_defun', MockIterator,
+        device_and_data_format(), defun=True)
 
   def benchmark_eager_train_datasets(self):
 
@@ -294,7 +299,8 @@ class ResNet50Benchmarks(tf.test.Benchmark):
       return tfe.Iterator(ds)
 
     self._benchmark_eager_train(
-        'eager_train_dataset', make_iterator, defun=False)
+        'eager_train_dataset', make_iterator,
+        device_and_data_format(), defun=False)
 
   def benchmark_eager_train_datasets_with_defun(self):
 
@@ -304,7 +310,8 @@ class ResNet50Benchmarks(tf.test.Benchmark):
       return tfe.Iterator(ds)
 
     self._benchmark_eager_train(
-        'eager_train_dataset_with_defun', make_iterator, defun=True)
+        'eager_train_dataset_with_defun', make_iterator,
+        device_and_data_format(), defun=True)
 
 
 if __name__ == '__main__':
diff --git a/tensorflow/python/eager/function.py b/tensorflow/python/eager/function.py
index 89257bb20a..b478b6b0db 100644
--- a/tensorflow/python/eager/function.py
+++ b/tensorflow/python/eager/function.py
@@ -23,6 +23,7 @@ import collections
 
 import numpy as np
 
+from tensorflow.core.framework import attr_value_pb2
 from tensorflow.core.framework import function_pb2
 from tensorflow.python import pywrap_tensorflow
 from tensorflow.python.eager import context
@@ -227,7 +228,7 @@ def _inference_name(n):
 class _EagerDefinedFunction(object):
   """Function object with the interface of tf _DefinedFunction."""
 
-  def __init__(self, name, graph, operations, inputs, outputs):
+  def __init__(self, name, graph, operations, inputs, outputs, attrs):
     """Initializes an eager defined function.
 
     Args:
@@ -237,6 +238,7 @@ class _EagerDefinedFunction(object):
         which will be in the function
       inputs: the tensors in the graph to be used as inputs to the function
       outputs: the tensors in the graph which will be outputs to the function
+      attrs: dict mapping names of attributes to their AttrValue values
     """
     fn = pywrap_tensorflow.TF_GraphToFunction_wrapper(
         graph._c_graph,  # pylint: disable=protected-access
@@ -248,6 +250,14 @@ class _EagerDefinedFunction(object):
         [],
         None,
         compat.as_str(""))
+
+    for name, attr_value in attrs.items():
+      serialized = attr_value.SerializeToString()
+      # TODO(iga): this creates and deletes a new TF_Status for every attr.
+      # It might be worth creating a convenient way to re-use status.
+      pywrap_tensorflow.TF_FunctionSetAttrValueProto(
+          fn, compat.as_str(name), serialized)
+
     # TODO(apassos) avoid creating a FunctionDef (specially to grab the
     # signature, but also in general it's nice not to depend on it.
     with c_api_util.tf_buffer() as buffer_:
@@ -289,25 +299,6 @@ def _flatten(sequence):
 
 class GraphModeFunction(object):
   """Callable object representing a graph-mode function.
-
-  Args:
-    name: str the name of the created function
-    input_placeholders: list of placeholder values (tensors) to feed when
-      calling the wrapped function.
-    extra_inputs: Tensor inputs this function definition closed over which
-      are passed as arguments. Need to track so gradients are supported
-      correctly.
-    graph: the Graph from which the operations will be pulled. Used as
-      a context when computing gradients.
-    operations: the subset of Operations in the graph used in the function
-      definition.
-    outputs: a flat list of the Tensors in the graph used as outputs to the
-      function
-    func_outputs: a possibly nested python object which will be returned by
-      this function. The Tensors in this structure will be replaced by their
-      corresponding values in outputs.
-    output_shapes: List of shapes of all tensors in outputs
-    variables: (optional) List of variables to watch during function execution.
   """
 
   def __init__(self,
@@ -319,9 +310,36 @@ class GraphModeFunction(object):
                outputs,
                func_outputs,
                output_shapes,
-               variables=None):
+               variables=None,
+               attrs=None):
+    """Initialize a GraphModeFunction.
+
+    Args:
+      name: str the name of the created function
+      input_placeholders: list of placeholder values (tensors) to feed when
+        calling the wrapped function.
+      extra_inputs: Tensor inputs this function definition closed over which
+        are passed as arguments. Need to track so gradients are supported
+        correctly.
+      graph: the Graph from which the operations will be pulled. Used as
+        a context when computing gradients.
+      operations: the subset of Operations in the graph used in the function
+        definition.
+      outputs: a flat list of the Tensors in the graph used as outputs to the
+        function
+      func_outputs: a possibly nested python object which will be returned by
+        this function. The Tensors in this structure will be replaced by their
+        corresponding values in outputs.
+      output_shapes: List of shapes of all tensors in outputs
+      variables: (optional) List of variables to watch during function
+        execution.
+      attrs: (optional) dict mapping names of attributes to their AttrValue
+        values. Attributes in `attrs` will be included in this function's
+        definition.
+    """
+    self._attrs = attrs or {}
     defined_function = _EagerDefinedFunction(
-        name, graph, operations, input_placeholders, outputs)
+        name, graph, operations, input_placeholders, outputs, self._attrs)
     if len(input_placeholders) != len(defined_function.signature.input_arg):
       raise ValueError("Internal error: invalid lengths. %s %s" % (
           len(input_placeholders), len(defined_function.signature.input_arg)))
@@ -374,7 +392,7 @@ class GraphModeFunction(object):
     forward_name = _forward_name(self._func_name)
     self._forward_fdef = _EagerDefinedFunction(
         forward_name, self._graph, self._ops, self._input_placeholders,
-        filtered_outputs + captures)
+        filtered_outputs + captures, self._attrs)
     all_inputs = self._out_grad_placeholders + captures
     # Excluding input ops from the body as we do not intend to execute these
     # operations when the function is executed.
@@ -388,7 +406,7 @@ class GraphModeFunction(object):
     bname = _backward_name(self._func_name)
     self._backward_function = GraphModeFunction(
         bname, all_inputs, [], self._graph, function_def_ops,
-        backward_outputs, in_gradients, output_shapes)
+        backward_outputs, in_gradients, output_shapes, attrs=self._attrs)
 
   def _backprop_call(self, args):
     """Calls the wrapped function and records the result on a tape."""
@@ -562,7 +580,7 @@ def _get_defun_inputs(args):
   return nest.pack_sequence_as(args, ret)
 
 
-def _defun_internal(name, func, args, kwds):
+def _defun_internal(name, func, compiled, args, kwds):
   """Defines and returns graph-mode version of func."""
   graph_key = ops.get_default_graph()._graph_key  # pylint: disable=protected-access
   with context.graph_mode():
@@ -627,9 +645,14 @@ def _defun_internal(name, func, args, kwds):
     for f in tmp_graph._functions.values():  # pylint: disable=protected-access
       # TODO(ashankar): What about the gradient registry?
       _register(f._c_func.func)  # pylint: disable=protected-access
+
+  attrs = {}
+  if compiled:
+    attrs["_XlaCompile"] = attr_value_pb2.AttrValue(b=True)
+
   return GraphModeFunction(
       fname, all_inputs, extra_inputs, tmp_graph, operations, func_def_outputs,
-      func_outputs, output_shapes, variables)
+      func_outputs, output_shapes, variables, attrs)
 
 
 # Defun uses this instead of Tensor as a cache key. Using dtype because
@@ -671,7 +694,7 @@ def _register(fn):
 
 
 # TODO(apassos): better error messages for non-hashable arguments.
-def named_defun(func, name):
+def named_defun(func, name, compiled=False):
   """Defines a function with a given name.
 
   See the documentation for `defun` for more information on the semantics of the
@@ -680,6 +703,7 @@ def named_defun(func, name):
   Args:
     func: the function to be wrapped.
     name: the name given to it.
+    compiled: if true, the framework will attempt to compile func with XLA.
 
   Returns:
     the wrapped function.
@@ -696,13 +720,13 @@ def named_defun(func, name):
 
     if cache_key not in arguments_to_functions:
       arguments_to_functions[cache_key] = _defun_internal(
-          name, func, args, kwds)
+          name, func, compiled, args, kwds)
     return arguments_to_functions[cache_key](*args)
 
   return decorated
 
 
-def defun(func):
+def defun(func=None, compiled=False):
   """Decorator to compile func into graph_mode.
 
   `defun` converts a function that constructs a TensorFlow graph into a function
@@ -745,18 +769,45 @@ def defun(func):
   ```
 
   Args:
-    func: function to be compiled.
+    func: function to be compiled. If `func` is None, returns a
+      decorator that can be invoked with a single argument - `func`. The
+      end result is equivalent to providing all the arguments up front.
+      In other words, defun(compiled=True)(func) is equivalent to
+      defun(func, compiled=True). The former allows the following use case:
+        @tfe.defun(compiled=True)
+        def foo(...):
+          ...
+    compiled: If True, an attempt to compile `func` with XLA will be made.
+      If it fails, function will be run normally. Experimental.
+      Currently, supported only for execution on TPUs.
 
   Returns:
-     A callable that will execute the compiled function (and return zero
-     or more `tf.Tensor` objects).
+     If `func` is not None, returns callable that will execute the compiled
+     function (and return zero or more `tf.Tensor` objects).
+     If `func` is None, returns a decorator that, when invoked with a single
+     `func` argument, returns a callable equivalent to the case above.
   """
   # TODO(apassos): deal with captured global state. Deal with control flow.
-  try:
-    name = func.__name__
-  except AttributeError:
-    name = "function"
-  return tf_decorator.make_decorator(func, named_defun(func, name))
+  def decorated(function):
+    try:
+      name = function.__name__
+    except AttributeError:
+      name = "function"
+    return tf_decorator.make_decorator(
+        function, named_defun(function, name, compiled=compiled))
+
+  # This code path is for the `foo = tfe.defun(foo, ...)` use case
+  if func is not None:
+    return decorated(func)
+
+  # This code path is for the
+  #
+  # @tfe.defun(...)
+  # def foo(...):
+  #    ...
+  #
+  # use case, which is equivalent to `foo = tfe.defun(...)(foo)`
+  return decorated
 
 
 def make_defun_op(func, *args, **kwds):
@@ -808,7 +859,7 @@ def make_defun_op(func, *args, **kwds):
   name = func.__name__
   if any(isinstance(x, ops.EagerTensor) for x in kwds.values()):
     raise ValueError("Tensor keyword arguments are not supported.")
-  return _defun_internal(name, func, args, kwds)
+  return _defun_internal(name, func, False, args, kwds)
 
 
 class AutomaticControlDependencies(object):
-- 
GitLab


From d14a530533a049bb4096d1789c626f7c3f3e1d83 Mon Sep 17 00:00:00 2001
From: Yifei Feng <yifeif@google.com>
Date: Tue, 8 May 2018 17:16:45 -0700
Subject: [PATCH 390/395] Hardcode EndpointSpec deprecated input to False for
 now after cl/195718061.

---
 tensorflow/java/src/gen/cc/op_specs.cc | 3 +--
 tensorflow/java/src/gen/cc/op_specs.h  | 6 +++---
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensorflow/java/src/gen/cc/op_specs.cc b/tensorflow/java/src/gen/cc/op_specs.cc
index 081062ceaf..4bcfc7fe01 100644
--- a/tensorflow/java/src/gen/cc/op_specs.cc
+++ b/tensorflow/java/src/gen/cc/op_specs.cc
@@ -382,8 +382,7 @@ EndpointSpec CreateEndpoint(const OpDef& op_def, const ApiDef& api_def,
   return EndpointSpec(package,
       name,
       Javadoc::Create(ParseDocumentation(api_def.summary()))
-          .details(ParseDocumentation(api_def.description())),
-      endpoint_def.deprecation_version() > 0);
+          .details(ParseDocumentation(api_def.description())));
 }
 
 }  // namespace
diff --git a/tensorflow/java/src/gen/cc/op_specs.h b/tensorflow/java/src/gen/cc/op_specs.h
index 81582ea207..034cf636ed 100644
--- a/tensorflow/java/src/gen/cc/op_specs.h
+++ b/tensorflow/java/src/gen/cc/op_specs.h
@@ -34,11 +34,11 @@ class EndpointSpec {
   // package: package of this endpoint (from which also derives its package)
   // name: name of this endpoint class
   // javadoc: the endpoint class documentation
-  // deprecated: true if this endpoint is now deprecated
+  // TODO(annarev): hardcode depcreated to false until deprecated is possible
   EndpointSpec(const string& package, const string& name,
-      const Javadoc& javadoc, bool deprecated)
+      const Javadoc& javadoc)
     : package_(package), name_(name), javadoc_(javadoc),
-      deprecated_(deprecated) {}
+      deprecated_(false) {}
 
   const string& package() const { return package_; }
   const string& name() const { return name_; }
-- 
GitLab


From 1f03f829285ca0fbd47a99350e9f5d99aa10e9b9 Mon Sep 17 00:00:00 2001
From: Yifei Feng <1192265+yifeif@users.noreply.github.com>
Date: Tue, 8 May 2018 17:35:21 -0700
Subject: [PATCH 391/395] Switch to use str instead of number for colab_url

Fix nightly failure:
  File "tensorflow/tools/ci_build/update_version.py", line 253, in colab_url
    version_string = "%d.%d.%d" % (version.major, version.minor, version.patch)
TypeError: %d format: a number is required, not str
---
 tensorflow/tools/ci_build/update_version.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/tools/ci_build/update_version.py b/tensorflow/tools/ci_build/update_version.py
index 9ddb219048..00bfcfd49b 100755
--- a/tensorflow/tools/ci_build/update_version.py
+++ b/tensorflow/tools/ci_build/update_version.py
@@ -250,7 +250,7 @@ def update_md_files(old_version, new_version):
 
   # Update any links to colab notebooks.
   def colab_url(version):
-    version_string = "%d.%d.%d" % (version.major, version.minor, version.patch)
+    version_string = "%s.%s.%s" % (version.major, version.minor, version.patch)
     prefix = "https://colab.research.google.com/github/tensorflow/models/blob/r"
     return prefix + version_string + "/"
 
-- 
GitLab


From c317afd07eb11abe416080cdced9ec00198dbbb0 Mon Sep 17 00:00:00 2001
From: Yong Tang <yong.tang.github@outlook.com>
Date: Wed, 9 May 2018 08:55:19 -0700
Subject: [PATCH 392/395] Enable test case for float64 with conv1d (#19179)

The float64 for conv2d support has been added to tensorflow
in e3468b56d323783fdfb79fa2d6c24effc58bcaa9. (Thanks brianwa84!)
Since conv1d implementation invokes conv2d, the float64 support
for conv1d is supported now as well.

This fix adds the test case for float64 support of conv1d and
removes the TODO.

This fix fixes 19175.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
---
 tensorflow/python/kernel_tests/conv1d_test.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensorflow/python/kernel_tests/conv1d_test.py b/tensorflow/python/kernel_tests/conv1d_test.py
index e2e6205911..fcba456004 100644
--- a/tensorflow/python/kernel_tests/conv1d_test.py
+++ b/tensorflow/python/kernel_tests/conv1d_test.py
@@ -31,9 +31,7 @@ class Conv1DTest(test.TestCase):
 
   def testBasic(self):
     """Test that argument passing to conv1d is handled properly."""
-    # TODO(yongtang): dtypes.float64 can only be enabled once conv2d support
-    # dtypes.float64, as conv1d implicitly calls conv2d after expand_dims.
-    for dtype in [dtypes.float16, dtypes.float32]:
+    for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]:
       x = constant_op.constant([1, 2, 3, 4], dtype=dtype)
       x = array_ops.expand_dims(x, 0)  # Add batch dimension
       x = array_ops.expand_dims(x, 2)  # And depth dimension
-- 
GitLab


From 76e8a4ec287c11d5b1286244d1821994640dbecf Mon Sep 17 00:00:00 2001
From: ctiijima <ctiijima@us.ibm.com>
Date: Wed, 9 May 2018 09:50:48 -0700
Subject: [PATCH 393/395] Grammar fixes for Programmers guide FAQ (#19170)

---
 tensorflow/docs_src/programmers_guide/faq.md | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/tensorflow/docs_src/programmers_guide/faq.md b/tensorflow/docs_src/programmers_guide/faq.md
index 51c1a1e032..b6291a9ffa 100644
--- a/tensorflow/docs_src/programmers_guide/faq.md
+++ b/tensorflow/docs_src/programmers_guide/faq.md
@@ -72,7 +72,7 @@ tensors in the execution of a step.
 
 If `t` is a @{tf.Tensor} object,
 @{tf.Tensor.eval} is shorthand for
-@{tf.Session.run} (where `sess` is the
+@{tf.Session.run}, where `sess` is the
 current @{tf.get_default_session}. The
 two following snippets of code are equivalent:
 
@@ -101,9 +101,8 @@ sessions, it may be more straightforward to make explicit calls to
 Sessions can own resources, such as
 @{tf.Variable},
 @{tf.QueueBase}, and
-@{tf.ReaderBase}; and these resources can use
-a significant amount of memory. These resources (and the associated memory) are
-released when the session is closed, by calling
+@{tf.ReaderBase}. These resources can sometimes use
+a significant amount of memory, and can be released when the session is closed by calling
 @{tf.Session.close}.
 
 The intermediate tensors that are created as part of a call to
@@ -137,7 +136,7 @@ TensorFlow also has a
 to help build support for more client languages.  We invite contributions of new
 language bindings.
 
-Bindings for various other languages (such as [C#](https://github.com/migueldeicaza/TensorFlowSharp), [Julia](https://github.com/malmaud/TensorFlow.jl), [Ruby](https://github.com/somaticio/tensorflow.rb) and [Scala](https://github.com/eaplatanios/tensorflow_scala)) created and supported by the opensource community build on top of the C API supported by the TensorFlow maintainers.
+Bindings for various other languages (such as [C#](https://github.com/migueldeicaza/TensorFlowSharp), [Julia](https://github.com/malmaud/TensorFlow.jl), [Ruby](https://github.com/somaticio/tensorflow.rb) and [Scala](https://github.com/eaplatanios/tensorflow_scala)) created and supported by the open source community build on top of the C API supported by the TensorFlow maintainers.
 
 #### Does TensorFlow make use of all the devices (GPUs and CPUs) available on my machine?
 
@@ -210,8 +209,8 @@ a new tensor with a different dynamic shape.
 
 #### How do I build a graph that works with variable batch sizes?
 
-It is often useful to build a graph that works with variable batch sizes, for
-example so that the same code can be used for (mini-)batch training, and
+It is often useful to build a graph that works with variable batch sizes 
+so that the same code can be used for (mini-)batch training, and
 single-instance inference. The resulting graph can be
 @{tf.Graph.as_graph_def$saved as a protocol buffer}
 and
@@ -260,7 +259,7 @@ See the how-to documentation for
 There are three main options for dealing with data in a custom format.
 
 The easiest option is to write parsing code in Python that transforms the data
-into a numpy array. Then use @{tf.data.Dataset.from_tensor_slices} to
+into a numpy array. Then, use @{tf.data.Dataset.from_tensor_slices} to
 create an input pipeline from the in-memory data.
 
 If your data doesn't fit in memory, try doing the parsing in the Dataset
@@ -274,7 +273,7 @@ If your data is not easily parsable with the built-in TensorFlow operations,
 consider converting it, offline, to a format that is easily parsable, such
 as @{tf.python_io.TFRecordWriter$`TFRecord`} format.
 
-The more efficient method to customize the parsing behavior is to
+The most efficient method to customize the parsing behavior is to
 @{$adding_an_op$add a new op written in C++} that parses your
 data format. The @{$new_data_formats$guide to handling new data formats} has
 more information about the steps for doing this.
-- 
GitLab


From baeb356fbf209bd8ef325704fa9bd22e6f2a0887 Mon Sep 17 00:00:00 2001
From: Letian Feng <letian.feng@hotmail.com>
Date: Wed, 9 May 2018 18:50:57 +0200
Subject: [PATCH 394/395] Fix 2 typos in documents (#19177)

* fix minor typo in doc: tf.layer to tf.layers

* removed a duplicated line
---
 tensorflow/docs_src/programmers_guide/variables.md | 2 +-
 tensorflow/docs_src/tutorials/layers.md            | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensorflow/docs_src/programmers_guide/variables.md b/tensorflow/docs_src/programmers_guide/variables.md
index e8cf771155..cd8c4b5b9a 100644
--- a/tensorflow/docs_src/programmers_guide/variables.md
+++ b/tensorflow/docs_src/programmers_guide/variables.md
@@ -237,7 +237,7 @@ TensorFlow supports two ways of sharing variables:
 While code which explicitly passes variables around is very clear, it is
 sometimes convenient to write TensorFlow functions that implicitly use
 variables in their implementations. Most of the functional layers from
-`tf.layer` use this approach, as well as all `tf.metrics`, and a few other
+`tf.layers` use this approach, as well as all `tf.metrics`, and a few other
 library utilities.
 
 Variable scopes allow you to control variable reuse when calling functions which
diff --git a/tensorflow/docs_src/tutorials/layers.md b/tensorflow/docs_src/tutorials/layers.md
index 37cd2bb139..496b1e4da9 100644
--- a/tensorflow/docs_src/tutorials/layers.md
+++ b/tensorflow/docs_src/tutorials/layers.md
@@ -209,7 +209,6 @@ for two-dimensional image data expect input tensors to have a shape of
 *   _`channels`_. Number of color channels in the example images. For color
     images, the number of channels is 3 (red, green, blue). For monochrome
     images, there is just 1 channel (black).
-*   _`image_height`_. Height of the example images.
 *   _`data_format`_. A string, one of `channels_last` (default) or `channels_first`.
       `channels_last` corresponds to inputs with shape
       `(batch, ..., channels)` while `channels_first` corresponds to
-- 
GitLab


From 4fb125264c5394c9e4295ed437adb1d9711bd456 Mon Sep 17 00:00:00 2001
From: AG Ramesh <ag.ramesh@intel.com>
Date: Wed, 9 May 2018 09:51:10 -0700
Subject: [PATCH 395/395] [INTEL MKL] Fixes a failure in
 //tensorflow/python/profiler:model_analyzer_test. (#19152)

* Modified  testComplexCodeView test

 Modified  testComplexCodeView to look for lower total_float_ops. The value of total_float_ops is lower when Tensorflow is compiled with Intel MKL.

* Added code to check if MKL is enabled

* Fixed Pylint errors
---
 tensorflow/python/profiler/model_analyzer_test.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/profiler/model_analyzer_test.py b/tensorflow/python/profiler/model_analyzer_test.py
index 75580fc630..9e49188c1e 100644
--- a/tensorflow/python/profiler/model_analyzer_test.py
+++ b/tensorflow/python/profiler/model_analyzer_test.py
@@ -232,7 +232,12 @@ class PrintModelAnalysisTest(test.TestCase):
 
         self.assertLess(0, tfprof_node.total_exec_micros)
         self.assertEqual(2844, tfprof_node.total_parameters)
-        self.assertLess(145660, tfprof_node.total_float_ops)
+        #The graph is modifed when MKL is enabled,total_float_ops will
+        #be different
+        if test_util.IsMklEnabled():
+          self.assertLess(101600, tfprof_node.total_float_ops)
+        else:
+          self.assertLess(145660, tfprof_node.total_float_ops)
         self.assertEqual(8, len(tfprof_node.children))
         self.assertEqual('_TFProfRoot', tfprof_node.name)
         self.assertEqual(
-- 
GitLab